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:
+17
-1
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
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]:
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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
|
||||
|
||||
Reference in New Issue
Block a user