src/document/ingest_v2.py

"""Unstructured document ingestion v2 — parse, store figures, embed.

UNSTRUCTURED_V2_PLAN U2. Replaces Go's v1 path (Tesseract / DOCX XML / TXT →
fixed-size chunks) for `pdf`, `docx` and `txt` uploads made through the v2
document routes. Tabular files never come here; they stay Go's.

    original in Supabase ──► Parsing v2 (PaddleOCR-VL or Mistral OCR, per
                              DOCUMENT_PARSER_BACKEND; TXT skips OCR)
                              ├── figures ──► SupabaseAssetSink  (private bucket)
                              └── chunks  ──► embeddings ──► langchain_pg_embedding
                                                              collection "documents"

What a stored chunk looks like, and why (plan §4.1):

- `document` = the chunk's heading line + its verbatim text, figure markers
  `![](asset://<asset_id>)` left where the document put them. The marker is the
  only image reference that is ever persisted; the browser gets a URL from the
  image-token endpoint at render time.
- `cmetadata` = Go's embedder shape (verified on all 20 live rows, 2026-09-23)
  — `user_id`, `source_type`, `updated_at`, `data{document_id, filename,
  file_type, page_label, chunk_index}` — so the retriever, `_normalize_chunks` and
  traceability read v2 rows unchanged. Added: `parser_version` (retrieval is
  v2-only, U-D6), `chunk_id`, and `assets` — the figures whose marker is in THIS
  chunk, with their `storage_key`. `assets` sits outside `data` on purpose:
  nothing that renders `data` to a prompt or an API response ever sees a key.

Rows are written with raw SQL mirroring Go's insert
(`internal/documents/service.go:1226`) rather than LangChain's PGVector: this
table is the older uuid-`id` + `custom_id` schema, which PGVector mis-types (see
`retrieval/document.py`'s module note).

Failure is LOUD on the document row (`status='failed'` + `error_message`) — a
user must see that a file did not index. Only individual figure uploads degrade
silently (the sink is never-throw), because one figure is not worth a document.
"""

from __future__ import annotations

import asyncio
import json
import math
import tempfile
import uuid
from concurrent.futures import ThreadPoolExecutor
from dataclasses import dataclass, field
from datetime import datetime, timedelta, timezone
from pathlib import Path
from typing import Any

from langchain.text_splitter import RecursiveCharacterTextSplitter
from sqlalchemy import text

from src.middlewares.logging import get_logger

logger = get_logger("document_ingest_v2")

PARSER_VERSION = "v2"
UNSTRUCTURED_TYPES = ("pdf", "docx", "txt")
# Which OCR parser processing uses — server-side only, set by
# DOCUMENT_PARSER_BACKEND (added 2026-09-24). PaddleOCR takes PDF and images only.
PARSER_BACKENDS = ("paddleocr", "mistral")
COLLECTION_NAME = "documents"

# Mirrors Go's processAsync budget (`context.WithTimeout(…, 10*time.Minute)`).
PROCESS_TIMEOUT_SECONDS = 600

# Guard, not a chunker. Parsed v2 chunks measured 2026-09-23 on three documents
# (BUMA / Komatsu / Open Pit): max 7,846 chars, zero above 8,000 — so this never
# fires on them. It exists so one unstructured wall of text cannot exceed the
# embedding model's input limit and fail the whole document.
_MAX_EMBED_CHARS = 8000
# TXT has no layout to parse, so it is chunked by paragraphs to roughly the
# measured median size of a parsed chunk.
_TXT_CHUNK_CHARS = 1500
_TXT_OVERLAP = 150

_ASSET_MARKER = "asset://"

_JAKARTA_TZ = timezone(timedelta(hours=7))

# Parsing blocks (Mistral upload + OCR + normalisation), so it runs off the event
# loop in its own small pool — the #49 `dbexec` precedent, so a long parse never
# competes with the shared default pool.
_EXECUTOR = ThreadPoolExecutor(max_workers=2, thread_name_prefix="docv2")


@dataclass
class EmbedChunk:
    """One row to embed and store."""

    text: str
    page_label: int
    chunk_index: int
    chunk_id: str
    assets: list[dict[str, str]] = field(default_factory=list)
    # "<parser>/<model>", e.g. "paddleocr-vl/PaddleOCR-VL-1.6"; None for TXT.
    parser: str | None = None


# ---------------------------------------------------------------------------
# Chunk building (pure — unit-tested without a model or a database)
# ---------------------------------------------------------------------------


def _figure_assets(chunk: Any) -> dict[str, str]:
    """asset_id → storage_key for this chunk's STORED figures.

    Tables have no bytes, and a figure whose upload failed has no key; neither
    can ever be signed, so neither is recorded.
    """
    out: dict[str, str] = {}
    for a in getattr(chunk, "assets", None) or []:
        # "chart" too since PaddleOCR (2026-09-24): its chart blocks are stored images.
        if getattr(a, "kind", None) in {"figure", "chart", "diagram"} and getattr(
            a, "storage_key", None
        ):
            out[a.asset_id] = a.storage_key
    return out


def _split(body: str) -> list[str]:
    if len(body) <= _MAX_EMBED_CHARS:
        return [body]
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=_MAX_EMBED_CHARS,
        chunk_overlap=_MAX_EMBED_CHARS // 20,
        length_function=len,
    )
    return splitter.split_text(body)


def chunks_from_artifact(artifact: Any) -> list[EmbedChunk]:
    """ParsedDocument → rows. Heading line + verbatim text; empty chunks dropped.

    A figure is attached to a row only when that row's text carries its marker,
    so an over-long chunk split in two keeps each figure with the half it sits in.
    """
    out: list[EmbedChunk] = []
    name = getattr(artifact, "parser_name", None)
    parser = f"{name}/{getattr(artifact, 'parser_version', None)}" if name else None
    for chunk in artifact.chunks:
        heading = (chunk.heading or "").strip()
        body_text = (chunk.text or "").strip()
        if not body_text:
            continue
        body = f"{heading}\n{body_text}" if heading else body_text
        stored = _figure_assets(chunk)
        for piece in _split(body):
            assets = [
                {"asset_id": aid, "storage_key": key}
                for aid, key in stored.items()
                if f"{_ASSET_MARKER}{aid}" in piece
            ]
            out.append(
                EmbedChunk(
                    text=piece,
                    page_label=chunk.page_idx + 1,
                    chunk_index=len(out),
                    chunk_id=chunk.chunk_id,
                    assets=assets,
                    parser=parser,
                )
            )
    return out


def chunks_from_text(document_id: str, raw: str) -> list[EmbedChunk]:
    """TXT → rows. No layout, no figures, one page."""
    splitter = RecursiveCharacterTextSplitter(
        chunk_size=_TXT_CHUNK_CHARS, chunk_overlap=_TXT_OVERLAP, length_function=len
    )
    pieces = [p.strip() for p in splitter.split_text(raw) if p.strip()]
    return [
        EmbedChunk(
            text=p,
            page_label=1,
            chunk_index=i,
            chunk_id=f"{document_id}::{i:04d}",
        )
        for i, p in enumerate(pieces)
    ]


def build_metadata(document: Any, chunk: EmbedChunk) -> dict[str, Any]:
    meta = {
        "user_id": document.user_id,
        "source_type": "document",
        "updated_at": datetime.now(_JAKARTA_TZ).isoformat(),
        "data": {
            "document_id": document.id,
            "filename": document.filename,
            "file_type": document.file_type,
            "page_label": chunk.page_label,
            "chunk_index": chunk.chunk_index,
        },
        "parser_version": PARSER_VERSION,
        "chunk_id": chunk.chunk_id,
        "assets": chunk.assets,
    }
    # Additive (2026-09-24): which OCR produced this row. Outside `data`, like
    # `assets`, so nothing that renders `data` into a prompt sees it.
    if chunk.parser:
        meta["parser"] = chunk.parser
    return meta


# ---------------------------------------------------------------------------
# The service
# ---------------------------------------------------------------------------


def parser_backend() -> str:
    """The configured parser, validated — an unknown value fails loudly."""
    from src.config.settings import settings

    backend = (settings.document_parser_backend or "").strip().lower()
    if backend not in PARSER_BACKENDS:
        raise ValueError(
            f"DOCUMENT_PARSER_BACKEND={backend!r} is not one of {', '.join(PARSER_BACKENDS)}"
        )
    return backend


def _parse_sync(source: Path, doc_id: str, sink: Any, backend: str) -> Any:
    from src.config.settings import settings
    from src.knowledge_parsing.config import PipelineConfig
    from src.knowledge_parsing.run import parse_one

    cfg = PipelineConfig(
        input_dir=source.parent,
        backend=backend,
        paddle_model=settings.paddleocr_model,
        paddle_job_url=settings.paddleocr_job_url,
    )
    return parse_one(source, cfg, doc_id, asset_sink=sink)


def _vector_literal(vec: list[float]) -> str:
    if not all(math.isfinite(v) for v in vec):
        raise ValueError("Embedding vector contains NaN or Infinity values.")
    return "[" + ",".join(str(v) for v in vec) + "]"


class DocumentIngestV2:
    """Process one uploaded `documents` row end to end."""

    def __init__(self, storage: Any = None, embeddings: Any = None, session_factory: Any = None):
        # All three are injectable so tests run with no bucket, no model, no DB.
        self._storage = storage
        self._embeddings = embeddings
        self._session_factory = session_factory

    def _get_storage(self):
        if self._storage is None:
            from src.storage.object_storage.supabase_s3 import object_storage

            self._storage = object_storage
        return self._storage

    def _get_embeddings(self):
        if self._embeddings is None:
            # The SAME client the retriever embeds queries with, so stored and
            # query vectors always come from one deployment.
            from src.retrieval.document import _get_embeddings

            self._embeddings = _get_embeddings()
        return self._embeddings

    def _sessions(self):
        if self._session_factory is None:
            from src.db.postgres.connection import AsyncSessionLocal

            self._session_factory = AsyncSessionLocal
        return self._session_factory

    async def build_chunks(self, document: Any) -> list[EmbedChunk]:
        """Download the original and turn it into rows (parsing off-loop)."""
        from src.knowledge_ingest.asset_sink import SupabaseAssetSink

        file_type = (document.file_type or "").lower()
        backend = None if file_type == "txt" else parser_backend()
        # Refused before the download: PaddleOCR's API takes PDF and images only.
        if backend == "paddleocr" and file_type == "docx":
            raise ValueError(
                "DOCX is not supported by the PaddleOCR parser — upload the document "
                "as PDF, or set DOCUMENT_PARSER_BACKEND=mistral"
            )
        storage = self._get_storage()
        content = await storage.download_file(document.blob_name)

        if file_type == "txt":
            return chunks_from_text(document.id, content.decode("utf-8", errors="replace"))

        sink = SupabaseAssetSink(storage, document.user_id, document.id)
        with tempfile.TemporaryDirectory(prefix="docv2-") as tmp:
            # Named for the original so the artifact's `source_title` is readable;
            # identity is `documents.id`, passed explicitly.
            stem = Path(document.filename or "document").stem
            source = Path(tmp) / f"{stem}.{file_type}"
            source.write_bytes(content)
            loop = asyncio.get_running_loop()
            artifact = await loop.run_in_executor(
                _EXECUTOR, _parse_sync, source, document.id, sink, backend
            )
        return chunks_from_artifact(artifact)

    async def write_rows(self, document: Any, chunks: list[EmbedChunk]) -> None:
        """Embed and replace this document's v2 rows in one transaction.

        Deletes only `parser_version='v2'` rows for this document: a re-process
        replaces its own output and never touches Go's v1 rows (U-D6).
        """
        vectors = await self._get_embeddings().aembed_documents([c.text for c in chunks])
        async with self._sessions()() as session:
            async with session.begin():
                coll = await session.execute(
                    text("SELECT uuid FROM langchain_pg_collection WHERE name = :n LIMIT 1"),
                    {"n": COLLECTION_NAME},
                )
                collection_id = coll.scalar()
                if collection_id is None:
                    # Go creates it on its first embedding; it exists in dedorch
                    # (verified 2026-09-23). Creating a Go-owned row here is not
                    # this service's job — fail loudly instead.
                    raise RuntimeError(f"pgvector collection {COLLECTION_NAME!r} not found")
                await session.execute(
                    text(
                        "DELETE FROM langchain_pg_embedding "
                        "WHERE collection_id = :c "
                        "AND cmetadata->'data'->>'document_id' = :d "
                        "AND cmetadata->>'parser_version' = :v"
                    ),
                    {"c": collection_id, "d": document.id, "v": PARSER_VERSION},
                )
                for chunk, vec in zip(chunks, vectors, strict=True):
                    await session.execute(
                        text(
                            "INSERT INTO langchain_pg_embedding "
                            "(id, collection_id, embedding, document, cmetadata, custom_id) "
                            "VALUES (:id, :c, CAST(:e AS vector), :doc, CAST(:m AS jsonb), :cid)"
                        ),
                        {
                            "id": uuid.uuid4(),
                            "c": collection_id,
                            "e": _vector_literal(vec),
                            "doc": chunk.text,
                            "m": json.dumps(build_metadata(document, chunk), ensure_ascii=False),
                            "cid": f"{document.id}_{PARSER_VERSION}_{chunk.chunk_index}",
                        },
                    )

    async def _set_status(
        self, document_id: str, status: str, *, error: str | None = None, chunks: int | None = None
    ) -> None:
        async with self._sessions()() as session:
            async with session.begin():
                if status == "completed":
                    await session.execute(
                        text(
                            "UPDATE documents SET status = :s, processed_at = now(), "
                            "chunks_count = :n, error_message = NULL WHERE id = :id"
                        ),
                        {"s": status, "n": chunks, "id": document_id},
                    )
                else:
                    await session.execute(
                        text("UPDATE documents SET status = :s, error_message = :e WHERE id = :id"),
                        {"s": status, "e": error, "id": document_id},
                    )

    async def _invalidate_retrieval_cache(self, user_id: str) -> None:
        """Drop the user's cached retrievals so new chunks are visible at once
        (the v1 route did the same, `api/v1/document.py:117`). Never-throw: a
        stale cache only delays visibility by at most its 1h TTL."""
        try:
            from src.retrieval.router import retrieval_router

            await retrieval_router.invalidate_cache(user_id)
        except Exception as e:
            logger.warning(
                "retrieval cache invalidation failed",
                user_id=user_id,
                error=repr(e),
                degraded_seam="retrieval_cache_invalidate",
            )

    async def process(self, document: Any) -> None:
        """Background entry point. Never raises — the outcome lands on the row."""
        try:
            chunks = await asyncio.wait_for(
                self.build_chunks(document), timeout=PROCESS_TIMEOUT_SECONDS
            )
            if not chunks:
                raise ValueError("no text extracted from document")
            await self.write_rows(document, chunks)
            await self._set_status(document.id, "completed", chunks=len(chunks))
            await self._invalidate_retrieval_cache(document.user_id)
            logger.info(
                "document processed",
                document_id=document.id,
                chunks=len(chunks),
                figures=sum(len(c.assets) for c in chunks),
            )
        except Exception as e:
            logger.error("document processing failed", document_id=document.id, error=repr(e))
            try:
                await self._set_status(document.id, "failed", error=repr(e)[:500])
            except Exception as e2:
                logger.error(
                    "could not record failure",
                    document_id=document.id,
                    error=repr(e2),
                    degraded_seam="document_status",
                )


document_ingest_v2 = DocumentIngestV2()