Extraction: OCR + LLM

from __future__ import annotations
import io
import pdfplumber
import pytesseract
from pdf2image import convert_from_bytes


def text_from_native(file_bytes: bytes) -> str:
    parts = []
    with pdfplumber.open(io.BytesIO(file_bytes)) as pdf:
        for page in pdf.pages:
            t = page.extract_text() or ""
            parts.append(t)
    return "\n\n".join(parts)


def text_from_ocr(file_bytes: bytes) -> str:
    images = convert_from_bytes(file_bytes, dpi=200)
    parts = []
    for img in images:
        parts.append(pytesseract.image_to_string(img))
    return "\n\n".join(parts)


def extract_text(file_bytes: bytes) -> tuple[str, str]:
    """Return (text, source) where source is 'native' or 'ocr'."""
    native = text_from_native(file_bytes)
    if len(native.strip()) > 200:
        return native, "native"
    # Fall back to OCR for image-heavy docs
    return text_from_ocr(file_bytes), "ocr"

from __future__ import annotations
from pydantic import BaseModel, Field


class Citation(BaseModel):
    quote: str = Field(description="The exact text from the document supporting this value")
    page: int | None = Field(default=None, description="Page number (1-indexed) if known")


class FieldWithCitation(BaseModel):
    value: str | None
    confidence: float = Field(ge=0.0, le=1.0)
    citation: Citation | None = None


class RenewalTerms(BaseModel):
    auto_renew: bool
    term_length_months: int
    notice_period_days: int


class RenewalTermsWithCitation(BaseModel):
    value: RenewalTerms | None
    confidence: float = Field(ge=0.0, le=1.0)
    citation: Citation | None = None


class MSAExtraction(BaseModel):
    """Fields extracted from a Master Service Agreement."""
    parties: list[str] = Field(description="The contracting entities (e.g., ['Acme Corp', 'Customer Inc'])")
    parties_confidence: float = Field(ge=0.0, le=1.0)
    effective_date: FieldWithCitation
    term_length_months: FieldWithCitation
    renewal_terms: RenewalTermsWithCitation
    governing_law: FieldWithCitation

SYSTEM_PROMPT = """You are an expert at extracting structured information from Master Service Agreements (MSAs).

You will be given the text of an MSA. Extract the fields requested.

Rules:
1. If a field cannot be found in the document, return null for that field's value.
2. Do not guess. Do not fabricate. If you are uncertain, return null and set confidence below 0.7.
3. For every extracted value, return a citation: the exact quote from the document supporting it, and the page if known.
4. Confidence is your honest estimate of correctness: 1.0 means certain, 0.7 is the default for HITL routing, < 0.5 means very uncertain.
5. Normalize dates to ISO format (YYYY-MM-DD). If only a month/year is present, leave value as null.
"""

USER_TEMPLATE = """Extract structured fields from the following Master Service Agreement.

<document>
{text}
</document>

Return your answer as a single JSON object matching the schema."""

from __future__ import annotations
import json
from anthropic import Anthropic
from pydantic import ValidationError

from ci.config import settings
from ci.extraction.schemas import MSAExtraction


SYSTEM_PROMPT = "..."  # as above
USER_TEMPLATE = "..."  # as above

_client = Anthropic(api_key=settings.anthropic_api_key)


def extract_msa(text: str) -> MSAExtraction:
    """Extract MSA fields from document text."""
    tool = {
        "name": "extract_msa_fields",
        "description": "Extract structured MSA fields",
        "input_schema": MSAExtraction.model_json_schema(),
    }

    response = _client.messages.create(
        model=settings.extraction_model,
        max_tokens=settings.extraction_max_tokens,
        temperature=settings.extraction_temperature,
        system=SYSTEM_PROMPT,
        tools=[tool],
        tool_choice={"type": "tool", "name": "extract_msa_fields"},
        messages=[
            {"role": "user", "content": USER_TEMPLATE.format(text=text[:30000])}
            # 30k char truncation; LLMs handle longer but cost / latency adds up.
        ],
    )

    # Anthropic returns the tool_use block with structured input
    tool_block = next(b for b in response.content if b.type == "tool_use")
    try:
        return MSAExtraction.model_validate(tool_block.input)
    except ValidationError as e:
        # In production, raise a parse-or-retry; here we surface the error
        raise RuntimeError(f"Extraction schema validation failed: {e}") from e

from __future__ import annotations
from pathlib import Path

from ci.db import db_session
from ci.models import Document
from ci.extraction.ocr import extract_text
from ci.extraction.llm import extract_msa
from ci.extraction.schemas import MSAExtraction


def extract_document(document_id: str) -> MSAExtraction:
    """End-to-end: fetch document, OCR, LLM extract."""
    with db_session() as s:
        doc = s.get(Document, document_id)
        if not doc:
            raise ValueError(f"Document {document_id} not found")
        path = Path(doc.source_uri)
        if not path.exists():
            raise FileNotFoundError(f"Source file missing: {path}")
        doc_class = doc.doc_class

    file_bytes = path.read_bytes()
    text, source = extract_text(file_bytes)

    if doc_class == "master_service_agreement":
        return extract_msa(text)
    # Add per-class extractors as you grow the build.
    raise NotImplementedError(f"No extractor for doc_class={doc_class}")

from sqlalchemy import select

from ci.db import db_session
from ci.extraction import extract_document
from ci.models import Document


@app.command()
def extract(document_id: str | None = None) -> None:
    """Run extraction on a document (or all MSAs if none specified)."""
    with db_session() as s:
        if document_id:
            doc_ids = [document_id]
        else:
            doc_ids = list(
                s.scalars(select(Document.document_id)
                          .where(Document.doc_class == "master_service_agreement"))
            )

    for doc_id in doc_ids:
        console.print(f"[blue]Extracting {doc_id}...[/blue]")
        try:
            result = extract_document(doc_id)
            console.print(result.model_dump_json(indent=2))
        except Exception as e:
            console.print(f"[red]Failed: {e}[/red]")

ci extract
# Should print extracted fields per document.