src/knowledge_extraction/extract/client.py

"""LLM clients for the extraction stage — **the only place this pipeline spends
money.**

Two things this module is careful about:

- **Structured output is probed, not assumed.** `json_schema` needs a recent
  api_version and we cannot confirm from here what the resource exposes. The
  first call tries it; on rejection it falls back to `json_object` plus
  validate-and-retry, and records which mode actually applied.
- **Cached tokens are read from the API, never modelled.** Caching does not
  engage below the token floor, so an under-length prefix caches nothing.
  `usage.prompt_tokens_details.cached_tokens` is the only source of truth, and a
  cached price must never be reported without it.

All four branches route to the **nano** deployment (`__54n`). That is a recorded
decision, not an oversight: nano measured 0.75 schema-fill precision against a
0.80 line, and `rule`/`summary` — whose failure mode is least detectable, since
a plausible summary cannot be span-checked — run there too until a larger
deployment exists.
"""

from __future__ import annotations

import json
import time
from typing import Any

from ...config.settings import settings as app_settings
from ...middlewares.logging import get_logger
from ..models import Branch, CallUsage
from ..settings import TEMPERATURE

logger = get_logger("knowledge_extract_client")

MAX_RETRIES = 3


class LLMResult:
    def __init__(self, data: dict, usage: CallUsage, raw: str = ""):
        self.data = data
        self.usage = usage
        self.raw = raw


class AzureExtractor:
    """Real calls, real spend. Always dry-run before a corpus-scale run."""

    def __init__(self, client=None, deployment: str | None = None):
        self.deployment = deployment or app_settings.azureai_deployment_name_54n
        self._client = client or self._build_client()
        self._mode: str | None = None  # resolved on the first successful call

    @staticmethod
    def _build_client():
        from openai import AzureOpenAI

        endpoint = app_settings.azureai_endpoint_url_54n
        api_key = app_settings.azureai_api_key_54n
        if not endpoint or not api_key:
            raise RuntimeError(
                "azureai__endpoint__url__54n / azureai__api_key__54n are not set. "
                "Use the mock extractor to run without Azure."
            )
        return AzureOpenAI(
            azure_endpoint=endpoint,
            api_key=api_key,
            api_version=app_settings.azureai_api_version_54n,
        )

    def complete(
        self,
        branch: Branch,
        system_prompt: str,
        user_prompt: str,
        schema: dict,
        schema_name: str,
    ) -> LLMResult:
        messages = [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
        last_error: Exception | None = None

        for attempt in range(MAX_RETRIES):
            mode = self._mode or "json_schema"
            started = time.time()
            try:
                response = self._client.chat.completions.create(
                    model=self.deployment,
                    messages=messages,
                    temperature=TEMPERATURE,
                    response_format=self._response_format(mode, schema, schema_name),
                )
            except Exception as exc:
                if mode == "json_schema" and self._looks_unsupported(exc):
                    logger.info(
                        "json_schema unsupported — falling back to json_object",
                        error=repr(exc),
                    )
                    self._mode = "json_object"
                    continue
                last_error = exc
                logger.warning("call failed", branch=branch, attempt=attempt, error=repr(exc))
                time.sleep(2**attempt)
                continue

            self._mode = mode
            content = response.choices[0].message.content or "{}"
            try:
                data = json.loads(content)
            except json.JSONDecodeError as exc:
                last_error = exc
                logger.warning("unparseable JSON", branch=branch, attempt=attempt)
                continue

            usage = self._usage(response, branch, time.time() - started, attempt, mode)
            return LLMResult(data, usage, content)

        raise RuntimeError(f"{branch}: all {MAX_RETRIES} attempts failed: {last_error!r}")

    @staticmethod
    def _response_format(mode: str, schema: dict, schema_name: str) -> dict:
        if mode == "json_schema":
            return {
                "type": "json_schema",
                "json_schema": {"name": schema_name, "schema": schema, "strict": False},
            }
        return {"type": "json_object"}

    @staticmethod
    def _looks_unsupported(exc: Exception) -> bool:
        text = str(exc).lower()
        return any(
            s in text
            for s in ("response_format", "json_schema", "unsupported", "invalid_request")
        )

    def _usage(
        self, response: Any, branch: Branch, latency: float, retries: int, mode: str
    ) -> CallUsage:
        usage = getattr(response, "usage", None)
        details = getattr(usage, "prompt_tokens_details", None)
        # The ONLY source of truth for caching. Absent -> cached stays 0 and the
        # uncached regime is what gets reported.
        cached = int(getattr(details, "cached_tokens", 0) or 0) if details else 0
        return CallUsage(
            branch=branch,
            deployment=self.deployment,
            tier="nano",
            prompt_tokens=int(getattr(usage, "prompt_tokens", 0) or 0),
            cached_tokens=cached,
            completion_tokens=int(getattr(usage, "completion_tokens", 0) or 0),
            latency_s=round(latency, 3),
            retries=retries,
            structured_output_mode=mode,
            simulated=False,
        )


class MockExtractor:
    """No network, no spend. Every record it produces is stamped `simulated`.

    Exercises the wiring — schema validation, span checking, escalation,
    conflicts, diff, queue — without credentials. It is **not** a model-quality
    measurement and its output must never be reported as one.

    It abstains by default (returns null definitions), because abstention is the
    dominant real behaviour: on the reference document 56 of 66 entries had no
    definition. A mock that always answers would make the downstream stages look
    far better exercised than they are.
    """

    def __init__(self, responses: dict[str, dict] | None = None, deployment: str = "mock"):
        self.responses = responses or {}
        self.deployment = deployment
        self.calls: list[tuple[str, str]] = []

    def complete(
        self,
        branch: Branch,
        system_prompt: str,
        user_prompt: str,
        schema: dict,
        schema_name: str,
    ) -> LLMResult:
        self.calls.append((branch, user_prompt))
        data = self.responses.get(branch) or self._abstain(branch, user_prompt)
        usage = CallUsage(
            branch=branch,
            deployment=self.deployment,
            prompt_tokens=len(system_prompt) // 4 + len(user_prompt) // 4,
            completion_tokens=40,
            structured_output_mode="mock",
            simulated=True,
        )
        return LLMResult(data, usage, json.dumps(data))

    @staticmethod
    def _abstain(branch: Branch, user_prompt: str) -> dict:
        # Quote a real fragment so the span check has something locatable and is
        # genuinely exercised rather than trivially passed.
        span = ""
        if "EVIDENCE" in user_prompt:
            body = user_prompt.split("EVIDENCE", 1)[1]
            for line in body.splitlines():
                if line.strip() and not line.startswith("["):
                    span = line.strip()[:60]
                    break
        prov = {"section_no": None, "page": 1, "span": span}
        if branch == "glossary":
            term = "unknown"
            for line in user_prompt.splitlines():
                if line.startswith("CANDIDATE TERM:"):
                    term = line.split(":", 1)[1].strip()
                    break
            return {"term": term, "definition": None, "provenance": prov}
        if branch == "rule":
            return {"rule_id": "r_mock", "statement": None, "provenance": prov}
        if branch == "formula":
            return {"name": None, "formula_latex": None, "provenance": prov}
        return {"title": None, "summary_md": None, "provenance": prov}