chore(weave): Add db and server support for cached tokens

tests/trace/test_weave_client.py

@@ -1697,6 +1697,12 @@ def models(text):
             "pricing_level": "default",
             "pricing_level_id": "default",
             "created_by": "system",
+            "cache_read_input_tokens": 0,
+            "cache_creation_input_tokens": 0,
+            "cache_read_input_tokens_total_cost": 0.0,
+            "cache_creation_input_tokens_total_cost": 0.0,
+            "cache_read_input_token_cost": 0.0,
+            "cache_creation_input_token_cost": 0.0,
         }
     )
 
@@ -1716,6 +1722,12 @@ def models(text):
             "pricing_level": "default",
             "pricing_level_id": "default",
             "created_by": "system",
+            "cache_read_input_tokens": 0,
+            "cache_creation_input_tokens": 0,
+            "cache_read_input_tokens_total_cost": 0.0,
+            "cache_creation_input_tokens_total_cost": 0.0,
+            "cache_read_input_token_cost": 0.0,
+            "cache_creation_input_token_cost": 0.0,
         }
     )
 

tests/trace_server/costs/test_insert_costs.py

@@ -122,6 +122,8 @@ def test_insert_costs_into_db(self, mock_datetime, mock_uuid4):
                     "USD",
                     cost.get("output", 0),
                     "USD",
+                    cost.get("cache_read_input", 0),
+                    cost.get("cache_creation_input", 0),
                     "system",
                     created_at,
                 )
@@ -142,6 +144,8 @@ def test_insert_costs_into_db(self, mock_datetime, mock_uuid4):
                 "prompt_token_cost_unit",
                 "completion_token_cost",
                 "completion_token_cost_unit",
+                "cache_read_input_token_cost",
+                "cache_creation_input_token_cost",
                 "created_by",
                 "created_at",
             ],

tests/trace_server/test_opentelemetry.py

@@ -1494,6 +1494,12 @@ def test_opentelemetry_cost_calculation(self, client: weave_client.WeaveClient):
                 "pricing_level": "default",
                 "pricing_level_id": "default",
                 "created_by": "system",
+                "cache_read_input_tokens": 0,
+                "cache_creation_input_tokens": 0,
+                "cache_read_input_tokens_total_cost": 0,
+                "cache_creation_input_tokens_total_cost": 0,
+                "cache_read_input_token_cost": 0,
+                "cache_creation_input_token_cost": 0,
             }
         )
 

tests/trace_server/test_trace_usage.py

@@ -15,10 +15,16 @@
     "output_tokens",
     "requests",
     "total_tokens",
+    "cache_read_input_tokens",
+    "cache_creation_input_tokens",
     "prompt_tokens_total_cost",
     "completion_tokens_total_cost",
+    "cache_read_input_tokens_total_cost",
+    "cache_creation_input_tokens_total_cost",
     "prompt_token_cost",
     "completion_token_cost",
+    "cache_read_input_token_cost",
+    "cache_creation_input_token_cost",
     "prompt_token_cost_unit",
     "completion_token_cost_unit",
     "effective_date",
@@ -778,3 +784,83 @@ def test_calls_usage_handles_missing_usage(
     assert usage.completion_tokens == expected_root[1]
     assert usage.total_tokens == expected_root[2]
     assert usage.requests == expected_root[3]
+
+
+def test_aggregate_usage_with_cache_tokens_rolls_up() -> None:
+    """Cache token counts and costs are extracted, rolled up, and merged correctly."""
+    root_id = "root"
+    child_id = "child"
+
+    calls = [
+        _make_call(
+            root_id,
+            None,
+            _usage_summary(
+                {
+                    "claude-3.5-sonnet": {
+                        "input_tokens": 100,
+                        "output_tokens": 50,
+                        "cache_read_input_tokens": 80,
+                        "cache_creation_input_tokens": 20,
+                        "requests": 1,
+                    }
+                },
+                costs={
+                    "claude-3.5-sonnet": {
+                        "prompt_tokens_total_cost": 0.3,
+                        "completion_tokens_total_cost": 0.15,
+                        "cache_read_input_tokens_total_cost": 0.04,
+                        "cache_creation_input_tokens_total_cost": 0.05,
+                    }
+                },
+            ),
+        ),
+        _make_call(
+            child_id,
+            root_id,
+            _usage_summary(
+                {
+                    "claude-3.5-sonnet": {
+                        "input_tokens": 60,
+                        "output_tokens": 30,
+                        "cache_read_input_tokens": 40,
+                        "cache_creation_input_tokens": 10,
+                        "requests": 1,
+                    }
+                },
+                costs={
+                    "claude-3.5-sonnet": {
+                        "prompt_tokens_total_cost": 0.18,
+                        "completion_tokens_total_cost": 0.09,
+                        "cache_read_input_tokens_total_cost": 0.02,
+                        "cache_creation_input_tokens_total_cost": 0.0,
+                    }
+                },
+            ),
+        ),
+    ]
+
+    # Token counts roll up
+    result = usage_utils.aggregate_usage_with_descendants(calls, include_costs=True)
+    root_usage = result[root_id]["claude-3.5-sonnet"]
+    assert root_usage.prompt_tokens == 160
+    assert root_usage.completion_tokens == 80
+    assert root_usage.cache_read_input_tokens == 120
+    assert root_usage.cache_creation_input_tokens == 30
+    assert root_usage.requests == 2
+
+    # Costs roll up
+    assert root_usage.cache_read_input_tokens_total_cost == pytest.approx(0.06)
+    assert root_usage.cache_creation_input_tokens_total_cost == pytest.approx(0.05)
+
+    child_usage = result[child_id]["claude-3.5-sonnet"]
+    assert child_usage.cache_read_input_tokens == 40
+    assert child_usage.cache_creation_input_tokens == 10
+
+    # Without costs flag, cost fields are None
+    result_no_costs = usage_utils.aggregate_usage_with_descendants(
+        calls, include_costs=False
+    )
+    root_no_costs = result_no_costs[root_id]["claude-3.5-sonnet"]
+    assert root_no_costs.cache_read_input_tokens_total_cost is None
+    assert root_no_costs.cache_creation_input_tokens_total_cost is None

weave/trace_server/calls_query_builder/usage_query_builder.py

@@ -71,6 +71,8 @@ def _get_usage_metric_extraction_sql(metric: str, json_col: str) -> str:
             ifNull(toFloat64OrNull(JSONExtractRaw({json_col}, 'completion_tokens')), 0) +
             ifNull(toFloat64OrNull(JSONExtractRaw({json_col}, 'output_tokens')), 0)
         )"""
+    elif metric in {"cache_read_input_tokens", "cache_creation_input_tokens"}:
+        return f"ifNull(toFloat64OrNull(JSONExtractRaw({json_col}, '{metric}')), 0)"
     else:
         return f"toFloat64OrNull(JSONExtractRaw({json_col}, '{metric}'))"
 

weave/trace_server/clickhouse_trace_server_batched.py

@@ -1089,12 +1089,24 @@ def _get_prices_for_models(
 
         prices: dict[str, dict[str, float]] = {}
         for row in result.result_rows:
-            llm_id, prompt_cost, completion_cost = row
+            (
+                llm_id,
+                prompt_cost,
+                completion_cost,
+                cache_read_cost,
+                cache_creation_cost,
+            ) = row
             prices[llm_id] = {
                 "prompt_token_cost": float(prompt_cost) if prompt_cost else 0.0,
                 "completion_token_cost": float(completion_cost)
                 if completion_cost
                 else 0.0,
+                "cache_read_input_token_cost": float(cache_read_cost)
+                if cache_read_cost
+                else 0.0,
+                "cache_creation_input_token_cost": float(cache_creation_cost)
+                if cache_creation_cost
+                else 0.0,
             }
         return prices
 
@@ -1126,22 +1138,42 @@ def _compute_costs_for_buckets(
             model_prices = prices.get(model, {})
             prompt_cost = model_prices.get("prompt_token_cost", 0.0)
             completion_cost = model_prices.get("completion_token_cost", 0.0)
+            cache_read_cost = model_prices.get("cache_read_input_token_cost", 0.0)
+            cache_creation_cost = model_prices.get(
+                "cache_creation_input_token_cost", 0.0
+            )
 
             input_tokens = bucket.get("sum_input_tokens", 0) or 0
             output_tokens = bucket.get("sum_output_tokens", 0) or 0
+            cache_read_tokens = bucket.get("sum_cache_read_input_tokens", 0) or 0
+            cache_creation_tokens = (
+                bucket.get("sum_cache_creation_input_tokens", 0) or 0
+            )
+
+            # Subtract cache tokens from input: they are billed at cache
+            # rates, not the regular prompt rate.
+            net_input_tokens = (
+                input_tokens - cache_read_tokens - cache_creation_tokens
+            )
 
             if "input_cost" in requested_cost_metrics:
-                bucket["sum_input_cost"] = input_tokens * prompt_cost
+                bucket["sum_input_cost"] = net_input_tokens * prompt_cost
 
             if "output_cost" in requested_cost_metrics:
                 bucket["sum_output_cost"] = output_tokens * completion_cost
 
             if "total_cost" in requested_cost_metrics:
-                input_cost = bucket.get("sum_input_cost", input_tokens * prompt_cost)
+                input_cost = bucket.get(
+                    "sum_input_cost", net_input_tokens * prompt_cost
+                )
                 output_cost = bucket.get(
                     "sum_output_cost", output_tokens * completion_cost
                 )
-                bucket["sum_total_cost"] = input_cost + output_cost
+                cache_read_total = cache_read_tokens * cache_read_cost
+                cache_creation_total = cache_creation_tokens * cache_creation_cost
+                bucket["sum_total_cost"] = (
+                    input_cost + output_cost + cache_read_total + cache_creation_total
+                )
 
     def call_stats(self, req: tsi.CallStatsReq) -> tsi.CallStatsRes:
         """Return call statistics grouped by bucket with requested aggregations.
@@ -5742,6 +5774,8 @@ def cost_create(self, req: tsi.CostCreateReq) -> tsi.CostCreateRes:
                 ),
                 "prompt_token_cost": cost.prompt_token_cost,
                 "completion_token_cost": cost.completion_token_cost,
+                "cache_read_input_token_cost": cost.cache_read_input_token_cost,
+                "cache_creation_input_token_cost": cost.cache_creation_input_token_cost,
                 "prompt_token_cost_unit": cost.prompt_token_cost_unit,
                 "completion_token_cost_unit": cost.completion_token_cost_unit,
             }

weave/trace_server/costs/insert_costs.py

@@ -42,6 +42,8 @@ def get_current_costs(
 class CostDetails(TypedDict):
     input: float
     output: float
+    cache_read_input: float
+    cache_creation_input: float
     provider: str
     created_at: str
 
@@ -69,6 +71,8 @@ def insert_costs_into_db(client: Client, data: dict[str, list[CostDetails]]) ->
             provider_id = cost.get("provider", "default")
             input_token_cost = cost.get("input", 0)
             output_token_cost = cost.get("output", 0)
+            cache_read_input_token_cost = cost.get("cache_read_input", 0)
+            cache_creation_input_token_cost = cost.get("cache_creation_input", 0)
             date_str = cost.get(
                 "created_at", datetime.now().strftime("%Y-%m-%d %H:%M:%S")
             )
@@ -87,6 +91,8 @@ def insert_costs_into_db(client: Client, data: dict[str, list[CostDetails]]) ->
                     "USD",
                     output_token_cost,
                     "USD",
+                    cache_read_input_token_cost,
+                    cache_creation_input_token_cost,
                     "system",
                     created_at,
                 ),
@@ -106,6 +112,8 @@ def insert_costs_into_db(client: Client, data: dict[str, list[CostDetails]]) ->
             "prompt_token_cost_unit",
             "completion_token_cost",
             "completion_token_cost_unit",
+            "cache_read_input_token_cost",
+            "cache_creation_input_token_cost",
             "created_by",
             "created_at",
         ],

weave/trace_server/migrations/027_add_cache_token_costs.down.sql

@@ -0,0 +1,2 @@
+ALTER TABLE llm_token_prices DROP COLUMN IF EXISTS cache_read_input_token_cost;
+ALTER TABLE llm_token_prices DROP COLUMN IF EXISTS cache_creation_input_token_cost;

weave/trace_server/migrations/027_add_cache_token_costs.up.sql

@@ -0,0 +1,2 @@
+ALTER TABLE llm_token_prices ADD COLUMN IF NOT EXISTS cache_read_input_token_cost Float DEFAULT 0;
+ALTER TABLE llm_token_prices ADD COLUMN IF NOT EXISTS cache_creation_input_token_cost Float DEFAULT 0;

weave/trace_server/sqlite_trace_server.py

@@ -239,6 +239,12 @@ def _cost_usage_from_summary(
             "requests": _safe_int_for_costs(usage.get("requests")),
             # Match ClickHouse: keep total_tokens as-reported rather than deriving it.
             "total_tokens": _safe_int_for_costs(usage.get("total_tokens")),
+            "cache_read_input_tokens": _safe_int_for_costs(
+                usage.get("cache_read_input_tokens")
+            ),
+            "cache_creation_input_tokens": _safe_int_for_costs(
+                usage.get("cache_creation_input_tokens")
+            ),
         }
     return normalized_usage
 
@@ -437,6 +443,8 @@ def setup_tables(self) -> None:
                 effective_date TEXT NOT NULL,
                 prompt_token_cost REAL NOT NULL,
                 completion_token_cost REAL NOT NULL,
+                cache_read_input_token_cost REAL NOT NULL DEFAULT 0,
+                cache_creation_input_token_cost REAL NOT NULL DEFAULT 0,
                 prompt_token_cost_unit TEXT NOT NULL,
                 completion_token_cost_unit TEXT NOT NULL,
                 created_by TEXT NOT NULL,
@@ -671,6 +679,8 @@ def _ensure_default_costs(self, cursor: sqlite3.Cursor) -> bool:
                 row["effective_date"],
                 row["prompt_token_cost"],
                 row["completion_token_cost"],
+                row.get("cache_read_input_token_cost", 0),
+                row.get("cache_creation_input_token_cost", 0),
                 row["prompt_token_cost_unit"],
                 row["completion_token_cost_unit"],
                 row["created_by"],
@@ -689,11 +699,13 @@ def _ensure_default_costs(self, cursor: sqlite3.Cursor) -> bool:
                 effective_date,
                 prompt_token_cost,
                 completion_token_cost,
+                cache_read_input_token_cost,
+                cache_creation_input_token_cost,
                 prompt_token_cost_unit,
                 completion_token_cost_unit,
                 created_by,
                 created_at
-            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+            ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
             """,
             default_rows,
         )
@@ -767,6 +779,8 @@ def _apply_costs_to_calls(
                 effective_date,
                 prompt_token_cost,
                 completion_token_cost,
+                cache_read_input_token_cost,
+                cache_creation_input_token_cost,
                 prompt_token_cost_unit,
                 completion_token_cost_unit,
                 created_by,
@@ -800,6 +814,8 @@ def _apply_costs_to_calls(
                         "effective_date",
                         "prompt_token_cost",
                         "completion_token_cost",
+                        "cache_read_input_token_cost",
+                        "cache_creation_input_token_cost",
                         "prompt_token_cost_unit",
                         "completion_token_cost_unit",
                         "created_by",
@@ -836,18 +852,43 @@ def _apply_costs_to_calls(
 
                 prompt_cost = float(best_row["prompt_token_cost"] or 0.0)
                 completion_cost = float(best_row["completion_token_cost"] or 0.0)
+                cache_read_cost = float(
+                    best_row.get("cache_read_input_token_cost") or 0.0
+                )
+                cache_creation_cost = float(
+                    best_row.get("cache_creation_input_token_cost") or 0.0
+                )
                 prompt_tokens = usage["prompt_tokens"]
                 completion_tokens = usage["completion_tokens"]
+                cache_read_input_tokens = usage.get("cache_read_input_tokens", 0)
+                cache_creation_input_tokens = usage.get(
+                    "cache_creation_input_tokens", 0
+                )
 
                 call_costs[llm_id] = {
                     "prompt_tokens": prompt_tokens,
                     "completion_tokens": completion_tokens,
+                    "cache_read_input_tokens": cache_read_input_tokens,
+                    "cache_creation_input_tokens": cache_creation_input_tokens,
                     "requests": usage["requests"],
                     "total_tokens": usage["total_tokens"],
-                    "prompt_tokens_total_cost": prompt_tokens * prompt_cost,
+                    # Subtract cached tokens: they are billed at the cache
+                    # rate, not the regular input rate.
+                    "prompt_tokens_total_cost": (
+                        prompt_tokens
+                        - cache_read_input_tokens
+                        - cache_creation_input_tokens
+                    )
+                    * prompt_cost,
                     "completion_tokens_total_cost": completion_tokens * completion_cost,
+                    "cache_read_input_tokens_total_cost": cache_read_input_tokens
+                    * cache_read_cost,
+                    "cache_creation_input_tokens_total_cost": cache_creation_input_tokens
+                    * cache_creation_cost,
                     "prompt_token_cost": prompt_cost,
                     "completion_token_cost": completion_cost,
+                    "cache_read_input_token_cost": cache_read_cost,
+                    "cache_creation_input_token_cost": cache_creation_cost,
                     "prompt_token_cost_unit": best_row["prompt_token_cost_unit"],
                     "completion_token_cost_unit": best_row[
                         "completion_token_cost_unit"