diff --git a/NOTICE b/NOTICE index 8c6e112..a528e87 100644 --- a/NOTICE +++ b/NOTICE @@ -86,3 +86,19 @@ Octop voice https://github.com/TencentCloud/Octop Copyright (c) 2026 Octop Licensed under the MIT License. + +Jonathan Ai document search is adapted from Octop's knowledge base. PDF, +DOCX, Markdown, plain text, and HTML are split into passages and stored in +a per-user SQLite file under the user's home directory. Answers cite the +file name and the passage. Local embeddings use fastembed with the ONNX +build of BAAI/bge-small-en-v1.5 from Hugging Face, and each file is checked +against a pinned checksum before it is used. Hosts other than huggingface.co +are refused. When that model is not on the machine, search uses BM25 and +still cites the file and passage. Document text stays on this machine unless +the user selects the OpenAI embedding provider. Jonathan Ai does not vendor +Octop's OCR, spreadsheet parsers, or its download race against other hosts. + +Octop knowledge base +https://github.com/TencentCloud/Octop +Copyright (c) 2026 Octop +Licensed under the MIT License. diff --git a/pyproject.toml b/pyproject.toml index d8c3b5a..022b6b3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,10 @@ voice = [ "sounddevice>=0.4.6", "piper-tts>=1.2.0", ] +kb = [ + "fastembed>=0.4.2", + "pypdf>=4.2.0", +] [project.scripts] clawd = "src.cli:main" diff --git a/src/cli.py b/src/cli.py index 0068cc9..9c81d2d 100644 --- a/src/cli.py +++ b/src/cli.py @@ -35,6 +35,8 @@ def main(): clawd cron list List local scheduled jobs clawd cron run Run one scheduled job now clawd voice settings Show local speech settings + clawd kb add Index a local PDF, DOCX, MD, TXT, or HTML file + clawd kb search Search local documents and cite the passage clawd Start interactive REPL """ ) @@ -110,6 +112,16 @@ def main(): voice_parser = subparsers.add_parser('voice', help='Local speech, push-to-talk, and devices') voice_parser.add_argument('words', nargs='*', help='settings, devices, stt, tts, ptt, wake, say') + kb_parser = subparsers.add_parser('kb', help='Add, list, remove, or search the local document index') + kb_sub = kb_parser.add_subparsers(dest='kb_command') + kb_add = kb_sub.add_parser('add', help='Index a local document') + kb_add.add_argument('path') + kb_sub.add_parser('list', help='List indexed documents') + kb_remove = kb_sub.add_parser('remove', help='Remove one document by id') + kb_remove.add_argument('doc_id') + kb_search = kb_sub.add_parser('search', help='Search indexed documents') + kb_search.add_argument('query', nargs='+') + args = parser.parse_args() # Handle --version @@ -189,6 +201,8 @@ def main(): return handle_cron(args) elif args.command == 'voice': return handle_voice(args) + elif args.command == 'kb': + return handle_kb(args) # Default: start REPL return start_repl(stream=args.stream) @@ -308,6 +322,27 @@ def handle_voice(args) -> int: return 0 +def handle_kb(args) -> int: + """Add, list, remove, or search documents stored on this machine.""" + from src.knowledge.service import format_kb_command + + action = getattr(args, "kb_command", None) or "list" + if action == "add": + text = format_kb_command("add " + args.path) + elif action == "remove": + text = format_kb_command("remove " + args.doc_id) + elif action == "search": + text = format_kb_command("search " + " ".join(args.query)) + elif action == "list": + text = format_kb_command("list") + else: + text = format_kb_command("") + print(text) + if text.startswith("No document") or text.startswith("document not found"): + return 1 + return 0 + + def start_repl(stream: bool = False): """Start interactive REPL.""" from src.config import get_default_provider diff --git a/src/command_system/__init__.py b/src/command_system/__init__.py index 43bb88b..8d8963b 100644 --- a/src/command_system/__init__.py +++ b/src/command_system/__init__.py @@ -14,6 +14,7 @@ HELP_COMMAND, CRON_COMMAND, INIT_COMMAND, + KB_COMMAND, MEMORY_COMMAND, VOICE_COMMAND, SKILLS_COMMAND, @@ -94,6 +95,7 @@ "CONTEXT_COMMAND", "COMPACT_COMMAND", "CRON_COMMAND", + "KB_COMMAND", "MEMORY_COMMAND", "VOICE_COMMAND", "INIT_COMMAND", diff --git a/src/command_system/builtins.py b/src/command_system/builtins.py index 6ffefc7..09a3d9a 100644 --- a/src/command_system/builtins.py +++ b/src/command_system/builtins.py @@ -652,6 +652,22 @@ def voice_command_call(args: str, context: CommandContext) -> LocalCommandResult return LocalCommandResult(type="text", value=format_voice_command(args)) + +KB_COMMAND = LocalCommand( + name="kb", + description="Add, list, remove, or search local documents", + argument_hint="add | list | remove | search ", + supports_non_interactive=True, +) + + +def kb_command_call(args: str, context: CommandContext) -> LocalCommandResult: + """Search documents indexed on this machine.""" + del context + from ..knowledge.service import format_kb_command + + return LocalCommandResult(type="text", value=format_kb_command(args)) + INIT_COMMAND = PromptCommand( name="init", description="Initialize new CLAUDE.md file(s) and optional skills/hooks with codebase documentation", @@ -702,6 +718,8 @@ def execute_command_sync(cmd_name: str, args: str, context: CommandContext) -> t result = cron_command_call(args, context) elif cmd is VOICE_COMMAND: result = voice_command_call(args, context) + elif cmd is KB_COMMAND: + result = kb_command_call(args, context) else: return False, None, f"Command not implemented for sync execution: {cmd_name}" @@ -721,6 +739,7 @@ def execute_command_sync(cmd_name: str, args: str, context: CommandContext) -> t MEMORY_COMMAND.set_call(memory_command_call) CRON_COMMAND.set_call(cron_command_call) VOICE_COMMAND.set_call(voice_command_call) +KB_COMMAND.set_call(kb_command_call) def get_builtin_commands() -> list[Command]: @@ -736,6 +755,7 @@ def get_builtin_commands() -> list[Command]: MEMORY_COMMAND, CRON_COMMAND, VOICE_COMMAND, + KB_COMMAND, INIT_COMMAND, ] diff --git a/src/desktop/runtime.py b/src/desktop/runtime.py index b94cafe..3b2484e 100644 --- a/src/desktop/runtime.py +++ b/src/desktop/runtime.py @@ -1174,6 +1174,7 @@ def _format_palette(self) -> str: - `/cron` — list, run, enable, disable, or delete local schedules (`/cron run `) - Browser tool — local Chrome over CDP. Logins stay in the profile. Submits, purchases, and logins need approval. - `/voice` — local speech and speakers (`/voice settings`, `/voice say hello`) +- `/kb` — add, list, remove, or search local documents (`/kb search `) Type a message to run the existing agent loop. Destructive and network tools prompt for approval before they run. diff --git a/src/desktop/server.py b/src/desktop/server.py index 5822b3e..d79187b 100644 --- a/src/desktop/server.py +++ b/src/desktop/server.py @@ -243,6 +243,14 @@ def handle_api(self, method: str, path: str, query: dict[str, list[str]], body: from src.voice.service import push_to_talk_action return 200, push_to_talk_action(str(body.get("action") or "")) + if path == "/api/kb" and method == "GET": + from src.knowledge.service import knowledge_status + + return 200, knowledge_status() + if path == "/api/kb" and method == "POST": + from src.knowledge.service import update_knowledge_settings + + return 200, update_knowledge_settings(body) return 404, {"error": f"unknown route: {method} {path}"} def serve_forever(self) -> None: diff --git a/src/desktop/web/app.js b/src/desktop/web/app.js index 788e35c..2a3a65e 100644 --- a/src/desktop/web/app.js +++ b/src/desktop/web/app.js @@ -670,6 +670,16 @@ async function sendMessage(path = "/api/chat") { } } +async function loadKnowledgeSettings() { + const info = await api("/api/kb"); + if ($("kb-backend")) $("kb-backend").value = info.embedding_backend || "local"; + if ($("kb-model")) $("kb-model").value = info.model_id || ""; + const stay = info.documents_stay_local + ? "Documents stay on this machine." + : "Cloud embeddings are selected, so document text is sent when indexing."; + if ($("kb-status")) $("kb-status").textContent = `${stay} ${info.document_count || 0} documents.`; +} + function applyRoute() { const route = currentRoute(); document.body.dataset.route = route; @@ -679,7 +689,10 @@ function applyRoute() { const showSetup = route === "setup" || (route === "chat" && Boolean(state.status?.needs_setup)); $("setup-modal")?.classList.toggle("hidden", !showSetup); if (route === "library") refreshLibrary().catch(reportError); - if (route === "settings") loadVoiceSettings().catch(reportError); + if (route === "settings") { + loadVoiceSettings().catch(reportError); + loadKnowledgeSettings().catch(reportError); + } } async function refreshLibrary() { @@ -980,6 +993,21 @@ function bindUi() { }); $("open-settings").addEventListener("click", () => go("settings")); $("settings-close").addEventListener("click", () => go("chat")); + $("kb-save")?.addEventListener("click", guard(async () => { + const saved = await api("/api/kb", { + method: "POST", + body: JSON.stringify({ + embedding_backend: $("kb-backend").value, + model_id: $("kb-model").value, + openai_api_key: $("kb-key").value, + }), + }); + $("kb-key").value = ""; + const stay = saved.documents_stay_local + ? "Documents stay on this machine." + : "Cloud embeddings are selected, so document text is sent when indexing."; + $("kb-status").textContent = `${stay} ${saved.document_count || 0} documents.`; + })); $("settings-save").addEventListener("click", async () => { try { const provider = $("settings-provider").value; diff --git a/src/desktop/web/index.html b/src/desktop/web/index.html index 1566921..f14e181 100644 --- a/src/desktop/web/index.html +++ b/src/desktop/web/index.html @@ -194,6 +194,23 @@

Voice

+

Knowledge base

+

Documents stay on this machine unless a cloud embedding provider is explicitly chosen.

+ + + +
+ +

Source and updates: https://github.com/GoDeskio/Clawd-Code. Default local tree: ~/Jonathan/Jonathan-Ai.

diff --git a/src/knowledge/__init__.py b/src/knowledge/__init__.py new file mode 100644 index 0000000..ea8e943 --- /dev/null +++ b/src/knowledge/__init__.py @@ -0,0 +1,22 @@ +"""Local document knowledge base for Jonathan Ai. + +Documents are parsed, chunked, and stored in a per-user SQLite file on this +machine. Search uses a local ONNX embedding model when its Hugging Face +checksums match, and BM25 otherwise. Document text is sent to a cloud +embedding provider only when that provider is explicitly selected. + +The chunking and citation shape follow Octop's knowledge base. Jonathan Ai +does not vendor Octop's service, OCR, or mirror download race. + +Source: https://github.com/TencentCloud/Octop +Copyright (c) 2026 Octop +SPDX-License-Identifier: MIT +""" + +from src.knowledge.service import format_kb_command, knowledge_status, update_knowledge_settings + +__all__ = [ + "format_kb_command", + "knowledge_status", + "update_knowledge_settings", +] diff --git a/src/knowledge/bm25.py b/src/knowledge/bm25.py new file mode 100644 index 0000000..e0476d6 --- /dev/null +++ b/src/knowledge/bm25.py @@ -0,0 +1,49 @@ +"""Stdlib BM25 ranking used when local embeddings are not available.""" + +from __future__ import annotations + +import math +import re +from collections import Counter + +_TOKEN = re.compile(r"[a-z0-9]+") + + +def tokenize(text: str) -> list[str]: + return _TOKEN.findall(text.lower()) + + +def bm25_scores( + query: str, + documents: list[str], + *, + k1: float = 1.5, + b: float = 0.75, +) -> list[float]: + """Score each document against ``query``. Higher is a closer match.""" + if not documents: + return [] + query_terms = tokenize(query) + if not query_terms: + return [0.0] * len(documents) + tokenized = [tokenize(document) for document in documents] + lengths = [len(tokens) for tokens in tokenized] + average = sum(lengths) / len(lengths) if lengths else 0.0 + doc_count = len(tokenized) + document_frequency: Counter[str] = Counter() + for tokens in tokenized: + document_frequency.update(set(tokens)) + scores: list[float] = [] + for tokens, length in zip(tokenized, lengths): + frequency = Counter(tokens) + score = 0.0 + for term in query_terms: + seen = document_frequency.get(term, 0) + if seen == 0 or term not in frequency: + continue + idf = math.log(1.0 + (doc_count - seen + 0.5) / (seen + 0.5)) + term_frequency = frequency[term] + denominator = term_frequency + k1 * (1.0 - b + b * (length / average if average else 0.0)) + score += idf * (term_frequency * (k1 + 1.0)) / denominator + scores.append(score) + return scores diff --git a/src/knowledge/chunk.py b/src/knowledge/chunk.py new file mode 100644 index 0000000..8a81c36 --- /dev/null +++ b/src/knowledge/chunk.py @@ -0,0 +1,29 @@ +"""Overlapping character windows for knowledge documents. + +Adapted from Octop's knowledge chunker (MIT). The window size and overlap +match that splitter so a passage stays a stable citation target. + +Source: https://github.com/TencentCloud/Octop +Copyright (c) 2026 Octop +SPDX-License-Identifier: MIT +""" + +from __future__ import annotations + + +def chunk_text(text: str, *, size: int = 800, overlap: int = 120) -> list[str]: + """Split non-empty text into overlapping character windows.""" + if size <= 0: + raise ValueError("size must be positive") + if overlap < 0 or overlap >= size: + raise ValueError("overlap must be non-negative and smaller than size") + content = text.strip() + if not content: + return [] + step = size - overlap + chunks: list[str] = [] + for start in range(0, len(content), step): + chunks.append(content[start : start + size]) + if start + size >= len(content): + break + return chunks diff --git a/src/knowledge/embed.py b/src/knowledge/embed.py new file mode 100644 index 0000000..1668140 --- /dev/null +++ b/src/knowledge/embed.py @@ -0,0 +1,143 @@ +"""Local fastembed embeddings, with an explicit OpenAI cloud path. + +The local engine is imported only when a verified model directory is +embedded. Search falls back to BM25 when the extra or the model is absent. +Cloud embeddings run only after the caller has selected that provider, and +only against api.openai.com. +""" + +from __future__ import annotations + +import json +import os +import urllib.request +from typing import Any, Callable +from pathlib import Path + +from src.knowledge.models import ( + DEFAULT_MODEL_ID, + HF_HOST, + ChecksumError, + assert_cloud_embed_url, + verify_model_dir, +) + +CLOUD_EMBED_URL = "https://api.openai.com/v1/embeddings" +DEFAULT_CLOUD_MODEL = "text-embedding-3-small" +_LOCAL_EXTRA = 'Local embeddings need the optional knowledge extra: pip install "clawd-codex[kb]"' +_NOT_SENT = "cloud embedding key is not set; document text was not sent" + +Transport = Callable[[str, dict[str, Any], dict[str, str]], dict[str, Any]] +Embedder = Callable[[list[str]], list[list[float]]] + + +def embed_local( + texts: list[str], + model_dir: Path, + *, + embedder: Embedder | None = None, +) -> list[list[float]] | None: + """Embed with the injected engine or fastembed. + + Returns None when the local model is not on disk and no engine was + injected, so the caller can rank with BM25. A checksum mismatch still + raises after the bad file is removed. + """ + if not texts: + return [] + if embedder is not None: + return [_as_floats(row) for row in embedder(texts)] + try: + ready = verify_model_dir(model_dir) + except ChecksumError: + return None + if not ready: + return None + return _fastembed(texts, model_dir) + + +def embed_cloud( + texts: list[str], + *, + api_key: str, + model: str, + transport: Transport | None = None, + url: str = CLOUD_EMBED_URL, +) -> list[list[float]]: + """Send texts to the selected cloud embedder. Refuses other hosts first.""" + if not texts: + return [] + if not api_key.strip(): + raise RuntimeError(_NOT_SENT) + assert_cloud_embed_url(url) + sender = transport or _urllib_transport + payload = sender( + url, + {"model": model.strip() or DEFAULT_CLOUD_MODEL, "input": list(texts)}, + {"Authorization": f"Bearer {api_key}", "Content-Type": "application/json"}, + ) + return _vectors_from_response(payload, expected=len(texts)) + + +def _fastembed(texts: list[str], model_dir: Path) -> list[list[float]]: + try: + from fastembed import TextEmbedding + except ImportError as exc: + raise RuntimeError(_LOCAL_EXTRA) from exc + saved = { + "HF_ENDPOINT": os.environ.get("HF_ENDPOINT"), + "HF_HUB_OFFLINE": os.environ.get("HF_HUB_OFFLINE"), + } + os.environ["HF_ENDPOINT"] = f"https://{HF_HOST}" + os.environ["HF_HUB_OFFLINE"] = "1" + try: + engine = TextEmbedding( + model_name=DEFAULT_MODEL_ID, + specific_model_path=str(model_dir), + local_files_only=True, + ) + vectors = list(engine.embed(texts)) + finally: + for key, value in saved.items(): + if value is None: + os.environ.pop(key, None) + else: + os.environ[key] = value + return [_as_floats(row) for row in vectors] + + +def _as_floats(row: Any) -> list[float]: + values = row.tolist() if hasattr(row, "tolist") else list(row) + floats = [float(value) for value in values] + if not floats: + raise RuntimeError("embedding was empty") + return floats + + +def _vectors_from_response(payload: dict[str, Any], *, expected: int) -> list[list[float]]: + data = payload.get("data") + if not isinstance(data, list): + raise RuntimeError("embedding response has no data") + ordered = sorted( + data, + key=lambda item: int(item.get("index", 0)) if isinstance(item, dict) else 0, + ) + vectors = [_as_floats(item["embedding"]) for item in ordered if isinstance(item, dict)] + if len(vectors) != expected: + raise RuntimeError("embedding count mismatch") + return vectors + + +def _urllib_transport(url: str, body: dict[str, Any], headers: dict[str, str]) -> dict[str, Any]: + assert_cloud_embed_url(url) + request = urllib.request.Request( + url, + data=json.dumps(body).encode("utf-8"), + headers=headers, + method="POST", + ) + with urllib.request.urlopen(request, timeout=60) as response: + parsed = json.loads(response.read().decode("utf-8")) + if not isinstance(parsed, dict): + raise RuntimeError("embedding response was not an object") + return parsed diff --git a/src/knowledge/index.py b/src/knowledge/index.py new file mode 100644 index 0000000..88c8f18 --- /dev/null +++ b/src/knowledge/index.py @@ -0,0 +1,184 @@ +"""Per-user SQLite index for knowledge documents and passages.""" + +from __future__ import annotations + +import sqlite3 +import struct +from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path + + +@dataclass(frozen=True) +class DocumentRecord: + doc_id: str + filename: str + source_path: str + added_at: str + method: str + passages: int + + +@dataclass(frozen=True) +class ChunkRecord: + chunk_id: str + doc_id: str + filename: str + ordinal: int + text: str + embedding: tuple[float, ...] | None + + +_SCHEMA = """ +CREATE TABLE IF NOT EXISTS documents ( + doc_id TEXT PRIMARY KEY, + filename TEXT NOT NULL, + source_path TEXT NOT NULL, + added_at TEXT NOT NULL, + method TEXT NOT NULL +); + +CREATE TABLE IF NOT EXISTS chunks ( + chunk_id TEXT PRIMARY KEY, + doc_id TEXT NOT NULL, + ordinal INTEGER NOT NULL, + text TEXT NOT NULL, + embedding BLOB +); + +CREATE INDEX IF NOT EXISTS idx_chunks_doc ON chunks(doc_id); +CREATE INDEX IF NOT EXISTS idx_documents_path ON documents(source_path); +""" + + +def _now() -> str: + return datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ") + + +def _pack(vector: list[float] | None) -> bytes | None: + if not vector: + return None + return struct.pack(f"<{len(vector)}f", *[float(value) for value in vector]) + + +def _unpack(blob: bytes | memoryview | None) -> tuple[float, ...] | None: + if not blob: + return None + raw = bytes(blob) + count = len(raw) // 4 + if count <= 0 or count * 4 != len(raw): + return None + return struct.unpack(f"<{count}f", raw) + + +class KnowledgeIndex: + """One SQLite file for one user's documents.""" + + def __init__(self, path: Path) -> None: + self.path = path + self.path.parent.mkdir(parents=True, exist_ok=True) + with self._connect() as conn: + conn.executescript(_SCHEMA) + _restrict(self.path) + + def _connect(self) -> sqlite3.Connection: + conn = sqlite3.connect(self.path) + conn.row_factory = sqlite3.Row + return conn + + def replace_document( + self, + *, + doc_id: str, + filename: str, + source_path: str, + method: str, + texts: list[str], + embeddings: list[list[float]] | None, + ) -> None: + if embeddings is not None and len(embeddings) != len(texts): + raise ValueError("embeddings and passages must have the same length") + rows: list[tuple[str, str, int, str, bytes | None]] = [] + for ordinal, text in enumerate(texts): + vector = embeddings[ordinal] if embeddings is not None else None + rows.append((f"{doc_id}:{ordinal}", doc_id, ordinal, text, _pack(vector))) + with self._connect() as conn: + existing = conn.execute( + "SELECT doc_id FROM documents WHERE source_path = ?", + (source_path,), + ).fetchall() + for record in existing: + previous = str(record["doc_id"]) + conn.execute("DELETE FROM chunks WHERE doc_id = ?", (previous,)) + conn.execute("DELETE FROM documents WHERE doc_id = ?", (previous,)) + conn.execute( + "INSERT INTO documents(doc_id, filename, source_path, added_at, method) VALUES (?, ?, ?, ?, ?)", + (doc_id, filename, source_path, _now(), method), + ) + conn.executemany( + "INSERT INTO chunks(chunk_id, doc_id, ordinal, text, embedding) VALUES (?, ?, ?, ?, ?)", + rows, + ) + _restrict(self.path) + + def delete(self, doc_id: str) -> bool: + with self._connect() as conn: + found = conn.execute("SELECT 1 FROM documents WHERE doc_id = ?", (doc_id,)).fetchone() + if found is None: + return False + conn.execute("DELETE FROM chunks WHERE doc_id = ?", (doc_id,)) + conn.execute("DELETE FROM documents WHERE doc_id = ?", (doc_id,)) + return True + + def list_documents(self) -> list[DocumentRecord]: + with self._connect() as conn: + rows = conn.execute( + """ + SELECT d.doc_id, d.filename, d.source_path, d.added_at, d.method, + COUNT(c.chunk_id) AS passages + FROM documents d + LEFT JOIN chunks c ON c.doc_id = d.doc_id + GROUP BY d.doc_id + ORDER BY d.added_at, d.filename + """ + ).fetchall() + return [ + DocumentRecord( + doc_id=str(row["doc_id"]), + filename=str(row["filename"]), + source_path=str(row["source_path"]), + added_at=str(row["added_at"]), + method=str(row["method"]), + passages=int(row["passages"]), + ) + for row in rows + ] + + def chunks(self) -> list[ChunkRecord]: + with self._connect() as conn: + rows = conn.execute( + """ + SELECT c.chunk_id, c.doc_id, d.filename, c.ordinal, c.text, c.embedding + FROM chunks c + JOIN documents d ON d.doc_id = c.doc_id + ORDER BY d.filename, c.ordinal + """ + ).fetchall() + return [ + ChunkRecord( + chunk_id=str(row["chunk_id"]), + doc_id=str(row["doc_id"]), + filename=str(row["filename"]), + ordinal=int(row["ordinal"]), + text=str(row["text"]), + embedding=_unpack(row["embedding"]), + ) + for row in rows + ] + + +def _restrict(path: Path) -> None: + try: + path.chmod(0o600) + except OSError: + return diff --git a/src/knowledge/models.py b/src/knowledge/models.py new file mode 100644 index 0000000..071f88c --- /dev/null +++ b/src/knowledge/models.py @@ -0,0 +1,228 @@ +"""Hugging Face ONNX embedding files with checksum checks. + +The default model is the quantized ONNX build of BAAI/bge-small-en-v1.5 +published by Qdrant on Hugging Face, which is the model fastembed loads. +Weights are accepted only from huggingface.co. Mirror hosts are refused +before any download starts. A checksum mismatch deletes the file. + +Octop races an official Hugging Face download against hf-mirror and a COS +bucket. Jonathan Ai keeps the official host and the checksum, and does not +call those other hosts. + +Source: https://github.com/TencentCloud/Octop +Copyright (c) 2026 Octop +SPDX-License-Identifier: MIT +""" + +from __future__ import annotations + +import hashlib +import os +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Callable +from urllib.parse import urlsplit + +DEFAULT_MODEL_ID = "BAAI/bge-small-en-v1.5" +HF_REPO = "Qdrant/bge-small-en-v1.5-onnx-Q" +HF_HOST = "huggingface.co" + +_BLOCKED_HOSTS = ( + "hf-mirror.com", + "modelscope.cn", + "modelscope.com", + "npmmirror.com", + "gitee.com", + "ghproxy.com", + "aliyuncs.com", + "myqcloud.com", +) + + +class ModelSourceError(ValueError): + """The model URL is not an allowlisted Hugging Face file.""" + + +class ChecksumError(ValueError): + """A model file did not match the pinned Hugging Face checksum.""" + + +@dataclass(frozen=True) +class PinnedFile: + name: str + sha256: str + size: int + + +# sha256 values were read from Hugging Face on 2026-10-10. +# model_optimized.onnx is the Git LFS oid. The other files are content hashes. +PINNED_FILES: tuple[PinnedFile, ...] = ( + PinnedFile( + "model_optimized.onnx", + "51f1bd0addd6e859e42c2c8021a5e5461385bb676a649f4b269aa445449f2431", + 66465124, + ), + PinnedFile( + "config.json", + "13582bcf2effc85b7bf3d3f5532e686bc1c9ce86bb009d10f0ec33cbe92299dd", + 706, + ), + PinnedFile( + "ort_config.json", + "99881c45e073696289224931dd48694398bc6bcd1fe7cb7018bca1e0cc00e1fc", + 1272, + ), + PinnedFile( + "special_tokens_map.json", + "5d5b662e421ea9fac075174bb0688ee0d9431699900b90662acd44b2a350503a", + 695, + ), + PinnedFile( + "tokenizer_config.json", + "0b29c7bfc889e53b36d9dd3e686dd4300f6525110eaa98c76a5dafceb2029f53", + 1242, + ), + PinnedFile( + "tokenizer.json", + "d241a60d5e8f04cc1b2b3e9ef7a4921b27bf526d9f6050ab90f9267a1f9e5c66", + 711396, + ), + PinnedFile( + "vocab.txt", + "07eced375cec144d27c900241f3e339478dec958f92fddbc551f295c992038a3", + 231508, + ), +) + +_PINNED_BY_NAME = {item.name: item for item in PINNED_FILES} + + +def file_url(name: str) -> str: + if name not in _PINNED_BY_NAME: + raise ModelSourceError(f"refusing model file {name}") + return f"https://{HF_HOST}/{HF_REPO}/resolve/main/{name}" + + +def assert_model_url(url: str, *, redirect: bool = False) -> None: + """Refuse anything that is not the pinned Hugging Face model. + + Redirects may land on a huggingface.co CDN host. They still have to stay + on that host over HTTPS. The first request must name a pinned file. + """ + parsed = urlsplit(url) + host = (parsed.hostname or "").lower().rstrip(".") + if parsed.scheme != "https" or not host or parsed.username or parsed.password: + raise ModelSourceError(f"refusing model URL {url}") + if parsed.query or parsed.fragment: + raise ModelSourceError(f"refusing model URL {url}") + if host.endswith(".cn") or any(host == blocked or host.endswith("." + blocked) for blocked in _BLOCKED_HOSTS): + raise ModelSourceError(f"refusing model host {host}") + if host != HF_HOST and not host.endswith("." + HF_HOST): + raise ModelSourceError(f"refusing model host {host}") + if redirect: + return + if host != HF_HOST: + raise ModelSourceError(f"refusing model host {host}") + prefix = f"/{HF_REPO}/resolve/main/" + if not parsed.path.startswith(prefix): + raise ModelSourceError(f"refusing model path {parsed.path}") + name = parsed.path[len(prefix) :] + if name not in _PINNED_BY_NAME: + raise ModelSourceError(f"refusing model file {name}") + + +def assert_cloud_embed_url(url: str) -> None: + """Cloud document embeddings go only to the official OpenAI endpoint.""" + parsed = urlsplit(url) + host = (parsed.hostname or "").lower().rstrip(".") + if parsed.scheme != "https" or host != "api.openai.com" or parsed.username: + raise ModelSourceError(f"refusing embedding host {host or url}") + if host.endswith(".cn") or any(blocked in host for blocked in _BLOCKED_HOSTS): + raise ModelSourceError(f"refusing embedding host {host}") + if parsed.path.rstrip("/") != "/v1/embeddings": + raise ModelSourceError("refusing embedding path") + + +def sha256_bytes(data: bytes) -> str: + return hashlib.sha256(data).hexdigest() + + +def sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for block in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(block) + return digest.hexdigest() + + +def verify_model_dir(model_dir: Path) -> bool: + """Return True when every pinned file is present and matches. + + Missing files mean the model is not installed, so search can use BM25. + A file that exists and does not match is deleted, and this raises. + """ + saw_file = False + for item in PINNED_FILES: + path = model_dir / item.name + if not path.is_file(): + continue + saw_file = True + if path.stat().st_size != item.size or sha256_file(path) != item.sha256: + path.unlink() + raise ChecksumError(f"checksum mismatch for {item.name}; file discarded") + if not saw_file: + return False + return all((model_dir / item.name).is_file() for item in PINNED_FILES) + + +def store_verified_bytes(dest: Path, data: bytes, item: PinnedFile) -> None: + """Write ``data`` only when it matches ``item``. A mismatch leaves nothing.""" + if len(data) != item.size or sha256_bytes(data) != item.sha256: + if dest.exists(): + dest.unlink() + raise ChecksumError(f"checksum mismatch for {item.name}; file discarded") + dest.parent.mkdir(parents=True, exist_ok=True) + dest.write_bytes(data) + + +def read_model_bytes(url: str, fetch: Callable[[str], bytes]) -> bytes: + """Check the URL, then download. ``fetch`` is not called for a refused host.""" + assert_model_url(url) + return fetch(url) + + +def download_model(model_dir: Path, fetch: Callable[[str], bytes]) -> None: + """Download each pinned file through ``fetch`` and keep it only if it matches.""" + model_dir.mkdir(parents=True, exist_ok=True) + for item in PINNED_FILES: + data = read_model_bytes(file_url(item.name), fetch) + store_verified_bytes(model_dir / item.name, data, item) + + +def _redirect_handler() -> type[urllib.request.HTTPRedirectHandler]: + class _CheckedRedirect(urllib.request.HTTPRedirectHandler): + def redirect_request(self, req, fp, code, msg, headers, newurl): # type: ignore[no-untyped-def] + assert_model_url(newurl, redirect=True) + return urllib.request.HTTPRedirectHandler.redirect_request( + self, req, fp, code, msg, headers, newurl + ) + + return _CheckedRedirect + + +def default_fetch(url: str) -> bytes: + """Download one allowlisted model URL. Redirects must stay on Hugging Face.""" + assert_model_url(url) + previous_endpoint = os.environ.get("HF_ENDPOINT") + os.environ["HF_ENDPOINT"] = f"https://{HF_HOST}" + try: + request = urllib.request.Request(url, headers={"User-Agent": "jonathan-ai-kb"}) + opener = urllib.request.build_opener(_redirect_handler()) + with opener.open(request, timeout=120) as response: + return response.read() + finally: + if previous_endpoint is None: + os.environ.pop("HF_ENDPOINT", None) + else: + os.environ["HF_ENDPOINT"] = previous_endpoint diff --git a/src/knowledge/parse.py b/src/knowledge/parse.py new file mode 100644 index 0000000..2635f4f --- /dev/null +++ b/src/knowledge/parse.py @@ -0,0 +1,138 @@ +"""Extract text from PDF, DOCX, Markdown, plain text, and HTML. + +PDF reading uses pypdf from the optional knowledge extra. DOCX, Markdown, +text, and HTML use the Python standard library. Tests can pass a PDF reader +so the extra does not need to be installed. + +Adapted from Octop's document parser (MIT). Image OCR and spreadsheet +formats from that parser are not included. + +Source: https://github.com/TencentCloud/Octop +Copyright (c) 2026 Octop +SPDX-License-Identifier: MIT +""" + +from __future__ import annotations + +import zipfile +from html.parser import HTMLParser +from pathlib import Path +from typing import Callable +from xml.etree import ElementTree + +SUPPORTED_SUFFIXES = frozenset({".pdf", ".docx", ".md", ".markdown", ".txt", ".html", ".htm"}) + +_W_NS = "http://schemas.openxmlformats.org/wordprocessingml/2006/main" +_MC_NS = "http://schemas.openxmlformats.org/markup-compatibility/2006" +_PDF_EXTRA = 'PDF ingest needs the optional knowledge extra: pip install "clawd-codex[kb]"' + + +def parse_document(path: Path, *, pdf_reader: Callable[[Path], str] | None = None) -> str: + """Return searchable text for a supported local document.""" + suffix = path.suffix.lower() + if suffix not in SUPPORTED_SUFFIXES: + raise ValueError(f"unsupported knowledge document: {suffix or '(none)'}") + if suffix in {".md", ".markdown", ".txt"}: + return _read_text(path) + if suffix in {".html", ".htm"}: + return _html_text(_read_text(path)) + if suffix == ".pdf": + return _parse_pdf(path, pdf_reader) + return _parse_docx(path) + + +def _read_text(path: Path) -> str: + data = path.read_bytes() + text = data.decode("utf-8-sig", errors="replace") + return text.replace("\r\n", "\n").replace("\r", "\n") + + +def _parse_pdf(path: Path, pdf_reader: Callable[[Path], str] | None) -> str: + if pdf_reader is not None: + return pdf_reader(path) + try: + from pypdf import PdfReader + except ImportError as exc: + raise RuntimeError(_PDF_EXTRA) from exc + return "\n".join(page.extract_text() or "" for page in PdfReader(str(path)).pages) + + +def _parse_docx(path: Path) -> str: + try: + with zipfile.ZipFile(path) as archive: + xml_bytes = archive.read("word/document.xml") + except (zipfile.BadZipFile, KeyError) as exc: + raise ValueError(f"could not read DOCX: {path.name}") from exc + root = ElementTree.fromstring(xml_bytes) + paragraphs: list[str] = [] + + def walk(node: ElementTree.Element) -> None: + if node.tag == f"{{{_MC_NS}}}Fallback": + return + if node.tag == f"{{{_W_NS}}}p": + paragraphs.append(_paragraph_text(node)) + return + for child in list(node): + walk(child) + + walk(root) + return "\n".join(line for line in paragraphs if line.strip()) + + +def _paragraph_text(paragraph: ElementTree.Element) -> str: + parts: list[str] = [] + + def walk(node: ElementTree.Element) -> None: + if node.tag == f"{{{_MC_NS}}}Fallback": + return + if node.tag == f"{{{_W_NS}}}t" and node.text: + parts.append(node.text) + if node.tag == f"{{{_W_NS}}}tab": + parts.append("\t") + for child in list(node): + walk(child) + + walk(paragraph) + return "".join(parts).strip() + + +def _html_text(raw: str) -> str: + parser = _HTMLTextParser() + parser.feed(raw) + parser.close() + lines = [" ".join(line.split()) for line in parser.text().splitlines()] + return "\n".join(line for line in lines if line) + + +class _HTMLTextParser(HTMLParser): + _SKIP = frozenset({"script", "style", "noscript", "template"}) + _BLOCK = frozenset( + {"p", "div", "br", "li", "tr", "h1", "h2", "h3", "h4", "h5", "h6", "section", "article", "blockquote"} + ) + + def __init__(self) -> None: + super().__init__(convert_charrefs=True) + self._skip_depth = 0 + self._parts: list[str] = [] + + def handle_starttag(self, tag: str, attrs: list[tuple[str, str | None]]) -> None: + del attrs + name = tag.lower() + if name in self._SKIP: + self._skip_depth += 1 + elif name in self._BLOCK and self._skip_depth == 0: + self._parts.append("\n") + + def handle_endtag(self, tag: str) -> None: + name = tag.lower() + if name in self._SKIP and self._skip_depth: + self._skip_depth -= 1 + elif name in self._BLOCK and self._skip_depth == 0: + self._parts.append("\n") + + def handle_data(self, data: str) -> None: + if self._skip_depth == 0 and data: + self._parts.append(data) + + def text(self) -> str: + return "".join(self._parts) diff --git a/src/knowledge/service.py b/src/knowledge/service.py new file mode 100644 index 0000000..7efb0f7 --- /dev/null +++ b/src/knowledge/service.py @@ -0,0 +1,326 @@ +"""Add, list, remove, and search the local document knowledge base. + +Answers name the file and the passage they came from. Document text stays +in the per-user SQLite file unless the OpenAI embedding provider is selected. +""" + +from __future__ import annotations + +import math +import os +import re +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Callable + +from src.agent.memory import current_user_id +from src.knowledge.bm25 import bm25_scores +from src.knowledge.chunk import chunk_text +from src.knowledge.embed import Embedder, Transport, embed_cloud, embed_local +from src.knowledge.index import ChunkRecord, KnowledgeIndex +from src.knowledge.models import ChecksumError, verify_model_dir +from src.knowledge.parse import parse_document +from src.knowledge.settings import CLOUD_BACKEND, load_settings, public_view, update_settings +from src.secret_mask import mask_secrets + +_USER_ID_RE = re.compile(r"[^A-Za-z0-9_.-]") + +_USAGE = ( + "Usage: /kb add | /kb list | /kb remove | /kb search \n" + "Documents stay on this machine unless a cloud embedding provider is chosen." +) + + +@dataclass(frozen=True) +class PassageHit: + filename: str + passage: int + text: str + score: float + method: str + + +def knowledge_root(directory: Path | None = None) -> Path: + if directory is not None: + path = directory + else: + raw = os.environ.get("JONATHAN_KB_DIR", "").strip() + path = Path(raw) if raw else Path.home() / ".clawd" / "knowledge" + path.mkdir(parents=True, exist_ok=True) + try: + path.chmod(0o700) + except OSError: + pass + return path + + +class KnowledgeBase: + """Index for one user under one local directory.""" + + def __init__( + self, + directory: Path | None = None, + *, + user_id: str | None = None, + embedder: Embedder | None = None, + transport: Transport | None = None, + pdf_reader: Callable[[Path], str] | None = None, + ) -> None: + self.root = knowledge_root(directory) + self.user_id = _safe_user_id(user_id) + self.index = KnowledgeIndex(self.root / f"knowledge-{self.user_id}.sqlite") + self.settings_path = self.root / f"settings-{self.user_id}.json" + self.model_dir = self.root / "models" / "bge-small-en-v1.5" + self._embedder = embedder + self._transport = transport + self._pdf_reader = pdf_reader + + def status(self) -> dict[str, object]: + settings = load_settings(self.settings_path) + return public_view( + settings, + document_count=len(self.index.list_documents()), + model_ready=self._model_ready(), + ) + + def update(self, changes: dict[str, object]) -> dict[str, object]: + update_settings(self.settings_path, changes) + return self.status() + + def add_document(self, path: str | Path) -> str: + source = Path(path).expanduser() + if not source.is_file(): + raise ValueError(f"document not found: {source}") + resolved = str(source.resolve()) + text = mask_secrets(parse_document(source, pdf_reader=self._pdf_reader)) + passages = chunk_text(text) + if not passages: + raise ValueError(f"document has no text: {source.name}") + settings = load_settings(self.settings_path) + backend = settings["embedding_backend"] + embeddings: list[list[float]] | None = None + method = "bm25" + if backend == CLOUD_BACKEND: + embeddings = embed_cloud( + passages, + api_key=settings.get("openai_api_key", ""), + model=settings.get("cloud_model", ""), + transport=self._transport, + ) + method = "embedding" + else: + local = embed_local(passages, self.model_dir, embedder=self._embedder) + if local is not None: + embeddings = local + method = "embedding" + doc_id = uuid.uuid4().hex[:12] + self.index.replace_document( + doc_id=doc_id, + filename=source.name, + source_path=resolved, + method=method, + texts=passages, + embeddings=embeddings, + ) + if backend == CLOUD_BACKEND: + how = "OpenAI embeddings" + stay = "Document text was sent to api.openai.com because a cloud embedding provider is selected." + elif method == "embedding": + how = "local embeddings" + stay = "Document text stayed on this machine." + else: + how = "BM25" + stay = "Document text stayed on this machine." + return f"Added {source.name} ({len(passages)} passages, {doc_id}) with {how}. {stay}" + + def list_documents(self) -> str: + documents = self.index.list_documents() + if not documents: + return "No documents in the local knowledge base." + lines = ["Local documents:"] + for document in documents: + lines.append( + f"{document.doc_id} {document.filename} {document.passages} passages {document.method}" + ) + return "\n".join(lines) + + def remove_document(self, doc_id: str) -> str: + token = doc_id.strip() + if not token: + raise ValueError("remove needs a document id") + if self.index.delete(token): + return f"Removed {token}." + return f"No document {token}." + + def search(self, query: str, *, k: int = 5) -> str: + text = query.strip() + if not text: + raise ValueError("search needs a query") + hits = self._search(text, k=k) + if not hits: + return "No matching passage." + lines: list[str] = [] + sources: list[str] = [] + for hit in hits: + label = f"{hit.filename}, passage {hit.passage}" + lines.append(f"{label} ({hit.method})") + lines.append(hit.text.strip()) + lines.append("") + sources.append(f"- {label}") + lines.append("Sources:") + lines.extend(sources) + return "\n".join(lines).strip() + + def _search(self, query: str, *, k: int) -> list[PassageHit]: + chunks = self.index.chunks() + if not chunks or k <= 0: + return [] + settings = load_settings(self.settings_path) + vector_hits = self._embedding_hits(query, chunks, settings, k=k) + if vector_hits: + return vector_hits + return self._bm25_hits(query, chunks, k=k) + + def _embedding_hits( + self, + query: str, + chunks: list[ChunkRecord], + settings: dict[str, str], + *, + k: int, + ) -> list[PassageHit]: + stored = [chunk for chunk in chunks if chunk.embedding] + if not stored: + return [] + backend = settings.get("embedding_backend", "local") + if backend == CLOUD_BACKEND: + vectors = embed_cloud( + [query], + api_key=settings.get("openai_api_key", ""), + model=settings.get("cloud_model", ""), + transport=self._transport, + ) + else: + vectors = embed_local([query], self.model_dir, embedder=self._embedder) + if not vectors: + return [] + return _cosine_hits(vectors[0], stored, k=k) + + def _bm25_hits(self, query: str, chunks: list[ChunkRecord], *, k: int) -> list[PassageHit]: + scores = bm25_scores(query, [chunk.text for chunk in chunks]) + ranked = sorted( + ( + PassageHit( + filename=chunk.filename, + passage=chunk.ordinal + 1, + text=chunk.text, + score=score, + method="bm25", + ) + for chunk, score in zip(chunks, scores) + if score > 0 + ), + key=lambda hit: hit.score, + reverse=True, + ) + return ranked[:k] + + def _model_ready(self) -> bool: + try: + return verify_model_dir(self.model_dir) + except ChecksumError: + return False + + +def _safe_user_id(user_id: str | None) -> str: + raw = user_id if user_id else current_user_id() + cleaned = _USER_ID_RE.sub("_", raw).strip("._-") + return cleaned or "local" + + +def format_kb_command( + args: str, + *, + directory: Path | None = None, + user_id: str | None = None, + embedder: Embedder | None = None, + transport: Transport | None = None, + pdf_reader: Callable[[Path], str] | None = None, +) -> str: + """Run one /kb command and return text for the user.""" + base = KnowledgeBase( + directory, + user_id=user_id, + embedder=embedder, + transport=transport, + pdf_reader=pdf_reader, + ) + parts = args.strip().split(maxsplit=1) + if not parts: + return _USAGE + action = parts[0].lower() + rest = parts[1].strip() if len(parts) > 1 else "" + try: + if action == "add": + if not rest: + return _USAGE + return base.add_document(rest) + if action == "list": + return base.list_documents() + if action == "remove": + if not rest: + return _USAGE + return base.remove_document(rest) + if action == "search": + if not rest: + return _USAGE + return base.search(rest) + except (ValueError, RuntimeError, OSError) as exc: + return str(exc) + return _USAGE + + +def knowledge_status(directory: Path | None = None, *, user_id: str | None = None) -> dict[str, object]: + return KnowledgeBase(directory, user_id=user_id).status() + + +def update_knowledge_settings( + changes: dict[str, Any], + directory: Path | None = None, + *, + user_id: str | None = None, +) -> dict[str, object]: + clean = {key: value for key, value in changes.items() if key in { + "embedding_backend", + "model_id", + "cloud_model", + "openai_api_key", + }} + return KnowledgeBase(directory, user_id=user_id).update(clean) + + +def _cosine_hits(query: list[float], chunks: list[ChunkRecord], *, k: int) -> list[PassageHit]: + query_norm = math.sqrt(sum(value * value for value in query)) + if query_norm == 0: + return [] + hits: list[PassageHit] = [] + for chunk in chunks: + embedding = chunk.embedding + if not embedding or len(embedding) != len(query): + continue + norm = math.sqrt(sum(value * value for value in embedding)) + score = 0.0 if norm == 0 else sum(left * right for left, right in zip(query, embedding)) / (query_norm * norm) + if score <= 0: + continue + hits.append( + PassageHit( + filename=chunk.filename, + passage=chunk.ordinal + 1, + text=chunk.text, + score=score, + method="embedding", + ) + ) + hits.sort(key=lambda hit: hit.score, reverse=True) + return hits[:k] diff --git a/src/knowledge/settings.py b/src/knowledge/settings.py new file mode 100644 index 0000000..4632c52 --- /dev/null +++ b/src/knowledge/settings.py @@ -0,0 +1,92 @@ +"""Per-user knowledge settings. The public view never includes the API key.""" + +from __future__ import annotations + +import json +from pathlib import Path + +from src.knowledge.models import DEFAULT_MODEL_ID +from src.knowledge.embed import DEFAULT_CLOUD_MODEL + +LOCAL_BACKEND = "local" +CLOUD_BACKEND = "openai" +_BACKENDS = {LOCAL_BACKEND, CLOUD_BACKEND} + + +def default_settings() -> dict[str, str]: + return { + "embedding_backend": LOCAL_BACKEND, + "model_id": DEFAULT_MODEL_ID, + "cloud_model": DEFAULT_CLOUD_MODEL, + "openai_api_key": "", + } + + +def load_settings(path: Path) -> dict[str, str]: + settings = default_settings() + if not path.is_file(): + return settings + try: + raw = json.loads(path.read_text(encoding="utf-8")) + except (OSError, json.JSONDecodeError): + return settings + if not isinstance(raw, dict): + return settings + backend = str(raw.get("embedding_backend") or LOCAL_BACKEND).strip().lower() + if backend in _BACKENDS: + settings["embedding_backend"] = backend + model_id = str(raw.get("model_id") or "").strip() + if model_id: + settings["model_id"] = model_id + cloud_model = str(raw.get("cloud_model") or "").strip() + if cloud_model: + settings["cloud_model"] = cloud_model + key = raw.get("openai_api_key") + if isinstance(key, str): + settings["openai_api_key"] = key + return settings + + +def save_settings(path: Path, settings: dict[str, str]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(settings, indent=2) + "\n", encoding="utf-8") + try: + path.chmod(0o600) + except OSError: + return + + +def update_settings(path: Path, changes: dict[str, object]) -> dict[str, str]: + settings = load_settings(path) + if "embedding_backend" in changes and changes["embedding_backend"] is not None: + backend = str(changes["embedding_backend"]).strip().lower() + if backend not in _BACKENDS: + raise ValueError("embedding backend must be local or openai") + settings["embedding_backend"] = backend + if "model_id" in changes and changes["model_id"] is not None: + model_id = str(changes["model_id"]).strip() + if model_id: + settings["model_id"] = model_id + if "cloud_model" in changes and changes["cloud_model"] is not None: + cloud_model = str(changes["cloud_model"]).strip() + if cloud_model: + settings["cloud_model"] = cloud_model + if "openai_api_key" in changes and changes["openai_api_key"] is not None: + key = str(changes["openai_api_key"]) + if key.strip(): + settings["openai_api_key"] = key.strip() + save_settings(path, settings) + return settings + + +def public_view(settings: dict[str, str], *, document_count: int, model_ready: bool) -> dict[str, object]: + backend = settings.get("embedding_backend", LOCAL_BACKEND) + return { + "embedding_backend": backend, + "model_id": settings.get("model_id", DEFAULT_MODEL_ID), + "cloud_model": settings.get("cloud_model", DEFAULT_CLOUD_MODEL), + "documents_stay_local": backend != CLOUD_BACKEND, + "cloud_key_configured": bool(settings.get("openai_api_key")), + "model_ready": model_ready, + "document_count": document_count, + } diff --git a/tests/test_knowledge.py b/tests/test_knowledge.py new file mode 100644 index 0000000..a3d88e0 --- /dev/null +++ b/tests/test_knowledge.py @@ -0,0 +1,392 @@ +"""Local document knowledge base: parse, cite, and keep text on this machine.""" + +from __future__ import annotations + +import hashlib +import io +import json +import os +import sys +import tempfile +import unittest +import zipfile +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[1] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from src.command_system import create_command_context, execute_command_sync +from src.desktop.runtime import _HELP_TEXT +from src.desktop.server import DesktopServer +from src.knowledge.bm25 import bm25_scores +from src.knowledge.chunk import chunk_text +from src.knowledge.models import ( + ChecksumError, + ModelSourceError, + PinnedFile, + assert_cloud_embed_url, + assert_model_url, + file_url, + read_model_bytes, + store_verified_bytes, + verify_model_dir, +) +from src.knowledge.parse import parse_document +from src.knowledge.service import KnowledgeBase, format_kb_command + +SECRET = "sk-knowledgebaseSecretKey99" + + +def _docx(path: Path, body: str) -> None: + buffer = io.BytesIO() + with zipfile.ZipFile(buffer, "w") as archive: + archive.writestr("word/document.xml", body) + path.write_bytes(buffer.getvalue()) + + +class TestParseAndChunk(unittest.TestCase): + def test_markdown_text_html_and_docx(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + notes = root / "notes.md" + notes.write_text("# Otters\n\nOtters live in rivers.\n", encoding="utf-8") + plain = root / "plain.txt" + plain.write_text("Plain text about beavers.\n", encoding="utf-8") + page = root / "page.html" + page.write_text( + "

Visible otter habitat.

", + encoding="utf-8", + ) + document = root / "brief.docx" + _docx( + document, + """ + + + DOCX otter paragraph + + choice otter + fallback secretphrase + + + +""", + ) + self.assertIn("Otters live in rivers", parse_document(notes)) + self.assertIn("beavers", parse_document(plain)) + html_text = parse_document(page) + self.assertIn("Visible otter habitat", html_text) + self.assertNotIn("hiddenphrase", html_text) + self.assertNotIn("skipme", html_text) + docx_text = parse_document(document) + self.assertIn("DOCX otter paragraph", docx_text) + self.assertIn("choice otter", docx_text) + self.assertNotIn("fallback secretphrase", docx_text) + with self.assertRaises(ValueError): + parse_document(root / "skip.csv") + + def test_pdf_uses_injected_reader_without_the_extra(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "paper.pdf" + path.write_bytes(b"%PDF-1.4") + text = parse_document(path, pdf_reader=lambda source: f"pdf otter from {source.name}") + self.assertIn("pdf otter", text) + self.assertNotIn("pypdf", sys.modules) + + def test_chunks_overlap(self) -> None: + parts = chunk_text("a" * 800 + "b" * 100, size=800, overlap=120) + self.assertEqual(len(parts), 2) + self.assertTrue(parts[1].startswith("a" * 120)) + self.assertEqual(chunk_text(" "), []) + + +class TestModelSources(unittest.TestCase): + def test_official_url_is_allowed_and_mirrors_are_refused_before_fetch(self) -> None: + assert_model_url(file_url("model_optimized.onnx")) + calls: list[str] = [] + + def fetch(url: str) -> bytes: + calls.append(url) + return b"" + + for url in ( + "https://hf-mirror.com/Qdrant/bge-small-en-v1.5-onnx-Q/resolve/main/model_optimized.onnx", + "https://npmmirror.com/mirrors/bge/model_optimized.onnx", + "https://www.modelscope.cn/models/bge/model_optimized.onnx", + "https://example.cn/model_optimized.onnx", + "http://huggingface.co/Qdrant/bge-small-en-v1.5-onnx-Q/resolve/main/model_optimized.onnx", + "https://huggingface.co/other/repo/resolve/main/model_optimized.onnx", + ): + with self.assertRaises(ModelSourceError): + read_model_bytes(url, fetch) + self.assertEqual(calls, []) + + def test_cloud_embeddings_only_use_openai(self) -> None: + assert_cloud_embed_url("https://api.openai.com/v1/embeddings") + for url in ( + "https://hf-mirror.com/v1/embeddings", + "https://api.openai.com.evil.example/v1/embeddings", + "http://api.openai.com/v1/embeddings", + "https://example.cn/v1/embeddings", + ): + with self.assertRaises(ModelSourceError): + assert_cloud_embed_url(url) + + def test_checksum_mismatch_discards_the_file(self) -> None: + with tempfile.TemporaryDirectory() as tmp: + root = Path(tmp) + payload = b"not-the-model" + item = PinnedFile("weights.bin", hashlib.sha256(b"real-model").hexdigest(), len(b"real-model")) + dest = root / "weights.bin" + dest.write_bytes(b"stale") + with self.assertRaises(ChecksumError): + store_verified_bytes(dest, payload, item) + self.assertFalse(dest.exists()) + store_verified_bytes(dest, b"real-model", item) + self.assertEqual(dest.read_bytes(), b"real-model") + + model_dir = root / "model" + model_dir.mkdir() + self.assertFalse(verify_model_dir(model_dir)) + bad = model_dir / "model_optimized.onnx" + bad.write_bytes(b"corrupt-onnx") + with self.assertRaises(ChecksumError): + verify_model_dir(model_dir) + self.assertFalse(bad.exists()) + + +class _KbCase(unittest.TestCase): + def setUp(self) -> None: + self._tmp = tempfile.TemporaryDirectory() + self.addCleanup(self._tmp.cleanup) + self.root = Path(self._tmp.name) + self._old_dir = os.environ.get("JONATHAN_KB_DIR") + os.environ["JONATHAN_KB_DIR"] = str(self.root) + self.addCleanup(self._restore_env) + + def _restore_env(self) -> None: + if self._old_dir is None: + os.environ.pop("JONATHAN_KB_DIR", None) + else: + os.environ["JONATHAN_KB_DIR"] = self._old_dir + + def _write(self, name: str, text: str) -> Path: + path = self.root / name + path.write_text(text, encoding="utf-8") + return path + + +class TestKnowledgeBase(_KbCase): + def test_bm25_cites_file_and_passage_without_calling_cloud(self) -> None: + def transport(*_args: object) -> dict[str, object]: + raise AssertionError("cloud embedding was called") + + base = KnowledgeBase(self.root, user_id="reader", transport=transport) + notes = self._write("rivers.md", "Otters live in rivers and eat fish.\n") + other = self._write("forest.txt", "Pine trees grow in the forest.\n") + added = base.add_document(notes) + base.add_document(other) + self.assertIn("BM25", added) + self.assertIn("stayed on this machine", added) + later = "x" * 900 + " zebrapassage is the later habitat note." + base.add_document(self._write("later.md", later)) + found = base.search("zebrapassage") + self.assertIn("later.md, passage 2", found) + self.assertIn("(bm25)", found) + self.assertIn("Sources:", found) + otter = base.search("otters") + self.assertIn("rivers.md, passage 1", otter) + self.assertNotIn("forest.txt", otter) + self.assertIn("rivers.md", base.list_documents()) + self.assertTrue(base.status()["documents_stay_local"]) + self.assertEqual(base.status()["embedding_backend"], "local") + + def test_local_embeddings_cite_the_matching_passage(self) -> None: + def embedder(texts: list[str]) -> list[list[float]]: + return [[1.0, 0.0] if "otter" in text.lower() else [0.0, 1.0] for text in texts] + + base = KnowledgeBase(self.root, user_id="reader", embedder=embedder) + base.add_document(self._write("rivers.md", "Otters live in rivers.\n")) + base.add_document(self._write("forest.txt", "Pine trees grow here.\n")) + added = base.list_documents() + self.assertIn("embedding", added) + found = base.search("where do otters live") + self.assertIn("rivers.md, passage 1 (embedding)", found) + self.assertNotIn("forest.txt", found) + + def test_cloud_is_used_only_after_it_is_selected(self) -> None: + sent: list[tuple[str, object]] = [] + + def transport(url: str, body: dict[str, object], headers: dict[str, str]) -> dict[str, object]: + del headers + sent.append((url, body["input"])) + vectors = [] + for text in body["input"]: + vectors.append([1.0, 0.0] if "otter" in str(text).lower() else [0.0, 1.0]) + return {"data": [{"index": index, "embedding": vector} for index, vector in enumerate(vectors)]} + + base = KnowledgeBase(self.root, user_id="reader", transport=transport) + notes = self._write("rivers.md", "Otters live in rivers.\n") + base.add_document(notes) + self.assertEqual(sent, []) + base.update({"embedding_backend": "openai"}) + with self.assertRaises(RuntimeError) as caught: + base.add_document(self._write("again.md", "Otters again.\n")) + self.assertIn("was not sent", str(caught.exception)) + self.assertEqual(sent, []) + self.assertNotIn("again.md", base.list_documents()) + base.update({"openai_api_key": "sk-cloudKeyForKnowledge1234"}) + cloud = base.add_document(self._write("cloud.md", "Otters live in rivers.\n")) + self.assertIn("api.openai.com", cloud) + self.assertEqual(sent[0][0], "https://api.openai.com/v1/embeddings") + self.assertIn("Otters live in rivers.", sent[0][1]) + found = base.search("otters") + self.assertIn("cloud.md, passage 1", found) + self.assertIn("otters", str(sent[-1][1]).lower()) + public = base.status() + self.assertFalse(public["documents_stay_local"]) + self.assertNotIn("sk-cloudKeyForKnowledge1234", json.dumps(public)) + stored = base.settings_path.read_text(encoding="utf-8") + self.assertIn("sk-cloudKeyForKnowledge1234", stored) + self.assertEqual(base.settings_path.stat().st_mode & 0o777, 0o600) + + def test_secrets_are_masked_before_they_are_stored(self) -> None: + base = KnowledgeBase(self.root, user_id="reader") + base.add_document(self._write("secret.txt", f"The key is {SECRET} and otters swim.\n")) + blob = base.index.path.read_bytes() + self.assertNotIn(SECRET.encode("utf-8"), blob) + self.assertNotIn(b"knowledgebaseSecret", blob) + visible = base.search("otters") + self.assertNotIn(SECRET, visible) + self.assertNotIn("knowledgebaseSecret", visible) + + def test_users_do_not_share_an_index(self) -> None: + left = KnowledgeBase(self.root, user_id="alpha") + right = KnowledgeBase(self.root, user_id="beta") + left.add_document(self._write("alpha.md", "Alpha keeps otter notes.\n")) + self.assertIn("alpha.md", left.list_documents()) + self.assertIn("No documents", right.list_documents()) + self.assertNotEqual(left.index.path, right.index.path) + + def test_remove_and_replace_the_same_file(self) -> None: + base = KnowledgeBase(self.root, user_id="reader") + path = self._write("notes.md", "First otter note.\n") + base.add_document(path) + path.write_text("Second otter note.\n", encoding="utf-8") + base.add_document(path) + listed = base.list_documents() + self.assertEqual(listed.count("notes.md"), 1) + doc_id = listed.split()[2] + self.assertIn("Removed", base.remove_document(doc_id)) + self.assertIn("No documents", base.list_documents()) + self.assertIn("No document", base.remove_document(doc_id)) + + def test_pdf_docx_and_html_are_indexed(self) -> None: + base = KnowledgeBase( + self.root, + user_id="reader", + pdf_reader=lambda source: "PDF passage about otters.", + ) + pdf = self.root / "paper.pdf" + pdf.write_bytes(b"%PDF-1.4") + html = self._write("page.html", "

HTML passage about herons.

") + document = self.root / "brief.docx" + _docx( + document, + """ + + DOCX passage about beavers. + +""", + ) + base.add_document(pdf) + base.add_document(html) + base.add_document(document) + self.assertIn("paper.pdf, passage 1", base.search("otters")) + heron = base.search("herons") + self.assertIn("page.html, passage 1", heron) + self.assertNotIn("hiddenphrase", heron) + self.assertIn("brief.docx, passage 1", base.search("beavers")) + + def test_corrupt_model_file_falls_back_to_bm25(self) -> None: + base = KnowledgeBase(self.root, user_id="reader") + model = base.model_dir + model.mkdir(parents=True) + bad = model / "model_optimized.onnx" + bad.write_bytes(b"corrupt-onnx") + base.add_document(self._write("rivers.md", "Otters live in rivers.\n")) + found = base.search("otters") + self.assertIn("(bm25)", found) + self.assertFalse(bad.exists()) + self.assertNotIn("fastembed", sys.modules) + + +class TestCommandsAndSettings(_KbCase): + def test_kb_command_roundtrip(self) -> None: + path = self._write("rivers.md", "Otters live in rivers.\n") + context = create_command_context(self.root) + success, added, error = execute_command_sync("kb", f"add {path}", context) + self.assertTrue(success, error) + self.assertIn("BM25", added or "") + success, listed, error = execute_command_sync("kb", "list", context) + self.assertTrue(success, error) + self.assertIn("rivers.md", listed or "") + success, found, error = execute_command_sync("kb", "search otters", context) + self.assertTrue(success, error) + self.assertIn("rivers.md, passage 1", found or "") + doc_id = (listed or "").split()[2] + success, removed, error = execute_command_sync("kb", f"remove {doc_id}", context) + self.assertTrue(success, error) + self.assertIn("Removed", removed or "") + success, usage, error = execute_command_sync("kb", "explode", context) + self.assertTrue(success, error) + self.assertIn("Usage: /kb", usage or "") + self.assertIn("/kb", _HELP_TEXT) + + def test_settings_panel_and_api(self) -> None: + html = (ROOT / "src" / "desktop" / "web" / "index.html").read_text(encoding="utf-8") + script = (ROOT / "src" / "desktop" / "web" / "app.js").read_text(encoding="utf-8") + self.assertIn("Knowledge base", html) + self.assertIn("Documents stay on this machine", html) + self.assertIn("kb-backend", html) + self.assertIn("/api/kb", script) + server = DesktopServer(runtime=object()) + status, body = server.handle_api("GET", "/api/kb", {}, {}) + self.assertEqual(status, 200) + self.assertEqual(body["embedding_backend"], "local") + self.assertTrue(body["documents_stay_local"]) + self.assertNotIn("openai_api_key", body) + status, saved = server.handle_api( + "POST", + "/api/kb", + {}, + {"embedding_backend": "openai", "openai_api_key": SECRET}, + ) + self.assertEqual(status, 200) + self.assertTrue(saved["cloud_key_configured"]) + self.assertNotIn(SECRET, json.dumps(saved)) + with self.assertRaises(ValueError): + server.handle_api("POST", "/api/kb", {}, {"embedding_backend": "hf-mirror"}) + self.assertIn('kb = [', (ROOT / "pyproject.toml").read_text(encoding="utf-8")) + notice = (ROOT / "NOTICE").read_text(encoding="utf-8") + self.assertIn("https://github.com/TencentCloud/Octop", notice) + self.assertIn("knowledge base", notice.lower()) + + def test_format_helper_matches_the_command(self) -> None: + path = self._write("notes.md", "Otters live in rivers.\n") + text = format_kb_command(f"add {path}", directory=self.root, user_id="reader") + self.assertIn("notes.md", text) + self.assertIn("bm25", format_kb_command("list", directory=self.root, user_id="reader")) + + +class TestImportDoesNotPullOptionalEngines(unittest.TestCase): + def test_fastembed_and_pypdf_stay_unimported(self) -> None: + self.assertNotIn("fastembed", sys.modules) + self.assertNotIn("pypdf", sys.modules) + self.assertGreater(bm25_scores("otters", ["otters live here", "pine trees"])[0], 0) + + +if __name__ == "__main__": + unittest.main()