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>

api/routes/kb_public.py

@@ -252,6 +252,11 @@ async def ask(body: AskRequest, request: Request):
         })
     messages.append({"role": "user", "content": body.message})
 
+    # Collect the sources Shira actually searched through during this turn.
+    # The legal_kb tool appends to this list on every search_insurance_kb
+    # call; the UI uses it to render a side PDF viewer next to the answer.
+    sources_used: list[dict] = []
+
     try:
         text = await runner.run(
             system_prompt=system_prompt,
@@ -260,9 +265,20 @@ async def ask(body: AskRequest, request: Request):
             espocrm_url=os.environ.get("ESPOCRM_URL", ""),
             espocrm_api_key=os.environ.get("ESPOCRM_API_KEY", ""),
             allowed_toolsets=["legal"],
+            kb_sources_used=sources_used,
         )
     except Exception as e:
         logger.exception("[kb.ask] error")
         raise HTTPException(status_code=500, detail=str(e))
 
-    return {"text": text, "conversationId": body.conversation_id}
+    # Keep only the ones that actually have a PDF attached.
+    pdf_sources = [
+        s for s in sources_used
+        if (s.get("original_path") or "").lower().endswith(".pdf")
+    ]
+
+    return {
+        "text": text,
+        "conversationId": body.conversation_id,
+        "sources": pdf_sources,
+    }

api/services/agent_runner.py

@@ -58,6 +58,7 @@ class AgentRunner:
         depth: int = 0,
         blocked_tools: set[str] | None = None,
         allowed_toolsets: list[str] | None = None,
+        kb_sources_used: list[dict] | None = None,
     ) -> str:
         """Run a conversation with tool-calling loop. Returns the final text response.
 
@@ -82,7 +83,7 @@ class AgentRunner:
             register_document_tools(tools_registry, crm, case_id, context)
         if should_register_all or "legal" in (allowed_toolsets or []):
             register_legal_tools(tools_registry, crm, case_id, context)
-            register_legal_kb_tools(tools_registry)
+            register_legal_kb_tools(tools_registry, sources_used=kb_sources_used)
         if should_register_all or "signature" in (allowed_toolsets or []):
             register_signature_tools(tools_registry, crm, case_id, user_id, context)
 

api/services/kb/chunker.py

@@ -1,7 +1,7 @@
 """Text chunking — strategies per source kind.
 
 All strategies return list[dict] with keys:
-    chunk_index, heading_path, section_ref, content, token_count
+    chunk_index, heading_path, section_ref, content, token_count, page_number
 """
 
 from __future__ import annotations
@@ -10,6 +10,12 @@ import re
 from dataclasses import dataclass
 from typing import Literal
 
+# Page marker sentinel embedded by the PDF parser (api/services/kb/ingest.py
+# _page_marker). Format: \x00KB_PAGE:<N>\x00 where N is 1-indexed. The chunker
+# strips markers from chunk.content and records the last-seen page as the
+# chunk's page_number.
+_PAGE_MARKER_RE = re.compile(r"\x00KB_PAGE:(\d+)\x00")
+
 # Approx tokens per char for Hebrew/mixed text (voyage-multilingual-2 tokenizer).
 # Voyage tends to split Hebrew into 1.5-2 tokens per word; 0.55 tokens/char
 # is a conservative estimate that matches what the API actually reports.
@@ -69,23 +75,44 @@ class Chunk:
     section_ref: str
     content: str
     token_count: int
+    page_number: int | None = None
 
 
 def _count_tokens(text: str) -> int:
     return max(1, int(len(text) * _TOKENS_PER_CHAR))
 
 
+def _strip_page_markers(text: str) -> tuple[str, int | None]:
+    """Remove page markers from `text` and return the first page number seen.
+
+    The first-marker rule matches how users think about a citation ("this
+    paragraph starts on page 42") — subsequent markers on the same chunk
+    mean the paragraph spans pages, but the jump-to-PDF link should point
+    at the page where the quote begins.
+    """
+    first_page: int | None = None
+    m = _PAGE_MARKER_RE.search(text)
+    if m:
+        try:
+            first_page = int(m.group(1))
+        except ValueError:
+            first_page = None
+    return _PAGE_MARKER_RE.sub("", text), first_page
+
+
 def _as_dicts(chunks: list[Chunk]) -> list[dict]:
-    return [
+    result: list[dict] = []
-        {
+    for c in chunks:
+        clean, page = _strip_page_markers(c.content)
+        result.append({
             "chunk_index": c.chunk_index,
             "heading_path": c.heading_path or None,
             "section_ref": c.section_ref or None,
-            "content": c.content,
+            "content": clean.strip(),
             "token_count": c.token_count,
-        }
+            "page_number": page if page is not None else c.page_number,
-        for c in chunks
+        })
-    ]
+    return result
 
 
 def _split_oversized(chunks: list[Chunk]) -> list[Chunk]:

api/services/kb/ingest.py

@@ -37,6 +37,17 @@ def _coerce_date(value) -> _dt.date | None:
     return None
 
 
+# Sentinel markers we embed in the parser's output so the chunker can
+# attach a page number to each resulting chunk. The markers are stripped
+# from chunk.content before it's stored; only the page number is kept.
+_PAGE_MARKER_PREFIX = "\x00KB_PAGE:"
+_PAGE_MARKER_SUFFIX = "\x00"
+
+
+def _page_marker(page_index_zero_based: int) -> str:
+    return f"{_PAGE_MARKER_PREFIX}{page_index_zero_based + 1}{_PAGE_MARKER_SUFFIX}"
+
+
 def _parse_pdf(data: bytes) -> str:
     """Extract text from a PDF, preferring block-based parsing with row merging.
 
@@ -56,7 +67,7 @@ def _parse_pdf(data: bytes) -> str:
     row_tolerance = 4.0  # pixels; blocks within this y-range are the same line
 
     with fitz.open(stream=data, filetype="pdf") as doc:
-        for page in doc:
+        for page_idx, page in enumerate(doc):
             raw = page.get_text("blocks") or []
             textual = [
                 (float(b[0]), float(b[1]), float(b[2]), float(b[3]), b[4])
@@ -90,7 +101,11 @@ def _parse_pdf(data: bytes) -> str:
                 current_row.sort(key=lambda p: -p[0])
                 lines.append(" ".join(p[1] for p in current_row))
 
-            parts.append("\n\n".join(lines))
+            # Embed a page marker at the start of each page. The chunker
+            # treats markers as zero-width separators: they don't affect
+            # chunking boundaries, but the last marker seen before a chunk
+            # determines its page_number.
+            parts.append(_page_marker(page_idx) + "\n\n".join(lines))
     return "\n\n".join(p for p in parts if p).strip()
 
 
@@ -221,6 +236,7 @@ async def ingest_source(
                     c["content"],
                     vectors[i],
                     c["token_count"],
+                    c.get("page_number"),
                 )
                 for i, c in enumerate(chunks)
             ]
@@ -228,8 +244,8 @@ async def ingest_source(
                 """
                 INSERT INTO kb_chunk
                     (source_id, chunk_index, heading_path, section_ref,
-                     content, embedding, token_count)
+                     content, embedding, token_count, page_number)
-                VALUES ($1,$2,$3,$4,$5,$6,$7)
+                VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
                 """,
                 rows,
             )

api/services/kb/search.py

@@ -67,9 +67,9 @@ async def search(
       ) u GROUP BY id
     )
     SELECT
-      s.kind, s.title, s.identifier, s.source_url,
+      s.id AS source_id, s.kind, s.title, s.identifier, s.source_url,
-      s.published_at, s.effective_at,
+      s.published_at, s.effective_at, s.original_path,
-      c.heading_path, c.section_ref, c.content,
+      c.heading_path, c.section_ref, c.content, c.page_number,
       f.score
     FROM fused f
     JOIN kb_chunk c ON c.id = f.id

mcp_server/tools/legal_kb_tools.py

@@ -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")