feat(phase2): provider-agnostic LLM client + composite cache (P1b)
This commit is contained in:
@@ -0,0 +1,536 @@
|
||||
"""Provider-agnostic LLM client for the phase-2 decision layer (P1.1/P1.2).
|
||||
|
||||
One abstraction for every LLM call the simulator makes, so sweeps, the
|
||||
compiler and the per-turn showcase all share the same safeguards:
|
||||
|
||||
- **Composite cache** (decision D6): a ``ResponseCache`` keyed by
|
||||
``(model, params, prompt_version, schema_version, prompt hash)``. The
|
||||
``prompt_version`` (external pack version + per-prompt content hash, see
|
||||
``prompt.py``) makes a campaign reproducible after a prompt edit (section
|
||||
14). The cache is sharded (decision D10) and thread-safe, so parallel
|
||||
sweep workers can share one instance.
|
||||
- **Cost/token guards** (P1.2, section 13): a ``CostBudget`` enforces
|
||||
optional ceilings on provider calls, total tokens and estimated USD cost,
|
||||
raising a typed ``BudgetExceededError`` instead of silently continuing.
|
||||
The ledger counts every attempt plus tokens/cost of successful completions.
|
||||
- **Retry/backoff**: a transient ``ProviderError`` is retried with
|
||||
exponential backoff (injectable ``sleep`` for tests), then a typed
|
||||
``RetryExhaustedError`` is raised.
|
||||
- **Circuit breaker** (section 13): after ``failure_threshold`` consecutive
|
||||
failures the client refuses further calls (``CircuitOpenError``) for
|
||||
``open_refusals`` calls, then allows one probe attempt; a success closes
|
||||
the breaker, a failure reopens it.
|
||||
|
||||
``LLMClient`` is abstract and provider-agnostic (section 13): subclasses
|
||||
implement ``_call_once`` (one attempt, returning text + token usage).
|
||||
``FakeLLMClient`` serves a scripted sequence of responses offline (P1.1),
|
||||
with optional simulated transient failures - this is what drives P1.3's
|
||||
fully offline battle tests.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import hashlib
|
||||
import time
|
||||
import zlib
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from threading import RLock
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from pf1e_simulator.combat import OBSERVATION_SCHEMA_VERSION
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import Callable, Sequence
|
||||
|
||||
|
||||
class ClientError(RuntimeError):
|
||||
"""Base error for the LLM client layer."""
|
||||
|
||||
|
||||
class ProviderError(ClientError):
|
||||
"""A transient provider failure - safe to retry."""
|
||||
|
||||
|
||||
class RetryExhaustedError(ClientError):
|
||||
"""Every retry attempt failed; the last provider error is the cause."""
|
||||
|
||||
|
||||
class CircuitOpenError(ClientError):
|
||||
"""The circuit breaker is open: calls are refused, not attempted."""
|
||||
|
||||
|
||||
class BudgetExceededError(ClientError):
|
||||
"""A cost/token/call ceiling was hit (P1.2, section 13)."""
|
||||
|
||||
|
||||
class ScriptExhaustedError(ClientError):
|
||||
"""A scripted fake client has no more responses to serve."""
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClientParams:
|
||||
"""Sampling parameters; part of the composite cache key (D6)."""
|
||||
|
||||
temperature: float = 0.0
|
||||
max_tokens: int = 1024
|
||||
top_p: float = 1.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Pricing:
|
||||
"""USD per 1000 tokens, used to estimate the cost ledger."""
|
||||
|
||||
prompt_per_1k_usd: float = 0.0
|
||||
completion_per_1k_usd: float = 0.0
|
||||
|
||||
def cost_usd(self, usage: TokenUsage) -> float:
|
||||
return (
|
||||
usage.prompt_tokens / 1000 * self.prompt_per_1k_usd
|
||||
+ usage.completion_tokens / 1000 * self.completion_per_1k_usd
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CostBudget:
|
||||
"""Optional hard ceilings enforced by ``LLMClient`` (section 13).
|
||||
|
||||
Each field defaults to ``None`` = unlimited for that resource.
|
||||
"""
|
||||
|
||||
max_calls: int | None = None
|
||||
max_total_tokens: int | None = None
|
||||
max_cost_usd: float | None = None
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TokenUsage:
|
||||
"""Token counts of one completion (measured by the provider, or estimated)."""
|
||||
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
|
||||
@property
|
||||
def total(self) -> int:
|
||||
return self.prompt_tokens + self.completion_tokens
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Completion:
|
||||
"""One successful provider response."""
|
||||
|
||||
text: str
|
||||
usage: TokenUsage
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClientStats:
|
||||
"""Read-only snapshot of the client ledger (for reports and sweeps)."""
|
||||
|
||||
calls: int
|
||||
prompt_tokens: int
|
||||
completion_tokens: int
|
||||
cost_usd: float
|
||||
cache_hits: int
|
||||
cache_misses: int
|
||||
|
||||
|
||||
def estimate_tokens(text: str) -> int:
|
||||
"""Rough token estimate (chars / 4) used when the provider reports none."""
|
||||
return max(1, len(text) // 4)
|
||||
|
||||
|
||||
def response_cache_key(
|
||||
*,
|
||||
model: str,
|
||||
params: ClientParams,
|
||||
prompt_version: str,
|
||||
schema_version: int,
|
||||
prompt: str,
|
||||
) -> str:
|
||||
"""Composite cache key (D6): model, params, prompt version, schema, hash.
|
||||
|
||||
The prompt hash covers the full rendered prompt text, so two different
|
||||
observations never collide even when they share every other component.
|
||||
"""
|
||||
params_component = (
|
||||
f"t={params.temperature:g};mt={params.max_tokens};tp={params.top_p:g}"
|
||||
)
|
||||
digest = hashlib.sha256(prompt.encode("utf-8")).hexdigest()
|
||||
return (
|
||||
f"{model}|{params_component}|{prompt_version}"
|
||||
f"|schema={schema_version}|{digest}"
|
||||
)
|
||||
|
||||
|
||||
class ResponseCache:
|
||||
"""Thread-safe, sharded cache of completion texts (D6 + D10).
|
||||
|
||||
The key is the composite string from :func:`response_cache_key`; the
|
||||
value is the completion text. Shards are selected by ``zlib.crc32`` so
|
||||
parallel sweep workers (D10) never contend on one dict.
|
||||
"""
|
||||
|
||||
def __init__(self, shard_count: int = 8) -> None:
|
||||
if shard_count < 1:
|
||||
msg = f"shard_count must be >= 1, got {shard_count}"
|
||||
raise ValueError(msg)
|
||||
self._shards: tuple[dict[str, str], ...] = tuple({} for _ in range(shard_count))
|
||||
self._locks: tuple[RLock, ...] = tuple(RLock() for _ in range(shard_count))
|
||||
self._stat_lock = RLock()
|
||||
self._hits = 0
|
||||
self._misses = 0
|
||||
|
||||
def get(self, key: str) -> str | None:
|
||||
"""Return the cached completion for ``key``, or ``None``."""
|
||||
shard, lock = self._shard_for(key)
|
||||
with lock:
|
||||
value = shard.get(key)
|
||||
with self._stat_lock:
|
||||
if value is None:
|
||||
self._misses += 1
|
||||
else:
|
||||
self._hits += 1
|
||||
return value
|
||||
|
||||
def put(self, key: str, text: str) -> None:
|
||||
"""Store ``text`` under ``key`` (idempotent)."""
|
||||
shard, lock = self._shard_for(key)
|
||||
with lock:
|
||||
shard[key] = text
|
||||
|
||||
def clear(self) -> None:
|
||||
"""Drop every entry and reset the hit/miss counters (tests only)."""
|
||||
for shard, lock in zip(self._shards, self._locks, strict=True):
|
||||
with lock:
|
||||
shard.clear()
|
||||
with self._stat_lock:
|
||||
self._hits = 0
|
||||
self._misses = 0
|
||||
|
||||
@property
|
||||
def hits(self) -> int:
|
||||
with self._stat_lock:
|
||||
return self._hits
|
||||
|
||||
@property
|
||||
def misses(self) -> int:
|
||||
with self._stat_lock:
|
||||
return self._misses
|
||||
|
||||
def _shard_for(self, key: str) -> tuple[dict[str, str], RLock]:
|
||||
index = zlib.crc32(key.encode("utf-8")) % len(self._shards)
|
||||
return self._shards[index], self._locks[index]
|
||||
|
||||
|
||||
class CircuitBreaker:
|
||||
"""Trips after N consecutive failures; probes after K refusals.
|
||||
|
||||
Not thread-safe by itself: ``LLMClient`` serializes all access with its
|
||||
own lock.
|
||||
"""
|
||||
|
||||
def __init__(self, *, failure_threshold: int = 3, open_refusals: int = 2) -> None:
|
||||
if failure_threshold < 1:
|
||||
msg = f"failure_threshold must be >= 1, got {failure_threshold}"
|
||||
raise ValueError(msg)
|
||||
if open_refusals < 0:
|
||||
msg = f"open_refusals must be >= 0, got {open_refusals}"
|
||||
raise ValueError(msg)
|
||||
self.failure_threshold = failure_threshold
|
||||
self.open_refusals = open_refusals
|
||||
self._consecutive_failures = 0
|
||||
self._refusals = 0
|
||||
|
||||
@property
|
||||
def is_open(self) -> bool:
|
||||
return self._consecutive_failures >= self.failure_threshold
|
||||
|
||||
@property
|
||||
def probe_due(self) -> bool:
|
||||
"""True when an open breaker allows exactly one probe attempt."""
|
||||
return self.is_open and self._refusals >= self.open_refusals
|
||||
|
||||
@property
|
||||
def consecutive_failures(self) -> int:
|
||||
return self._consecutive_failures
|
||||
|
||||
def record_failure(self) -> None:
|
||||
self._consecutive_failures += 1
|
||||
self._refusals = 0
|
||||
|
||||
def record_refusal(self) -> None:
|
||||
self._refusals += 1
|
||||
|
||||
def close(self) -> None:
|
||||
self._consecutive_failures = 0
|
||||
self._refusals = 0
|
||||
|
||||
|
||||
class LLMClient(ABC):
|
||||
"""Provider-agnostic client: caching, budgets, retry/backoff, breaker.
|
||||
|
||||
Subclasses implement :meth:`_call_once` (one provider attempt). The
|
||||
public :meth:`complete` adds the shared safeguards around it.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt_version: str,
|
||||
schema_version: int = OBSERVATION_SCHEMA_VERSION,
|
||||
cache: ResponseCache | None = None,
|
||||
budget: CostBudget | None = None,
|
||||
pricing: Pricing | None = None,
|
||||
retries: int = 2,
|
||||
backoff_base_s: float = 0.25,
|
||||
sleep: Callable[[float], None] = time.sleep,
|
||||
breaker: CircuitBreaker | None = None,
|
||||
) -> None:
|
||||
if retries < 0:
|
||||
msg = f"retries must be >= 0, got {retries}"
|
||||
raise ValueError(msg)
|
||||
if backoff_base_s < 0:
|
||||
msg = f"backoff_base_s must be >= 0, got {backoff_base_s}"
|
||||
raise ValueError(msg)
|
||||
self.model = model
|
||||
self.prompt_version = prompt_version
|
||||
self.schema_version = schema_version
|
||||
# Cache is mandatory (section 13); a client without an explicit one
|
||||
# gets its own instance, so caching is never silently off.
|
||||
self._cache = cache if cache is not None else ResponseCache()
|
||||
self._budget = budget
|
||||
self._pricing = pricing
|
||||
self._retries = retries
|
||||
self._backoff_base_s = backoff_base_s
|
||||
self._sleep = sleep
|
||||
self._breaker = breaker if breaker is not None else CircuitBreaker()
|
||||
self._lock = RLock()
|
||||
self._calls = 0
|
||||
self._prompt_tokens = 0
|
||||
self._completion_tokens = 0
|
||||
self._cost_usd = 0.0
|
||||
self._cache_hits = 0
|
||||
self._cache_misses = 0
|
||||
|
||||
@abstractmethod
|
||||
def _call_once(self, prompt: str, *, params: ClientParams) -> Completion:
|
||||
"""One provider attempt; raise ``ProviderError`` on transient failure."""
|
||||
|
||||
def complete(
|
||||
self,
|
||||
prompt: str,
|
||||
*,
|
||||
params: ClientParams | None = None,
|
||||
) -> str:
|
||||
"""Complete ``prompt`` through the cache/budget/breaker safeguards.
|
||||
|
||||
Returns the completion text (from the cache when possible). Raises a
|
||||
typed ``ClientError`` subclass on budget, circuit or retry failure.
|
||||
"""
|
||||
parameters = params if params is not None else ClientParams()
|
||||
key = response_cache_key(
|
||||
model=self.model,
|
||||
params=parameters,
|
||||
prompt_version=self.prompt_version,
|
||||
schema_version=self.schema_version,
|
||||
prompt=prompt,
|
||||
)
|
||||
|
||||
with self._lock:
|
||||
self._check_call_budget()
|
||||
cached = self._cache.get(key)
|
||||
if cached is not None:
|
||||
self._cache_hits += 1
|
||||
return cached
|
||||
self._cache_misses += 1
|
||||
if self._breaker.is_open and not self._breaker.probe_due:
|
||||
self._breaker.record_refusal()
|
||||
msg = (
|
||||
f"circuit breaker open "
|
||||
f"({self._breaker.consecutive_failures} consecutive "
|
||||
f"failures); call refused"
|
||||
)
|
||||
raise CircuitOpenError(msg)
|
||||
|
||||
try:
|
||||
completion = self._complete_with_retry(prompt, parameters)
|
||||
except RetryExhaustedError:
|
||||
with self._lock:
|
||||
self._breaker.record_failure()
|
||||
raise
|
||||
|
||||
with self._lock:
|
||||
self._record_usage(completion.usage)
|
||||
self._breaker.close()
|
||||
self._cache.put(key, completion.text)
|
||||
return completion.text
|
||||
|
||||
def _complete_with_retry(
|
||||
self,
|
||||
prompt: str,
|
||||
params: ClientParams,
|
||||
) -> Completion:
|
||||
"""Attempt ``_call_once`` up to ``retries + 1`` times, backing off."""
|
||||
last_error: ProviderError | None = None
|
||||
for attempt in range(self._retries + 1):
|
||||
with self._lock:
|
||||
self._calls += 1
|
||||
try:
|
||||
return self._call_once(prompt, params=params)
|
||||
except ProviderError as exc:
|
||||
last_error = exc
|
||||
if attempt < self._retries:
|
||||
self._sleep(self._backoff_base_s * (2**attempt))
|
||||
msg = (
|
||||
f"provider failed after {self._retries + 1} attempt(s): {last_error}"
|
||||
)
|
||||
raise RetryExhaustedError(msg) from last_error
|
||||
|
||||
def _check_call_budget(self) -> None:
|
||||
if self._budget is None or self._budget.max_calls is None:
|
||||
return
|
||||
if self._calls >= self._budget.max_calls:
|
||||
msg = (
|
||||
f"call budget exceeded: {self._calls} >= "
|
||||
f"{self._budget.max_calls} calls"
|
||||
)
|
||||
raise BudgetExceededError(msg)
|
||||
|
||||
def _record_usage(self, usage: TokenUsage) -> None:
|
||||
"""Add a completion to the ledger, enforcing the token/cost ceilings."""
|
||||
if self._budget is not None:
|
||||
ceiling = self._budget.max_total_tokens
|
||||
if ceiling is not None:
|
||||
total_after = (
|
||||
self._prompt_tokens + self._completion_tokens + usage.total
|
||||
)
|
||||
if total_after > ceiling:
|
||||
msg = (
|
||||
f"token budget exceeded: {total_after} > {ceiling} "
|
||||
f"total tokens"
|
||||
)
|
||||
raise BudgetExceededError(msg)
|
||||
self._prompt_tokens += usage.prompt_tokens
|
||||
self._completion_tokens += usage.completion_tokens
|
||||
if self._pricing is not None:
|
||||
cost = self._pricing.cost_usd(usage)
|
||||
self._cost_usd += cost
|
||||
ceiling = None if self._budget is None else self._budget.max_cost_usd
|
||||
if ceiling is not None and self._cost_usd > ceiling:
|
||||
msg = (
|
||||
f"cost budget exceeded: {self._cost_usd:.6f} USD > "
|
||||
f"{ceiling:.6f} USD"
|
||||
)
|
||||
raise BudgetExceededError(msg)
|
||||
|
||||
@property
|
||||
def calls(self) -> int:
|
||||
with self._lock:
|
||||
return self._calls
|
||||
|
||||
@property
|
||||
def prompt_tokens(self) -> int:
|
||||
with self._lock:
|
||||
return self._prompt_tokens
|
||||
|
||||
@property
|
||||
def completion_tokens(self) -> int:
|
||||
with self._lock:
|
||||
return self._completion_tokens
|
||||
|
||||
@property
|
||||
def cost_usd(self) -> float:
|
||||
with self._lock:
|
||||
return self._cost_usd
|
||||
|
||||
@property
|
||||
def cache_hits(self) -> int:
|
||||
with self._lock:
|
||||
return self._cache_hits
|
||||
|
||||
@property
|
||||
def cache_misses(self) -> int:
|
||||
with self._lock:
|
||||
return self._cache_misses
|
||||
|
||||
@property
|
||||
def stats(self) -> ClientStats:
|
||||
with self._lock:
|
||||
return ClientStats(
|
||||
calls=self._calls,
|
||||
prompt_tokens=self._prompt_tokens,
|
||||
completion_tokens=self._completion_tokens,
|
||||
cost_usd=self._cost_usd,
|
||||
cache_hits=self._cache_hits,
|
||||
cache_misses=self._cache_misses,
|
||||
)
|
||||
|
||||
|
||||
class FakeLLMClient(LLMClient):
|
||||
"""Scripted, offline LLM client (P1.1) for tests and the showcase.
|
||||
|
||||
Serves ``script`` in order, one response per successful call. With
|
||||
``failures > 0`` the first ``failures`` attempts raise a simulated
|
||||
``ProviderError`` (to exercise retry/backoff and the breaker); scripts
|
||||
are only consumed by successes. A call past the script raises
|
||||
``ScriptExhaustedError`` (a test-authoring bug, not a transient outage).
|
||||
Token usage is estimated from character counts.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
script: Sequence[str],
|
||||
*,
|
||||
model: str = "fake",
|
||||
prompt_version: str = "fixture:fake:0",
|
||||
failures: int = 0,
|
||||
schema_version: int = OBSERVATION_SCHEMA_VERSION,
|
||||
cache: ResponseCache | None = None,
|
||||
budget: CostBudget | None = None,
|
||||
pricing: Pricing | None = None,
|
||||
retries: int = 2,
|
||||
backoff_base_s: float = 0.25,
|
||||
breaker: CircuitBreaker | None = None,
|
||||
) -> None:
|
||||
if failures < 0:
|
||||
msg = f"failures must be >= 0, got {failures}"
|
||||
raise ValueError(msg)
|
||||
super().__init__(
|
||||
model=model,
|
||||
prompt_version=prompt_version,
|
||||
schema_version=schema_version,
|
||||
cache=cache,
|
||||
budget=budget,
|
||||
pricing=pricing,
|
||||
retries=retries,
|
||||
backoff_base_s=backoff_base_s,
|
||||
sleep=lambda _: None, # tests must never actually sleep
|
||||
breaker=breaker,
|
||||
)
|
||||
self._script = tuple(script)
|
||||
self._failures = failures
|
||||
self._attempts = 0
|
||||
self._served = 0
|
||||
|
||||
def _call_once(self, prompt: str, *, params: ClientParams) -> Completion:
|
||||
del params # a scripted fake ignores sampling parameters
|
||||
with self._lock:
|
||||
self._attempts += 1
|
||||
if self._attempts <= self._failures:
|
||||
msg = f"simulated provider outage (attempt {self._attempts})"
|
||||
raise ProviderError(msg)
|
||||
if self._served >= len(self._script):
|
||||
msg = (
|
||||
f"fake script exhausted after {self._served} served "
|
||||
f"response(s) of {len(self._script)}"
|
||||
)
|
||||
raise ScriptExhaustedError(msg)
|
||||
text = self._script[self._served]
|
||||
self._served += 1
|
||||
usage = TokenUsage(
|
||||
prompt_tokens=estimate_tokens(prompt),
|
||||
completion_tokens=estimate_tokens(text),
|
||||
)
|
||||
return Completion(text=text, usage=usage)
|
||||
Reference in New Issue
Block a user