379 lines
13 KiB
Python
379 lines
13 KiB
Python
"""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
|