src/extractors/extractor.py

"""Main extraction orchestrator.

    from src.extractors import DocumentExtractor
    ex = DocumentExtractor()
    result, metrics = ex.extract(file_bytes, "invoice.pdf", doc_type="invoice")
    # result is ExtractionResult[Invoice]
    # metrics has tokens, cost, latency
"""
from __future__ import annotations

from typing import Any

from src.extractors.document_loader import LoadedDocument, load_document
from src.extractors.envelope import compute_overall_confidence, make_envelope
from src.extractors.openai_client import OpenAIExtractionClient
from src.extractors.prompts import get_prompt
from src.extractors.section_chunker import chunk_filing
from src.schemas import ExtractionResult
from src.schemas.registry import get_schema
from src.utils.cost_tracker import ExtractionMetrics
from src.utils.logging import logger


class DocumentExtractor:
    """Public entry point for extraction.

    Stateless-ish: holds a shared OpenAI client + default model, no per-request state.
    Safe to reuse across many extractions.
    """

    def __init__(
        self,
        client: OpenAIExtractionClient | None = None,
        default_model: str | None = None,
    ):
        self._client = client or OpenAIExtractionClient(model=default_model)

    # ------------------------------------------------------------------

    def extract(
        self,
        file_bytes: bytes,
        filename: str,
        doc_type: str,
        *,
        model_override: str | None = None,
        render_images: bool = True,
        reasoning_effort: str | None = None,
    ) -> tuple[ExtractionResult, ExtractionMetrics]:
        """Extract a doc into an ExtractionResult[T] plus per-call metrics.

        - `doc_type` looks up both the schema (via registry) and the prompt.
        - `model_override` lets the caller swap models for benchmarking.
        - `render_images=False` skips vision (text-only) for cheap extractions.
        """
        # 1. Look up schema + prompt.
        schema_cls = get_schema(doc_type)
        system_prompt = get_prompt(doc_type)
        envelope_cls = make_envelope(schema_cls)

        # 2. Load the document.
        loaded = load_document(file_bytes, filename, render_images=render_images)
        if loaded.source_type == "empty":
            raise ValueError(f"Could not load document {filename!r} (unknown or corrupt format).")

        # 3. Build the messages. Filings use a section-aware path — a full 10-K
        #    is ~150K tokens; we ship only cover + Item 8 + Item 1A to keep
        #    per-call cost + latency reasonable and reduce distractor text.
        if doc_type.strip().lower() == "filing":
            messages = self._build_filing_messages(system_prompt, loaded)
        else:
            messages = self._build_messages(system_prompt, loaded)

        # 4. Call the model.
        envelope, metrics = self._client.parse_structured(
            response_format=envelope_cls,
            messages=messages,
            model=model_override,
            reasoning_effort=reasoning_effort,
        )

        # 5. Wrap into ExtractionResult.
        result = ExtractionResult(
            document_type=doc_type,
            data=envelope.data,
            field_confidences=envelope.field_confidences,
            overall_confidence=compute_overall_confidence(envelope.field_confidences),
            warnings=envelope.warnings,
            raw_text_snippet=loaded.text[:2000] if loaded.text else None,
        )

        logger.info(
            f"Extracted {doc_type} from {filename}: "
            f"overall_confidence={result.overall_confidence:.2f}, "
            f"warnings={len(result.warnings)}, cost=${metrics.cost_usd:.5f}"
        )
        return result, metrics

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    @staticmethod
    def _build_messages(system_prompt: str, loaded: LoadedDocument) -> list[dict[str, Any]]:
        """Assemble the OpenAI messages list from the loaded document."""
        user_content: list[dict[str, Any]] = []

        if loaded.text:
            user_content.append(
                {
                    "type": "text",
                    "text": (
                        "Extract the structured data from this document. "
                        "The document text follows (and page images may also be attached):\n\n"
                        f"---BEGIN DOCUMENT TEXT---\n{loaded.text}\n---END DOCUMENT TEXT---"
                    ),
                }
            )
        else:
            user_content.append(
                {
                    "type": "text",
                    "text": (
                        "Extract the structured data from this document. "
                        "Only page images are provided (no text was extractable)."
                    ),
                }
            )

        for img_b64 in loaded.images_b64:
            user_content.append(
                {
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/png;base64,{img_b64}",
                        "detail": "high",
                    },
                }
            )

        return [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_content},
        ]

    # ------------------------------------------------------------------
    # Filing path
    # ------------------------------------------------------------------

    # Per-section byte caps. Real 10-K sections rarely exceed ~40 KB; 60 KB
    # gives us headroom for verbose filers (JPM's Item 1A runs long) while
    # keeping total prompt size ~30-40K tokens — well inside gpt-5-nano's cost
    # sweet spot.
    _FILING_COVER_BYTES = 6_000
    _FILING_ITEM_1A_BYTES = 60_000
    _FILING_ITEM_8_BYTES = 60_000

    def _build_filing_messages(
        self,
        system_prompt: str,
        loaded: LoadedDocument,
    ) -> list[dict[str, Any]]:
        """Message builder for the 10-K path.

        Slices the loaded text into cover + Item 8 (financials) + Item 1A
        (risk factors) and hands each to the model as a clearly-labeled block.
        Filings are text-first — vision isn't used here (10-K images are
        usually chart infographics, not extraction targets).
        """
        text = loaded.text or ""
        chunked = chunk_filing(text, cover_bytes=self._FILING_COVER_BYTES)

        cover = chunked.cover[: self._FILING_COVER_BYTES]
        item_8 = chunked.get_text("8", default="(Item 8 not found in this filing.)")[
            : self._FILING_ITEM_8_BYTES
        ]
        item_1a = chunked.get_text("1A", default="(Item 1A not found in this filing.)")[
            : self._FILING_ITEM_1A_BYTES
        ]

        logger.info(
            f"[filing] chunked: cover={len(cover)}B, "
            f"item_8={len(item_8)}B (present={chunked.has('8')}), "
            f"item_1a={len(item_1a)}B (present={chunked.has('1A')}), "
            f"total_items={len(chunked.item_ids)}"
        )

        user_text = (
            "Extract the structured filing data. Three relevant sections of the "
            "10-K are provided below. Do NOT hallucinate values from other "
            "sections not shown.\n\n"
            "---COVER SECTION---\n"
            f"{cover}\n"
            "---END COVER SECTION---\n\n"
            "---FINANCIAL SECTION (Item 8)---\n"
            f"{item_8}\n"
            "---END FINANCIAL SECTION---\n\n"
            "---RISK FACTORS SECTION (Item 1A)---\n"
            f"{item_1a}\n"
            "---END RISK FACTORS SECTION---"
        )
        return [
            {"role": "system", "content": system_prompt},
            {"role": "user",   "content": [{"type": "text", "text": user_text}]},
        ]