fix(telemetry): include Anthropic cache tokens in total token extraction

packages/rhesis/src/rhesis/telemetry/token_extraction.py

@@ -182,8 +182,26 @@ def extract_token_usage(usage: Union[Dict, Any]) -> Tuple[int, int, int]:
         ],
     )
 
+    # Anthropic cache tokens (billed separately but part of actual usage)
+    cache_creation_tokens = get_first_value(
+        usage,
+        [
+            "cache_creation_input_tokens",
+            "cacheCreationInputTokens",  # camelCase variant
+        ],
+    )
+    cache_read_tokens = get_first_value(
+        usage,
+        [
+            "cache_read_input_tokens",
+            "cacheReadInputTokens",  # camelCase variant
+        ],
+    )
+
     # Calculate total if not explicitly provided
-    if not total_tokens and (input_tokens or output_tokens):
-        total_tokens = input_tokens + output_tokens
+    if not total_tokens and (input_tokens or output_tokens or cache_creation_tokens or cache_read_tokens):
+        total_tokens = (
+            input_tokens + output_tokens + cache_creation_tokens + cache_read_tokens
+        )
 
     return input_tokens, output_tokens, total_tokens

tests/sdk/telemetry/test_token_extraction.py

@@ -0,0 +1,36 @@
+"""Tests for provider-agnostic token usage extraction."""
+
+from rhesis.telemetry.token_extraction import extract_token_usage
+
+
+class TestExtractTokenUsage:
+    def test_openai_format(self):
+        usage = {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}
+        assert extract_token_usage(usage) == (10, 20, 30)
+
+    def test_anthropic_format_with_cache_tokens(self):
+        usage = {
+            "input_tokens": 50,
+            "output_tokens": 20,
+            "cache_creation_input_tokens": 1000,
+            "cache_read_input_tokens": 4000,
+        }
+        input_tk, output_tk, total_tk = extract_token_usage(usage)
+        assert input_tk == 50
+        assert output_tk == 20
+        assert total_tk == 50 + 20 + 1000 + 4000
+
+    def test_anthropic_format_without_cache_tokens(self):
+        usage = {"input_tokens": 50, "output_tokens": 20}
+        assert extract_token_usage(usage) == (50, 20, 70)
+
+    def test_gemini_format(self):
+        usage = {"prompt_token_count": 15, "candidates_token_count": 25}
+        assert extract_token_usage(usage) == (15, 25, 40)
+
+    def test_none_returns_zeroes(self):
+        assert extract_token_usage(None) == (0, 0, 0)
+
+    def test_cache_only_tokens(self):
+        usage = {"cache_creation_input_tokens": 1000, "cache_read_input_tokens": 4000}
+        assert extract_token_usage(usage) == (0, 0, 5000)