feat(kb): track page_number per chunk + expose sources_used from /kb/ask

To let the browse/ask UIs jump straight to the right PDF page:

- ingest.py: embed \x00KB_PAGE:<N>\x00 sentinels at the start of each
  PDF page during _parse_pdf.
- chunker.py: _strip_page_markers removes the sentinels and records the
  first page seen per chunk as page_number. Chunk dataclass + _as_dicts
  propagate page_number through both chunk_statute and chunk_circular.
- ingest.py INSERT: new kb_chunk.page_number column (DB migration
  applied manually).
- search.py SQL: SELECT c.page_number + s.id + s.original_path so the
  rerank results carry them forward.
- legal_kb_tools.py: register_legal_kb_tools(sources_used=list) — the
  tool appends {source_id, title, kind, identifier, heading_path,
  section_ref, page_number, original_path} for each hit, deduped by
  (source_id, page_number).
- agent_runner.py: new kb_sources_used kwarg threads the list through
  to the tool registration.
- kb_public.py /kb/ask: collects sources_used during the run, filters
  to PDF-only sources, and returns them alongside text.

Chunks stored before this change have page_number = NULL; PDFs need to
be re-ingested to populate them. The TXT law source is unaffected.

Refs Task Master #2

Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-04-24 12:20:52 +00:00
parent 748c8f10c9
commit 4a3a8945af
6 changed files with 100 additions and 17 deletions
+24 -1
View File
@@ -47,7 +47,14 @@ def _format_hits(hits: list[dict]) -> str:
return "\n".join(lines).rstrip()
def register_legal_kb_tools(tools: dict) -> None:
def register_legal_kb_tools(tools: dict, sources_used: list[dict] | None = None) -> None:
"""Register the search_insurance_kb tool.
If `sources_used` is provided, each hit's (source_id, page_number,
section_ref, title, kind, original_path) is appended — deduplicated
by (source_id, page_number) — so the caller can build a "view PDF"
link list alongside the agent's text answer.
"""
async def search_insurance_kb(
query: str,
kind: str = "any",
@@ -58,6 +65,22 @@ def register_legal_kb_tools(tools: dict) -> None:
kind = "any"
top_k = max(1, min(int(top_k or 8), 15))
hits = await kb_search.search(query=query, kind=kind, top_k=top_k)
if sources_used is not None:
seen = {(s.get("source_id"), s.get("page_number")) for s in sources_used}
for h in hits:
key = (h.get("source_id"), h.get("page_number"))
if h.get("source_id") and key not in seen:
sources_used.append({
"source_id": h["source_id"],
"title": h.get("title"),
"kind": h.get("kind"),
"identifier": h.get("identifier"),
"heading_path": h.get("heading_path"),
"section_ref": h.get("section_ref"),
"page_number": h.get("page_number"),
"original_path": h.get("original_path"),
})
seen.add(key)
return _format_hits(hits)
except Exception as e:
logger.exception("[search_insurance_kb] failed")