backend/services/ingestion.py

"""Document ingestion pipeline."""

import logging
import uuid
from datetime import datetime, timezone
from typing import Dict, List

from config import settings
from db.faiss_client import FaissDB
from models.embedder import MiniCPMEmbedder
from models.ocr import MiniCPMVOCR
from utils.chunker import FinanceAwareChunker
from utils.liteparse_parser import parse_document
from utils.pdf_parser import extract_pdf_spatial_pages, render_page_image

logger = logging.getLogger(__name__)


class IngestionService:
    def __init__(
        self,
        embedder: MiniCPMEmbedder,
        ocr: MiniCPMVOCR,
        db: FaissDB,
    ):
        self.embedder = embedder
        self.ocr = ocr
        self.db = db
        self.chunker = FinanceAwareChunker()

    def _embed_texts(self, texts: List[str]) -> List[List[float]]:
        batch_size = settings.EMBED_BATCH_SIZE
        vectors: List[List[float]] = []
        total = len(texts)
        for start in range(0, total, batch_size):
            batch = texts[start : start + batch_size]
            logger.info(
                "Embedding batch %d–%d of %d",
                start + 1,
                min(start + len(batch), total),
                total,
            )
            vectors.extend(self.embedder.embed_documents(batch))
        return vectors

    def ingest_pdf(self, file_bytes: bytes, filename: str) -> Dict:
        doc_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc).isoformat()
        all_chunks: List[Dict] = []

        logger.info("Parsing %s ...", filename)
        parse_result = parse_document(file_bytes, filename, self.ocr)
        sparse_pages = {
            page_num
            for page_num, _, is_sparse in extract_pdf_spatial_pages(file_bytes)
            if is_sparse
        }
        logger.info(
            "Parsed %d pages (%d OCR pages) from %s",
            len(parse_result.pages),
            len(sparse_pages),
            filename,
        )

        chart_ocr_count = 0
        for parsed_page in parse_result.pages:
            page_text = parsed_page.text.strip()
            if not page_text:
                continue

            page_num = parsed_page.page_num
            source = "liteparse" if page_num in sparse_pages else "embedded"
            page_chunks = self.chunker.chunk(
                page_text, page_num=page_num, source=source
            )

            if (
                chart_ocr_count < settings.CHART_OCR_MAX_PAGES
                and self.chunker.should_extract_chart(page_text)
            ):
                try:
                    logger.info(
                        "Chart OCR page %d (%d/%d cap)",
                        page_num,
                        chart_ocr_count + 1,
                        settings.CHART_OCR_MAX_PAGES,
                    )
                    page_image = render_page_image(file_bytes, page_num)
                    chart_desc = self.ocr.describe_chart(page_image)
                    if chart_desc and len(chart_desc.strip()) > 20:
                        chart_chunks = self.chunker.chunk(
                            chart_desc,
                            page_num=page_num,
                            source="ocr_chart",
                            section_override="chart_data",
                        )
                        page_chunks.extend(chart_chunks)
                    chart_ocr_count += 1
                except Exception as e:
                    logger.warning(
                        "Chart extraction failed on page %d: %s", page_num, e
                    )

            for chunk in page_chunks:
                chunk["document_id"] = doc_id
                chunk["document_name"] = filename
                chunk["document_type"] = "pdf"
                chunk["page_number"] = page_num
                chunk["created_at"] = now

            all_chunks.extend(page_chunks)

        if not all_chunks:
            return {"document_id": doc_id, "chunks_ingested": 0, "filename": filename}

        texts = [c["text"] for c in all_chunks]
        logger.info("Embedding %d chunks from %s ...", len(texts), filename)
        vectors = self._embed_texts(texts)
        logger.info("Saving %d chunks to FAISS ...", len(all_chunks))
        self.db.upsert_chunks(all_chunks, vectors)
        logger.info(
            "Ingestion complete: %s (%d chunks, %d chart OCR pages)",
            filename,
            len(all_chunks),
            chart_ocr_count,
        )

        return {
            "document_id": doc_id,
            "chunks_ingested": len(all_chunks),
            "filename": filename,
        }

    def ingest_image(self, file_bytes: bytes, filename: str) -> Dict:
        doc_id = str(uuid.uuid4())
        now = datetime.now(timezone.utc).isoformat()

        logger.info("Parsing image %s ...", filename)
        parse_result = parse_document(file_bytes, filename, self.ocr)
        ocr_text = parse_result.text.strip()
        chunks = self.chunker.chunk(ocr_text, page_num=1, source="liteparse")

        try:
            chart_desc = self.ocr.describe_chart(file_bytes)
            if chart_desc and len(chart_desc.strip()) > 20:
                chart_chunks = self.chunker.chunk(
                    chart_desc,
                    page_num=1,
                    source="ocr_chart",
                    section_override="chart_data",
                )
                chunks.extend(chart_chunks)
        except Exception as e:
            logger.warning("Chart extraction failed: %s", e)

        for chunk in chunks:
            chunk["document_id"] = doc_id
            chunk["document_name"] = filename
            chunk["document_type"] = "image"
            chunk["page_number"] = 1
            chunk["created_at"] = now

        texts = [c["text"] for c in chunks]
        logger.info("Embedding %d chunks from %s ...", len(texts), filename)
        vectors = self._embed_texts(texts)
        self.db.upsert_chunks(chunks, vectors)

        return {
            "document_id": doc_id,
            "chunks_ingested": len(chunks),
            "filename": filename,
        }