""" knowledge_base.py — AIPA Knowledge Base Filesystem-first Markdown knowledge base with JIT ChromaDB embeddings for search. Source of truth: Markdown files in KB_PATH//.md ChromaDB cache: KB_PATH/.qwen3_index/ (rebuilt on-demand from filesystem) Embeddings: Qwen3-Embedding-0.6B via sentence-transformers (local, CPU/GPU). - Documents are embedded without a prefix (passage side). - Search queries are prefixed with a retrieval instruction (query side). This asymmetric setup follows the model's recommended usage and improves retrieval quality over symmetric embeddings. File format: # Title Body content... --- *Changelog* - 2026-04-03 10:30 — Agent — Entry message - 2026-04-04 09:15 — Agent — Entry message Both ChromaDB and the embedding model are lazy-initialised on first use. """ from __future__ import annotations import os import re import time from datetime import datetime, timezone from pathlib import Path from typing import Any import config as _config # --------------------------------------------------------------------------- # Internal state # --------------------------------------------------------------------------- _chroma_client: Any = None # chromadb.PersistentClient, lazy-init _embed_model: Any = None # SentenceTransformer, lazy-init # --------------------------------------------------------------------------- # Embedding — Qwen3-Embedding-0.6B # --------------------------------------------------------------------------- # Task instruction prepended to search queries only (asymmetric retrieval). # Documents are embedded without any prefix. _QUERY_INSTRUCTION = ( "Instruct: Given a search query, retrieve relevant knowledge base documents " "that answer the query or contain related information.\nQuery: " ) def _get_embed_model(): """Lazy-load Qwen3-Embedding-0.6B via sentence-transformers.""" global _embed_model if _embed_model is None: try: from sentence_transformers import SentenceTransformer except ImportError: raise RuntimeError( "sentence-transformers is not installed. " "Run: pip install sentence-transformers" ) _embed_model = SentenceTransformer("Qwen/Qwen3-Embedding-0.6B") return _embed_model class _QwenDocumentEF: """ ChromaDB embedding function for the document (indexing) side. No instruction prefix — used when upserting KB documents into the index. """ @staticmethod def name() -> str: return "qwen3-embedding-0.6b" def __call__(self, input: list[str]) -> list[list[float]]: return ( _get_embed_model() .encode(input, normalize_embeddings=True, show_progress_bar=False) .tolist() ) def _embed_query(query: str) -> list[float]: """ Embed a search query with the Qwen3 retrieval instruction prefix. Returns a normalized 1024-dim float list ready to pass to ChromaDB as query_embeddings. """ return ( _get_embed_model() .encode( _QUERY_INSTRUCTION + query, normalize_embeddings=True, show_progress_bar=False, ) .tolist() ) # --------------------------------------------------------------------------- # Path and text helpers # --------------------------------------------------------------------------- # Separate dir from the old all-MiniLM index to avoid dimension mismatch. _CHROMA_DIR = ".qwen3_index" _CHANGELOG_SEP = "\n\n---\n*Changelog*\n" def _safe_filename(filename: str) -> str: """ Normalise a filename or relative subpath within a collection. Allows forward slashes for subdirectory paths (e.g. 'research/topic.md'). Rejects backslashes and any '..' components to prevent traversal. Ensures the path ends with .md. Raises ValueError on invalid input. """ if not filename: raise ValueError("filename must not be empty.") if "\\" in filename: raise ValueError("filename must use forward slashes, not backslashes.") # Reject any path component that is '..' parts = Path(filename).parts if ".." in parts: raise ValueError(f"Path traversal not allowed in filename: {filename!r}") if not filename.endswith(".md"): filename = filename + ".md" return filename def _col_path(collection: str) -> Path: """Absolute path to the collection directory.""" name = collection or _config.KB_DEFAULT_COLLECTION return _config.KB_PATH / name def _split_body_changelog(text: str) -> tuple[str, list[str]]: """ Split a document into (body, changelog_lines). Returns (full_text, []) if no changelog section exists. """ if _CHANGELOG_SEP in text: body, cl_block = text.split(_CHANGELOG_SEP, 1) lines = [ln for ln in cl_block.splitlines() if ln.strip()] return body, lines return text, [] def _build_document(body: str, changelog_lines: list[str]) -> str: """Reassemble body + changelog into the canonical on-disk format.""" if not changelog_lines: return body cl_block = "\n".join(changelog_lines) return body + _CHANGELOG_SEP + cl_block + "\n" def _now_str() -> str: """Current local datetime as 'YYYY-MM-DD HH:MM'.""" return datetime.now().strftime("%Y-%m-%d %H:%M") # --------------------------------------------------------------------------- # ChromaDB helpers # --------------------------------------------------------------------------- def _get_chroma_client(): global _chroma_client if _chroma_client is None: try: import chromadb except ImportError: raise RuntimeError("chromadb is not installed. Run: pip install chromadb") index_path = _config.KB_PATH / _CHROMA_DIR index_path.mkdir(parents=True, exist_ok=True) _chroma_client = chromadb.PersistentClient(path=str(index_path)) return _chroma_client def _get_chroma_collection(collection: str): """Get or create a ChromaDB collection using Qwen3 document embeddings.""" name = collection or _config.KB_DEFAULT_COLLECTION return _get_chroma_client().get_or_create_collection( name=name, metadata={"hnsw:space": "cosine"}, embedding_function=_QwenDocumentEF(), ) def _jit_sync(collection: str) -> None: """ Synchronise the ChromaDB index for a collection with the filesystem. - New or modified .md files are upserted. - Files deleted from disk are removed from the index. - Uses file mtime stored as ChromaDB metadata to skip unchanged files. """ col_dir = _col_path(collection) if not col_dir.exists(): return chroma_col = _get_chroma_collection(collection) # Map id → mtime for all currently indexed documents all_ids_result = chroma_col.get(include=["metadatas"]) indexed: dict[str, float] = {} if all_ids_result and all_ids_result["ids"]: for doc_id, meta in zip(all_ids_result["ids"], all_ids_result["metadatas"]): indexed[doc_id] = float((meta or {}).get("mtime", 0)) # Walk filesystem recursively (supports subdirectories within a collection) on_disk: set[str] = set() for md_file in col_dir.rglob("*.md"): # Use relative path as doc_id to avoid collisions across subdirs doc_id = str(md_file.relative_to(col_dir)) on_disk.add(doc_id) mtime = md_file.stat().st_mtime if indexed.get(doc_id, -1) >= mtime: continue # up to date text = md_file.read_text(encoding="utf-8") body, _ = _split_body_changelog(text) chroma_col.upsert( ids=[doc_id], documents=[body], metadatas=[{"source": str(md_file), "mtime": mtime, "collection": collection or _config.KB_DEFAULT_COLLECTION}], ) # Remove stale entries stale = set(indexed) - on_disk if stale: chroma_col.delete(ids=list(stale)) # --------------------------------------------------------------------------- # Public API # --------------------------------------------------------------------------- def kb_write( filename: str, content: str, collection: str = "", agent: str = "", changelog_entry: str = "", ) -> dict: """ Write or overwrite a Markdown KB document. If the file already exists its changelog is preserved and a new entry appended (unless changelog_entry is empty). If the file is new the changelog is initialised with a "Created document" entry. Returns {"ok": True, "path": str, "collection": str}. """ filename = _safe_filename(filename) col_name = collection or _config.KB_DEFAULT_COLLECTION col_dir = _col_path(col_name) path = col_dir / filename path.parent.mkdir(parents=True, exist_ok=True) # Preserve existing changelog existing_changelog: list[str] = [] if path.exists(): existing_text = path.read_text(encoding="utf-8") _, existing_changelog = _split_body_changelog(existing_text) # Build new changelog entry agent_label = agent or "system" ts = _now_str() if changelog_entry: entry_text = f"- {ts} — {agent_label} — {changelog_entry}" elif not path.exists(): entry_text = f"- {ts} — {agent_label} — Created document" else: entry_text = f"- {ts} — {agent_label} — Updated document" new_changelog = existing_changelog + [entry_text] doc = _build_document(content, new_changelog) path.write_text(doc, encoding="utf-8") return {"ok": True, "path": str(path), "collection": col_name} def kb_read(filename: str, collection: str = "") -> dict: """ Read a KB document. Returns {"filename": str, "body": str, "changelog": list[str], "collection": str}. """ filename = _safe_filename(filename) col_name = collection or _config.KB_DEFAULT_COLLECTION path = _col_path(col_name) / filename if not path.exists(): return {"error": f"Document not found: {filename!r} in collection {col_name!r}"} text = path.read_text(encoding="utf-8") body, changelog = _split_body_changelog(text) return { "filename": filename, "body": body, "changelog": changelog, "collection": col_name, } def kb_search( query: str, n_results: int = 5, collection: str = "", ) -> list[dict]: """ Semantic search within a collection. Syncs the ChromaDB index from disk before querying. Returns results ordered by similarity (closest first, lower distance = more similar). """ col_name = collection or _config.KB_DEFAULT_COLLECTION _jit_sync(col_name) chroma_col = _get_chroma_collection(col_name) count = chroma_col.count() if count == 0: return [] results = chroma_col.query( query_embeddings=[_embed_query(query)], n_results=min(n_results, count), include=["documents", "metadatas", "distances"], ) return [ { "filename": results["ids"][0][i], "snippet": results["documents"][0][i][:500], "collection": col_name, "distance": round(results["distances"][0][i], 4), } for i in range(len(results["ids"][0])) ] def kb_list(collection: str = "") -> list[dict]: """ List all Markdown files in a collection with their title and modification time. Title is extracted from the first `# Heading` line; falls back to filename stem. """ col_name = collection or _config.KB_DEFAULT_COLLECTION col_dir = _col_path(col_name) if not col_dir.exists(): return [] results = [] for md_file in sorted(col_dir.rglob("*.md")): rel = str(md_file.relative_to(col_dir)) text = md_file.read_text(encoding="utf-8") title = md_file.stem # default for line in text.splitlines(): stripped = line.strip() if stripped.startswith("# "): title = stripped[2:].strip() break results.append({ "filename": rel, "title": title, "collection": col_name, "modified": datetime.fromtimestamp(md_file.stat().st_mtime).strftime("%Y-%m-%d %H:%M"), }) return results def kb_delete(filename: str, collection: str = "") -> dict: """ Delete a Markdown KB document. The ChromaDB entry is removed on the next JIT sync. Returns {"ok": True, "deleted": filename, "collection": col_name}. """ filename = _safe_filename(filename) col_name = collection or _config.KB_DEFAULT_COLLECTION path = _col_path(col_name) / filename if not path.exists(): return {"error": f"Document not found: {filename!r} in collection {col_name!r}"} path.unlink() return {"ok": True, "deleted": filename, "collection": col_name} def kb_list_collections() -> list[dict]: """ List all collections (subdirectories of KB_PATH, excluding hidden dirs). Returns name and document count for each. """ if not _config.KB_PATH.exists(): return [] results = [] for entry in sorted(_config.KB_PATH.iterdir()): if not entry.is_dir() or entry.name.startswith("."): continue count = len(list(entry.rglob("*.md"))) results.append({"name": entry.name, "count": count}) return results