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


# 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_first_pages(data: bytes, filename: str, n: int = 3) -> str:
    """Plain-text extraction of the first ~N pages, used by the
    AI classifier on upload (api/services/kb/classifier.py).

    Faster and lighter than `parse(...)` because:
      • Uses `get_text("text")` (line-by-line, no block sorting / RTL
        merging). The classifier is robust to interleaved columns —
        it just needs the document title, citation header, and a
        paragraph or two of body to decide kind/subject/dates.
      • Caps at N pages so a 600-page ספר הליקויים doesn't burn
        seconds of parsing every time it's re-uploaded.

    Returns empty string if the file isn't a PDF/DOCX/TXT or has no
    extractable text. Caller (classifier) treats empty as fail-soft.
    """
    name = (filename or "").lower()
    try:
        if name.endswith(".pdf"):
            with fitz.open(stream=data, filetype="pdf") as doc:
                pages = []
                for i in range(min(n, len(doc))):
                    pages.append(doc[i].get_text("text") or "")
                return "\n\n".join(p.strip() for p in pages if p.strip())
        if name.endswith(".docx"):
            # DOCX has no real "page" concept; first 1500 chars of body
            # is a reasonable proxy for "first 3 pages worth".
            return _parse_docx(data)[:6000]
        if name.endswith(".txt") or name.endswith(".md"):
            return _parse_text(data)[:6000]
    except Exception as e:
        logger.warning("[kb.ingest] parse_first_pages failed for %s: %s", filename, e)
    return ""


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", "caselaw", "tool"],
    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,
    topic_id: int | None = None,
    description: str | None = None,
    ai_classified: bool = False,
) -> 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}")

    # Phase 1 made topic_id NOT NULL on kb_source. Legacy callers (admin
    # /ingest, /scan-inbox) don't pass it yet — fall back to topic_id=1
    # (ביטוח לאומי) so the existing inbox cron keeps working unchanged.
    if topic_id is None:
        topic_id = 1

    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).
            ai_ts = "now()" if ai_classified else "NULL"
            row = await conn.fetchrow(
                f"""
                INSERT INTO kb_source
                    (kind, title, identifier, published_at, effective_at,
                     source_url, original_path, checksum, topic_id,
                     description, ai_classified_at)
                VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,{ai_ts})
                RETURNING id
                """,
                kind, title, identifier,
                _coerce_date(published_at), _coerce_date(effective_at),
                source_url, original_path, checksum, topic_id,
                description,
            )
            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},
    }