feat(agent): grep-based completeness sweep for the query agent

openkb/agent/query.py

@@ -8,6 +8,7 @@
 from agents import ToolOutputImage, ToolOutputText
 from openkb.agent.tools import (
     get_wiki_page_content,
+    grep_wiki_files,
     read_wiki_file,
     read_wiki_image,
     write_kb_file,
@@ -38,7 +39,19 @@
      ranges to help you target. Never fetch the whole document.
 6. Source content may reference images (e.g. ![image](sources/images/doc/file.png)).
    Use the get_image tool to view them when needed.
-7. Synthesize a clear, concise, well-cited answer grounded in wiki content.
+7. COMPLETENESS SWEEP (do this before finalizing): the summary layer is
+   lossy, so before you commit to an answer, call grep_wiki for the
+   question's salient terms — proper nouns, technical terms, numbers, key
+   entities — plus any claim you asserted in your draft that you have not
+   yet seen on a wiki page. Because grep is lexical (not semantic), try a
+   few term variants: acronym and expansion, singular/plural, close
+   synonyms. For any matching page you have NOT already read, read_file it
+   and fold in relevant content. If grep surfaces a claim that contradicts
+   your draft, note both claims with their citations rather than silently
+   choosing one. Do at most 3 grep rounds (a round = one concept and its
+   variants); stop once a round surfaces no new page. grep_wiki is a check,
+   not the primary search — index.md and summaries still come first.
+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.
@@ -87,12 +100,36 @@ def get_image(image_path: str) -> ToolOutputImage | ToolOutputText:
             return ToolOutputImage(image_url=result["image_url"])
         return ToolOutputText(text=result["text"])
 
+    @function_tool
+    def grep_wiki(pattern: str, ignore_case: bool = True, fixed_string: bool = False) -> str:
+        """Lexically grep the wiki's markdown for a pattern.
+
+        Use this as a FINAL completeness check, after you have drafted an
+        answer from index.md / summaries / concepts / entities. It searches
+        every wiki .md file (including short-doc sources/) for the literal
+        terms of the question — catching details the summaries compressed
+        away, pages you never opened, or contradicting mentions. It does NOT
+        search long-document page content (use get_page_content for that).
+
+        Returns up to 50 matches as 'relative/path.md:LINE: text'. Feed any
+        new path into read_file. Try a few term variants (acronym/expansion,
+        singular/plural, synonyms) — this is lexical, not semantic.
+
+        Args:
+            pattern: Search pattern (regex by default).
+            ignore_case: Case-insensitive (default True).
+            fixed_string: Treat pattern as a literal string, not a regex.
+        """
+        return grep_wiki_files(
+            pattern, wiki_root, ignore_case=ignore_case, fixed_string=fixed_string,
+        )
+
     from agents.model_settings import ModelSettings
 
     return Agent(
         name="wiki-query",
         instructions=instructions,
-        tools=[read_file, get_page_content, get_image],
+        tools=[read_file, get_page_content, get_image, grep_wiki],
         model=f"litellm/{model}",
         model_settings=ModelSettings(parallel_tool_calls=False),
     )

openkb/agent/tools.py

@@ -8,8 +8,14 @@
 
 import contextlib
 import json as _json
+import shutil
+import subprocess
 from pathlib import Path
 
+# grep_wiki_files tuning
+_GREP_MAX_LINES = 50
+_GREP_TIMEOUT_S = 10
+
 
 def list_wiki_files(directory: str, wiki_root: str) -> str:
     """List all Markdown files in a wiki subdirectory.
@@ -54,6 +60,100 @@ def read_wiki_file(path: str, wiki_root: str) -> str:
     return full_path.read_text(encoding="utf-8")
 
 
+def grep_wiki_files(
+    pattern: str,
+    wiki_root: str,
+    *,
+    ignore_case: bool = True,
+    fixed_string: bool = False,
+) -> str:
+    """Lexically search the wiki's markdown layer for ``pattern``.
+
+    A completeness sweep: shells out to ripgrep (preferred) or grep
+    (fallback) over every ``*.md`` file under *wiki_root* — summaries,
+    concepts, entities, explorations, ``index.md``, and short-doc
+    ``sources/*.md``. Long-doc per-page ``*.json`` (PageIndex's domain) and
+    ``log.md`` bookkeeping are excluded.
+
+    Args:
+        pattern: Search pattern. Regex by default; literal when
+            *fixed_string* is True.
+        wiki_root: Absolute path to the wiki root directory.
+        ignore_case: Case-insensitive match (default True).
+        fixed_string: Treat *pattern* as a literal string, not a regex.
+
+    Returns:
+        Up to :data:`_GREP_MAX_LINES` matches as ``relative/path.md:LINE: text``
+        lines, plus a truncation notice if capped. On no match / missing
+        binary / timeout / error, returns an explicit message string. Never
+        raises and never invokes a shell (``shell=False``), so a hostile
+        *pattern* cannot inject commands.
+    """
+    root = Path(wiki_root).resolve()
+    if not root.exists():
+        return f"Wiki root not found: {wiki_root}"
+
+    rg = shutil.which("rg")
+    grep = shutil.which("grep")
+
+    if rg:
+        # --no-ignore: the wiki dir is often gitignored; without this rg
+        # silently returns zero matches inside a real OpenKB checkout.
+        cmd = [
+            rg, "--line-number", "--no-heading", "--color", "never",
+            "--no-ignore", "-g", "*.md", "-g", "!log.md",
+        ]
+        if ignore_case:
+            cmd.append("-i")
+        if fixed_string:
+            cmd.append("-F")
+        cmd += ["-e", pattern, str(root)]
+    elif grep:
+        cmd = [grep, "-rn", "--include=*.md", "--exclude-dir=images"]
+        if ignore_case:
+            cmd.append("-i")
+        if fixed_string:
+            cmd.append("-F")
+        cmd += ["-e", pattern, str(root)]
+    else:
+        return "grep unavailable on this system."
+
+    try:
+        proc = subprocess.run(
+            cmd, capture_output=True, text=True,
+            timeout=_GREP_TIMEOUT_S, check=False,
+        )
+    except subprocess.TimeoutExpired:
+        return "grep timed out; narrow the pattern."
+
+    # rg/grep convention: 0 = matches, 1 = no matches, >=2 = real error.
+    if proc.returncode >= 2:
+        stderr_lines = (proc.stderr or "").strip().splitlines()
+        first = stderr_lines[0] if stderr_lines else "unknown error"
+        return f"grep error: {first}."
+
+    prefix = str(root) + "/"
+    results: list[str] = []
+    for line in proc.stdout.splitlines():
+        if not line.strip():
+            continue
+        rel = line[len(prefix):] if line.startswith(prefix) else line
+        path_part = rel.split(":", 1)[0]
+        # Defensive: grep --include=*.md still matches log.md; drop it.
+        if path_part == "log.md" or path_part.endswith("/log.md"):
+            continue
+        results.append(rel)
+
+    if not results:
+        return f"No matches for {pattern}."
+
+    truncated = len(results) > _GREP_MAX_LINES
+    out = "\n".join(results[:_GREP_MAX_LINES])
+    if truncated:
+        out += "\n… more matches; narrow the pattern."
+    return out
+
+
 def parse_pages(pages: str) -> list[int]:
     """Parse a page specification string into a sorted, deduplicated list of page numbers.
 

tests/test_grep.py

@@ -0,0 +1,186 @@
+"""Tests for openkb.agent.tools.grep_wiki_files — lexical wiki search."""
+from __future__ import annotations
+
+from openkb.agent.tools import grep_wiki_files
+
+
+def _wiki(tmp_path):
+    """Build a minimal wiki/ tree and return its root as a string."""
+    root = tmp_path / "wiki"
+    (root / "summaries").mkdir(parents=True)
+    (root / "concepts").mkdir(parents=True)
+    (root / "entities").mkdir(parents=True)
+    (root / "sources" / "images").mkdir(parents=True)
+    (root / "summaries" / "paper.md").write_text(
+        "# Paper\nThe transformer architecture uses self-attention.\n",
+        encoding="utf-8",
+    )
+    (root / "concepts" / "attention.md").write_text(
+        "# Attention\nScaled dot-product Attention is central.\n",
+        encoding="utf-8",
+    )
+    (root / "sources" / "note.md").write_text(
+        "Short note: the lottery ticket hypothesis appears here only.\n",
+        encoding="utf-8",
+    )
+    # Long-doc per-page JSON — must NEVER be grepped.
+    (root / "sources" / "book.json").write_text(
+        '[{"page": 1, "text": "transformer secret in json"}]\n',
+        encoding="utf-8",
+    )
+    # Bookkeeping — must NEVER be grepped.
+    (root / "log.md").write_text(
+        "# Operations Log\n## [2026-01-01] ingest | transformer\n",
+        encoding="utf-8",
+    )
+    return str(root)
+
+
+def test_finds_match_in_summaries(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("self-attention", wiki)
+    assert "summaries/paper.md:" in out
+    assert "self-attention" in out
+
+
+def test_finds_match_in_short_source_md(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("lottery ticket", wiki)
+    assert "sources/note.md:" in out
+
+
+def test_excludes_long_doc_json(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("transformer", wiki)
+    assert "book.json" not in out
+
+
+def test_excludes_log_md(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("transformer", wiki)
+    assert "log.md" not in out
+
+
+def test_case_insensitive_by_default(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("TRANSFORMER", wiki)
+    assert "summaries/paper.md:" in out
+
+
+def test_case_sensitive_when_disabled(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("TRANSFORMER", wiki, ignore_case=False)
+    assert out == "No matches for TRANSFORMER."
+
+
+def test_fixed_string_treats_regex_literally(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("self.attention", wiki, fixed_string=True)
+    # literal "self.attention" does not appear (text has "self-attention")
+    assert out == "No matches for self.attention."
+    # but as a regex, "." matches the hyphen
+    out2 = grep_wiki_files("self.attention", wiki, fixed_string=False)
+    assert "summaries/paper.md:" in out2
+
+
+def test_no_match_returns_message(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("nonexistentterm12345", wiki)
+    assert out == "No matches for nonexistentterm12345."
+
+
+def test_paths_are_relative_to_wiki_root(tmp_path):
+    wiki = _wiki(tmp_path)
+    out = grep_wiki_files("self-attention", wiki)
+    assert wiki not in out
+    assert out.splitlines()[0].startswith("summaries/")
+
+
+def test_result_cap_and_truncation_notice(tmp_path):
+    wiki = _wiki(tmp_path)
+    root = tmp_path / "wiki" / "summaries"
+    big = "\n".join(f"line {i} needle" for i in range(60))
+    (root / "big.md").write_text(big + "\n", encoding="utf-8")
+    out = grep_wiki_files("needle", wiki)
+    lines = out.splitlines()
+    assert lines[-1] == "… more matches; narrow the pattern."
+    assert len(lines) == 51
+
+
+def test_shell_metacharacters_do_not_execute(tmp_path):
+    wiki = _wiki(tmp_path)
+    sentinel = tmp_path / "pwned"
+    grep_wiki_files("; touch " + str(sentinel), wiki)
+    assert not sentinel.exists()
+
+
+def test_rg_branch_builds_expected_command(tmp_path, monkeypatch):
+    import subprocess as _sp
+    import openkb.agent.tools as tools_mod
+
+    wiki = _wiki(tmp_path)
+    monkeypatch.setattr(tools_mod.shutil, "which", lambda name: "/usr/bin/rg" if name == "rg" else None)
+
+    captured = {}
+
+    class _FakeProc:
+        returncode = 0
+        stdout = ""
+        stderr = ""
+
+    def _fake_run(cmd, *args, **kwargs):
+        captured["cmd"] = cmd
+        captured["shell"] = kwargs.get("shell", False)
+        return _FakeProc()
+
+    monkeypatch.setattr(tools_mod.subprocess, "run", _fake_run)
+
+    tools_mod.grep_wiki_files("needle", wiki, ignore_case=True, fixed_string=True)
+
+    cmd = captured["cmd"]
+    # rg binary chosen, shell never used
+    assert cmd[0] == "/usr/bin/rg"
+    assert captured["shell"] is False
+    # load-bearing flags present
+    assert "--no-ignore" in cmd
+    assert "--line-number" in cmd
+    assert "--no-heading" in cmd
+    assert "-i" in cmd          # ignore_case
+    assert "-F" in cmd          # fixed_string
+    # md include + log.md exclude globs
+    assert cmd[cmd.index("-g") :].count("-g") == 2 or cmd.count("-g") == 2
+    assert "*.md" in cmd
+    assert "!log.md" in cmd
+    # pattern passed via -e as a separate argv (injection-safe), root last
+    assert cmd[-3] == "-e"
+    assert cmd[-2] == "needle"
+    assert cmd[-1] == str(__import__("pathlib").Path(wiki).resolve())
+
+
+def test_rg_branch_omits_flags_when_disabled(tmp_path, monkeypatch):
+    import openkb.agent.tools as tools_mod
+
+    wiki = _wiki(tmp_path)
+    monkeypatch.setattr(tools_mod.shutil, "which", lambda name: "/usr/bin/rg" if name == "rg" else None)
+    captured = {}
+
+    class _FakeProc:
+        returncode = 1
+        stdout = ""
+        stderr = ""
+
+    monkeypatch.setattr(tools_mod.subprocess, "run", lambda cmd, *a, **k: (captured.__setitem__("cmd", cmd) or _FakeProc()))
+
+    tools_mod.grep_wiki_files("needle", wiki, ignore_case=False, fixed_string=False)
+    cmd = captured["cmd"]
+    assert "-i" not in cmd
+    assert "-F" not in cmd
+
+
+def test_finds_match_in_entities(tmp_path):
+    wiki = _wiki(tmp_path)
+    (tmp_path / "wiki" / "entities" / "vaswani.md").write_text(
+        "# Vaswani\nAshish Vaswani is a lead author.\n", encoding="utf-8",
+    )
+    out = grep_wiki_files("Ashish Vaswani", wiki)
+    assert "entities/vaswani.md:" in out

tests/test_query.py

@@ -16,16 +16,17 @@ 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 "read_file" in names
         assert "get_page_content" in names
         assert "get_image" in names
+        assert "grep_wiki" in names
 
     def test_instructions_mention_get_page_content(self, tmp_path):
         agent = build_query_agent(str(tmp_path), "gpt-4o-mini")