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:
    parts: list[str] = []
    with fitz.open(stream=data, filetype="pdf") as doc:
        for page in doc:
            parts.append(page.get_text("text"))
    return "\n\n".join(parts).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},
    }