chaim
4a3a8945af
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>
284 lines
10 KiB
Python
284 lines
10 KiB
Python
"""Ingest pipeline: bytes → parse → chunk → embed → upsert."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import datetime as _dt
|
|
import hashlib
|
|
import io
|
|
import logging
|
|
from typing import Literal
|
|
|
|
import fitz # pymupdf
|
|
|
|
from api.services.kb import voyage
|
|
from api.services.kb.chunker import chunk
|
|
from api.services.kb.db import get_pool
|
|
|
|
logger = logging.getLogger("shira.kb.ingest")
|
|
|
|
|
|
class IngestError(RuntimeError):
|
|
pass
|
|
|
|
|
|
def _coerce_date(value) -> _dt.date | None:
|
|
"""Accept date objects, 'YYYY-MM-DD' strings, or None; return date|None."""
|
|
if value is None or value == "":
|
|
return None
|
|
if isinstance(value, _dt.date) and not isinstance(value, _dt.datetime):
|
|
return value
|
|
if isinstance(value, _dt.datetime):
|
|
return value.date()
|
|
if isinstance(value, str):
|
|
try:
|
|
return _dt.date.fromisoformat(value[:10])
|
|
except ValueError:
|
|
return 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.
|
|
|
|
Plain `get_text("text")` interleaves columns and floating labels for
|
|
multi-column Hebrew PDFs (e.g. ספר הליקויים) — words come out split to
|
|
single letters on separate lines. Using `get_text("blocks")` and sorting
|
|
by vertical position keeps paragraphs intact for well-formatted pages.
|
|
|
|
Some PDFs (ספר הליקויים's title/preface pages) emit each word as its
|
|
own block. Sorting blocks by y-position alone leaves them one-word-per-
|
|
line. To fix that, we merge blocks that share the same visual row
|
|
(|Δy| ≤ row_tolerance) into a single line joined by spaces, right-to-
|
|
left. That makes prose pages readable without affecting well-formed
|
|
content pages.
|
|
"""
|
|
parts: list[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_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])
|
|
for b in raw
|
|
if isinstance(b, tuple) and len(b) >= 5 and isinstance(b[4], str) and b[4].strip()
|
|
]
|
|
if not textual:
|
|
continue
|
|
|
|
# Sort by y (top→bottom), then by -x so RTL rows read right→left
|
|
textual.sort(key=lambda t: (t[1], -t[0]))
|
|
|
|
# Group into rows by y proximity, then join each row with a space.
|
|
lines: list[str] = []
|
|
current_row: list[tuple[float, str]] = []
|
|
current_y: float | None = None
|
|
for x0, y0, x1, y1, text in textual:
|
|
text = text.strip()
|
|
if not text:
|
|
continue
|
|
if current_y is None or abs(y0 - current_y) <= row_tolerance:
|
|
current_row.append((x0, text))
|
|
current_y = y0 if current_y is None else current_y
|
|
else:
|
|
# Flush previous row
|
|
current_row.sort(key=lambda p: -p[0]) # RTL
|
|
lines.append(" ".join(p[1] for p in current_row))
|
|
current_row = [(x0, text)]
|
|
current_y = y0
|
|
if current_row:
|
|
current_row.sort(key=lambda p: -p[0])
|
|
lines.append(" ".join(p[1] for p in current_row))
|
|
|
|
# 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()
|
|
|
|
|
|
def _parse_docx(data: bytes) -> str:
|
|
from docx import Document # local import — keeps startup light
|
|
doc = Document(io.BytesIO(data))
|
|
return "\n\n".join(p.text for p in doc.paragraphs if p.text.strip())
|
|
|
|
|
|
def _parse_text(data: bytes) -> str:
|
|
return data.decode("utf-8", errors="replace").strip()
|
|
|
|
|
|
def parse(data: bytes, filename: str) -> str:
|
|
name = (filename or "").lower()
|
|
if name.endswith(".pdf"):
|
|
return _parse_pdf(data)
|
|
if name.endswith(".docx"):
|
|
return _parse_docx(data)
|
|
if name.endswith((".txt", ".md")):
|
|
return _parse_text(data)
|
|
# Fallback: try PDF then text
|
|
try:
|
|
return _parse_pdf(data)
|
|
except Exception:
|
|
return _parse_text(data)
|
|
|
|
|
|
async def ingest_source(
|
|
*,
|
|
kind: Literal["law", "regulation", "circular"],
|
|
title: str,
|
|
identifier: str | None,
|
|
content: bytes,
|
|
filename: str,
|
|
published_at: str | None = None,
|
|
effective_at: str | None = None,
|
|
source_url: str | None = None,
|
|
original_path: str | None = None,
|
|
) -> dict:
|
|
"""Parse, chunk, embed, and upsert a source document.
|
|
|
|
Returns: {source_id, chunks_created, was_updated, checksum}
|
|
"""
|
|
text = parse(content, filename)
|
|
if not text:
|
|
raise IngestError(f"no text extracted from {filename}")
|
|
|
|
checksum = hashlib.sha256(text.encode("utf-8")).hexdigest()
|
|
chunks = chunk(text, kind)
|
|
if not chunks:
|
|
raise IngestError("chunker produced zero chunks")
|
|
|
|
logger.info(
|
|
"[kb.ingest] kind=%s title=%r len=%d chunks=%d checksum=%s",
|
|
kind, title[:60], len(text), len(chunks), checksum[:12],
|
|
)
|
|
|
|
pool = await get_pool()
|
|
async with pool.acquire() as conn:
|
|
async with conn.transaction():
|
|
# Match a prior source by (kind + identifier) when identifier is provided,
|
|
# otherwise fall back to (kind + title). Without this, two different
|
|
# sources with NULL identifier would both compare equal and the second
|
|
# ingest would incorrectly supersede the first.
|
|
if identifier:
|
|
existing = await conn.fetchrow(
|
|
"""
|
|
SELECT id, checksum FROM kb_source
|
|
WHERE kind = $1 AND identifier = $2 AND superseded_by IS NULL
|
|
ORDER BY created_at DESC LIMIT 1
|
|
""",
|
|
kind, identifier,
|
|
)
|
|
else:
|
|
existing = await conn.fetchrow(
|
|
"""
|
|
SELECT id, checksum FROM kb_source
|
|
WHERE kind = $1 AND title = $2 AND identifier IS NULL AND superseded_by IS NULL
|
|
ORDER BY created_at DESC LIMIT 1
|
|
""",
|
|
kind, title,
|
|
)
|
|
was_updated = False
|
|
|
|
if existing and existing["checksum"] == checksum:
|
|
logger.info("[kb.ingest] skip — identical checksum, source_id=%s", existing["id"])
|
|
return {
|
|
"source_id": existing["id"],
|
|
"chunks_created": 0,
|
|
"was_updated": False,
|
|
"checksum": checksum,
|
|
}
|
|
|
|
# Upsert source (new row per checksum — supersession is explicit).
|
|
row = await conn.fetchrow(
|
|
"""
|
|
INSERT INTO kb_source
|
|
(kind, title, identifier, published_at, effective_at,
|
|
source_url, original_path, checksum)
|
|
VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
|
|
RETURNING id
|
|
""",
|
|
kind, title, identifier,
|
|
_coerce_date(published_at), _coerce_date(effective_at),
|
|
source_url, original_path, checksum,
|
|
)
|
|
source_id = row["id"]
|
|
|
|
# If a previous source exists, mark it superseded.
|
|
if existing:
|
|
await conn.execute(
|
|
"UPDATE kb_source SET superseded_by = $1 WHERE id = $2",
|
|
source_id, existing["id"],
|
|
)
|
|
was_updated = True
|
|
|
|
# Embed in batches via Voyage, then insert chunks.
|
|
texts = [c["content"] for c in chunks]
|
|
vectors = await voyage.embed(texts, input_type="document")
|
|
|
|
rows = [
|
|
(
|
|
source_id,
|
|
c["chunk_index"],
|
|
c["heading_path"],
|
|
c["section_ref"],
|
|
c["content"],
|
|
vectors[i],
|
|
c["token_count"],
|
|
c.get("page_number"),
|
|
)
|
|
for i, c in enumerate(chunks)
|
|
]
|
|
await conn.executemany(
|
|
"""
|
|
INSERT INTO kb_chunk
|
|
(source_id, chunk_index, heading_path, section_ref,
|
|
content, embedding, token_count, page_number)
|
|
VALUES ($1,$2,$3,$4,$5,$6,$7,$8)
|
|
""",
|
|
rows,
|
|
)
|
|
|
|
return {
|
|
"source_id": source_id,
|
|
"chunks_created": len(chunks),
|
|
"was_updated": was_updated,
|
|
"checksum": checksum,
|
|
}
|
|
|
|
|
|
async def stats() -> dict:
|
|
pool = await get_pool()
|
|
async with pool.acquire() as conn:
|
|
row = await conn.fetchrow(
|
|
"""
|
|
SELECT
|
|
(SELECT COUNT(*) FROM kb_source) AS total_sources,
|
|
(SELECT COUNT(*) FROM kb_source WHERE superseded_by IS NULL) AS active_sources,
|
|
(SELECT COUNT(*) FROM kb_chunk) AS total_chunks,
|
|
(SELECT COUNT(*) FROM kb_chunk WHERE embedding IS NOT NULL) AS embedded_chunks
|
|
"""
|
|
)
|
|
by_kind = await conn.fetch(
|
|
"""
|
|
SELECT kind, COUNT(*) AS n
|
|
FROM kb_source WHERE superseded_by IS NULL
|
|
GROUP BY kind ORDER BY kind
|
|
"""
|
|
)
|
|
return {
|
|
**dict(row),
|
|
"active_by_kind": {r["kind"]: r["n"] for r in by_kind},
|
|
}
|