From 941ed7b934c75be8bb02565861db0f95badb758c Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 16:25:44 +0200 Subject: [PATCH] feat(agent): add hybrid BM25 search_wiki tool to query/chat agent Adds a dependency-free BM25 full-text index (openkb/fulltext_index.py) over concepts/entities/summaries pages, exposed as a new search_wiki tool alongside index.md-driven navigation in build_query_agent. Additive hybrid retrieval: surfaces pages whose one-line index summary omits a buried detail, without replacing existing navigation. Resolves #233. --- README.md | 2 + openkb/agent/query.py | 29 +++++- openkb/agent/tools.py | 30 ++++++ openkb/fulltext_index.py | 177 +++++++++++++++++++++++++++++++++++ tests/test_agent_tools.py | 39 ++++++++ tests/test_fulltext_index.py | 103 ++++++++++++++++++++ tests/test_query.py | 3 +- 7 files changed, 378 insertions(+), 5 deletions(-) create mode 100644 openkb/fulltext_index.py create mode 100644 tests/test_fulltext_index.py diff --git a/README.md b/README.md index 988bebda0..0a5dbce82 100644 --- a/README.md +++ b/README.md @@ -207,6 +207,8 @@ A "generator" reads from the compiled wiki and produces something usable: an ans `openkb query "..."` answers a single question with a grounded, cited answer from your wiki. `openkb chat` is interactive, an ongoing multi-turn session over the same wiki (`--resume`, `--list`, `--delete` to manage sessions). → Walked through with real saved output in **[`examples/commands/`](examples/commands/)** (query) and **[`examples/chat/`](examples/chat/)** (chat). +Retrieval is hybrid: the agent primarily navigates via `index.md`'s one-line summaries, and additionally has a `search_wiki` tool — a dependency-free BM25 full-text search over `concepts/`, `entities/`, and `summaries/` — for surfacing pages whose index summary doesn't mention a specific buried detail. It's additive, not a replacement, so recall can only improve over index-only navigation. + Inside a chat, type `/` to access slash commands (Tab to complete).
diff --git a/openkb/agent/query.py b/openkb/agent/query.py index da1a939ef..e5cd79cd5 100644 --- a/openkb/agent/query.py +++ b/openkb/agent/query.py @@ -14,6 +14,9 @@ read_wiki_image, write_kb_file, ) +from openkb.agent.tools import ( + search_wiki as search_wiki_impl, +) from openkb.config import LlmCredentialBundle, resolve_model_settings from openkb.schema import get_agents_md @@ -33,18 +36,23 @@ 3. Read concept pages (concepts/) for cross-document synthesis. 4. For "who/what is X" questions about a specific named person, organization, place, or product, read the matching page in entities/ first. -5. When you need detailed source document content, each summary page has a +5. If index.md's one-line summaries don't surface a specific detail you + need (a niche term, an exact figure, a buried fact), use + search_wiki(query) — a keyword-level full-text search over + concepts/entities/summaries. This is a hybrid fallback: use it in + addition to, not instead of, index.md navigation. +6. When you need detailed source document content, each summary page has a `full_text` frontmatter field with the path to the original document content: - Short documents (doc_type: short): read_file with that path. - PageIndex documents (doc_type: pageindex): use get_page_content(doc_name, pages) with tight page ranges. The summary shows document tree structure with page ranges to help you target. Never fetch the whole document. -6. Source content may reference images. Short-doc .md pages link them +7. Source content may reference images. Short-doc .md pages link them note-relative (e.g. ![image](images/doc/file.png), resolved from wiki/sources/); long-doc JSON page metadata lists them wiki-root-relative (e.g. sources/images/doc/file.png). Pass either form as seen to the get_image tool — it accepts both. -7. Synthesize a clear, concise, well-cited answer grounded in wiki content. +8. Synthesize a clear, concise, well-cited answer grounded in wiki content. Answer based only on wiki content. Be concise. Before each tool call, output one short sentence explaining the reason. @@ -83,6 +91,19 @@ def get_page_content(doc_name: str, pages: str) -> str: """ return get_wiki_page_content(doc_name, pages, wiki_root) + @function_tool + def search_wiki(query: str) -> str: + """Full-text (BM25) keyword search over concepts/entities/summaries. + + Hybrid fallback for when index.md's one-line summaries don't surface + a specific buried detail (a niche term, an exact figure, a fact). + Use in addition to, not instead of, index.md navigation. + + Args: + query: Free-text search query (keywords or a natural-language question). + """ + return search_wiki_impl(query, wiki_root) + @function_tool def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: """View an image from the wiki. @@ -117,7 +138,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=[read_file, get_page_content, search_wiki, get_image], model=f"litellm/{model}", model_settings=ModelSettings(**model_settings), ) diff --git a/openkb/agent/tools.py b/openkb/agent/tools.py index eedd388de..a4fa3ad91 100644 --- a/openkb/agent/tools.py +++ b/openkb/agent/tools.py @@ -135,6 +135,36 @@ def get_wiki_page_content(doc_name: str, pages: str, wiki_root: str) -> str: return "\n\n".join(parts) + "\n\n" +def search_wiki(query: str, wiki_root: str, top_k: int = 5) -> str: + """Full-text (BM25) search over concepts/entities/summaries wiki pages. + + Hybrid retrieval helper: complements index.md-driven navigation by + surfacing pages whose one-line index summary doesn't mention a specific + buried detail the query is looking for (a niche term, a figure, an exact + fact). Additive — use alongside, not instead of, index.md navigation. + + Args: + query: Free-text search query (keywords or a natural-language question). + wiki_root: Absolute path to the wiki root directory. + top_k: Maximum number of ranked results to return. + + Returns: + A formatted, ranked list of page hits (wikilink, title, snippet), or + a message indicating no matches were found. + """ + from openkb.fulltext_index import WikiFullTextIndex + + hits = WikiFullTextIndex(wiki_root).search(query, top_k=top_k) + if not hits: + return "No matching pages found." + + lines = [] + for i, hit in enumerate(hits, start=1): + wikilink = hit.path[:-3] if hit.path.endswith(".md") else hit.path + lines.append(f"{i}. [[{wikilink}]] — {hit.title} (score: {hit.score})\n {hit.snippet}") + return "\n".join(lines) + + _MIME_TYPES = { ".png": "image/png", ".jpg": "image/jpeg", diff --git a/openkb/fulltext_index.py b/openkb/fulltext_index.py new file mode 100644 index 000000000..5c7852df9 --- /dev/null +++ b/openkb/fulltext_index.py @@ -0,0 +1,177 @@ +"""Dependency-free BM25 full-text index over compiled wiki pages. + +Hybrid retrieval: the query/chat agent's primary search strategy is +``index.md`` navigation (one-line summaries pointing at pages to read). That +strategy loses recall for details buried deep in a page body that the +one-liner doesn't mention. This module adds an additive, keyword-level +fallback — a BM25 index over the same compiled pages — exposed to the agent +as the ``search_wiki`` tool (see ``openkb.agent.tools.search_wiki``). It is a +union with index-driven navigation, not a replacement, so recall can only +improve relative to index-only navigation, never regress. + +No new dependency: OpenKB pins dependencies exactly and vets each one +deliberately (see ``pyproject.toml``), and BM25 over a few hundred wiki pages +is cheap enough in pure Python that a search-library dependency (e.g. Whoosh) +isn't warranted. +""" + +from __future__ import annotations + +import math +import re +from dataclasses import dataclass +from pathlib import Path + +from openkb.schema import PAGE_CONTENT_DIRS + +_TOKEN_RE = re.compile(r"[a-z0-9]+") + +# Standard BM25 hyperparameters (Robertson/Sparck-Jones defaults). +_K1 = 1.5 +_B = 0.75 + +_SNIPPET_RADIUS = 80 # characters of context on each side of the first match + + +def _tokenize(text: str) -> list[str]: + """Lowercase, alphanumeric-only tokenization (no stemming).""" + return _TOKEN_RE.findall(text.lower()) + + +def _extract_title(text: str) -> str | None: + """Return the first ``# heading`` line's text, or ``None``.""" + for line in text.splitlines(): + stripped = line.strip() + if stripped.startswith("# "): + return stripped[2:].strip() + return None + + +def _make_snippet(text: str, query_terms: list[str]) -> str: + """Return a short excerpt around the first query-term match in *text*.""" + lowered = text.lower() + match_pos = -1 + for term in query_terms: + pos = lowered.find(term) + if pos != -1 and (match_pos == -1 or pos < match_pos): + match_pos = pos + if match_pos == -1: + collapsed = " ".join(text.split()) + truncated = collapsed[: _SNIPPET_RADIUS * 2] + suffix = "…" if len(collapsed) > _SNIPPET_RADIUS * 2 else "" + return truncated + suffix + + start = max(0, match_pos - _SNIPPET_RADIUS) + end = min(len(text), match_pos + _SNIPPET_RADIUS) + collapsed = " ".join(text[start:end].split()) + prefix = "…" if start > 0 else "" + suffix = "…" if end < len(text) else "" + return f"{prefix}{collapsed}{suffix}" + + +@dataclass(frozen=True) +class SearchHit: + """A single BM25 search result over a wiki page.""" + + path: str # wiki-root-relative, e.g. "concepts/attention.md" + title: str + score: float + snippet: str + + +@dataclass(frozen=True) +class _IndexedPage: + path: str + title: str + text: str + tokens: list[str] + + +class WikiFullTextIndex: + """In-memory BM25 index over :data:`PAGE_CONTENT_DIRS` wiki pages. + + Rebuilt fresh on construction — cheap enough at the wiki sizes this + pattern targets (hundreds of pages); no on-disk cache or incremental + update is needed. + """ + + def __init__(self, wiki_root: str | Path) -> None: + self._wiki_root = Path(wiki_root).resolve() + self._pages: list[_IndexedPage] = [] + self._df: dict[str, int] = {} + self._avgdl = 0.0 + self._build() + + def _build(self) -> None: + for subdir in PAGE_CONTENT_DIRS: + target = self._wiki_root / subdir + if not target.is_dir(): + continue + for md_file in sorted(target.glob("*.md")): + text = md_file.read_text(encoding="utf-8") + tokens = _tokenize(text) + if not tokens: + continue + title = _extract_title(text) or md_file.stem + path = f"{subdir}/{md_file.name}" + self._pages.append(_IndexedPage(path=path, title=title, text=text, tokens=tokens)) + + if not self._pages: + return + + self._avgdl = sum(len(page.tokens) for page in self._pages) / len(self._pages) + for page in self._pages: + for term in set(page.tokens): + self._df[term] = self._df.get(term, 0) + 1 + + def _idf(self, term: str) -> float: + n = len(self._pages) + df = self._df.get(term, 0) + # +1 smoothing keeps idf non-negative even for very common terms. + return math.log((n - df + 0.5) / (df + 0.5) + 1) + + def _score(self, query_terms: list[str], page: _IndexedPage) -> float: + dl = len(page.tokens) + tf: dict[str, int] = {} + for term in page.tokens: + tf[term] = tf.get(term, 0) + 1 + + score = 0.0 + for term in query_terms: + f = tf.get(term, 0) + if f == 0: + continue + idf = self._idf(term) + numerator = f * (_K1 + 1) + denominator = f + _K1 * (1 - _B + _B * dl / self._avgdl) + score += idf * (numerator / denominator) + return score + + def search(self, query: str, top_k: int = 5) -> list[SearchHit]: + """Return the ``top_k`` highest-scoring pages for *query* (BM25). + + Args: + query: Free-text search query (keywords or a question). + top_k: Maximum number of results to return. + + Returns: + Ranked hits, highest score first. Empty if the query has no + tokens or the index has no pages. + """ + query_terms = _tokenize(query) + if not query_terms or not self._pages: + return [] + + scored = [(self._score(query_terms, page), page) for page in self._pages] + scored = [(score, page) for score, page in scored if score > 0] + scored.sort(key=lambda item: item[0], reverse=True) + + return [ + SearchHit( + path=page.path, + title=page.title, + score=round(score, 3), + snippet=_make_snippet(page.text, query_terms), + ) + for score, page in scored[:top_k] + ] diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 283a5a8b0..9c0463629 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -9,6 +9,7 @@ parse_pages, read_wiki_file, read_wiki_image, + search_wiki, write_wiki_file, ) @@ -320,3 +321,41 @@ def test_artifact_event_none_for_non_output_zone(): def test_artifact_event_none_for_bad_json(): assert artifact_event_from_write("write_file", "not json", "Written: output/x.html") is None + + +# --------------------------------------------------------------------------- +# search_wiki +# --------------------------------------------------------------------------- + + +class TestSearchWiki: + def test_finds_matching_page(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "cnn.md").write_text( + "# Convolutional Neural Networks\n\nDropout regularization prevents overfitting." + ) + + result = search_wiki("dropout regularization", wiki_root) + + assert "[[concepts/cnn]]" in result + assert "Convolutional Neural Networks" in result + + def test_no_matches_returns_message(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "cnn.md").write_text("# CNN\n\nSomething else entirely.") + + result = search_wiki("nonexistent_keyword_xyz", wiki_root) + + assert result == "No matching pages found." + + def test_respects_top_k(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "entities").mkdir() + for i in range(5): + (tmp_path / "entities" / f"e{i}.md").write_text(f"# Entity {i}\n\nkeyword {i}.") + + result = search_wiki("keyword", wiki_root, top_k=2) + + assert result.count("[[entities/") == 2 diff --git a/tests/test_fulltext_index.py b/tests/test_fulltext_index.py new file mode 100644 index 000000000..750b7498b --- /dev/null +++ b/tests/test_fulltext_index.py @@ -0,0 +1,103 @@ +"""Tests for openkb.fulltext_index (BM25 hybrid search).""" + +from __future__ import annotations + +from openkb.fulltext_index import WikiFullTextIndex + + +def _write(tmp_path, subdir, name, text): + directory = tmp_path / subdir + directory.mkdir(parents=True, exist_ok=True) + (directory / name).write_text(text, encoding="utf-8") + + +class TestWikiFullTextIndex: + def test_empty_wiki_returns_no_hits(self, tmp_path): + index = WikiFullTextIndex(str(tmp_path)) + assert index.search("anything") == [] + + def test_finds_page_by_keyword_in_body(self, tmp_path): + _write( + tmp_path, + "concepts", + "cnn.md", + "# Convolutional Neural Networks\n\nAlexNet popularized ReLU activations " + "and dropout regularization for large-scale image classification.", + ) + _write( + tmp_path, + "concepts", + "unrelated.md", + "# Gardening\n\nTomatoes need plenty of sunlight and water.", + ) + + hits = WikiFullTextIndex(str(tmp_path)).search("dropout regularization") + + assert len(hits) == 1 + assert hits[0].path == "concepts/cnn.md" + assert hits[0].title == "Convolutional Neural Networks" + assert hits[0].score > 0 + + def test_ranks_more_relevant_page_higher(self, tmp_path): + _write( + tmp_path, + "concepts", + "on-topic.md", + "# Topic\n\nAlexNet AlexNet AlexNet training data criticism bias bias.", + ) + _write( + tmp_path, + "concepts", + "off-topic.md", + "# Other\n\nA single passing mention of AlexNet in an unrelated paragraph " + "about something else entirely, padded with filler words to change length.", + ) + + hits = WikiFullTextIndex(str(tmp_path)).search("AlexNet bias") + + assert [hit.path for hit in hits[:1]] == ["concepts/on-topic.md"] + + def test_respects_top_k(self, tmp_path): + for i in range(10): + _write(tmp_path, "entities", f"e{i}.md", f"# Entity {i}\n\nkeyword appears here {i}.") + + hits = WikiFullTextIndex(str(tmp_path)).search("keyword", top_k=3) + + assert len(hits) == 3 + + def test_only_indexes_page_content_dirs(self, tmp_path): + _write(tmp_path, "sources", "raw.md", "# Raw\n\nkeyword raw source content.") + _write(tmp_path, "concepts", "c.md", "# Concept\n\nkeyword concept content.") + + hits = WikiFullTextIndex(str(tmp_path)).search("keyword") + + assert [hit.path for hit in hits] == ["concepts/c.md"] + + def test_falls_back_to_filename_when_no_heading(self, tmp_path): + _write(tmp_path, "summaries", "no-heading.md", "keyword content without a heading line.") + + hits = WikiFullTextIndex(str(tmp_path)).search("keyword") + + assert hits[0].title == "no-heading" + + def test_no_query_tokens_returns_no_hits(self, tmp_path): + _write(tmp_path, "concepts", "c.md", "# Concept\n\nkeyword concept content.") + + hits = WikiFullTextIndex(str(tmp_path)).search(" ") + + assert hits == [] + + def test_snippet_contains_context_around_match(self, tmp_path): + _write( + tmp_path, + "concepts", + "c.md", + "# Concept\n\n" + + ("padding " * 40) + + "the exact fee is five hundred dollars" + + (" more" * 40), + ) + + hits = WikiFullTextIndex(str(tmp_path)).search("fee") + + assert "fee" in hits[0].snippet.lower() diff --git a/tests/test_query.py b/tests/test_query.py index ecaceabd9..a720ccce3 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -19,13 +19,14 @@ def test_agent_name(self, tmp_path): def test_agent_has_three_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 "read_file" in names assert "get_page_content" in names + assert "search_wiki" in names assert "get_image" in names def test_instructions_mention_get_page_content(self, tmp_path):