GTI_GERENCIADOR_FISCAL/fiscal/central/app/ai/llm_provider.py

"""Real LLM provider (OpenAI-compatible) for the NCM semantic evaluation.

Contract (mission REAL LLM EVALUATION V1):
- The API key lives ONLY in the Hub environment (AI_API_KEY). Never committed, logged,
  written to artifacts, sent to the Agent or Firebird.
- The gold answer is never part of the prompt: items carry only description/brand/unit/
  segment-weak-context + the official top-K hybrid candidates with hierarchical context.
- Strict structured output: per item {work_item_id, selected_ncm, confidence, abstain,
  reason_short}; selected_ncm must be in the candidate list or null (abstain).
- Batching is safe only when the response is valid JSON with 1:1 work_item_id matching,
  no omitted items, no cross-product mixing. Verified progressively (5/10/20).
"""
import hashlib
import json
import re
from dataclasses import dataclass

PROMPT_VERSION_LLM = "llm-ncm-eval-v1"

_SYSTEM_PROMPT = (
    "You are a Brazilian customs classification assistant. For each product below, choose the "
    "most likely official NCM code from the provided candidate list, or abstain. "
    "Rules: selected_ncm MUST be one of the candidate codes for that exact item, or null with "
    "abstain=true when evidence is insufficient. Never invent a code. confidence is your calibrated "
    "probability in [0,1]. Reply with ONLY a JSON array, one object per item, in the same order, "
    'each: {"work_item_id": str, "selected_ncm": str|null, "confidence": float, "abstain": bool, "reason_short": str}.'
)


@dataclass
class LlmEvalItem:
    work_item_id: str
    description: str
    normalized_description: str
    brand: str | None
    unit: str | None
    segment: str
    candidates: list  # [(code, official_description, hybrid_score, official_score, kb_score, tenant_count)]


class LlmAnswerError(ValueError):
    pass


def item_prompt_payload(item: LlmEvalItem) -> dict:
    return {
        "work_item_id": item.work_item_id,
        "description": item.description,
        "normalized_description": item.normalized_description,
        "brand": item.brand,
        "unit": item.unit,
        "segment_weak_context": item.segment,
        "official_candidates": [
            {"ncm": code, "official_description": text}
            for code, text, *_ in item.candidates
        ],
    }


def candidate_set_hash(item: LlmEvalItem) -> str:
    codes = [code for code, *_ in item.candidates]
    return hashlib.sha256("|".join(codes).encode()).hexdigest()


def input_fingerprint(item: LlmEvalItem) -> str:
    payload = {**item_prompt_payload(item), "system": _SYSTEM_PROMPT}
    return hashlib.sha256(json.dumps(payload, sort_keys=True, ensure_ascii=False).encode()).hexdigest()


def validate_answer(answer: dict, item: LlmEvalItem) -> dict:
    """Fail-closed validator: schema, work_item_id, candidate-list confinement."""
    if not isinstance(answer, dict):
        raise LlmAnswerError("answer must be an object")
    if str(answer.get("work_item_id")) != str(item.work_item_id):
        raise LlmAnswerError(f"work_item_id mismatch: {answer.get('work_item_id')!r} != {item.work_item_id!r}")
    abstain = bool(answer.get("abstain"))
    selected = answer.get("selected_ncm")
    selected = None if selected in (None, "") else str(selected)
    try:
        confidence = float(answer.get("confidence"))
    except (TypeError, ValueError):
        raise LlmAnswerError("confidence must be numeric")
    if not 0.0 <= confidence <= 1.0:
        raise LlmAnswerError("confidence out of [0,1]")
    codes = {code for code, *_ in item.candidates}
    if abstain:
        if selected is not None:
            raise LlmAnswerError("abstain must carry selected_ncm=null")
    else:
        if selected is None:
            raise LlmAnswerError("non-abstain requires selected_ncm")
        if selected not in codes:
            raise LlmAnswerError(f"selected_ncm {selected!r} is outside the candidate list")
    return {"work_item_id": item.work_item_id, "selected_ncm": selected, "confidence": confidence, "abstain": abstain, "reason_short": str(answer.get("reason_short") or "")[:200]}


class OpenAICompatibleProvider:
    """Talks to any OpenAI-compatible /chat/completions endpoint. `transport` is
    injectable for tests; production default uses httpx with the key from settings only."""

    name = "openai-compatible"
    prompt_version = PROMPT_VERSION_LLM

    def __init__(self, api_key: str, model: str, base_url: str = "https://api.openai.com/v1", transport=None, timeout: float = 120.0):
        self.api_key = api_key
        self.model = model
        self.base_url = base_url.rstrip("/")
        self.timeout = timeout
        self._transport = transport or self._httpx_transport
        self.http_calls = 0
        self.retries = 0
        self.usage = {"input_tokens": 0, "output_tokens": 0, "cached_tokens": 0}

    def _httpx_transport(self, payload: dict) -> dict:
        import httpx
        response = httpx.post(f"{self.base_url}/chat/completions", json=payload, headers={"Authorization": f"Bearer {self.api_key}"}, timeout=self.timeout)
        response.raise_for_status()
        return response.json()

    def list_models(self) -> list:
        """GET /models with the account key — validates auth/base URL/model availability
        without spending completion tokens. Never raises with the key attached."""
        import httpx
        response = httpx.get(f"{self.base_url}/models", headers={"Authorization": f"Bearer {self.api_key}"}, timeout=self.timeout)
        response.raise_for_status()
        data = response.json()
        return sorted(str(item.get("id")) for item in data.get("data", []) if item.get("id"))

    def _request(self, messages: list, max_retries: int = 2) -> dict:
        payload = {"model": self.model, "messages": messages, "temperature": 0}
        last_error = None
        for attempt in range(max_retries + 1):
            self.http_calls += 1
            if attempt:
                self.retries += 1
            try:
                data = self._transport(payload)
                usage = data.get("usage") or {}
                self.usage["input_tokens"] += int(usage.get("prompt_tokens") or 0)
                self.usage["output_tokens"] += int(usage.get("completion_tokens") or 0)
                cached = usage.get("prompt_tokens_details") or {}
                self.usage["cached_tokens"] += int(cached.get("cached_tokens") or 0)
                return data
            except Exception as error:  # transport-level failure only; JSON contract errors are not retried
                last_error = error
        raise RuntimeError(f"LLM transport failed after {max_retries + 1} attempts: {last_error}")

    def select_batch(self, items: list[LlmEvalItem]) -> list[dict]:
        """One HTTP call for the whole batch. Raises LlmAnswerError on any 1:1/schema violation
        so the harness can shrink the batch size."""
        user_content = json.dumps([item_prompt_payload(item) for item in items], ensure_ascii=False)
        data = self._request([{"role": "system", "content": _SYSTEM_PROMPT}, {"role": "user", "content": user_content}])
        content = data["choices"][0]["message"]["content"]
        content = re.sub(r"^```(?:json)?|```$", "", content.strip(), flags=re.MULTILINE).strip()
        answers = json.loads(content)
        if not isinstance(answers, list) or len(answers) != len(items):
            raise LlmAnswerError(f"expected {len(items)} answers, got {len(answers) if isinstance(answers, list) else type(answers).__name__}")
        return [validate_answer(answer, item) for answer, item in zip(answers, items)]