diff --git a/pyproject.toml b/pyproject.toml index 3cc2f38..1e302b6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -56,6 +56,8 @@ allowed-confusables = ["σ"] "src/pf1e_simulator/combat.py" = ["PLR0913"] # runner.py: run_encounter needs 6 params (spec, runs, seed, policy, spell_registry, capture_first_transcript) "src/pf1e_simulator/runner.py" = ["PLR0913"] +# client.py: LLMClient/FakeLLMClient need the full safeguard knobs (cache, budget, pricing, retry, breaker) +"src/pf1e_simulator/client.py" = ["PLR0913"] # ── Basedpyright ────────────────────────────────────────────────────────────── [tool.basedpyright] diff --git a/src/pf1e_simulator/client.py b/src/pf1e_simulator/client.py new file mode 100644 index 0000000..75b2adf --- /dev/null +++ b/src/pf1e_simulator/client.py @@ -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) diff --git a/tests/test_phase2_client.py b/tests/test_phase2_client.py new file mode 100644 index 0000000..ffd3d42 --- /dev/null +++ b/tests/test_phase2_client.py @@ -0,0 +1,378 @@ +"""Phase-2 LLM client layer (docs/phase-2 P1.1/P1.2, decisions D6/D10). + +``client.py`` wraps every LLM call with the shared safeguards: composite +cache (D6: model, params, prompt_version, schema_version, prompt hash), +cost/token budgets + circuit breaker (P1.2, section 13), retry/backoff on +transient provider failures, and ``FakeLLMClient`` (P1.1) to drive battles +fully offline. These tests pin the ledger, the budgets and the breaker +behaviour without any network. +""" + +from __future__ import annotations + +import pytest + +from pf1e_simulator.client import ( + BudgetExceededError, + CircuitBreaker, + CircuitOpenError, + ClientParams, + ClientStats, + Completion, + CostBudget, + FakeLLMClient, + LLMClient, + Pricing, + ResponseCache, + RetryExhaustedError, + ScriptExhaustedError, + TokenUsage, + estimate_tokens, + response_cache_key, +) + +_PARAMS = ClientParams() +_PROMPT = "observation rendered prompt v1" + + +class _UnusedClient(LLMClient): + """Concrete client whose provider is never invoked (validation tests).""" + + def _call_once(self, prompt: str, *, params: ClientParams) -> Completion: + del prompt, params # validation-only: the provider never runs + msg = "provider must never be called" + raise AssertionError(msg) + + +# ── composite cache key (D6) ─────────────────────────────────────────────── + + +def test_cache_key_is_composite_and_deterministic() -> None: + """Each D6 component changes the key; identical inputs are stable.""" + base = response_cache_key( + model="m", + params=_PARAMS, + prompt_version="pack:1.0.0:nl:abc", + schema_version=1, + prompt=_PROMPT, + ) + assert ( + response_cache_key( + model="m", + params=_PARAMS, + prompt_version="pack:1.0.0:nl:abc", + schema_version=1, + prompt=_PROMPT, + ) + == base + ) + assert ( + response_cache_key( + model="other", + params=_PARAMS, + prompt_version="pack:1.0.0:nl:abc", + schema_version=1, + prompt=_PROMPT, + ) + != base + ) + assert ( + response_cache_key( + model="m", + params=ClientParams(temperature=0.7), + prompt_version="pack:1.0.0:nl:abc", + schema_version=1, + prompt=_PROMPT, + ) + != base + ) + assert ( + response_cache_key( + model="m", + params=_PARAMS, + prompt_version="pack:9.9.9:nl:xyz", + schema_version=1, + prompt=_PROMPT, + ) + != base + ) + assert ( + response_cache_key( + model="m", + params=_PARAMS, + prompt_version="pack:1.0.0:nl:abc", + schema_version=2, + prompt=_PROMPT, + ) + != base + ) + assert ( + response_cache_key( + model="m", + params=_PARAMS, + prompt_version="pack:1.0.0:nl:abc", + schema_version=1, + prompt="a different observation", + ) + != base + ) + + +def test_response_cache_round_trip_and_counters() -> None: + """get/put round-trip; hits and misses are counted separately.""" + cache = ResponseCache(shard_count=4) + key = response_cache_key( + model="m", + params=_PARAMS, + prompt_version="pack:1.0.0:nl:abc", + schema_version=1, + prompt=_PROMPT, + ) + assert cache.get(key) is None + assert cache.hits == 0 + assert cache.misses == 1 + cache.put(key, "action json 1") + assert cache.get(key) == "action json 1" + assert cache.hits == 1 + assert cache.misses == 1 + + +def test_response_cache_clear_drops_entries_and_counters() -> None: + """clear() empties every shard and resets the counters (tests only).""" + cache = ResponseCache(shard_count=2) + key = response_cache_key( + model="m", + params=_PARAMS, + prompt_version="p:1:nl:a", + schema_version=1, + prompt=_PROMPT, + ) + cache.put(key, "action json 1") + assert cache.get(key) == "action json 1" + cache.clear() + assert cache.get(key) is None + assert cache.hits == 0 + assert cache.misses == 1 + + +def test_response_cache_requires_positive_shard_count() -> None: + """A zero shard count is a data error, not a silent default.""" + with pytest.raises(ValueError, match="shard_count"): + ResponseCache(shard_count=0) + + +# ── FakeLLMClient: scripted offline behaviour (P1.1) ─────────────────────── + + +def test_fake_serves_script_in_order() -> None: + """The fake returns scripted responses one per successful call.""" + client = FakeLLMClient(script=("action a", "action b")) + assert client.complete("observation one") == "action a" + assert client.complete("observation two") == "action b" + assert client.stats.calls == 2 + + +def test_fake_script_exhaustion_raises_typed_error() -> None: + """A call past the script is a test-authoring bug, loudly reported.""" + client = FakeLLMClient(script=("action a",)) + assert client.complete("observation one") == "action a" + with pytest.raises(ScriptExhaustedError, match="script exhausted"): + client.complete("observation two") + + +def test_fake_reports_estimated_token_usage() -> None: + """Token usage is estimated from character counts (no provider metric).""" + client = FakeLLMClient(script=("0123456789",)) + client.complete(_PROMPT) + stats = client.stats + assert stats.prompt_tokens == estimate_tokens(_PROMPT) + assert stats.completion_tokens == estimate_tokens("0123456789") + assert isinstance(stats, ClientStats) + + +# ── cache integration on the client (D6) ─────────────────────────────────── + + +def test_client_cache_hit_consumes_no_script() -> None: + """A cached response is returned without a provider call.""" + client = FakeLLMClient(script=("action a", "action b")) + first = client.complete(_PROMPT) + second = client.complete(_PROMPT) + assert first == second == "action a" + assert client.calls == 1 # the second call came from the cache + assert client.cache_misses == 1 + assert client.cache_hits == 1 + + +def test_client_cache_miss_on_different_prompt() -> None: + """A different prompt is a different cache key: script advances.""" + client = FakeLLMClient(script=("action a", "action b")) + assert client.complete(_PROMPT) == "action a" + assert client.complete("another observation") == "action b" + assert client.calls == 2 + + +def test_client_cache_respects_prompt_version() -> None: + """The composite key includes prompt_version: packs do not collide.""" + cache = ResponseCache() + a = FakeLLMClient( + script=("action a",), + prompt_version="pack:1.0.0:nl:aaa", + cache=cache, + ) + b = FakeLLMClient( + script=("action b",), + prompt_version="pack:1.0.0:nl:bbb", + cache=cache, + ) + assert a.complete(_PROMPT) == "action a" + assert b.complete(_PROMPT) == "action b" + assert a.calls == 1 + assert b.calls == 1 + + +def test_client_cache_respects_model() -> None: + """The composite key includes the model name.""" + cache = ResponseCache() + a = FakeLLMClient(script=("action a",), model="model-a", cache=cache) + b = FakeLLMClient(script=("action b",), model="model-b", cache=cache) + assert a.complete(_PROMPT) == "action a" + assert b.complete(_PROMPT) == "action b" + assert a.calls == 1 + assert b.calls == 1 + + +# ── retry/backoff (P1.2) ─────────────────────────────────────────────────── + + +def test_retry_recovers_after_transient_failures() -> None: + """ProviderError is retried; the script is only consumed by successes.""" + client = FakeLLMClient( + script=("action a",), + failures=2, + retries=3, + ) + assert client.complete(_PROMPT) == "action a" + assert client.calls == 3 # two failures + one success + + +def test_retry_exhausted_after_too_many_failures() -> None: + """Past the retry budget, a typed RetryExhaustedError is raised.""" + client = FakeLLMClient(script=("action a",), failures=5, retries=2) + with pytest.raises(RetryExhaustedError, match="after 3 attempt"): + client.complete(_PROMPT) + + +def test_retry_counter_requires_positive_value() -> None: + """A negative retry count is a data error.""" + with pytest.raises(ValueError, match="retries"): + FakeLLMClient(script=("a",), retries=-1) + + +# ── circuit breaker (P1.2, section 13) ───────────────────────────────────── + + +def test_breaker_opens_refuses_probes_and_recovers() -> None: + """Open after N failures: refusals, then a probe that closes on success.""" + breaker = CircuitBreaker(failure_threshold=2, open_refusals=1) + client = FakeLLMClient( + script=("action a",), + failures=2, # needs two attempts before real service + retries=0, + breaker=breaker, + ) + with pytest.raises(RetryExhaustedError): + client.complete(_PROMPT) # failure 1: not open yet + with pytest.raises(RetryExhaustedError): + client.complete(_PROMPT) # failure 2: threshold reached, opens + assert breaker.is_open + with pytest.raises(CircuitOpenError, match="circuit breaker open"): + client.complete(_PROMPT) # refused + assert client.complete(_PROMPT) == "action a" # probe succeeds now + assert not breaker.is_open + + +def test_breaker_probe_failure_reopens() -> None: + """A failed probe reopens the breaker for another refusal window.""" + breaker = CircuitBreaker(failure_threshold=1, open_refusals=1) + client = FakeLLMClient( + script=("action a",), + failures=10, # every attempt fails, probe included + retries=0, + breaker=breaker, + ) + with pytest.raises(RetryExhaustedError): + client.complete(_PROMPT) + assert breaker.is_open + with pytest.raises(CircuitOpenError): + client.complete(_PROMPT) + with pytest.raises(RetryExhaustedError): + client.complete(_PROMPT) # the probe fails + assert breaker.is_open + with pytest.raises(CircuitOpenError): + client.complete(_PROMPT) # reopened: refuses again + + +def test_breaker_validation() -> None: + """Breaker thresholds must be sensible numbers.""" + with pytest.raises(ValueError, match="failure_threshold"): + CircuitBreaker(failure_threshold=0) + with pytest.raises(ValueError, match="open_refusals"): + CircuitBreaker(open_refusals=-1) + + +# ── cost/token budgets (P1.2) ────────────────────────────────────────────── + + +def test_call_budget_ceiling() -> None: + """max_calls refuses further calls with a typed error.""" + client = FakeLLMClient( + script=("a", "b", "c"), + budget=CostBudget(max_calls=1), + ) + assert client.complete(_PROMPT) == "a" + with pytest.raises(BudgetExceededError, match="call budget exceeded"): + client.complete("a different observation") + + +def test_token_budget_ceiling() -> None: + """A completion over max_total_tokens is rejected and not cached.""" + client = FakeLLMClient( + script=("x" * 400,), # ~100 estimated completion tokens + budget=CostBudget(max_total_tokens=50), + ) + with pytest.raises(BudgetExceededError, match="token budget exceeded"): + client.complete(_PROMPT) + + +def test_cost_budget_ceiling_and_ledger() -> None: + """Estimated USD cost is accumulated and capped by max_cost_usd.""" + pricing = Pricing(prompt_per_1k_usd=1.0, completion_per_1k_usd=1.0) + long_prompt = "p" * 2000 # ~500 prompt tokens ≈ 0.5 USD + client = FakeLLMClient( + script=("c" * 2000, "c" * 2000), # ~500 completion tokens each ≈ 0.5 USD + pricing=pricing, + budget=CostBudget(max_cost_usd=1.0), + ) + assert client.complete(long_prompt) == "c" * 2000 + assert client.cost_usd == pytest.approx(1.0, abs=1e-6) + with pytest.raises(BudgetExceededError, match="cost budget exceeded"): + client.complete("q" * 2000) + assert client.cost_usd == pytest.approx(2.0, abs=1e-6) + + +def test_backoff_must_be_positive() -> None: + """A negative backoff is a data error (validation-only client).""" + with pytest.raises(ValueError, match="backoff_base_s"): + _UnusedClient( + model="m", + prompt_version="p:1:nl:a", + backoff_base_s=-1.0, + ) + + +def test_token_usage_total() -> None: + """TokenUsage.total is the sum of both counters.""" + usage = TokenUsage(prompt_tokens=3, completion_tokens=4) + assert usage.total == 7 + assert estimate_tokens("abcd") >= 1