spb/wpai0002_uqo

1"""Anthropic Claude adapter: pinned model id, logged request id / usage / stop reason / cost, exponential2backoff, refusal handling, and Message Batches support (50 % price). Prices come from configs/rewriting/pricing.json3and are never assumed. Thinking is adaptive by default on Claude 5.x models; we only set ``output_config.effort``.4"""5from __future__ import annotations67import json8import time9from pathlib import Path1011from htr.rewrite.provider import Provider, ProviderResult, Timer1213ROOT = Path(__file__).resolve().parents[3]141516def load_pricing(path: Path | None = None) -> dict:17    p = path or ROOT / "configs/rewriting/pricing.json"18    return json.loads(Path(p).read_text())192021def estimate_cost(model: str, usage: dict, pricing: dict, batch: bool = False) -> float | None:22    m = pricing["models"].get(model)23    if not m:24        return None25    disc = pricing.get("batch_discount", 0.5) if batch else 1.026    inp = usage.get("input_tokens", 0) or 027    out = usage.get("output_tokens", 0) or 028    cr = usage.get("cache_read_input_tokens", 0) or 029    cw = usage.get("cache_creation_input_tokens", 0) or 030    return disc * (inp * m["input"] + cw * m["input"] * 1.25 + cr * m.get("cache_read", m["input"] * 0.1) + out * m["output"]) / 1e6313233class AnthropicProvider(Provider):34    name = "anthropic"3536    def __init__(self, model: str, effort: str = "low", pricing: dict | None = None, max_retries: int = 6, api_key: str | None = None):37        import anthropic3839        self.model = model40        self.effort = effort41        self.pricing = pricing or load_pricing()42        if model not in self.pricing["models"]:43            raise ValueError(f"model {model!r} missing from configs/rewriting/pricing.json — add verified prices first")44        self.client = anthropic.Anthropic(api_key=api_key, max_retries=2)45        self.max_retries = max_retries46        self._anthropic = anthropic4748    # ---- live, single request --------------------------------------------------------------------------49    def rewrite(self, system: str, messages: list[dict], max_tokens: int, params: dict) -> ProviderResult:50        a = self._anthropic51        for attempt in range(self.max_retries + 1):52            try:53                with Timer() as t:54                    resp = self.client.messages.with_raw_response.create(55                        model=self.model, max_tokens=max_tokens, system=system, messages=messages, output_config={"effort": self.effort}56                    )57                msg = resp.parse()58                req_id = resp.headers.get("request-id")59                return self._to_result(msg, req_id, t.elapsed, batch=False)60            except a.RateLimitError:61                time.sleep(min(90.0, 5.0 * 2**attempt))62            except (a.APIConnectionError, a.InternalServerError, a.APIStatusError) as e:  # noqa: PERF20363                if isinstance(e, a.APIStatusError) and e.status_code < 500 and e.status_code not in (408, 409, 429):64                    return ProviderResult(None, "error", self.model, error=f"{type(e).__name__}: {str(e)[:300]}")65                time.sleep(min(90.0, 5.0 * 2**attempt))66        return ProviderResult(None, "error", self.model, error="retries exhausted")6768    def _to_result(self, msg, req_id: str | None, latency: float | None, batch: bool) -> ProviderResult:69        usage = {k: getattr(msg.usage, k, None) for k in ("input_tokens", "output_tokens", "cache_read_input_tokens", "cache_creation_input_tokens")}70        text = "".join(b.text for b in msg.content if getattr(b, "type", None) == "text").strip()71        if msg.stop_reason == "refusal":72            status = "refused"73        elif msg.stop_reason == "max_tokens":74            status = "truncated"75        elif not text:76            status = "empty"77        else:78            status = "ok"79        return ProviderResult(text or None, status, msg.model, provider_request_id=req_id or getattr(msg, "id", None), usage=usage, cost_usd=estimate_cost(self.model, usage, self.pricing, batch), stop_reason=msg.stop_reason, latency_s=latency)8081    # ---- Message Batches -------------------------------------------------------------------------------------82    def supports_batch(self) -> bool:83        return True8485    def submit_batch(self, items: list[tuple[str, str, list[dict], int]]) -> str:86        """items: (custom_id, system, messages, max_tokens). Returns batch id."""87        from anthropic.types.message_create_params import MessageCreateParamsNonStreaming88        from anthropic.types.messages.batch_create_params import Request8990        reqs = [Request(custom_id=cid, params=MessageCreateParamsNonStreaming(model=self.model, max_tokens=mt, system=sy, messages=ms, output_config={"effort": self.effort})) for cid, sy, ms, mt in items]91        b = self.client.messages.batches.create(requests=reqs)92        return b.id9394    def batch_status(self, batch_id: str) -> dict:95        b = self.client.messages.batches.retrieve(batch_id)96        return {"id": b.id, "status": b.processing_status, "counts": b.request_counts.model_dump() if hasattr(b.request_counts, "model_dump") else dict(b.request_counts)}9798    def batch_results(self, batch_id: str) -> dict[str, ProviderResult]:99        out: dict[str, ProviderResult] = {}100        for r in self.client.messages.batches.results(batch_id):101            rt = r.result.type102            if rt == "succeeded":103                out[r.custom_id] = self._to_result(r.result.message, getattr(r.result.message, "id", None), None, batch=True)104            elif rt == "errored":105                err = r.result.error106                out[r.custom_id] = ProviderResult(None, "error", self.model, error=f"batch:{getattr(err, 'type', 'error')}")107            else:108                out[r.custom_id] = ProviderResult(None, "error", self.model, error=f"batch:{rt}")109        return out110111    def cancel_batch(self, batch_id: str) -> None:112        self.client.messages.batches.cancel(batch_id)113