From 941ed7b934c75be8bb02565861db0f95badb758c Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 28 Aug 2026 16:25:44 +0200 Subject: [PATCH 1/3] 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): From 0025ea06cdf6591203ebbd3b82ea0ec6be683d7a Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 11 Sep 2026 11:33:05 +0200 Subject: [PATCH 2/3] feat(search): add tiered BM25 search (briefs/summaries/sources) and taxonomy accessors - fulltext_index.py: extract shared _BM25Scorer from WikiFullTextIndex (no behavior change), add Locator (line/page) on SearchHit, add TieredWikiSearch with three independent tiers over summaries/ (briefs + full body) and sources/ (whole-file .md + per-page PageIndex .json, never the whole long doc as one BM25 unit). - frontmatter.py: add resolve_description()/body_only() shared helpers (kept separate from agent.compiler._resolve_description, which is under active unrelated development). - agent/tools.py: add list_taxonomy_items()/get_taxonomy_item() for semantic browsing of persisted concepts/entities (pending candidates in PendingTopicsStore are structurally excluded). - No wiring into CLI/MCP/query-agent yet (follow-up PRs); WikiFullTextIndex and agent.tools.search_wiki keep their existing signature/behavior. --- openkb/agent/tools.py | 104 ++++++++++++ openkb/frontmatter.py | 26 +++ openkb/fulltext_index.py | 307 +++++++++++++++++++++++++++++++---- tests/test_agent_tools.py | 118 ++++++++++++++ tests/test_fulltext_index.py | 153 ++++++++++++++++- 5 files changed, 678 insertions(+), 30 deletions(-) diff --git a/openkb/agent/tools.py b/openkb/agent/tools.py index a4fa3ad91..6d0d45fd9 100644 --- a/openkb/agent/tools.py +++ b/openkb/agent/tools.py @@ -9,10 +9,17 @@ import contextlib import json as _json +from dataclasses import dataclass from pathlib import Path, PurePosixPath +from typing import Literal +from openkb import frontmatter from openkb.locks import atomic_write_text +# Maps a taxonomy "kind" to its wiki subdirectory. Single source of truth for +# list_taxonomy_items/get_taxonomy_item below. +_TAXONOMY_DIRS: dict[str, str] = {"concept": "concepts", "entity": "entities"} + def list_wiki_files(directory: str, wiki_root: str) -> str: """List all Markdown files in a wiki subdirectory. @@ -135,6 +142,103 @@ def get_wiki_page_content(doc_name: str, pages: str, wiki_root: str) -> str: return "\n\n".join(parts) + "\n\n" +@dataclass(frozen=True) +class TaxonomyItem: + """One persisted concept or entity page (never a pending candidate). + + ``PendingTopicsStore`` (see ``openkb.pending``) buffers not-yet-paged + concept/entity candidates separately from the compiled ``.md`` pages + under ``concepts/``/``entities/`` — this dataclass, and + :func:`list_taxonomy_items`, only ever surface the latter, so a caller + never sees an in-progress candidate as if it were a real page. + """ + + kind: Literal["concept", "entity"] + slug: str + path: str # wiki-root-relative, e.g. "concepts/attention.md" + brief: str + # Entity type (e.g. "person", "organization"); always None for concepts. + type: str | None = None + + +def list_taxonomy_items(wiki_root: str, kind: str | None = None) -> list[TaxonomyItem]: + """List persisted concept and/or entity pages with their one-line briefs. + + Intended as the first step of the search strategy: browse this compact, + semantically-scannable list and let the caller (an LLM) pick the + relevant slug(s) by meaning — this is deliberately not a keyword search + (see ``search_wiki`` for that, over summaries/sources only). + + Args: + wiki_root: Absolute path to the wiki root directory. + kind: Restrict to ``"concept"`` or ``"entity"``; ``None`` returns both. + + Returns: + Items sorted by kind, then slug. Empty list if the KB has neither + directory yet or both are empty. + + Raises: + ValueError: *kind* is neither ``None``, ``"concept"``, nor ``"entity"``. + """ + root = Path(wiki_root).resolve() + kinds = [kind] if kind else ["concept", "entity"] + for k in kinds: + if k not in _TAXONOMY_DIRS: + raise ValueError(f"Unknown kind {k!r}; expected 'concept' or 'entity'.") + + items: list[TaxonomyItem] = [] + for k in kinds: + directory = root / _TAXONOMY_DIRS[k] + if not directory.is_dir(): + continue + for md_file in sorted(directory.glob("*.md")): + text = md_file.read_text(encoding="utf-8") + fm = frontmatter.parse(text) + brief = frontmatter.resolve_description(fm) + etype = None + if k == "entity": + etype = str(fm.get("type") or "").strip().lower() or "other" + items.append( + TaxonomyItem( + kind=k, # type: ignore[arg-type] # validated against _TAXONOMY_DIRS above + slug=md_file.stem, + path=f"{_TAXONOMY_DIRS[k]}/{md_file.name}", + brief=brief, + type=etype, + ) + ) + return items + + +def get_taxonomy_item(slug: str, wiki_root: str, kind: str | None = None) -> str: + """Read a persisted concept or entity page's full Markdown content. + + Args: + slug: Page slug (filename without ``.md``), e.g. ``"attention"``. + wiki_root: Absolute path to the wiki root directory. + kind: ``"concept"`` or ``"entity"`` to disambiguate a same-named + slug; ``None`` checks ``concepts/`` first, then ``entities/``. + + Returns: + Full file content, or a "not found" message if no match exists in + the requested (or either) directory. + + Raises: + ValueError: *kind* is neither ``None``, ``"concept"``, nor ``"entity"``. + """ + root = Path(wiki_root).resolve() + kinds = [kind] if kind else ["concept", "entity"] + for k in kinds: + if k not in _TAXONOMY_DIRS: + raise ValueError(f"Unknown kind {k!r}; expected 'concept' or 'entity'.") + + for k in kinds: + path = (root / _TAXONOMY_DIRS[k] / f"{slug}.md").resolve() + if path.is_relative_to(root) and path.exists(): + return path.read_text(encoding="utf-8") + return f"Taxonomy item not found: {slug}" + + def search_wiki(query: str, wiki_root: str, top_k: int = 5) -> str: """Full-text (BM25) search over concepts/entities/summaries wiki pages. diff --git a/openkb/frontmatter.py b/openkb/frontmatter.py index 34c504bac..9143f2f92 100644 --- a/openkb/frontmatter.py +++ b/openkb/frontmatter.py @@ -110,3 +110,29 @@ def set_line(fm_block: str, key: str, value: str) -> str: def drop_line(fm_block: str, key: str) -> str: """Remove any ``key:`` line from a frontmatter block (no-op if absent).""" return re.sub(rf"^{re.escape(key)}:.*\n?", "", fm_block, flags=re.MULTILINE) + + +def resolve_description(fm: dict) -> str: + """Return a non-empty description string from a parsed frontmatter dict. + + Checks ``description`` first, then the legacy ``brief`` key (pre-migration + pages). Returns an empty string when neither key holds a non-blank value. + Mirrors ``agent.compiler._resolve_description`` — kept as a separate, + dependency-free copy here so callers outside the compiler (search/taxonomy + tooling) don't need to import from ``agent.compiler``, which is under + active, unrelated development. + """ + for key in ("description", "brief"): + v = fm.get(key) + if isinstance(v, str) and v.strip(): + return v.strip() + return "" + + +def body_only(text: str) -> str: + """Return *text* with any leading YAML frontmatter block removed. + + Returns *text* unchanged when it has no well-formed frontmatter. + """ + parts = split(text) + return parts[1] if parts is not None else text diff --git a/openkb/fulltext_index.py b/openkb/fulltext_index.py index 5c7852df9..b55cd6ef4 100644 --- a/openkb/fulltext_index.py +++ b/openkb/fulltext_index.py @@ -9,6 +9,30 @@ union with index-driven navigation, not a replacement, so recall can only improve relative to index-only navigation, never regress. +Concepts and entities are deliberately excluded from full-text search (see +:class:`TieredWikiSearch` below) — they are found by semantic browsing +(``list_taxonomy_items``/``get_taxonomy_item`` in ``agent.tools``), not +keyword search, so :class:`WikiFullTextIndex` (kept for backward +compatibility with the original single-tier ``search_wiki`` tool) and +:class:`TieredWikiSearch` cover different, non-overlapping surfaces: + +- :class:`WikiFullTextIndex` — the original combined BM25 index over + ``concepts/`` + ``entities/`` + ``summaries/`` (:data:`PAGE_CONTENT_DIRS`). +- :class:`TieredWikiSearch` — three independent BM25 tiers, each scoped to a + different part of a document's lifecycle so a query only "wastes" recall + budget on the granularity it's actually likely to match at: + 1. ``briefs`` — one-line ``description``/``brief`` frontmatter per + ``summaries/*.md`` (same short text ``index.md`` shows). High precision, + low recall — good for on-topic queries, filters out incidental + word-frequency noise from long documents. + 2. ``summaries`` — full body of ``summaries/*.md``. Higher recall for + specific terms/figures the one-liner omits. + 3. ``sources`` — raw ``sources/*.md`` (whole file) and ``sources/*.json`` + PageIndex documents (indexed **per page**, not per document, so a hit + can point at an exact page via a :class:`Locator` instead of forcing a + re-score over an entire long document). Covers details that never make + it into a summary at all (creation dates, authors, exact field names). + 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) @@ -17,11 +41,14 @@ from __future__ import annotations +import json as _json import math import re from dataclasses import dataclass from pathlib import Path +from typing import Literal +from openkb import frontmatter from openkb.schema import PAGE_CONTENT_DIRS _TOKEN_RE = re.compile(r"[a-z0-9]+") @@ -32,6 +59,9 @@ _SNIPPET_RADIUS = 80 # characters of context on each side of the first match +# Valid `scope` values for TieredWikiSearch.search() — one BM25 tier each. +TIERED_SCOPES = ("briefs", "summaries", "sources") + def _tokenize(text: str) -> list[str]: """Lowercase, alphanumeric-only tokenization (no stemming).""" @@ -69,14 +99,30 @@ def _make_snippet(text: str, query_terms: list[str]) -> str: return f"{prefix}{collapsed}{suffix}" +@dataclass(frozen=True) +class Locator: + """Points at a specific location within a hit's page for a follow-up read. + + ``kind="line"``: 1-based line number within a Markdown file (computed at + query time from the first query-term match). ``kind="page"``: 1-based + PageIndex page number within a long-doc ``sources/*.json`` array — fixed + at indexing time (one page = one BM25 "document"), and directly usable + with ``get_page_content(doc_name, pages=str(value))``. + """ + + kind: Literal["line", "page"] + value: int + + @dataclass(frozen=True) class SearchHit: - """A single BM25 search result over a wiki page.""" + """A single BM25 search result over a wiki page (or a PageIndex page).""" path: str # wiki-root-relative, e.g. "concepts/attention.md" title: str score: float snippet: str + locator: Locator | None = None @dataclass(frozen=True) @@ -85,42 +131,42 @@ class _IndexedPage: title: str text: str tokens: list[str] + # Fixed at indexing time for pseudo-documents that are inherently + # page-scoped (one PageIndex page = one _IndexedPage); None otherwise, in + # which case _BM25Scorer computes a "line" Locator at query time instead. + fixed_locator: Locator | None = None -class WikiFullTextIndex: - """In-memory BM25 index over :data:`PAGE_CONTENT_DIRS` wiki pages. +def _find_line_locator(text: str, query_terms: list[str]) -> Locator | None: + """Return a 1-based ``line`` Locator for the first query-term match line. - 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. + Returns ``None`` if no line contains any query term (can happen when the + match is only visible after tokenization, e.g. across punctuation). """ + for line_no, line in enumerate(text.splitlines(), start=1): + lowered = line.lower() + if any(term in lowered for term in query_terms): + return Locator(kind="line", value=line_no) + return None - def __init__(self, wiki_root: str | Path) -> None: - self._wiki_root = Path(wiki_root).resolve() - self._pages: list[_IndexedPage] = [] + +class _BM25Scorer: + """Pure BM25 ranking (Robertson/Sparck-Jones) over a fixed page list. + + Extracted from the original :class:`WikiFullTextIndex` so the same + scoring math is shared between the legacy combined index and + :class:`TieredWikiSearch`'s three independent tiers, without duplicating + the formula. No I/O — callers build the ``pages`` list. + """ + + def __init__(self, pages: list[_IndexedPage]) -> None: + self._pages = pages 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: + if not pages: return - - self._avgdl = sum(len(page.tokens) for page in self._pages) / len(self._pages) - for page in self._pages: + self._avgdl = sum(len(page.tokens) for page in pages) / len(pages) + for page in pages: for term in set(page.tokens): self._df[term] = self._df.get(term, 0) + 1 @@ -172,6 +218,209 @@ def search(self, query: str, top_k: int = 5) -> list[SearchHit]: title=page.title, score=round(score, 3), snippet=_make_snippet(page.text, query_terms), + locator=page.fixed_locator or _find_line_locator(page.text, query_terms), ) for score, page in scored[:top_k] ] + + +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. Kept for backward compatibility with the original + (PR #234) single-tier ``search_wiki`` tool — new callers should prefer + :class:`TieredWikiSearch`, which separates concepts/entities (browsed via + ``list_taxonomy_items``, not indexed here) from summaries/sources. + """ + + def __init__(self, wiki_root: str | Path) -> None: + self._wiki_root = Path(wiki_root).resolve() + self._pages: list[_IndexedPage] = _build_pages_from_dirs(self._wiki_root, PAGE_CONTENT_DIRS) + self._scorer = _BM25Scorer(self._pages) + + 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. + """ + return self._scorer.search(query, top_k=top_k) + + +def _build_pages_from_dirs(wiki_root: Path, subdirs: tuple[str, ...]) -> list[_IndexedPage]: + """Index every ``*.md`` file's full text under each of *subdirs*.""" + pages: list[_IndexedPage] = [] + for subdir in subdirs: + target = 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}" + pages.append(_IndexedPage(path=path, title=title, text=text, tokens=tokens)) + return pages + + +def _build_brief_pages(wiki_root: Path) -> list[_IndexedPage]: + """One pseudo-document per ``summaries/*.md``, text = its one-line brief. + + Uses the ``description``/legacy ``brief`` frontmatter field — the same + short text ``index.md``'s ``## Documents`` section shows — not the full + body. Pages without a resolvable brief are skipped (nothing to index). + """ + summaries_dir = wiki_root / "summaries" + if not summaries_dir.is_dir(): + return [] + pages: list[_IndexedPage] = [] + for md_file in sorted(summaries_dir.glob("*.md")): + text = md_file.read_text(encoding="utf-8") + brief = frontmatter.resolve_description(frontmatter.parse(text)) + tokens = _tokenize(brief) + if not tokens: + continue + title = _extract_title(text) or md_file.stem + pages.append( + _IndexedPage(path=f"summaries/{md_file.name}", title=title, text=brief, tokens=tokens) + ) + return pages + + +def _build_summary_pages(wiki_root: Path) -> list[_IndexedPage]: + """One document per ``summaries/*.md``, text = full body (no frontmatter).""" + summaries_dir = wiki_root / "summaries" + if not summaries_dir.is_dir(): + return [] + pages: list[_IndexedPage] = [] + for md_file in sorted(summaries_dir.glob("*.md")): + text = md_file.read_text(encoding="utf-8") + body = frontmatter.body_only(text) + tokens = _tokenize(body) + if not tokens: + continue + title = _extract_title(text) or md_file.stem + pages.append( + _IndexedPage(path=f"summaries/{md_file.name}", title=title, text=body, tokens=tokens) + ) + return pages + + +def _build_source_pages(wiki_root: Path) -> list[_IndexedPage]: + """Sources tier: ``sources/*.md`` (whole file) + ``sources/*.json`` (per page). + + A PageIndex ``sources/*.json`` document is a JSON array of + ``{"page": int, "content": str, ...}`` objects (see + ``agent.tools.get_wiki_page_content``). Each page is indexed as its own + ``_IndexedPage`` with a fixed ``page`` :class:`Locator` — never the whole + document as one BM25 unit — so a hit points at an exact page instead of + diluting the score across a potentially very long document, and so the + locator is directly usable with ``get_page_content(doc_name, pages=...)``. + """ + sources_dir = wiki_root / "sources" + if not sources_dir.is_dir(): + return [] + pages: list[_IndexedPage] = [] + for src_file in sorted(sources_dir.iterdir()): + if src_file.suffix == ".md": + text = src_file.read_text(encoding="utf-8") + tokens = _tokenize(text) + if not tokens: + continue + title = _extract_title(text) or src_file.stem + pages.append( + _IndexedPage(path=f"sources/{src_file.name}", title=title, text=text, tokens=tokens) + ) + elif src_file.suffix == ".json": + pages.extend(_index_pageindex_source(src_file)) + return pages + + +def _index_pageindex_source(src_file: Path) -> list[_IndexedPage]: + """Return one ``_IndexedPage`` per page of a PageIndex ``sources/*.json`` doc. + + Tolerant of malformed/foreign JSON (skips, doesn't raise) — a hand-edited + or unexpected file under ``sources/`` shouldn't break indexing of the rest + of the KB. + """ + try: + data = _json.loads(src_file.read_text(encoding="utf-8")) + except (_json.JSONDecodeError, OSError, UnicodeDecodeError): + return [] + if not isinstance(data, list): + return [] + + pages: list[_IndexedPage] = [] + for entry in data: + if not isinstance(entry, dict): + continue + page_num = entry.get("page") + content = entry.get("content", "") + if not isinstance(page_num, int) or not isinstance(content, str): + continue + tokens = _tokenize(content) + if not tokens: + continue + pages.append( + _IndexedPage( + path=f"sources/{src_file.name}", + title=f"{src_file.stem} (page {page_num})", + text=content, + tokens=tokens, + fixed_locator=Locator(kind="page", value=page_num), + ) + ) + return pages + + +class TieredWikiSearch: + """Three independent BM25 tiers over ``summaries/`` and ``sources/``. + + Concepts and entities are intentionally out of scope here — they are + browsed semantically via ``list_taxonomy_items``/``get_taxonomy_item`` + (``agent.tools``), not keyword-searched. Rebuilt fresh on construction, + same no-cache rationale as :class:`WikiFullTextIndex` (see module + docstring); cheap at the wiki sizes this pattern targets. + """ + + def __init__(self, wiki_root: str | Path) -> None: + wiki_root = Path(wiki_root).resolve() + self._scorers: dict[str, _BM25Scorer] = { + "briefs": _BM25Scorer(_build_brief_pages(wiki_root)), + "summaries": _BM25Scorer(_build_summary_pages(wiki_root)), + "sources": _BM25Scorer(_build_source_pages(wiki_root)), + } + + def search( + self, query: str, scope: list[str] | None = None, top_k: int = 5 + ) -> dict[str, list[SearchHit]]: + """Search one or more tiers; returns ``{tier_name: [SearchHit, ...]}``. + + Args: + query: Free-text search query (keywords or a question). + scope: Subset of :data:`TIERED_SCOPES` to search; ``None`` + searches all three tiers. + top_k: Maximum ranked results to return per tier. + + Returns: + One entry per searched tier (only the requested/valid tiers are + present as keys — never an empty-list placeholder for tiers the + caller didn't ask for). + + Raises: + ValueError: *scope* contains a name outside :data:`TIERED_SCOPES`. + """ + tiers = scope if scope else list(TIERED_SCOPES) + invalid = [t for t in tiers if t not in TIERED_SCOPES] + if invalid: + raise ValueError(f"Unknown scope(s) {invalid}; expected any of {TIERED_SCOPES}.") + return {tier: self._scorers[tier].search(query, top_k=top_k) for tier in tiers} diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index 9c0463629..dba78644c 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -3,8 +3,11 @@ from __future__ import annotations from openkb.agent.tools import ( + TaxonomyItem, artifact_event_from_write, + get_taxonomy_item, get_wiki_page_content, + list_taxonomy_items, list_wiki_files, parse_pages, read_wiki_file, @@ -359,3 +362,118 @@ def test_respects_top_k(self, tmp_path): result = search_wiki("keyword", wiki_root, top_k=2) assert result.count("[[entities/") == 2 + + +# --------------------------------------------------------------------------- +# list_taxonomy_items / get_taxonomy_item +# --------------------------------------------------------------------------- + + +class TestListTaxonomyItems: + def test_lists_concepts_and_entities_by_default(self, tmp_path): + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "attention.md").write_text( + '---\ndescription: "How attention works"\n---\n\n# Attention\n\nBody.' + ) + (tmp_path / "entities").mkdir() + (tmp_path / "entities" / "acme.md").write_text( + '---\ntype: organization\ndescription: "A company"\n---\n\n# Acme\n\nBody.' + ) + + items = list_taxonomy_items(str(tmp_path)) + + assert len(items) == 2 + by_slug = {i.slug: i for i in items} + assert by_slug["attention"].kind == "concept" + assert by_slug["attention"].brief == "How attention works" + assert by_slug["attention"].type is None + assert by_slug["acme"].kind == "entity" + assert by_slug["acme"].type == "organization" + assert by_slug["acme"].brief == "A company" + + def test_kind_filter_restricts_to_one_directory(self, tmp_path): + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "c.md").write_text("# C\n\nBody.") + (tmp_path / "entities").mkdir() + (tmp_path / "entities" / "e.md").write_text("# E\n\nBody.") + + items = list_taxonomy_items(str(tmp_path), kind="concept") + + assert len(items) == 1 + assert items[0].kind == "concept" + + def test_legacy_brief_key_resolves(self, tmp_path): + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "c.md").write_text('---\nbrief: "legacy brief"\n---\n\n# C\n\nX.') + + items = list_taxonomy_items(str(tmp_path)) + + assert items[0].brief == "legacy brief" + + def test_missing_directories_return_empty_list(self, tmp_path): + assert list_taxonomy_items(str(tmp_path)) == [] + + def test_no_frontmatter_yields_empty_brief(self, tmp_path): + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "c.md").write_text("# C\n\nNo frontmatter here.") + + items = list_taxonomy_items(str(tmp_path)) + + assert items[0].brief == "" + + def test_invalid_kind_raises_value_error(self, tmp_path): + import pytest + + with pytest.raises(ValueError, match="Unknown kind"): + list_taxonomy_items(str(tmp_path), kind="document") + + def test_items_are_taxonomy_item_instances(self, tmp_path): + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "c.md").write_text("# C\n\nBody.") + + items = list_taxonomy_items(str(tmp_path)) + + assert isinstance(items[0], TaxonomyItem) + + +class TestGetTaxonomyItem: + def test_reads_concept_page(self, tmp_path): + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "attention.md").write_text("# Attention\n\nFull content here.") + + result = get_taxonomy_item("attention", str(tmp_path)) + + assert "Full content here." in result + + def test_kind_disambiguates_same_slug(self, tmp_path): + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "acme.md").write_text("# Acme concept") + (tmp_path / "entities").mkdir() + (tmp_path / "entities" / "acme.md").write_text("# Acme entity") + + assert "concept" in get_taxonomy_item("acme", str(tmp_path), kind="concept") + assert "entity" in get_taxonomy_item("acme", str(tmp_path), kind="entity") + + def test_without_kind_checks_concepts_before_entities(self, tmp_path): + (tmp_path / "entities").mkdir() + (tmp_path / "entities" / "acme.md").write_text("# Acme entity only") + + result = get_taxonomy_item("acme", str(tmp_path)) + + assert "Acme entity only" in result + + def test_not_found_returns_message(self, tmp_path): + result = get_taxonomy_item("nonexistent", str(tmp_path)) + + assert result == "Taxonomy item not found: nonexistent" + + def test_invalid_kind_raises_value_error(self, tmp_path): + import pytest + + with pytest.raises(ValueError, match="Unknown kind"): + get_taxonomy_item("slug", str(tmp_path), kind="document") + + def test_path_traversal_is_rejected(self, tmp_path): + result = get_taxonomy_item("../../etc/passwd", str(tmp_path)) + + assert result == "Taxonomy item not found: ../../etc/passwd" diff --git a/tests/test_fulltext_index.py b/tests/test_fulltext_index.py index 750b7498b..8ed48e2a2 100644 --- a/tests/test_fulltext_index.py +++ b/tests/test_fulltext_index.py @@ -2,7 +2,9 @@ from __future__ import annotations -from openkb.fulltext_index import WikiFullTextIndex +import json + +from openkb.fulltext_index import Locator, TieredWikiSearch, WikiFullTextIndex def _write(tmp_path, subdir, name, text): @@ -101,3 +103,152 @@ def test_snippet_contains_context_around_match(self, tmp_path): hits = WikiFullTextIndex(str(tmp_path)).search("fee") assert "fee" in hits[0].snippet.lower() + + +class TestTieredWikiSearchBriefs: + def test_matches_brief_frontmatter_not_body(self, tmp_path): + _write( + tmp_path, + "summaries", + "doc-a.md", + '---\ndescription: "Salesforce Case Management overview"\n---\n\n' + "# Doc A\n\nUnrelated body text about something else entirely.", + ) + + result = TieredWikiSearch(str(tmp_path)).search("case management", scope=["briefs"]) + + assert len(result["briefs"]) == 1 + assert result["briefs"][0].path == "summaries/doc-a.md" + + def test_legacy_brief_key_still_resolves(self, tmp_path): + _write( + tmp_path, + "summaries", + "doc-a.md", + '---\nbrief: "legacy field name lookup notes"\n---\n\n# Doc A\n\nBody.', + ) + + result = TieredWikiSearch(str(tmp_path)).search("field name lookup", scope=["briefs"]) + + assert len(result["briefs"]) == 1 + + def test_no_brief_frontmatter_yields_no_hit(self, tmp_path): + _write(tmp_path, "summaries", "doc-a.md", "# Doc A\n\nkeyword body text, no frontmatter.") + + result = TieredWikiSearch(str(tmp_path)).search("keyword", scope=["briefs"]) + + assert result["briefs"] == [] + + +class TestTieredWikiSearchSummaries: + def test_matches_full_body_not_just_brief(self, tmp_path): + _write( + tmp_path, + "summaries", + "doc-a.md", + '---\ndescription: "General overview"\n---\n\n' + "# Doc A\n\nDetails about custom_field_xyz appear only here.", + ) + + result = TieredWikiSearch(str(tmp_path)).search("custom_field_xyz", scope=["summaries"]) + + assert len(result["summaries"]) == 1 + assert result["summaries"][0].locator is not None + assert result["summaries"][0].locator.kind == "line" + + def test_frontmatter_block_itself_is_not_indexed(self, tmp_path): + _write( + tmp_path, + "summaries", + "doc-a.md", + '---\ndescription: "uniquefrontmatterterm should not match body search"\n---\n\n' + "# Doc A\n\nUnrelated body.", + ) + + result = TieredWikiSearch(str(tmp_path)).search( + "uniquefrontmatterterm", scope=["summaries"] + ) + + assert result["summaries"] == [] + + +class TestTieredWikiSearchSources: + def test_short_source_doc_gets_line_locator(self, tmp_path): + _write( + tmp_path, + "sources", + "notes.md", + "Line one.\nLine two.\nAuthor: Jane Doe, created 2024-03-15.\nLine four.", + ) + + result = TieredWikiSearch(str(tmp_path)).search("Jane Doe", scope=["sources"]) + + assert len(result["sources"]) == 1 + hit = result["sources"][0] + assert hit.path == "sources/notes.md" + assert hit.locator == Locator(kind="line", value=3) + + def test_pageindex_json_hit_gets_page_locator_not_whole_document(self, tmp_path): + pages = [ + {"page": 1, "content": "Introduction, nothing special here."}, + {"page": 2, "content": "The field_xyz default value is 42."}, + {"page": 3, "content": "Conclusion, also nothing special."}, + ] + sources_dir = tmp_path / "sources" + sources_dir.mkdir(parents=True) + (sources_dir / "long-doc.json").write_text(json.dumps(pages), encoding="utf-8") + + result = TieredWikiSearch(str(tmp_path)).search("field_xyz", scope=["sources"]) + + assert len(result["sources"]) == 1 + hit = result["sources"][0] + assert hit.path == "sources/long-doc.json" + assert hit.locator == Locator(kind="page", value=2) + + def test_malformed_json_source_is_skipped_not_raised(self, tmp_path): + sources_dir = tmp_path / "sources" + sources_dir.mkdir(parents=True) + (sources_dir / "broken.json").write_text("{not valid json", encoding="utf-8") + + result = TieredWikiSearch(str(tmp_path)).search("anything", scope=["sources"]) + + assert result["sources"] == [] + + +class TestTieredWikiSearchScope: + def test_default_scope_searches_all_three_tiers(self, tmp_path): + _write( + tmp_path, + "summaries", + "doc.md", + '---\ndescription: "keyword brief"\n---\n\n# Doc\n\nkeyword body.', + ) + _write(tmp_path, "sources", "doc.md", "keyword raw source.") + + result = TieredWikiSearch(str(tmp_path)).search("keyword") + + assert set(result.keys()) == {"briefs", "summaries", "sources"} + assert len(result["briefs"]) == 1 + assert len(result["summaries"]) == 1 + assert len(result["sources"]) == 1 + + def test_concepts_and_entities_are_never_searched(self, tmp_path): + _write(tmp_path, "concepts", "c.md", "# Concept\n\nkeyword concept content.") + _write(tmp_path, "entities", "e.md", "# Entity\n\nkeyword entity content.") + + result = TieredWikiSearch(str(tmp_path)).search("keyword") + + assert result["briefs"] == [] + assert result["summaries"] == [] + assert result["sources"] == [] + + def test_invalid_scope_raises_value_error(self, tmp_path): + import pytest + + with pytest.raises(ValueError, match="Unknown scope"): + TieredWikiSearch(str(tmp_path)).search("keyword", scope=["not-a-real-tier"]) + + def test_empty_wiki_returns_empty_lists_for_all_tiers(self, tmp_path): + result = TieredWikiSearch(str(tmp_path)).search("anything") + + assert result == {"briefs": [], "summaries": [], "sources": []} From ff7a333207e4a51a66ad50ad1f555c99b4164268 Mon Sep 17 00:00:00 2001 From: Sebastian Braun Date: Fri, 11 Sep 2026 11:47:17 +0200 Subject: [PATCH 3/3] feat(cli,agent): expose tiered search + taxonomy via CLI and wire into query/chat agent - cli.py: new 'openkb list-taxonomy [--kind concept|entity] [--json]' and 'openkb search [--scope briefs,summaries,sources] [--top-k N] [--json]' commands. - agent/tools.py: search_wiki now searches the new tiered briefs/summaries/sources index instead of the old combined concepts+entities+summaries index (concepts/entities move to the new list_taxonomy tool - semantic browsing, not keyword search); new list_taxonomy() text-formatting wrapper over list_taxonomy_items(). - agent/query.py (+ chat.py via tool inheritance): wires list_taxonomy and the retiered search_wiki in as agent tools; search strategy instructions updated to browse taxonomy first, then use scope-restricted search_wiki as a keyword fallback. - README.md: updated hybrid-retrieval paragraph and command table. - Intentional behavior change to agent.tools.search_wiki (scope param, concepts/entities out of scope, output grouped by tier) - safe since #234/#259, which introduced it, are not yet merged upstream; existing tests updated to the new contract. --- README.md | 4 +- openkb/agent/query.py | 74 ++++++++++++++++------- openkb/agent/tools.py | 83 ++++++++++++++++++++------ openkb/cli.py | 122 ++++++++++++++++++++++++++++++++++++++ tests/test_agent_tools.py | 100 ++++++++++++++++++++++++++++--- tests/test_query.py | 5 +- 6 files changed, 337 insertions(+), 51 deletions(-) diff --git a/README.md b/README.md index 0a5dbce82..30e8966cf 100644 --- a/README.md +++ b/README.md @@ -185,6 +185,8 @@ OpenKB commands fall into two layers: the **wiki foundation** (compile + manage | ------------------------------------------------------------------ | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ | | openkb remove <doc> | Remove a document and clean up its wiki pages, images, registry, and PageIndex state (`--dry-run` to preview, `--keep-raw` / `--keep-empty` to retain artifacts) | | openkb recompile [<doc>] [--all] | Re-run the compile pipeline on already-indexed docs without re-indexing. Regenerates summaries and rewrites concept pages; manual edits are overwritten (`--dry-run` to preview, `--refresh-schema` to also update `wiki/AGENTS.md`) | +| openkb list-taxonomy [--kind concept|entity] | List persisted concept/entity pages with their one-line briefs — semantic browsing, not keyword search (`--json` for scripting) | +| openkb search "query" [--scope briefs,summaries,sources] | Tiered BM25 full-text search over `summaries/`/`sources/` (never `concepts/`/`entities/` — use `list-taxonomy` for those); `--json` for scripting | | openkb feedback ["msg"] | File feedback by opening a prefilled GitHub issue (`--type bug/feature/question` to tag it) |
@@ -207,7 +209,7 @@ 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. +Retrieval is hybrid: the agent primarily navigates via `index.md`'s one-line summaries, and additionally has `list_taxonomy`/`get_taxonomy_item` (semantic browsing of `concepts/`/`entities/` pages by their one-line briefs — not a keyword search) and a tiered `search_wiki` tool — a dependency-free BM25 full-text search, in three independent tiers over `summaries/` briefs, full `summaries/` bodies, and `sources/` (including per-page indexing of long PageIndex documents) — for surfacing details a summary omits (an exact term, an author, a creation date). It's additive, not a replacement, so recall can only improve over index-only navigation. The same search/browse capability is available outside the agent via `openkb list-taxonomy` and `openkb search` (see `openkb --help`). Inside a chat, type `/` to access slash commands (Tab to complete). diff --git a/openkb/agent/query.py b/openkb/agent/query.py index e5cd79cd5..d7e613e9e 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 ( + list_taxonomy as list_taxonomy_impl, +) from openkb.agent.tools import ( search_wiki as search_wiki_impl, ) @@ -28,31 +31,40 @@ {schema_md} ## Search strategy -1. Read index.md to see all documents and concepts with brief summaries. - Each document is marked (short) or (pageindex) to indicate its type. +1. Read index.md to see all documents 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 - `full_text` frontmatter field to the source (see step 4). -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. 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 to the source (see step 5). +3. For concepts (cross-document synthesis) and entities ("who/what is X" + questions about a specific named person, organization, place, or + product), call list_taxonomy first — it's a compact, one-line-per-item + browse list, not a keyword search. Pick the slug(s) that match the + question's meaning by their brief, then read_file the matching + concepts/.md or entities/.md. +4. If index.md's one-line summaries and list_taxonomy don't surface a + specific detail you need (a niche term, an exact figure, an + author/creation-date only present in a raw source), use + search_wiki(query, scope) — a tiered, keyword-level full-text search + over summaries/sources only (concepts/entities are step 3's job, never + search_wiki's). This is a hybrid fallback: use it in addition to, not + instead of, index.md/list_taxonomy navigation. Narrow scope to + ["sources"] when you specifically need a source-only detail (an exact + field name, an author, a date) that a generated summary would likely + omit; leave scope unset to search all tiers. +5. 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. -7. Source content may reference images. Short-doc .md pages link them + ranges to help you target. Never fetch the whole document. A search_wiki + hit with a "page" locator names the exact page to fetch. +6. 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. -8. Synthesize a clear, concise, well-cited answer grounded in wiki content. +7. 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. @@ -92,17 +104,35 @@ 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. + def list_taxonomy(kind: str | None = None) -> str: + """List persisted concept/entity pages with one-line briefs (semantic browsing). + + Call this first for concept/entity questions and pick a slug by + meaning — this is a browse list, not a keyword search. Follow up + with read_file on the matching concepts/.md or + entities/.md to get the full page. + + Args: + kind: "concept" or "entity" to restrict the list; omit for both. + """ + return list_taxonomy_impl(wiki_root, kind=kind) + + @function_tool + def search_wiki(query: str, scope: list[str] | None = None) -> str: + """Tiered full-text (BM25) keyword search over summaries/sources. - 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. + Hybrid fallback for when index.md's one-line summaries and + list_taxonomy don't surface a specific buried detail (a niche term, + an exact figure, a fact only present in a raw source). Never + searches concepts/entities — use list_taxonomy for those. Use in + addition to, not instead of, index.md/list_taxonomy navigation. Args: query: Free-text search query (keywords or a natural-language question). + scope: Restrict to a subset of "briefs", "summaries", "sources"; + omit to search all three tiers. """ - return search_wiki_impl(query, wiki_root) + return search_wiki_impl(query, wiki_root, scope=scope) @function_tool def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: @@ -138,7 +168,7 @@ def get_image(image_path: str) -> ToolOutputImage | ToolOutputText: return Agent( name="wiki-query", instructions=instructions, - tools=[read_file, get_page_content, search_wiki, get_image], + tools=[read_file, get_page_content, list_taxonomy, 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 6d0d45fd9..62c90d52e 100644 --- a/openkb/agent/tools.py +++ b/openkb/agent/tools.py @@ -239,36 +239,83 @@ def get_taxonomy_item(slug: str, wiki_root: str, kind: str | None = None) -> str return f"Taxonomy item not found: {slug}" -def search_wiki(query: str, wiki_root: str, top_k: int = 5) -> str: - """Full-text (BM25) search over concepts/entities/summaries wiki pages. +def list_taxonomy(wiki_root: str, kind: str | None = None) -> str: + """Agent-facing text listing of persisted concept/entity 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. + Thin formatting wrapper around :func:`list_taxonomy_items` for use as an + LLM tool (see ``agent.query.build_query_agent``): one line per item with + its wikilink, entity type (if any), and one-line brief, so an LLM can + scan the whole taxonomy cheaply and pick a slug by meaning before calling + ``read_file`` on the matching page. 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. + kind: Restrict to ``"concept"`` or ``"entity"``; ``None`` returns both. Returns: - A formatted, ranked list of page hits (wikilink, title, snippet), or - a message indicating no matches were found. + One ``- [[path]] (type) — brief`` line per item, or a message if + none exist. """ - from openkb.fulltext_index import WikiFullTextIndex - - hits = WikiFullTextIndex(wiki_root).search(query, top_k=top_k) - if not hits: - return "No matching pages found." + items = list_taxonomy_items(wiki_root, kind=kind) + if not items: + return "No concepts or entities 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}") + for item in items: + wikilink = item.path[:-3] if item.path.endswith(".md") else item.path + type_suffix = f" ({item.type})" if item.type else "" + brief_suffix = f" — {item.brief}" if item.brief else "" + lines.append(f"- [[{wikilink}]]{type_suffix}{brief_suffix}") return "\n".join(lines) +def search_wiki(query: str, wiki_root: str, scope: list[str] | None = None, top_k: int = 5) -> str: + """Tiered full-text (BM25) search over summaries/sources wiki pages. + + Hybrid retrieval helper: complements index.md/``list_taxonomy`` navigation + by surfacing pages whose one-line brief 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/``list_taxonomy`` + navigation. Concepts/entities are never covered here — see + ``list_taxonomy``/``get_taxonomy_item`` for those (semantic browsing, not + keyword search). + + Args: + query: Free-text search query (keywords or a natural-language question). + wiki_root: Absolute path to the wiki root directory. + scope: Restrict to a subset of ``fulltext_index.TIERED_SCOPES`` + (``"briefs"``, ``"summaries"``, ``"sources"``); ``None`` searches + all three. + top_k: Maximum number of ranked results to return per tier. + + Returns: + Ranked hits grouped by tier (wikilink, locator if any, title, + snippet), or a message indicating no matches were found, or an error + message if *scope* contains an invalid tier name. + """ + from openkb.fulltext_index import TIERED_SCOPES, TieredWikiSearch + + try: + results = TieredWikiSearch(wiki_root).search(query, scope=scope, top_k=top_k) + except ValueError as exc: + return str(exc) + + sections = [] + for tier in TIERED_SCOPES: + hits = results.get(tier) + if not hits: + continue + lines = [f"## {tier}"] + for i, hit in enumerate(hits, start=1): + wikilink = hit.path[:-3] if hit.path.endswith(".md") else hit.path + locator = f" [{hit.locator.kind} {hit.locator.value}]" if hit.locator else "" + lines.append( + f"{i}. [[{wikilink}]]{locator} — {hit.title} (score: {hit.score})\n {hit.snippet}" + ) + sections.append("\n".join(lines)) + return "\n\n".join(sections) if sections else "No matching pages found." + + _MIME_TYPES = { ".png": "image/png", ".jpg": "image/jpeg", diff --git a/openkb/cli.py b/openkb/cli.py index c9e543181..c9c11de0e 100644 --- a/openkb/cli.py +++ b/openkb/cli.py @@ -2630,6 +2630,128 @@ def list_cmd(ctx): print_list(kb_dir) +def _taxonomy_items_to_json(items) -> list[dict]: + """Convert ``TaxonomyItem`` dataclasses to plain JSON-serializable dicts.""" + return [ + {"kind": i.kind, "slug": i.slug, "path": i.path, "brief": i.brief, "type": i.type} + for i in items + ] + + +@cli.command(name="list-taxonomy") +@click.option( + "--kind", + type=click.Choice(["concept", "entity"]), + default=None, + help="Restrict to concepts or entities (default: both).", +) +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +@click.pass_context +@_with_kb_lock(exclusive=False) +def list_taxonomy_cmd(ctx, kind, as_json): + """List persisted concept/entity pages with their one-line briefs. + + Intended for semantic browsing (external agents/scripts pick a slug by + meaning), not keyword search — see ``openkb search`` for that. Never + includes not-yet-paged pending candidates, only committed ``.md`` pages. + """ + 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.agent.tools import list_taxonomy_items + + items = list_taxonomy_items(str(kb_dir / "wiki"), kind=kind) + + if as_json: + click.echo(json.dumps(_taxonomy_items_to_json(items), ensure_ascii=False, indent=2)) + return + + if not items: + click.echo("No concepts or entities found.") + return + for item in items: + type_suffix = f" ({item.type})" if item.type else "" + brief_suffix = f" — {item.brief}" if item.brief else "" + click.echo(f"[{item.kind}] {item.slug}{type_suffix}{brief_suffix}") + + +def _search_results_to_json(results: dict) -> dict: + """Convert ``{tier: [SearchHit, ...]}`` to plain JSON-serializable dicts.""" + return { + tier: [ + { + "path": hit.path, + "title": hit.title, + "score": hit.score, + "snippet": hit.snippet, + "locator": ( + {"kind": hit.locator.kind, "value": hit.locator.value} if hit.locator else None + ), + } + for hit in hits + ] + for tier, hits in results.items() + } + + +@cli.command(name="search") +@click.argument("query") +@click.option( + "--scope", + default=None, + help="Comma-separated subset of briefs,summaries,sources (default: all three).", +) +@click.option("--top-k", default=5, show_default=True, help="Max ranked results per tier.") +@click.option("--json", "as_json", is_flag=True, default=False, help="Output as JSON.") +@click.pass_context +@_with_kb_lock(exclusive=False) +def search_cmd(ctx, query, scope, top_k, as_json): + """Full-text (BM25) search over summaries/sources, tier by tier. + + Concepts/entities are not covered — use ``openkb list-taxonomy`` for + those (semantic browsing, not keyword search). Each tier is scored and + ranked independently: ``briefs`` (one-line document summaries), rich + ``summaries`` (full document-summary text), and ``sources`` (raw source + files, with a page/line locator pointing at the exact hit location). + """ + 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.fulltext_index import TieredWikiSearch + + scope_list = [s.strip() for s in scope.split(",") if s.strip()] if scope else None + try: + results = TieredWikiSearch(str(kb_dir / "wiki")).search( + query, scope=scope_list, top_k=top_k + ) + except ValueError as exc: + click.echo(str(exc)) + ctx.exit(1) + return + + if as_json: + click.echo(json.dumps(_search_results_to_json(results), ensure_ascii=False, indent=2)) + return + + any_hits = False + for tier in ("briefs", "summaries", "sources"): + hits = results.get(tier) + if not hits: + continue + any_hits = True + click.echo(f"\n=== {tier} ===") + for i, hit in enumerate(hits, start=1): + locator_suffix = f" [{hit.locator.kind} {hit.locator.value}]" if hit.locator else "" + click.echo(f"{i}. {hit.path}{locator_suffix} — {hit.title} (score: {hit.score})") + click.echo(f" {hit.snippet}") + if not any_hits: + click.echo("No matching pages found.") + + def print_status(kb_dir: Path) -> None: """Print knowledge base status. Usable from CLI and chat REPL.""" wiki_dir = kb_dir / "wiki" diff --git a/tests/test_agent_tools.py b/tests/test_agent_tools.py index dba78644c..58e2a95a0 100644 --- a/tests/test_agent_tools.py +++ b/tests/test_agent_tools.py @@ -7,6 +7,7 @@ artifact_event_from_write, get_taxonomy_item, get_wiki_page_content, + list_taxonomy, list_taxonomy_items, list_wiki_files, parse_pages, @@ -332,7 +333,20 @@ def test_artifact_event_none_for_bad_json(): class TestSearchWiki: - def test_finds_matching_page(self, tmp_path): + def test_finds_matching_page_in_sources(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "sources").mkdir() + (tmp_path / "sources" / "cnn.md").write_text( + "# Convolutional Neural Networks\n\nDropout regularization prevents overfitting." + ) + + result = search_wiki("dropout regularization", wiki_root) + + assert "[[sources/cnn]]" in result + assert "Convolutional Neural Networks" in result + assert "## sources" in result + + def test_concepts_and_entities_are_not_searched(self, tmp_path): wiki_root = str(tmp_path) (tmp_path / "concepts").mkdir() (tmp_path / "concepts" / "cnn.md").write_text( @@ -341,13 +355,12 @@ def test_finds_matching_page(self, tmp_path): result = search_wiki("dropout regularization", wiki_root) - assert "[[concepts/cnn]]" in result - assert "Convolutional Neural Networks" in result + assert result == "No matching pages found." 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.") + (tmp_path / "sources").mkdir() + (tmp_path / "sources" / "cnn.md").write_text("# CNN\n\nSomething else entirely.") result = search_wiki("nonexistent_keyword_xyz", wiki_root) @@ -355,13 +368,42 @@ def test_no_matches_returns_message(self, tmp_path): def test_respects_top_k(self, tmp_path): wiki_root = str(tmp_path) - (tmp_path / "entities").mkdir() + (tmp_path / "sources").mkdir() for i in range(5): - (tmp_path / "entities" / f"e{i}.md").write_text(f"# Entity {i}\n\nkeyword {i}.") + (tmp_path / "sources" / 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 + assert result.count("[[sources/") == 2 + + def test_scope_restricts_to_requested_tiers(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "summaries").mkdir() + (tmp_path / "summaries" / "doc.md").write_text( + '---\ndescription: "keyword brief"\n---\n\n# Doc\n\nkeyword body.' + ) + (tmp_path / "sources").mkdir() + (tmp_path / "sources" / "doc.md").write_text("keyword raw source.") + + result = search_wiki("keyword", wiki_root, scope=["sources"]) + + assert "## sources" in result + assert "## briefs" not in result + assert "## summaries" not in result + + def test_invalid_scope_returns_error_message(self, tmp_path): + result = search_wiki("keyword", str(tmp_path), scope=["not-a-real-tier"]) + + assert "Unknown scope" in result + + def test_result_includes_locator_for_source_hit(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "sources").mkdir() + (tmp_path / "sources" / "notes.md").write_text("Line one.\nkeyword on line two.") + + result = search_wiki("keyword", wiki_root, scope=["sources"]) + + assert "[line 2]" in result # --------------------------------------------------------------------------- @@ -477,3 +519,45 @@ def test_path_traversal_is_rejected(self, tmp_path): result = get_taxonomy_item("../../etc/passwd", str(tmp_path)) assert result == "Taxonomy item not found: ../../etc/passwd" + + +# --------------------------------------------------------------------------- +# list_taxonomy (agent-facing text formatter over list_taxonomy_items) +# --------------------------------------------------------------------------- + + +class TestListTaxonomy: + def test_formats_concepts_and_entities_as_wikilinks(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "attention.md").write_text( + '---\ndescription: "How attention works"\n---\n\n# Attention\n\nBody.' + ) + (tmp_path / "entities").mkdir() + (tmp_path / "entities" / "acme.md").write_text( + '---\ntype: organization\ndescription: "A company"\n---\n\n# Acme\n\nBody.' + ) + + result = list_taxonomy(wiki_root) + + assert "[[concepts/attention]]" in result + assert "How attention works" in result + assert "[[entities/acme]] (organization)" in result + assert "A company" in result + + def test_kind_filter(self, tmp_path): + wiki_root = str(tmp_path) + (tmp_path / "concepts").mkdir() + (tmp_path / "concepts" / "c.md").write_text("# C\n\nBody.") + (tmp_path / "entities").mkdir() + (tmp_path / "entities" / "e.md").write_text("# E\n\nBody.") + + result = list_taxonomy(wiki_root, kind="concept") + + assert "concepts/c" in result + assert "entities/e" not in result + + def test_empty_taxonomy_returns_message(self, tmp_path): + result = list_taxonomy(str(tmp_path)) + + assert result == "No concepts or entities found." diff --git a/tests/test_query.py b/tests/test_query.py index a720ccce3..e13b55729 100644 --- a/tests/test_query.py +++ b/tests/test_query.py @@ -17,15 +17,16 @@ 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_five_tools(self, tmp_path): agent = build_query_agent(str(tmp_path), "gpt-4o-mini") - assert len(agent.tools) == 4 + assert len(agent.tools) == 5 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 "list_taxonomy" in names assert "search_wiki" in names assert "get_image" in names