140 lines
6.0 KiB
Python
140 lines
6.0 KiB
Python
"""Provider abstraction with strict structured output and a hard budget guard.
|
|
|
|
The fiscal engine is never coupled to a vendor: it talks to AISemanticProvider.
|
|
V1 ships BaselineProvider (deterministic, zero-cost) used by the evaluation
|
|
harness and as the no-credential fallback; real providers subclass the same
|
|
contract and are only invoked with AI_ENABLED=true and credentials in the Hub.
|
|
"""
|
|
from dataclasses import dataclass, field
|
|
from datetime import datetime, timezone
|
|
|
|
from ..classification import ENGINE_VERSION
|
|
from .fingerprint import overlap_score
|
|
|
|
PROMPT_VERSION = "ai-semantic-v1"
|
|
|
|
|
|
class SemanticOutputError(ValueError):
|
|
"""Raised when a provider answer violates the strict structured-output schema."""
|
|
|
|
|
|
@dataclass
|
|
class StructuredResult:
|
|
task_type: str
|
|
input_fingerprint: str
|
|
selected_candidate: str # candidate code, or "NONE"
|
|
confidence: float
|
|
reason: str
|
|
candidate_ids_seen: list[str]
|
|
model: str
|
|
prompt_version: str
|
|
created_at: str
|
|
dataset_versions: dict = field(default_factory=dict)
|
|
engine_version: str = ENGINE_VERSION
|
|
|
|
def to_dict(self):
|
|
return {
|
|
"task_type": self.task_type,
|
|
"input_fingerprint": self.input_fingerprint,
|
|
"selected_candidate": self.selected_candidate,
|
|
"confidence": self.confidence,
|
|
"reason": self.reason,
|
|
"candidate_ids_seen": self.candidate_ids_seen,
|
|
"model": self.model,
|
|
"prompt_version": self.prompt_version,
|
|
"created_at": self.created_at,
|
|
"dataset_versions": self.dataset_versions,
|
|
"engine_version": self.engine_version,
|
|
}
|
|
|
|
|
|
def result_from_dict(data) -> StructuredResult:
|
|
"""Fail-closed schema check: the selection must be NONE or one of the seen candidates."""
|
|
result = StructuredResult(
|
|
task_type=str(data["task_type"]),
|
|
input_fingerprint=str(data["input_fingerprint"]),
|
|
selected_candidate=str(data["selected_candidate"]),
|
|
confidence=float(data["confidence"]),
|
|
reason=str(data["reason"]),
|
|
candidate_ids_seen=[str(item) for item in data["candidate_ids_seen"]],
|
|
model=str(data["model"]),
|
|
prompt_version=str(data["prompt_version"]),
|
|
created_at=str(data["created_at"]),
|
|
dataset_versions=dict(data.get("dataset_versions") or {}),
|
|
)
|
|
if result.selected_candidate != "NONE" and result.selected_candidate not in result.candidate_ids_seen:
|
|
raise SemanticOutputError(f"selected_candidate {result.selected_candidate!r} is not among the official candidates")
|
|
if not 0.0 <= result.confidence <= 1.0:
|
|
raise SemanticOutputError("confidence must be within [0,1]")
|
|
return result
|
|
|
|
|
|
class AISemanticProvider:
|
|
"""Contract every provider must honour. `select` receives ONLY official candidates."""
|
|
name = "abstract"
|
|
model = "abstract"
|
|
|
|
def select(self, work_item, candidates) -> StructuredResult:
|
|
raise NotImplementedError
|
|
|
|
def estimate(self, work_items) -> dict:
|
|
raise NotImplementedError
|
|
|
|
|
|
def _score(text: str, official: str) -> float:
|
|
return overlap_score(text, official)
|
|
|
|
|
|
class BaselineProvider(AISemanticProvider):
|
|
"""Deterministic zero-cost baseline: highest token-overlap candidate, or NONE.
|
|
Used for the evaluation pipeline and as the mandatory no-credential fallback."""
|
|
|
|
name = "baseline"
|
|
model = "baseline-deterministic-v1"
|
|
|
|
def select(self, work_item, candidates, dataset_versions=None) -> StructuredResult:
|
|
description = work_item.payload.get("description", "") if hasattr(work_item, "payload") else ""
|
|
scored = [(code, official, _score(description, official)) for code, official in [(str(c[0]), c[1]) for c in candidates]]
|
|
best = max(scored, key=lambda row: row[2]) if scored else None
|
|
selected, _, score = best if best else ("NONE", "", 0.0)
|
|
if not scored or score <= 0.0:
|
|
selected = "NONE"
|
|
return StructuredResult(
|
|
task_type=work_item.task_type,
|
|
input_fingerprint=work_item.dedup_key,
|
|
selected_candidate=selected,
|
|
confidence=round(score, 4),
|
|
reason="baseline token-overlap top-1" if selected != "NONE" else "no token overlap with any official candidate",
|
|
candidate_ids_seen=[str(c[0]) for c in candidates],
|
|
model=self.model,
|
|
prompt_version=PROMPT_VERSION,
|
|
created_at=datetime.now(timezone.utc).isoformat(),
|
|
dataset_versions=dict(dataset_versions or {}),
|
|
)
|
|
|
|
def estimate(self, work_items) -> dict:
|
|
return {"input_tokens": 0, "output_tokens": 0, "estimated_cost": 0.0}
|
|
|
|
|
|
class BudgetGuard:
|
|
"""Refuses execution when the estimated cost exceeds AI_MAX_RUN_COST_USD.
|
|
Never falls back to a more expensive model. Prices are operator-configured — when
|
|
absent, cost is UNKNOWN: the residual run is refused, but an explicitly limited
|
|
benchmark (e.g. the 65-item GOLD eval) may opt in with allow_unknown_cost."""
|
|
|
|
def __init__(self, max_run_cost_usd: float, price_in_per_m: float | None, price_out_per_m: float | None):
|
|
self.max_run_cost_usd = float(max_run_cost_usd)
|
|
self.price_in_per_m = None if price_in_per_m is None else float(price_in_per_m)
|
|
self.price_out_per_m = None if price_out_per_m is None else float(price_out_per_m)
|
|
|
|
def estimate(self, items_in_tokens: int, items_out_tokens: int) -> dict:
|
|
if self.price_in_per_m is None or self.price_out_per_m is None:
|
|
return {"input_tokens": items_in_tokens, "output_tokens": items_out_tokens, "estimated_cost": None, "cost_unknown": True}
|
|
cost = (items_in_tokens / 1_000_000) * self.price_in_per_m + (items_out_tokens / 1_000_000) * self.price_out_per_m
|
|
return {"input_tokens": items_in_tokens, "output_tokens": items_out_tokens, "estimated_cost": round(cost, 6)}
|
|
|
|
def authorize(self, plan: dict, allow_unknown_cost: bool = False) -> bool:
|
|
if plan.get("estimated_cost") is None:
|
|
return allow_unknown_cost
|
|
return plan["estimated_cost"] <= self.max_run_cost_usd
|