"""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}