DealDocumentScreening/src/contract_check/worker_prescreen/extractor_hybrid.py

"""Hybrid two-stage orchestrator for prescreen metadata extraction.

Stage 1 — HeuristicExtractor (sync, via to_thread): deterministic, zero cost.
Stage 2 — LLMPrescreenExtractor: runs only when the kill-switch is off AND the
heuristic confidence is below `prescreen_llm_fallback_threshold`.

Failure semantics (plan §3.4): an LLM fallback error never fails the message —
the heuristic meta is kept, `extractor_version="heuristic-v2"`, and the error
is recorded on the result for `prescreen_results.auto_findings` plus the
`prescreen_fallback_runs_total{outcome="failed"}` counter.
"""

from __future__ import annotations

import asyncio
from typing import Any

from src.contract_check.core.config import get_settings
from src.contract_check.core.llm.factory import build_llm_provider
from src.contract_check.core.llm.port import LLMProvider
from src.contract_check.core.logging import get_logger
from src.contract_check.core.metrics import prescreen_extraction_stage, prescreen_fallback_runs
from src.contract_check.worker_prescreen.extractor import (
    ExtractionResult,
    PrescreenContractMeta,
    _score_confidence,
)
from src.contract_check.worker_prescreen.extractor_heuristic import (
    EXTRACTOR_VERSION as HEURISTIC_VERSION,
)
from src.contract_check.worker_prescreen.extractor_heuristic import HeuristicExtractor
from src.contract_check.worker_prescreen.extractor_llm import (
    EXTRACTOR_VERSION as HYBRID_LLM_VERSION,
)
from src.contract_check.worker_prescreen.extractor_llm import LLMPrescreenExtractor

log = get_logger(__name__)

_SCALAR_FIELDS = (
    "contract_type",
    "party_a",
    "party_b",
    "total_amount",
    "currency",
    "start_date",
    "end_date",
)
_BOOL_FIELDS = (
    "has_penalty_clause",
    "has_termination_clause",
    "has_arbitration",
)


def _merge(heuristic: PrescreenContractMeta, llm: PrescreenContractMeta) -> PrescreenContractMeta:
    """LLM values override heuristic Nones; boolean clauses OR-merge.

    Booleans from both stages are presence-checks — false positives are cheap,
    false negatives route wrong (plan §3.3).
    """
    update: dict[str, Any] = {}
    for field in _SCALAR_FIELDS:
        llm_value = getattr(llm, field)
        if llm_value is not None:
            update[field] = llm_value
    for field in _BOOL_FIELDS:
        h_value, l_value = getattr(heuristic, field), getattr(llm, field)
        if h_value is None:
            if l_value is not None:
                update[field] = l_value
        elif l_value is not None:
            update[field] = bool(h_value or l_value)

    merged = heuristic.model_copy(update=update)
    merged.confidence_score = _score_confidence(merged)
    return merged


class HybridMetaExtractor:
    """Default MetadataExtractor: heuristic first, LLM only when confidence is low."""

    def __init__(
        self,
        *,
        provider: LLMProvider | None = None,
        fallback_enabled: bool | None = None,
        fallback_threshold: float | None = None,
        heuristic: HeuristicExtractor | None = None,
    ) -> None:
        settings = get_settings()
        self._heuristic = heuristic or HeuristicExtractor()
        self._fallback_enabled = (
            settings.prescreen_llm_fallback_enabled
            if fallback_enabled is None
            else fallback_enabled
        )
        self._fallback_threshold = (
            settings.prescreen_llm_fallback_threshold
            if fallback_threshold is None
            else fallback_threshold
        )
        self._provider = provider
        self._owns_provider = provider is None
        self._llm: LLMPrescreenExtractor | None = None

    async def _llm_extractor(self) -> LLMPrescreenExtractor:
        if self._llm is None:
            if self._provider is None:
                self._provider = build_llm_provider(get_settings())
            self._llm = LLMPrescreenExtractor(self._provider)
        return self._llm

    async def aclose(self) -> None:
        if self._provider is not None and self._owns_provider:
            await self._provider.aclose()

    async def extract(self, text: str) -> ExtractionResult:
        with prescreen_extraction_stage.labels(stage="heuristic").time():
            meta = await asyncio.to_thread(self._heuristic.extract, text)

        if not self._fallback_enabled:
            prescreen_fallback_runs.labels(outcome="disabled").inc()
            return ExtractionResult(meta=meta, extractor_version=HEURISTIC_VERSION)

        if meta.confidence_score >= self._fallback_threshold:
            prescreen_fallback_runs.labels(outcome="skipped").inc()
            return ExtractionResult(meta=meta, extractor_version=HEURISTIC_VERSION)

        try:
            llm_extractor = await self._llm_extractor()
            with prescreen_extraction_stage.labels(stage="llm").time():
                llm_meta = await llm_extractor.extract(text)
        except Exception as exc:  # noqa: BLE001 — never fail the message (plan §3.4)
            prescreen_fallback_runs.labels(outcome="failed").inc()
            error = f"{type(exc).__name__}: {exc}"
            log.warning(
                "prescreen_llm_fallback_failed",
                error=error,
                heuristic_confidence=meta.confidence_score,
            )
            return ExtractionResult(
                meta=meta,
                extractor_version=HEURISTIC_VERSION,
                llm_fallback_error=error,
            )

        prescreen_fallback_runs.labels(outcome="used").inc()
        merged = _merge(meta, llm_meta)
        log.info(
            "prescreen_llm_fallback_merged",
            heuristic_confidence=meta.confidence_score,
            merged_confidence=merged.confidence_score,
        )
        return ExtractionResult(meta=merged, extractor_version=HYBRID_LLM_VERSION)


__all__ = ["HybridMetaExtractor"]