"""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