app/extraction/extractor.py

"""Citation-grounded, schema-constrained section extraction (STEP 3/4/8).

Pipeline for one section:

    section-scoped chunks
      -> deterministic LLM extraction (temperature=0, top_p=1, seed, JSON mode)
      -> Pydantic schema validation (closed enums, mandatory evidence)
      -> citation/source-span validation (drop unsupported claims)
      -> contradiction audit (resolve conflicts, collapse duplicates)
      -> SectionExtraction (with confidence + dropped-claim audit trail)

The LLM is the only non-deterministic component, and it is pinned as hard as
the API allows. Everything after it is pure, deterministic Python.

This module is additive: it does not alter the existing generation path. Callers
opt in (see ``settings.enable_citation_extraction``). It degrades safely to an
empty extraction when no API key is configured, so import and unit tests never
require network access.
"""

from __future__ import annotations

import json
import logging

from app.config import settings
from app.extraction.citation_validator import validate_findings
from app.extraction.contradiction import audit_contradictions
from app.extraction.domain_scope import GENERAL, classify_section, scope_chunks
from app.extraction.prompts import (
    EXTRACTION_SYSTEM_PROMPT,
    EXTRACTION_USER_TEMPLATE,
    build_chunks_block,
)
from app.extraction.output_validator import should_abstain
from app.extraction.schemas import (
    AtomicClaim,
    ClaimEvidence,
    ClaimType,
    ClaimVerification,
    ConditionRating,
    SectionExtraction,
    SupportLevel,
    SurveyFinding,
)

logger = logging.getLogger(__name__)

# Fixed seed for reproducible decoding across identical inputs.
_EXTRACTION_SEED = 7


def _chunk_tuples(chunks: list[object]) -> list[tuple[str, str, str | None]]:
    """Adapt SearchResult-like rows into (chunk_id, text, label) tuples."""
    out: list[tuple[str, str, str | None]] = []
    for c in chunks:
        cid = getattr(c, "chunk_id", None)
        text = getattr(c, "text", None)
        if not cid or not text:
            continue
        label = getattr(c, "section_title", None)
        out.append((str(cid), str(text), label))
    return out


def _parse_findings(raw_json: str, section: str) -> list[SurveyFinding]:
    """Parse model JSON into validated SurveyFinding models, skipping malformed rows."""
    try:
        data = json.loads(raw_json)
    except (json.JSONDecodeError, TypeError):
        logger.warning("extractor: model returned non-JSON for section=%s", section)
        return []
    rows = data.get("findings", []) if isinstance(data, dict) else []
    findings: list[SurveyFinding] = []
    for row in rows:
        if not isinstance(row, dict):
            continue
        row.setdefault("section", section)
        try:
            findings.append(SurveyFinding.model_validate(row))
        except Exception as exc:  # malformed row -> skip, never fabricate
            logger.debug("extractor: skipped malformed finding: %s", exc)
    return findings


async def extract_section(
    *,
    section: str,
    chunks: list[object],
    tenant_id: str | None = None,
    drop_partial: bool = False,
    domain: str | None = None,
    model: str | None = None,
) -> SectionExtraction:
    """Extract verified, contradiction-free findings for one section.

    When ``domain`` is provided (STEP 6), chunks clearly belonging to a
    different domain are excluded before extraction as defense-in-depth — the
    caller should still scope retrieval, but this guarantees the extractor never
    sees cross-domain evidence. The supplied (scoped) pool is the ONLY
    admissible evidence.

    Returns a :class:`SectionExtraction`. With no API key (or no chunks), returns
    an empty extraction rather than raising.
    """
    if domain and domain != GENERAL:
        chunks = scope_chunks(domain, chunks)
    pool = _chunk_tuples(chunks)
    if not pool:
        return SectionExtraction(section=section)
    if not (settings.openai_api_key or "").strip():
        logger.info("extractor: no API key; returning empty extraction for %s", section)
        return SectionExtraction(section=section)

    user_prompt = EXTRACTION_USER_TEMPLATE.format(
        section=section,
        chunks_block=build_chunks_block(pool),
    )

    from app.llm.openai_chat import chat_completions_create

    raw = await chat_completions_create(
        messages=[
            {"role": "system", "content": EXTRACTION_SYSTEM_PROMPT},
            {"role": "user", "content": user_prompt},
        ],
        model=model or settings.chat_model,
        max_tokens=settings.extraction_max_tokens,
        temperature=0.0,
        top_p=1.0,
        seed=_EXTRACTION_SEED,
        response_format={"type": "json_object"},
        phase="extraction",
        section_id=section,
        tenant_id=tenant_id,
    )

    findings = _parse_findings(raw, section)

    # Hard citation gate: drop everything not grounded in the supplied pool.
    kept, dropped = validate_findings(findings, pool, drop_partial=drop_partial)

    # Contradiction/duplicate audit on the survivors.
    resolved, contradictions = audit_contradictions(kept)

    logger.info(
        "extract_section=%s extracted=%d kept=%d dropped=%d contradictions=%d",
        section, len(findings), len(resolved), len(dropped), len(contradictions),
    )

    return SectionExtraction(
        section=section,
        findings=resolved,
        contradictions=contradictions,
        dropped_claims=dropped,
    )


_SUPPORT_CONFIDENCE: dict[SupportLevel, float] = {
    SupportLevel.SUPPORTED: 1.0,
    SupportLevel.PARTIAL: 0.5,
    SupportLevel.NOT_FOUND: 0.0,
}


def findings_to_atomic_claims(
    findings: list[SurveyFinding],
    *,
    min_confidence: float = 1.0,
) -> list[AtomicClaim]:
    """Decompose grounded findings into atomic, evidence-bound claim records.

    Each finding yields at most two atomic claims: a ``condition_rating`` claim
    (only when the source stated a real CR1/CR2/CR3) and an ``observation``
    claim for the finding text. Every claim is bound to the finding's first
    valid evidence span and then passed through the abstention gate
    (:func:`should_abstain`) — any claim that introduces forbidden/fabricated
    language or falls below ``min_confidence`` is dropped (RETURN NOTHING).

    Pure and deterministic: no network, no synthesis. Findings are assumed to
    have already passed citation validation.
    """
    claims: list[AtomicClaim] = []
    real_ratings = {ConditionRating.CR1, ConditionRating.CR2, ConditionRating.CR3}
    for f in findings:
        if not f.evidence:
            continue
        span = f.evidence[0]
        evidence_text = " \n ".join(e.text for e in f.evidence)
        confidence = _SUPPORT_CONFIDENCE.get(f.support, 0.0)
        page = span.page if span.page is not None else 0
        ev = ClaimEvidence(
            text=span.text,
            page=page,
            section=span.section_label or "",
            chunk_id=span.chunk_id,
        )

        if f.condition_rating in real_ratings:
            rating_claim = f"{f.element} condition rating is {f.condition_rating.value}"
            if not should_abstain(
                rating_claim, evidence_text,
                min_confidence=min_confidence, confidence=confidence,
            ):
                claims.append(AtomicClaim(
                    claim=rating_claim,
                    claim_type=ClaimType.CONDITION_RATING,
                    evidence=ev,
                    verification=ClaimVerification(
                        supported=True, contradiction_detected=False, confidence=confidence,
                    ),
                ))

        if not should_abstain(
            f.finding, evidence_text,
            min_confidence=min_confidence, confidence=confidence,
        ):
            claims.append(AtomicClaim(
                claim=f.finding,
                claim_type=ClaimType.OBSERVATION,
                evidence=ev,
                verification=ClaimVerification(
                    supported=True, contradiction_detected=False, confidence=confidence,
                ),
            ))
    return claims


async def audit_section_grounding(
    *,
    section_name: str,
    template_id: str | None,
    chunks: list[object],
    tenant_id: str | None = None,
) -> SectionExtraction:
    """High-level, non-destructive grounding audit for a generated section.

    Resolves the section's domain, scopes the retrieved evidence to that domain,
    and runs citation-grounded extraction + contradiction audit. Intended to be
    called alongside generation (behind ``settings.enable_citation_extraction``)
    to produce a traceability artifact — it never mutates generated prose.
    """
    domain = classify_section(section_name, template_id)
    return await extract_section(
        section=section_name or (template_id or "section"),
        chunks=chunks,
        tenant_id=tenant_id,
        domain=domain,
    )