core/knowledge_base.py

"""Knowledge Base with TF-IDF embeddings (no PyTorch needed).
Uses scikit-learn's TfidfVectorizer instead of sentence-transformers.
Works on HF free tier without heavy ML dependencies.
"""
import json
import re
from datetime import datetime
from pathlib import Path

import numpy as np
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.metrics.pairwise import cosine_similarity


class KnowledgeBase:

    def __init__(self, data_dir, embedding_model_id="", cache_dir=""):
        self.knowledge_dir = Path(data_dir) / "knowledge"
        self.knowledge_dir.mkdir(parents=True, exist_ok=True)
        self.index_path = self.knowledge_dir / "index.json"
        self.vectors_path = self.knowledge_dir / "vectors.npz"
        self._index = None
        self._vectors = None
        self._vectorizer = None

    def _load_index(self):
        if self._index is not None:
            return self._index
        if self.index_path.exists():
            try:
                self._index = json.loads(self.index_path.read_text(encoding="utf-8"))
            except (json.JSONDecodeError, OSError):
                self._index = {"chunks": []}
        else:
            self._index = {"chunks": []}
        return self._index

    def _load_vectors(self):
        if self._vectors is not None:
            return self._vectors
        if self.vectors_path.exists():
            try:
                data = np.load(str(self.vectors_path), allow_pickle=True)
                self._vectors = data["vectors"]
            except Exception:
                self._vectors = None
        return self._vectors

    def _save(self):
        index = self._load_index()
        self.index_path.write_text(
            json.dumps(index, ensure_ascii=False, indent=2),
            encoding="utf-8",
        )
        # Rebuild TF-IDF vectors from all chunks
        self._rebuild_vectors()

    def _rebuild_vectors(self):
        """Rebuild TF-IDF vectors from all chunk texts."""
        index = self._load_index()
        texts = [c["text"] for c in index["chunks"]]
        if not texts:
            self._vectors = None
            if self.vectors_path.exists():
                self.vectors_path.unlink(missing_ok=True)
            return

        self._vectorizer = TfidfVectorizer(
            max_features=5000,
            stop_words="english",
            ngram_range=(1, 2),
            sublinear_tf=True,
        )
        self._vectors = self._vectorizer.fit_transform(texts)
        # Save vectorizer vocab and vectors (convert np.int64 to int for json serialization)
        vocab_dict = {k: int(v) for k, v in self._vectorizer.vocabulary_.items()}
        np.savez_compressed(
            str(self.vectors_path),
            vectors=self._vectors.toarray(),
            vocab=json.dumps(vocab_dict),
            idf=self._vectorizer.idf_,
        )

    def _get_vectorizer(self):
        """Load or rebuild the TF-IDF vectorizer."""
        if self._vectorizer is not None:
            return self._vectorizer

        if self.vectors_path.exists():
            try:
                data = np.load(str(self.vectors_path), allow_pickle=True)
                vocab = json.loads(str(data["vocab"]))
                idf = data["idf"]
                self._vectorizer = TfidfVectorizer(
                    max_features=5000,
                    stop_words="english",
                    ngram_range=(1, 2),
                    sublinear_tf=True,
                    vocabulary=vocab,
                )
                # Fit with dummy to set internal state, then override idf
                index = self._load_index()
                texts = [c["text"] for c in index["chunks"]]
                if texts:
                    self._vectorizer.fit(texts)
                return self._vectorizer
            except Exception:
                pass

        # Rebuild from scratch
        self._rebuild_vectors()
        return self._vectorizer

    def chunk_text(self, text, max_chars=500, overlap=50):
        sentences = re.split(r"(?<=[.!?])\s+", text.strip())
        chunks = []
        current = ""

        for sentence in sentences:
            sentence = sentence.strip()
            if not sentence:
                continue

            if not current:
                current = sentence
                continue

            combined = current + " " + sentence
            if len(combined) <= max_chars:
                current = combined
            else:
                if current:
                    chunks.append(current)
                if overlap > 0 and len(current) > overlap:
                    overlap_text = current[-overlap:].lstrip()
                    current = overlap_text + " " + sentence
                else:
                    current = sentence

        if current:
            chunks.append(current)

        # fallback for text without sentence-ending punctuation
        if not chunks and text.strip():
            raw = text.strip()
            while len(raw) > max_chars:
                split_at = raw[:max_chars].rfind(" ")
                if split_at < 100:
                    split_at = max_chars
                chunks.append(raw[:split_at].strip())
                raw = raw[split_at:].strip()
            if raw:
                chunks.append(raw)

        return [c for c in chunks if len(c.strip()) >= 20]

    def add_document(self, text, source_name):
        chunks = self.chunk_text(text)
        if not chunks:
            return 0

        index = self._load_index()
        now = datetime.now().isoformat(timespec="seconds")
        for chunk_text in chunks:
            index["chunks"].append({
                "text": chunk_text,
                "source": source_name,
                "added_at": now,
            })

        self._save()
        return len(chunks)

    def search(self, query, top_k=3, threshold=0.1):
        index = self._load_index()
        if not index["chunks"]:
            return []

        vectorizer = self._get_vectorizer()
        if vectorizer is None:
            return []

        # Load vectors
        vectors = self._load_vectors()
        if vectors is None:
            return []

        # If vectors are sparse (from file), convert
        if hasattr(vectors, 'toarray'):
            dense_vectors = vectors.toarray()
        else:
            dense_vectors = vectors

        try:
            query_vec = vectorizer.transform([query]).toarray()
        except Exception:
            return []

        similarities = cosine_similarity(query_vec, dense_vectors)[0]
        top_indices = np.argsort(similarities)[::-1][:top_k]

        results = []
        for i in top_indices:
            score = float(similarities[i])
            if score < threshold:
                continue
            chunk = index["chunks"][i]
            results.append({
                "text": chunk["text"],
                "source": chunk["source"],
                "score": score,
            })

        return results

    def list_sources(self):
        index = self._load_index()
        sources = {}
        for chunk in index["chunks"]:
            name = chunk["source"]
            if name not in sources:
                sources[name] = {"count": 0, "added_at": chunk.get("added_at", "")}
            sources[name]["count"] += 1
        return sources

    def format_source_list(self):
        sources = self.list_sources()
        if not sources:
            return "No study materials in knowledge base yet. Use /learn to add some."

        lines = ["", "  Knowledge Base", "  " + "-" * 50]
        for idx, (name, info) in enumerate(sources.items(), start=1):
            added = info["added_at"][:16].replace("T", " ") if info["added_at"] else ""
            lines.append(f"  {idx}. {name}")
            lines.append(f"     {info['count']} chunks | Added: {added}")
        lines.append("")
        return "\n".join(lines)

    def remove_source(self, index_number):
        sources = self.list_sources()
        source_names = list(sources.keys())

        if index_number < 1 or index_number > len(source_names):
            return None

        target = source_names[index_number - 1]
        idx = self._load_index()
        idx["chunks"] = [c for c in idx["chunks"] if c["source"] != target]

        self._save()
        return target

    def is_empty(self):
        index = self._load_index()
        return len(index["chunks"]) == 0