396 lines
13 KiB
Python
396 lines
13 KiB
Python
"""
|
|
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/<collection>/<filename>.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
|