Merge pull request #27 from SpowZy/feat/cache-aware-pricing

feat: natively support Anthropic cache-aware pricing (#26)

Author: achris7

src/floe_guard/guard.py

@@ -281,6 +281,9 @@ def settle(
         *,
         reserved: float = 0.0,
         price: ManualPrice | None = None,
+        cache_creation_input_tokens: int = 0,
+        cache_creation_input_tokens_1h: int = 0,
+        cache_read_input_tokens: int = 0,
         label: str | None = None,
     ) -> float:
         """Release a reservation and record the actual cost. Concurrency-safe.
@@ -316,7 +319,14 @@ def settle(
             return 0.0
 
         try:
-            cost = price_tokens(priced, prompt_tokens, completion_tokens)
+            cost = price_tokens(
+                priced,
+                prompt_tokens,
+                completion_tokens,
+                cache_creation_input_tokens=cache_creation_input_tokens,
+                cache_creation_input_tokens_1h=cache_creation_input_tokens_1h,
+                cache_read_input_tokens=cache_read_input_tokens,
+            )
         except Exception:
             # price_tokens can raise (e.g. non-finite token counts). Release the
             # in-flight hold before propagating so _reserved doesn't leak and
@@ -356,6 +366,9 @@ def record(
         completion_tokens: int,
         *,
         price: ManualPrice | None = None,
+        cache_creation_input_tokens: int = 0,
+        cache_creation_input_tokens_1h: int = 0,
+        cache_read_input_tokens: int = 0,
         label: str | None = None,
     ) -> float:
         """Price one response's tokens offline and add the cost to the total.
@@ -365,7 +378,15 @@ def record(
         docstring): warn + raise (default), or warn + skip accrual.
         """
         return self.settle(
-            model, prompt_tokens, completion_tokens, reserved=0.0, price=price, label=label
+            model,
+            prompt_tokens,
+            completion_tokens,
+            reserved=0.0,
+            price=price,
+            cache_creation_input_tokens=cache_creation_input_tokens,
+            cache_creation_input_tokens_1h=cache_creation_input_tokens_1h,
+            cache_read_input_tokens=cache_read_input_tokens,
+            label=label,
         )
 
     def record_tool(self, tool: str, cost_usd: float, *, label: str | None = None) -> float:

src/floe_guard/integrations/anthropic.py

@@ -24,10 +24,8 @@
 from ..guard import BudgetGuard
 from ..pricing import resolve_price
 
-# Anthropic prompt-cache pricing multipliers, relative to the base input rate:
-# 5-minute cache writes bill at ~1.25x, cache reads at ~0.1x. See _usage_from.
-_CACHE_WRITE_WEIGHT = 1.25
-_CACHE_READ_WEIGHT = 0.1
+# Anthropic prompt-cache pricing is now handled natively in the core (guard.py / pricing.py).
+# We extract the buckets and pass them through.
 
 
 def _model_from(kwargs: dict[str, Any], response: Any) -> str:
@@ -41,38 +39,46 @@ def _model_from(kwargs: dict[str, Any], response: Any) -> str:
     return str(model or "")
 
 
-def _usage_from(response: Any) -> tuple[int, int]:
-    """Map an Anthropic message's usage onto (prompt_tokens, completion_tokens).
+def _usage_from(response: Any) -> tuple[int, int, int, int, int]:
+    """Map an Anthropic message's usage onto token buckets.
 
-    Anthropic reports ``usage.input_tokens`` / ``usage.output_tokens``; the guard
-    settles on the OpenAI-style ``(prompt_tokens, completion_tokens)`` pair, so
-    input maps to prompt and output to completion.
-
-    Prompt caching: ``input_tokens`` counts only the *fresh* prompt tokens —
-    cached tokens are reported separately as ``cache_creation_input_tokens``
-    (billed ~1.25x base input) and ``cache_read_input_tokens`` (~0.1x). Dropping
-    them would under-meter a cached call, and a budget guard must never
-    under-count. Since the guard prices with a flat per-model input rate, we fold
-    the cache tokens into the prompt bucket at their *effective* weight so the
-    metered cost approximates the real spend (an estimate, not exact cache-aware
-    pricing — safe-direction).
+    Anthropic reports ``usage.input_tokens`` / ``usage.output_tokens``.
+    Prompt caching tokens are reported as ``cache_creation_input_tokens``
+    and ``cache_read_input_tokens``. We return all five buckets so the
+    core pricing engine can apply exact cost multipliers.
     """
     usage = getattr(response, "usage", None)
     if usage is None and isinstance(response, dict):
         usage = response.get("usage")
     if usage is None:
-        return 0, 0
+        return 0, 0, 0, 0, 0
     get = usage.get if isinstance(usage, dict) else lambda k, d=0: getattr(usage, k, d)
     input_tokens = int(get("input_tokens", 0) or 0)
     output_tokens = int(get("output_tokens", 0) or 0)
-    cache_write = int(get("cache_creation_input_tokens", 0) or 0)
+    cache_write_total = int(get("cache_creation_input_tokens", 0) or 0)
     cache_read = int(get("cache_read_input_tokens", 0) or 0)
-    prompt_tokens = (
-        input_tokens
-        + round(cache_write * _CACHE_WRITE_WEIGHT)
-        + round(cache_read * _CACHE_READ_WEIGHT)
-    )
-    return prompt_tokens, output_tokens
+
+    # Resolve TTL buckets from `cache_creation` if present
+    cache_creation = get("cache_creation", None)
+    cache_write_5m = 0
+    cache_write_1h = 0
+    if cache_creation is not None:
+        get_cc = (
+            cache_creation.get
+            if isinstance(cache_creation, dict)
+            else lambda k, d=0: getattr(cache_creation, k, d)
+        )
+        cache_write_5m = int(get_cc("ephemeral_5m_input_tokens", 0) or 0)
+        cache_write_1h = int(get_cc("ephemeral_1h_input_tokens", 0) or 0)
+
+    # Fallback/defense-in-depth: if there are missing or future buckets, allocate
+    # any leftover cache_creation tokens to the most expensive known bucket (1h) so
+    # the guard over-counts rather than under-meters spend it can't fully attribute.
+    leftover = cache_write_total - (cache_write_5m + cache_write_1h)
+    if leftover > 0:
+        cache_write_1h += leftover
+
+    return input_tokens, output_tokens, cache_write_5m, cache_write_1h, cache_read
 
 
 def _settle_model(guard: BudgetGuard, kwargs: dict[str, Any], response: Any) -> str:
@@ -103,16 +109,36 @@ def _record_response(
     if not isinstance(kwargs, dict):
         kwargs = {}
     model = _settle_model(guard, kwargs, response)
-    prompt_tokens, completion_tokens = _usage_from(response)
-    if prompt_tokens <= 0 and completion_tokens <= 0:
+    (
+        prompt_tokens,
+        completion_tokens,
+        cache_creation_5m,
+        cache_creation_1h,
+        cache_read,
+    ) = _usage_from(response)
+    if (
+        prompt_tokens <= 0
+        and completion_tokens <= 0
+        and cache_creation_5m <= 0
+        and cache_creation_1h <= 0
+        and cache_read <= 0
+    ):
         # No tokens spent — free the reservation.
         guard.release(reserved)
         return
     # There IS spend to account for. Route it through settle() even when the
     # model id is missing, so the guard's policy applies (fail-closed → warn +
     # raise; fail-open → warn + skip) rather than letting a completed call go
     # unmetered and skew the next check().
-    guard.settle(model, prompt_tokens, completion_tokens, reserved=reserved)
+    guard.settle(
+        model,
+        prompt_tokens,
+        completion_tokens,
+        reserved=reserved,
+        cache_creation_input_tokens=cache_creation_5m,
+        cache_creation_input_tokens_1h=cache_creation_1h,
+        cache_read_input_tokens=cache_read,
+    )
 
 
 def _reject_streaming(kwargs: dict[str, Any]) -> None:

src/floe_guard/pricing.py

@@ -132,6 +132,13 @@ def resolve_price(
     return None
 
 
+# Anthropic prompt-cache pricing multipliers: creation (5m) is 1.25x base input,
+# creation (1h) is 2.0x base input, read is 0.1x base input.
+_CACHE_CREATION_MULTIPLIER = 1.25
+_CACHE_CREATION_1H_MULTIPLIER = 2.00
+_CACHE_READ_MULTIPLIER = 0.10
+
+
 def _both_finite(a: Any, b: Any) -> bool:
     return (
         isinstance(a, (int, float))
@@ -141,11 +148,33 @@ def _both_finite(a: Any, b: Any) -> bool:
     )
 
 
-def price_tokens(priced: PricedModel, prompt_tokens: int, completion_tokens: int) -> float:
+def price_tokens(
+    priced: PricedModel,
+    prompt_tokens: int,
+    completion_tokens: int,
+    *,
+    cache_creation_input_tokens: int = 0,
+    cache_read_input_tokens: int = 0,
+    cache_creation_input_tokens_1h: int = 0,
+) -> float:
     """USD cost for token usage. Negative counts are clamped to zero."""
     p = max(0, prompt_tokens)
     c = max(0, completion_tokens)
-    cost = p * priced.input_cost_per_token + c * priced.output_cost_per_token
+    cc = max(0, cache_creation_input_tokens)
+    cc_1h = max(0, cache_creation_input_tokens_1h)
+    cr = max(0, cache_read_input_tokens)
+
+    cache_creation_cost = cc * priced.input_cost_per_token * _CACHE_CREATION_MULTIPLIER
+    cache_creation_1h_cost = cc_1h * priced.input_cost_per_token * _CACHE_CREATION_1H_MULTIPLIER
+    cache_read_cost = cr * priced.input_cost_per_token * _CACHE_READ_MULTIPLIER
+
+    cost = (
+        (p * priced.input_cost_per_token)
+        + (c * priced.output_cost_per_token)
+        + cache_creation_cost
+        + cache_creation_1h_cost
+        + cache_read_cost
+    )
     if not math.isfinite(cost):
         # Defense-in-depth: resolve_price already guarantees finite rates.
         raise ValueError("Non-finite LLM cost — pricing entry is invalid")

tests/test_anthropic_adapter.py

@@ -65,24 +65,55 @@ def __init__(self, messages: object) -> None:
 _MODEL = "claude-3-7-sonnet-20250219"  # present in the bundled cost map
 
 
-def test_usage_maps_input_output_to_prompt_completion() -> None:
-    # The Anthropic-specific bit: input_tokens -> prompt, output_tokens -> completion.
-    assert _usage_from(_Response(_MODEL, _Usage(5, 7))) == (5, 7)
-    assert _usage_from({"usage": {"input_tokens": 5, "output_tokens": 7}}) == (5, 7)
+def test_usage_maps_input_output_and_cache_buckets() -> None:
+    # Anthropic-specific: input_tokens -> prompt, output_tokens -> completion, plus cache buckets
+    assert _usage_from(_Response(_MODEL, _Usage(5, 7))) == (5, 7, 0, 0, 0)
+    assert _usage_from({"usage": {"input_tokens": 5, "output_tokens": 7}}) == (5, 7, 0, 0, 0)
 
 
-def test_usage_folds_prompt_cache_tokens() -> None:
-    # Cached calls must not be under-metered: cache write ~1.25x, read ~0.1x,
-    # folded into the prompt bucket. 100 + round(200*1.25) + round(1000*0.1).
+def test_usage_extracts_prompt_cache_tokens() -> None:
+    # Cached calls must return buckets unmodified so pricing engine can multiply them.
     usage = {
         "input_tokens": 100,
         "output_tokens": 50,
         "cache_creation_input_tokens": 200,
         "cache_read_input_tokens": 1000,
     }
-    assert _usage_from({"usage": usage}) == (100 + 250 + 100, 50)
+    # Direct fallback: no per-TTL cache_creation breakdown, so the leftover
+    # cache_creation tokens default to the most expensive known bucket (1h) —
+    # defense-in-depth so the guard over-counts rather than under-meters.
+    assert _usage_from({"usage": usage}) == (100, 50, 0, 200, 1000)
+
+    # Detailed cache_creation breakdown with 5m and 1h TTLs
+    usage_with_breakdown = {
+        "input_tokens": 100,
+        "output_tokens": 50,
+        "cache_creation_input_tokens": 300,
+        "cache_read_input_tokens": 1000,
+        "cache_creation": {
+            "ephemeral_5m_input_tokens": 100,
+            "ephemeral_1h_input_tokens": 200,
+        }
+    }
+    assert _usage_from({"usage": usage_with_breakdown}) == (100, 50, 100, 200, 1000)
+
+    # Detailed cache_creation breakdown with leftover tokens mapped to the 1h
+    # bucket (defense-in-depth: leftover goes to the most expensive known bucket).
+    usage_with_leftover = {
+        "input_tokens": 100,
+        "output_tokens": 50,
+        "cache_creation_input_tokens": 300,
+        "cache_read_input_tokens": 1000,
+        "cache_creation": {
+            "ephemeral_5m_input_tokens": 50,
+            "ephemeral_1h_input_tokens": 200,
+        }
+    }
+    # 5m=50, 1h=200 → leftover 50 → 1h, so 5m=50, 1h=250.
+    assert _usage_from({"usage": usage_with_leftover}) == (100, 50, 50, 250, 1000)
+
     # No cache fields → unchanged (the object path getattr-defaults to 0).
-    assert _usage_from(_Response(_MODEL, _Usage(5, 7))) == (5, 7)
+    assert _usage_from(_Response(_MODEL, _Usage(5, 7))) == (5, 7, 0, 0, 0)
 
 
 def test_model_from_prefers_response_then_kwargs() -> None:

tests/test_pricing.py

@@ -142,3 +142,34 @@ def test_price_tokens_math() -> None:
 def test_price_tokens_clamps_negative_counts() -> None:
     priced = PricedModel(input_cost_per_token=1e-6, output_cost_per_token=2e-6, source="cost_map")
     assert price_tokens(priced, -50, -50) == 0.0
+
+
+def test_price_tokens_with_prompt_caching_math() -> None:
+    priced = PricedModel(input_cost_per_token=1e-6, output_cost_per_token=2e-6, source="cost_map")
+    # Base input price: 1e-6
+    # 5-minute write (1.25x): 100 * 1e-6 * 1.25 = 0.000125
+    # 1-hour write (2.0x): 200 * 1e-6 * 2.0 = 0.0004
+    # Read (0.1x): 1000 * 1e-6 * 0.1 = 0.0001
+    # Regular prompt (0 tokens), Completion (0 tokens)
+    expected_cost = 0.000125 + 0.0004 + 0.0001
+    cost = price_tokens(
+        priced,
+        prompt_tokens=0,
+        completion_tokens=0,
+        cache_creation_input_tokens=100,
+        cache_read_input_tokens=1000,
+        cache_creation_input_tokens_1h=200,
+    )
+    assert cost == pytest.approx(expected_cost)
+
+
+def test_price_tokens_caching_constants() -> None:
+    from floe_guard.pricing import (
+        _CACHE_CREATION_1H_MULTIPLIER,
+        _CACHE_CREATION_MULTIPLIER,
+        _CACHE_READ_MULTIPLIER,
+    )
+    assert _CACHE_CREATION_MULTIPLIER == 1.25
+    assert _CACHE_CREATION_1H_MULTIPLIER == 2.00
+    assert _CACHE_READ_MULTIPLIER == 0.10
+