From c4d67f244129fa9faede21a8bf8b1a8d46bae2ea Mon Sep 17 00:00:00 2001 From: runmanfang Date: Sun, 25 Oct 2026 12:56:18 +0800 Subject: [PATCH] feat(query): add FTS5 keyword search as hybrid retrieval pre-filter MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The query agent navigated the wiki purely by reading index.md, then summaries, concepts, and entities — several LLM turns of latency before any content is seen, and it handles "which page mentions X" lookups poorly. Add a zero-dependency SQLite FTS5 index (bm25() ranking) as a fast first-pass retrieval layer: - openkb/fts.py: builds {kb_dir}/.openkb/fts.db from the markdown under summaries/, concepts/, entities/, sources/, plus per-page rows for long-document JSON. CJK runs are indexed as overlapping bigrams so Chinese queries match properly with the stock unicode61 tokenizer — no external segmenter needed. Mixed-script tokens become adjacency phrases; FTS5 syntax characters are stripped from query tokens so pasted punctuation can never break MATCH. Frontmatter is stripped to keep wikilink lists out of snippets, and snippets are highlighted with marker-wrapped terms. The index refreshes incrementally by (mtime, size) on every search, so freshly added files are searchable without a manual rebuild. - Query agent: new `search` tool, first in the tool list, and the search-strategy prompt now starts with it, falling back to index.md navigation on no hits. Chat sessions inherit it automatically via build_chat_agent(). FTS failures degrade gracefully to navigation instead of breaking the agent loop. - CLI: `openkb fts rebuild` and `openkb fts search` for full rebuilds and quick debugging lookups. On a real ~1000-file wiki: a full build takes ~3s, and entity-style questions that previously took ~2.5 min / 6 agent turns now answer in ~25s / 2 tool calls (search -> read_file). --- openkb/agent/query.py | 30 ++++- openkb/cli.py | 38 ++++++ openkb/fts.py | 301 ++++++++++++++++++++++++++++++++++++++++++ tests/test_fts.py | 143 ++++++++++++++++++++ tests/test_query.py | 5 +- 5 files changed, 513 insertions(+), 4 deletions(-) create mode 100644 openkb/fts.py create mode 100644 tests/test_fts.py diff --git a/openkb/agent/query.py b/openkb/agent/query.py index da1a939e..f41d77f1 100644 --- a/openkb/agent/query.py +++ b/openkb/agent/query.py @@ -25,7 +25,12 @@ {schema_md} ## Search strategy -1. Read index.md to see all documents and concepts with brief summaries. +1. Use the `search` tool with the question's key terms (Chinese and/or + English) first — it is a BM25-ranked full-text index over the ENTIRE wiki + and returns candidate pages with snippets in milliseconds. Hits on long + documents include page numbers you can pass straight to get_page_content. + If `search` returns no hits or the question is about overall KB scope, + read index.md to see all documents and concepts with brief summaries. Each document is marked (short) or (pageindex) to indicate its type. 2. Read relevant summary pages (summaries/) for document overviews. Summaries may omit details — if you need more, follow the summary's @@ -83,6 +88,27 @@ def get_page_content(doc_name: str, pages: str) -> str: """ return get_wiki_page_content(doc_name, pages, wiki_root) + @function_tool + def search(query: str, limit: int = 8) -> str: + """Full-text keyword search over the entire wiki (BM25-ranked FTS5). + + Fast first step: returns candidate pages with snippets. Supports + Chinese and English terms (CJK is bigram-indexed). Multi-word queries + are ANDed (falls back to OR when nothing matches). Then verify by + reading the actual pages with read_file / get_page_content. + + Args: + query: Keywords, e.g. 'ETA 川普 豁免' or 'transformer attention'. + limit: Max hits (default 8). + """ + from openkb.fts import format_hits, search as fts_search + + kb_dir = Path(wiki_root).parent + try: + return format_hits(fts_search(kb_dir, query, limit=limit)) + except Exception as exc: # never break the agent loop on FTS issues + return f"Full-text search unavailable ({exc}). Navigate via index.md instead." + @function_tool def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: """View an image from the wiki. @@ -117,7 +143,7 @@ def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: return Agent( name="wiki-query", instructions=instructions, - tools=[read_file, get_page_content, get_image], + tools=[search, read_file, get_page_content, get_image], model=f"litellm/{model}", model_settings=ModelSettings(**model_settings), ) diff --git a/openkb/cli.py b/openkb/cli.py index c9e54318..edc958b0 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -1216,6 +1216,44 @@ def save_exploration(kb_dir: Path, question: str, answer: str) -> Path | None: return explore_path +@cli.group() +def fts(): + """FTS5 full-text index (hybrid retrieval pre-filter for queries).""" + + +@fts.command(name="rebuild") +@click.pass_context +def fts_rebuild(ctx): + """Rebuild the full-text index from scratch.""" + kb_dir = _find_kb_dir(ctx.obj.get("kb_dir_override")) + if kb_dir is None: + click.echo("No knowledge base found. Run `openkb init` first.") + return + from openkb.fts import rebuild_index + + t0 = time.time() + files_seen, rows = rebuild_index(kb_dir) + click.echo( + f"FTS index rebuilt: {files_seen} files, {rows} rows, " + f"{time.time() - t0:.1f}s -> {kb_dir / '.openkb' / 'fts.db'}" + ) + + +@fts.command(name="search") +@click.argument("fts_query") +@click.option("--limit", default=8, show_default=True, help="Max hits.") +@click.pass_context +def fts_search_cmd(ctx, fts_query, limit): + """BM25 keyword search over the wiki (debug / quick lookup).""" + kb_dir = _find_kb_dir(ctx.obj.get("kb_dir_override")) + if kb_dir is None: + click.echo("No knowledge base found. Run `openkb init` first.") + return + from openkb.fts import format_hits, search as fts_search + + click.echo(format_hits(fts_search(kb_dir, fts_query, limit=limit))) + + @cli.command() @click.argument("question") @click.option("--save", is_flag=True, default=False, help="Save the answer to wiki/explorations/.") diff --git a/openkb/fts.py b/openkb/fts.py new file mode 100644 index 00000000..bc253e9c --- /dev/null +++ b/openkb/fts.py @@ -0,0 +1,301 @@ +"""FTS5 hybrid retrieval for OpenKB: a fast BM25-ranked keyword pre-filter +over the whole wiki, complementing the query agent's navigation loop. + +The query agent's default strategy (read index.md -> summaries -> concepts -> +sources) is thorough but costs several LLM turns before any content is seen. +This module gives it a millisecond keyword index to jump straight to +candidate pages; the agent still verifies by reading the actual pages. + +Design notes: + +- Storage: ``{kb_dir}/.openkb/fts.db`` (SQLite FTS5, ``bm25()`` ranking). +- CJK: FTS5's unicode61 tokenizer treats a run of CJK ideographs as ONE + token, so "川普政府" would only match as a whole. We transform CJK runs + into overlapping bigrams ("川普 普政 政府") before indexing and apply the + same transform to query terms (classic CJK bigram trick, no external + segmenter dependency). +- Long docs (``sources/*.json`` page arrays) are indexed one row per page, + so a hit tells the agent exactly which ``pages=`` to pass to + ``get_page_content``. +- Freshness: incremental update keyed on (mtime, size) per file, checked on + every search; a missing db triggers a full build. ``openkb fts rebuild`` + forces a rebuild. +""" + +from __future__ import annotations + +import json +import re +import sqlite3 +import time +from pathlib import Path + +# CJK Unified Ideographs (+ ext A, compat). Runs of these get bigram-ised. +_CJK_RE = re.compile(r"[\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff〇]+") + +_INDEX_DIRS = ("summaries", "concepts", "entities", "sources", "explorations") +_KIND = { + "summaries": "summary", + "concepts": "concept", + "entities": "entity", + "sources": "source", + "explorations": "exploration", +} +_FRONTMATTER_RE = re.compile(r"\A---\n.*?\n---\n", re.DOTALL) +_RAW_CAP = 30000 # chars of original text kept per row for snippets + + +def fts_db_path(kb_dir: Path) -> Path: + return kb_dir / ".openkb" / "fts.db" + + +def _bigramize(text: str) -> str: + """Replace each CJK run with space-joined overlapping bigrams.""" + + def repl(m: re.Match) -> str: + s = m.group(0) + if len(s) == 1: + return s + return " ".join(s[i : i + 2] for i in range(len(s) - 1)) + + return _CJK_RE.sub(repl, text) + + +def _connect(db_path: Path) -> sqlite3.Connection: + con = sqlite3.connect(str(db_path)) + con.execute( + "CREATE VIRTUAL TABLE IF NOT EXISTS fts USING fts5(" + "content, raw UNINDEXED, kind UNINDEXED, path UNINDEXED, page UNINDEXED," + " tokenize='unicode61')" + ) + con.execute( + "CREATE TABLE IF NOT EXISTS files(" + "path TEXT PRIMARY KEY, mtime REAL, size INTEGER)" + ) + return con + + +def _iter_targets(wiki_root: Path): + """Yield (rel_path, kind, page, text) for every indexable unit.""" + for sub in _INDEX_DIRS: + d = wiki_root / sub + if not d.is_dir(): + continue + for f in sorted(d.glob("*.md")): + try: + text = f.read_text(encoding="utf-8", errors="replace") + except OSError: + continue + # Drop YAML frontmatter: wikilink lists there drown real mentions + # in snippets and pollute ranking. + text = _FRONTMATTER_RE.sub("", text) + yield (f"{sub}/{f.name}", _KIND[sub], None, text) + src = wiki_root / "sources" + if src.is_dir(): + for j in sorted(src.glob("*.json")): + try: + data = json.loads(j.read_text(encoding="utf-8", errors="replace")) + except (OSError, ValueError): + continue + if not isinstance(data, list): + continue + for entry in data: + if isinstance(entry, dict) and "page" in entry: + yield ( + f"sources/{j.name}", + "page", + entry["page"], + str(entry.get("content", "")), + ) + + +def _row_for(rel: str, kind: str, page, text: str): + raw = text[:_RAW_CAP] + return (_bigramize(raw), raw, kind, rel, page) + + +def _sync(con: sqlite3.Connection, wiki_root: Path) -> tuple[int, int]: + """Incremental update: returns (files_seen, rows_written).""" + cur = con.execute("SELECT path, mtime, size FROM files") + known = {p: (m, s) for p, m, s in cur.fetchall()} + seen: dict[str, list] = {} + files_seen = 0 + rows_written = 0 + + # Group rows by file so each file is replaced atomically-ish. + for rel, kind, page, text in _iter_targets(wiki_root): + seen.setdefault(rel, []).append((rel, kind, page, text)) + + wiki_files = {rel: (wiki_root / rel) for rel in seen} + stale = [p for p in known if p not in seen] + + for rel, rows in seen.items(): + f = wiki_files[rel] + try: + st = f.stat() + except OSError: + continue + files_seen += 1 + if known.get(rel) == (st.st_mtime, st.st_size): + continue + con.execute("DELETE FROM fts WHERE path = ?", (rel,)) + con.executemany( + "INSERT INTO fts (content, raw, kind, path, page) VALUES (?,?,?,?,?)", + [_row_for(rel, kind, page, text) for rel, kind, page, text in rows], + ) + con.execute( + "INSERT OR REPLACE INTO files (path, mtime, size) VALUES (?,?,?)", + (rel, st.st_mtime, st.st_size), + ) + rows_written += len(rows) + + for rel in stale: + con.execute("DELETE FROM fts WHERE path = ?", (rel,)) + con.execute("DELETE FROM files WHERE path = ?", (rel,)) + + con.commit() + return files_seen, rows_written + + +def ensure_index(kb_dir: Path) -> Path: + """Build the index if missing, else incrementally refresh it.""" + wiki_root = kb_dir / "wiki" + db_path = fts_db_path(kb_dir) + db_path.parent.mkdir(parents=True, exist_ok=True) + con = _connect(db_path) + try: + _sync(con, wiki_root) + finally: + con.close() + return db_path + + +def rebuild_index(kb_dir: Path) -> tuple[int, int]: + """Drop and rebuild the index. Returns (files_seen, rows_written).""" + wiki_root = kb_dir / "wiki" + db_path = fts_db_path(kb_dir) + db_path.parent.mkdir(parents=True, exist_ok=True) + con = _connect(db_path) + try: + con.execute("DROP TABLE IF EXISTS fts") + con.execute("DROP TABLE IF EXISTS files") + con.commit() + con.close() + con = _connect(db_path) + stats = _sync(con, wiki_root) + finally: + con.close() + return stats + + +def _fts_expr(query: str) -> tuple[str, str, list[str]]: + """Build (AND_expr, OR_expr, raw_terms) from a free-text query.""" + clauses_and: list[str] = [] + clauses_or: list[str] = [] + raw_terms: list[str] = [] + for tok in query.split(): + # Drop anything that is FTS5 syntax (quotes, parens, stars, colons...) + # so pasted punctuation can never break the MATCH expression. + tok = re.sub(r"[^\w一-鿿豈-﫿]+", "", tok) + if not tok: + continue + raw_terms.append(tok) + subs: list[str] = [] + pos = 0 + for m in _CJK_RE.finditer(tok): + if m.start() > pos: + subs.append(("latin", tok[pos : m.start()])) + subs.append(("cjk", m.group(0))) + pos = m.end() + if pos < len(tok): + subs.append(("latin", tok[pos:])) + + if len(subs) > 1: + # Mixed token (e.g. "ETA公司"): adjacency phrase, no prefix star. + parts = [] + for kind, s in subs: + parts.append(_bigramize(s) if kind == "cjk" else s) + clauses_and.append('"' + " ".join(parts) + '"') + clauses_or.append('"' + " ".join(parts) + '"') + elif subs: + kind, s = subs[0] + if kind == "cjk": + expr = '"' + _bigramize(s) + '"' + clauses_and.append(expr) + clauses_or.append(expr) + else: + clauses_and.append(s + "*") + clauses_or.append(s + "*") + return " AND ".join(clauses_and), " OR ".join(clauses_or), raw_terms + + +def _snippet(raw: str, terms: list[str], width: int = 90) -> str: + low = raw.lower() + best = -1 + for t in terms: + i = low.find(t.lower()) + if i >= 0 and (best < 0 or i < best): + best = i + flat = re.sub(r"\s+", " ", raw) + if best < 0: + return flat[: width * 2] + ("…" if len(flat) > width * 2 else "") + start = max(0, best - width) + end = min(len(flat), best + width) + # `best` indexes into raw; recompute on flat text approximately. + i = flat.lower().find(terms[0].lower()) if terms else -1 + if i >= 0: + start = max(0, i - width) + end = min(len(flat), i + width) + return ("…" if start > 0 else "") + flat[start:end] + ("…" if end < len(flat) else "") + + +def search(kb_dir: Path, query: str, limit: int = 8) -> list[dict]: + """BM25-ranked keyword search. Returns [{kind, path, page, snippet, score}].""" + db_path = ensure_index(kb_dir) + and_expr, or_expr, terms = _fts_expr(query) + if not and_expr: + return [] + con = sqlite3.connect(str(db_path)) + try: + rows = con.execute( + "SELECT kind, path, page, raw, bm25(fts) FROM fts" + " WHERE fts MATCH ? ORDER BY bm25(fts) LIMIT ?", + (and_expr, limit), + ).fetchall() + if not rows and or_expr and or_expr != and_expr: + rows = con.execute( + "SELECT kind, path, page, raw, bm25(fts) FROM fts" + " WHERE fts MATCH ? ORDER BY bm25(fts) LIMIT ?", + (or_expr, limit), + ).fetchall() + finally: + con.close() + return [ + { + "kind": k, + "path": p, + "page": pg, + "snippet": _snippet(raw or "", terms), + "score": round(score, 3), + } + for k, p, pg, raw, score in rows + ] + + +def format_hits(hits: list[dict]) -> str: + """Render hits as the text returned by the agent's `search` tool.""" + if not hits: + return ( + "No full-text hits. Fall back to navigation: read index.md, " + "then summaries/concepts/entities." + ) + lines = [] + for i, h in enumerate(hits, 1): + loc = h["path"] + (f" (page {h['page']})" if h["page"] is not None else "") + lines.append(f"[{i}] {h['kind']} | {loc}\n {h['snippet']}") + lines.append( + "\nNext: read_file() the .md hits; for 'page' hits call " + "get_page_content(doc_name, pages) with doc_name = path minus " + "'sources/' and '.json'." + ) + return "\n".join(lines) diff --git a/tests/test_fts.py b/tests/test_fts.py new file mode 100644 index 00000000..a28c9ffa --- /dev/null +++ b/tests/test_fts.py @@ -0,0 +1,143 @@ +"""Unit tests for openkb.fts (FTS5 hybrid retrieval pre-filter). + +All offline: builds tiny temporary wikis and exercises the index directly, +no LLM involved. +""" + +import json + +from openkb.fts import ( + _bigramize, + _fts_expr, + ensure_index, + fts_db_path, + format_hits, + rebuild_index, + search, +) + + +def _make_kb(tmp_path): + kb_dir = tmp_path / "kb" + for sub in ("summaries", "concepts", "entities", "sources"): + (kb_dir / "wiki" / sub).mkdir(parents=True, exist_ok=True) + return kb_dir + + +class TestBigramize: + def test_cjk_run_becomes_overlapping_bigrams(self): + assert _bigramize("川普政府") == "川普 普政 政府" + + def test_single_cjk_char_kept(self): + assert _bigramize("中") == "中" + + def test_latin_untouched(self): + assert _bigramize("DeepSeek v3") == "DeepSeek v3" + + def test_mixed_transforms_cjk_run_in_place(self): + assert _bigramize("ETA公司") == "ETA公司" + + +class TestFtsExpr: + def test_cjk_terms_become_phrases(self): + and_expr, or_expr, terms = _fts_expr("川普 政府") + assert '"川普"' in and_expr and '"政府"' in and_expr + assert " AND " in and_expr and " OR " in or_expr + assert terms == ["川普", "政府"] + + def test_latin_gets_prefix_star(self): + and_expr, _, _ = _fts_expr("transformer") + assert and_expr == "transformer*" + + def test_mixed_token_is_adjacency_phrase(self): + and_expr, _, _ = _fts_expr("ETA公司") + assert and_expr == '"ETA 公司"' + + def test_special_chars_stripped(self): + and_expr, _, _ = _fts_expr('"豁免"(2025)') + assert and_expr == '"豁免 2025"' + + +class TestIndexAndSearch: + def _seed(self, kb_dir): + (kb_dir / "wiki" / "entities" / "feras.md").write_text( + "---\nsee also: [[concepts/eta]]\n---\n" + "# Feras Khatib\n费拉斯·哈提布是 ETA 董事长，曾在美国海军情报部门服役。\n", + encoding="utf-8", + ) + (kb_dir / "wiki" / "summaries" / "deepseek-loi.md").write_text( + "# DeepSeek 意向书\nETA 谋求 DeepSeek 技术在北美与欧盟的独家许可授权。\n", + encoding="utf-8", + ) + (kb_dir / "wiki" / "sources" / "long-doc.json").write_text( + json.dumps( + [ + {"page": 1, "content": "封面与目录。"}, + {"page": 7, "content": "卢特尼克提出芯片豁免的五项要求。"}, + ] + ), + encoding="utf-8", + ) + + def test_rebuild_and_cjk_search(self, tmp_path): + kb_dir = _make_kb(tmp_path) + self._seed(kb_dir) + files, rows = rebuild_index(kb_dir) + assert files == 3 and rows == 4 # 2 md rows + 2 page rows + + hits = search(kb_dir, "董事长") + assert hits and hits[0]["path"] == "entities/feras.md" + assert hits[0]["kind"] == "entity" + assert "费拉斯" in hits[0]["snippet"] + + def test_frontmatter_stripped_from_snippets(self, tmp_path): + kb_dir = _make_kb(tmp_path) + self._seed(kb_dir) + rebuild_index(kb_dir) + hits = search(kb_dir, "eta") + assert hits + assert "see also" not in hits[0]["snippet"] + + def test_long_doc_page_hit(self, tmp_path): + kb_dir = _make_kb(tmp_path) + self._seed(kb_dir) + rebuild_index(kb_dir) + hits = search(kb_dir, "豁免 五项") + assert hits + page_hit = next(h for h in hits if h["kind"] == "page") + assert page_hit["path"] == "sources/long-doc.json" and page_hit["page"] == 7 + + def test_multi_term_and_with_or_fallback(self, tmp_path): + kb_dir = _make_kb(tmp_path) + self._seed(kb_dir) + rebuild_index(kb_dir) + # AND hits: both terms in one doc + assert search(kb_dir, "DeepSeek 独家许可")[0]["path"] == "summaries/deepseek-loi.md" + # No AND hit anywhere -> OR fallback still returns something + hits = search(kb_dir, "独家许可 卢特尼克") + assert len(hits) >= 2 + + def test_no_hit_message(self, tmp_path): + kb_dir = _make_kb(tmp_path) + self._seed(kb_dir) + rebuild_index(kb_dir) + out = format_hits(search(kb_dir, "zqxwv 火星鲱鱼")) + assert "No full-text hits" in out + + def test_incremental_update_picks_up_new_file(self, tmp_path): + kb_dir = _make_kb(tmp_path) + self._seed(kb_dir) + rebuild_index(kb_dir) + (kb_dir / "wiki" / "entities" / "jester.md").write_text( + "# Norman Jester\n环球交换电信负责人，在墨西哥创建全国性电信运营商。\n", + encoding="utf-8", + ) + # ensure_index (not rebuild) must notice the new file + hits = search(kb_dir, "墨西哥") + assert any(h["path"] == "entities/jester.md" for h in hits) + + def test_ensure_index_creates_db(self, tmp_path): + kb_dir = _make_kb(tmp_path) + self._seed(kb_dir) + db = ensure_index(kb_dir) + assert db == fts_db_path(kb_dir) and db.exists() diff --git a/tests/test_query.py b/tests/test_query.py index ecaceabd..4ff5ba88 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -17,13 +17,14 @@ def test_agent_name(self, tmp_path): agent = build_query_agent(str(tmp_path), "gpt-4o-mini") assert agent.name == "wiki-query" - def test_agent_has_three_tools(self, tmp_path): + def test_agent_has_four_tools(self, tmp_path): agent = build_query_agent(str(tmp_path), "gpt-4o-mini") - assert len(agent.tools) == 3 + assert len(agent.tools) == 4 def test_agent_tool_names(self, tmp_path): agent = build_query_agent(str(tmp_path), "gpt-4o-mini") names = {t.name for t in agent.tools} + assert "search" in names assert "read_file" in names assert "get_page_content" in names assert "get_image" in names