From 4a3a8945af34133ec74fa69925d0d868d7b1732d Mon Sep 17 00:00:00 2001 From: Chaim Date: Fri, 24 Apr 2026 12:20:52 +0000 Subject: [PATCH] feat(kb): track page_number per chunk + expose sources_used from /kb/ask MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit To let the browse/ask UIs jump straight to the right PDF page: - ingest.py: embed \x00KB_PAGE:\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) --- api/routes/kb_public.py | 18 ++++++++++++- api/services/agent_runner.py | 3 ++- api/services/kb/chunker.py | 41 +++++++++++++++++++++++++----- api/services/kb/ingest.py | 24 ++++++++++++++--- api/services/kb/search.py | 6 ++--- mcp_server/tools/legal_kb_tools.py | 25 +++++++++++++++++- 6 files changed, 100 insertions(+), 17 deletions(-) diff --git a/api/routes/kb_public.py b/api/routes/kb_public.py index f3573e6..c253735 100644 --- a/api/routes/kb_public.py +++ b/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, + } diff --git a/api/services/agent_runner.py b/api/services/agent_runner.py index 62a03f8..96ea90d 100644 --- a/api/services/agent_runner.py +++ b/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) diff --git a/api/services/kb/chunker.py b/api/services/kb/chunker.py index 2b02e92..d8f3284 100644 --- a/api/services/kb/chunker.py +++ b/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:\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, - } - for c in chunks - ] + "page_number": page if page is not None else c.page_number, + }) + return result def _split_oversized(chunks: list[Chunk]) -> list[Chunk]: diff --git a/api/services/kb/ingest.py b/api/services/kb/ingest.py index fb30709..2314c81 100644 --- a/api/services/kb/ingest.py +++ b/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) - VALUES ($1,$2,$3,$4,$5,$6,$7) + content, embedding, token_count, page_number) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8) """, rows, ) diff --git a/api/services/kb/search.py b/api/services/kb/search.py index 1549ab1..82fe751 100644 --- a/api/services/kb/search.py +++ b/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.published_at, s.effective_at, - c.heading_path, c.section_ref, c.content, + s.id AS source_id, s.kind, s.title, s.identifier, s.source_url, + s.published_at, s.effective_at, s.original_path, + 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 diff --git a/mcp_server/tools/legal_kb_tools.py b/mcp_server/tools/legal_kb_tools.py index 2da7dad..146c816 100644 --- a/mcp_server/tools/legal_kb_tools.py +++ b/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")