Spaces:
Sleeping
Sleeping
File size: 4,184 Bytes
465b477 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | """Provider-agnostic LLM request/response types and the model spec read from config."""
from __future__ import annotations
import time
from dataclasses import dataclass, field
from typing import Any, Literal, Protocol
from pydantic import BaseModel, Field
class ModelSpec(BaseModel):
"""One named model endpoint. Prices are USD per 1M tokens and only used for reporting."""
provider: Literal["anthropic", "openai_compat"]
model: str
base_url: str | None = None
api_key_env: str | None = None
max_tokens: int = 4096
# None means "do not send": several current Claude models reject sampling params.
temperature: float | None = None
effort: Literal["low", "medium", "high", "xhigh", "max"] | None = None
json_mode: Literal["schema", "object", "none"] = "schema"
price_in_per_mtok: float = 0.0
price_out_per_mtok: float = 0.0
max_concurrency: int = 4
timeout_s: float = 300.0
extra: dict[str, Any] = Field(default_factory=dict)
@dataclass
class LLMRequest:
system: str
user: str
json_schema: dict[str, Any] | None = None
schema_name: str = "output"
max_tokens: int | None = None
def cache_payload(self, spec: ModelSpec) -> dict[str, Any]:
return {
"provider": spec.provider,
"model": spec.model,
"base_url": spec.base_url,
"temperature": spec.temperature,
"effort": spec.effort,
"json_mode": spec.json_mode,
"max_tokens": self.max_tokens or spec.max_tokens,
"system": self.system,
"user": self.user,
"json_schema": self.json_schema,
}
@dataclass
class LLMResponse:
text: str
input_tokens: int
output_tokens: int
latency_s: float
model: str
stop_reason: str | None = None
cached: bool = False
cache_key: str | None = None
class LLMClient(Protocol):
spec: ModelSpec
def complete(self, request: LLMRequest) -> LLMResponse: ...
class LLMError(RuntimeError):
"""Raised when a provider call fails after retries or returns unusable output."""
@dataclass
class CallRecord:
stage: str
model: str
input_tokens: int
output_tokens: int
latency_s: float
cached: bool
cache_key: str | None
cost_usd: float
@dataclass
class UsageMeter:
"""Collects every model call made while answering one question."""
calls: list[CallRecord] = field(default_factory=list)
started: float = field(default_factory=time.perf_counter)
def record(self, stage: str, spec: ModelSpec, resp: LLMResponse) -> None:
cost = (resp.input_tokens * spec.price_in_per_mtok + resp.output_tokens * spec.price_out_per_mtok) / 1_000_000
self.calls.append(
CallRecord(
stage=stage,
model=spec.model,
input_tokens=resp.input_tokens,
output_tokens=resp.output_tokens,
latency_s=resp.latency_s,
cached=resp.cached,
cache_key=resp.cache_key,
cost_usd=cost,
)
)
def summary(self) -> dict[str, Any]:
by_stage: dict[str, dict[str, float]] = {}
for c in self.calls:
s = by_stage.setdefault(c.stage, {"calls": 0, "input_tokens": 0, "output_tokens": 0, "cost_usd": 0.0})
s["calls"] += 1
s["input_tokens"] += c.input_tokens
s["output_tokens"] += c.output_tokens
s["cost_usd"] += c.cost_usd
return {
"calls": len(self.calls),
"input_tokens": sum(c.input_tokens for c in self.calls),
"output_tokens": sum(c.output_tokens for c in self.calls),
# Latency is the provider latency summed over calls (sequential-equivalent);
# cached calls report the latency of the original call.
"llm_latency_s": round(sum(c.latency_s for c in self.calls), 3),
"cost_usd": round(sum(c.cost_usd for c in self.calls), 6),
"cache_hits": sum(1 for c in self.calls if c.cached),
"by_stage": by_stage,
"cache_keys": [c.cache_key for c in self.calls],
}
|