shira-hermes/api/services/kb/ingest.py

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


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 in 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))

            parts.append("\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"],
                )
                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)
                VALUES ($1,$2,$3,$4,$5,$6,$7)
                """,
                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},
    }