Merge pull request #139 from Takki-cd/feat/add-cache-tokens

feat: add fine-grained cache token tracking across 5 LLM providers

Author: JDOxygen

oxygent/schemas/usage.py

@@ -2,7 +2,7 @@
 
 from enum import Enum
 
-from pydantic import BaseModel, Field, computed_field
+from pydantic import BaseModel, Field, model_validator
 
 
 class EstimationMethod(str, Enum):
@@ -19,23 +19,29 @@ class TokenUsage(BaseModel):
     Attributes:
         input_tokens: Number of input (prompt) tokens.
         output_tokens: Number of output (completion) tokens.
-        cached_tokens: Number of input tokens served from cache (e.g. OpenAI
-            prompt_tokens_details.cached_tokens, Anthropic cache_read_input_tokens).
-        reasoning_tokens: Number of tokens used for reasoning/thinking
-            (e.g. OpenAI completion_tokens_details.reasoning_tokens).
+        cached_input_tokens: Number of input tokens served from cache read hit.
+            (e.g. OpenAI prompt_tokens_details.cached_tokens,
+            Anthropic cache_read_input_tokens, DeepSeek prompt_cache_hit_tokens,
+            Gemini cachedContentTokenCount).
+        cache_creation_input_tokens: Number of input tokens written to cache.
+            (e.g. Anthropic cache_creation_input_tokens).
+        reasoning_tokens: Number of tokens used for reasoning/thinking.
         model_name: Name of the LLM model used.
         estimation_method: Method used for token counting.
     """
 
     input_tokens: int = Field(default=0, ge=0)
     output_tokens: int = Field(default=0, ge=0)
-    cached_tokens: int = Field(default=0, ge=0)
+    total_tokens: int = Field(default=0, ge=0)
+    cached_input_tokens: int = Field(default=0, ge=0)
+    cache_creation_input_tokens: int = Field(default=0, ge=0)
     reasoning_tokens: int = Field(default=0, ge=0)
     model_name: str = ""
     estimation_method: EstimationMethod = EstimationMethod.EXACT
 
-    @computed_field
-    @property
-    def total_tokens(self) -> int:
-        """Total tokens (input + output), always computed."""
-        return self.input_tokens + self.output_tokens
+    @model_validator(mode="after")
+    def _compute_total(self) -> "TokenUsage":
+        """Use API-provided total_tokens when available, fallback to input+output."""
+        if self.total_tokens == 0:
+            self.total_tokens = self.input_tokens + self.output_tokens
+        return self

oxygent/utils/token_utils.py

@@ -80,16 +80,35 @@ def build_token_usage(
     u = _to_dict(usage_data)
 
     # --- input / output tokens ---
-    # OpenAI-family: prompt_tokens; Gemini native: promptTokenCount
-    input_tokens = u.get("prompt_tokens") or u.get("promptTokenCount", 0)
-    output_tokens = u.get("completion_tokens") or u.get("candidatesTokenCount", 0)
+    # OpenAI: prompt_tokens; Gemini native: promptTokenCount; Anthropic: input_tokens
+    input_tokens = u.get("prompt_tokens") or u.get("promptTokenCount") or u.get("input_tokens", 0)
+    output_tokens = u.get("completion_tokens") or u.get("candidatesTokenCount") or u.get("output_tokens", 0)
+
+    # --- total tokens (API-provided) ---
+    # Gemini: total = prompt + completion + thoughts (thoughts NOT included in completion)
+    # OpenAI: total = prompt + completion (reasoning already included in completion)
+    total_tokens = u.get("total_tokens") or 0
 
     # --- cached tokens ---
-    cached_tokens = 0
+    # cached_input_tokens: cache read hits (chain fallback)
+    # 1. prompt_tokens_details.cached_tokens        -> OpenAI / Doubao / Alibaba Cloud
+    # 2. top-level cache_read_input_tokens          -> Anthropic
+    # 3. top-level prompt_cache_hit_tokens          -> DeepSeek
+    # 4. top-level cachedContentTokenCount          -> Gemini native
+    cached_input_tokens = 0
     prompt_details = u.get("prompt_tokens_details")
     if prompt_details:
         d = _to_dict(prompt_details)
-        cached_tokens = d.get("cached_tokens") or d.get("cache_read_input_tokens", 0)
+        cached_input_tokens = d.get("cached_tokens") or 0
+    if not cached_input_tokens:
+        cached_input_tokens = u.get("cache_read_input_tokens") or 0
+    if not cached_input_tokens:
+        cached_input_tokens = u.get("prompt_cache_hit_tokens") or 0
+    if not cached_input_tokens:
+        cached_input_tokens = u.get("cachedContentTokenCount") or 0
+
+    # cache_creation_input_tokens: cache writes (Anthropic / Alibaba Cloud)
+    cache_creation_input_tokens = u.get("cache_creation_input_tokens") or 0
 
     # --- reasoning / thinking tokens ---
     # Priority: completion_tokens_details.reasoning_tokens (OpenAI/Doubao)
@@ -110,7 +129,9 @@ def build_token_usage(
     return TokenUsage(
         input_tokens=input_tokens,
         output_tokens=output_tokens,
-        cached_tokens=cached_tokens,
+        total_tokens=total_tokens,
+        cached_input_tokens=cached_input_tokens,
+        cache_creation_input_tokens=cache_creation_input_tokens,
         reasoning_tokens=reasoning_tokens,
         model_name=model_name,
         estimation_method=EstimationMethod.EXACT,
@@ -138,7 +159,8 @@ def aggregate_token_usage(oxy_request: Any, usage: TokenUsage) -> None:
             "total_tokens": 0,
             "input_tokens": 0,
             "output_tokens": 0,
-            "cached_tokens": 0,
+            "cached_input_tokens": 0,
+            "cache_creation_input_tokens": 0,
             "reasoning_tokens": 0,
             "request_count": 0,
             "by_model": {},
@@ -148,7 +170,8 @@ def aggregate_token_usage(oxy_request: Any, usage: TokenUsage) -> None:
     tm["total_tokens"] += usage.total_tokens
     tm["input_tokens"] += usage.input_tokens
     tm["output_tokens"] += usage.output_tokens
-    tm["cached_tokens"] += usage.cached_tokens
+    tm["cached_input_tokens"] += usage.cached_input_tokens
+    tm["cache_creation_input_tokens"] += usage.cache_creation_input_tokens
     tm["reasoning_tokens"] += usage.reasoning_tokens
     tm["request_count"] += 1
 
@@ -158,7 +181,8 @@ def aggregate_token_usage(oxy_request: Any, usage: TokenUsage) -> None:
             "total_tokens": 0,
             "input_tokens": 0,
             "output_tokens": 0,
-            "cached_tokens": 0,
+            "cached_input_tokens": 0,
+            "cache_creation_input_tokens": 0,
             "reasoning_tokens": 0,
             "request_count": 0,
         }
@@ -167,7 +191,8 @@ def aggregate_token_usage(oxy_request: Any, usage: TokenUsage) -> None:
     mm["total_tokens"] += usage.total_tokens
     mm["input_tokens"] += usage.input_tokens
     mm["output_tokens"] += usage.output_tokens
-    mm["cached_tokens"] += usage.cached_tokens
+    mm["cached_input_tokens"] += usage.cached_input_tokens
+    mm["cache_creation_input_tokens"] += usage.cache_creation_input_tokens
     mm["reasoning_tokens"] += usage.reasoning_tokens
     mm["request_count"] += 1
 

tests/unittest/test_token_metering.py

@@ -169,6 +169,24 @@ def test_gemini_openai_compat_format(self):
         assert usage.output_tokens == 9
         assert usage.reasoning_tokens == 26
 
+    def test_gemini_total_tokens_includes_thinking(self):
+        """Gemini total_tokens includes thoughts_tokens (differs from input+output)."""
+        usage = build_token_usage(
+            {
+                "prompt_tokens": 7,
+                "completion_tokens": 10,
+                "thoughts_tokens": 21,
+                "total_tokens": 38,
+            },
+            [],
+            "",
+            "gemini-2.5-flash",
+        )
+        assert usage.input_tokens == 7
+        assert usage.output_tokens == 10
+        assert usage.reasoning_tokens == 21
+        assert usage.total_tokens == 38  # API total includes thinking
+
     def test_gemini_native_format(self):
         """Gemini native: camelCase fields from usageMetadata."""
         usage = build_token_usage(
@@ -203,7 +221,7 @@ def test_doubao_format(self):
         assert usage.output_tokens == 27
 
     def test_cached_tokens(self):
-        """Extract cached_tokens from prompt_tokens_details."""
+        """Extract cached_input_tokens from prompt_tokens_details."""
         usage = build_token_usage(
             {
                 "prompt_tokens": 500,
@@ -214,7 +232,59 @@ def test_cached_tokens(self):
             "",
             "gpt-4o",
         )
-        assert usage.cached_tokens == 300
+        assert usage.cached_input_tokens == 300
+        assert usage.cache_creation_input_tokens == 0
+
+    def test_anthropic_format(self):
+        """Anthropic: top-level cache_read/cache_creation, no prompt_tokens_details."""
+        usage = build_token_usage(
+            {
+                "input_tokens": 50,
+                "output_tokens": 100,
+                "cache_read_input_tokens": 100000,
+                "cache_creation_input_tokens": 248,
+            },
+            [],
+            "",
+            "claude-sonnet-4-20250514",
+        )
+        assert usage.input_tokens == 50
+        assert usage.output_tokens == 100
+        assert usage.cached_input_tokens == 100000
+        assert usage.cache_creation_input_tokens == 248
+
+    def test_deepseek_cache_format(self):
+        """DeepSeek: top-level prompt_cache_hit_tokens."""
+        usage = build_token_usage(
+            {
+                "prompt_tokens": 405,
+                "total_tokens": 415,
+                "completion_tokens": 10,
+                "prompt_tokens_details": None,
+                "prompt_cache_hit_tokens": 380,
+                "reasoning_tokens": 42,
+            },
+            [],
+            "",
+            "deepseek-r1",
+        )
+        assert usage.cached_input_tokens == 380
+        assert usage.cache_creation_input_tokens == 0
+
+    def test_gemini_native_cache_format(self):
+        """Gemini native: cachedContentTokenCount。"""
+        usage = build_token_usage(
+            {
+                "promptTokenCount": 200,
+                "candidatesTokenCount": 15,
+                "thoughtsTokenCount": 30,
+                "cachedContentTokenCount": 150,
+            },
+            [],
+            "",
+            "gemini-2.5-flash",
+        )
+        assert usage.cached_input_tokens == 150
 
     def test_fallback_to_estimator(self):
         usage = build_token_usage(
@@ -284,30 +354,38 @@ def test_multi_model_accumulation(self):
 
     def test_cached_and_reasoning_accumulation(self):
         req = MagicMock(shared_data={})
-        aggregate_token_usage(
-            req,
-            TokenUsage(
-                input_tokens=100,
-                output_tokens=50,
-                cached_tokens=60,
-                reasoning_tokens=20,
-                model_name="deepseek-r1",
-            ),
+        u1 = build_token_usage(
+            {
+                "prompt_tokens": 100,
+                "completion_tokens": 50,
+                "prompt_tokens_details": {"cached_tokens": 60},
+                "cache_creation_input_tokens": 10,
+                "reasoning_tokens": 20,
+            },
+            [],
+            "",
+            "deepseek-r1",
         )
-        aggregate_token_usage(
-            req,
-            TokenUsage(
-                input_tokens=80,
-                output_tokens=30,
-                cached_tokens=40,
-                reasoning_tokens=10,
-                model_name="deepseek-r1",
-            ),
+        u2 = build_token_usage(
+            {
+                "prompt_tokens": 80,
+                "completion_tokens": 30,
+                "prompt_tokens_details": {"cached_tokens": 40},
+                "cache_creation_input_tokens": 5,
+                "reasoning_tokens": 10,
+            },
+            [],
+            "",
+            "deepseek-r1",
         )
+        aggregate_token_usage(req, u1)
+        aggregate_token_usage(req, u2)
         m = req.shared_data["_metrics"]["token_usage"]
-        assert m["cached_tokens"] == 100
+        assert m["cached_input_tokens"] == 100
+        assert m["cache_creation_input_tokens"] == 15
         assert m["reasoning_tokens"] == 30
-        assert m["by_model"]["deepseek-r1"]["cached_tokens"] == 100
+        assert m["by_model"]["deepseek-r1"]["cached_input_tokens"] == 100
+        assert m["by_model"]["deepseek-r1"]["cache_creation_input_tokens"] == 15
         assert m["by_model"]["deepseek-r1"]["reasoning_tokens"] == 30
 
     def test_disabled(self):
@@ -347,3 +425,28 @@ async def test_aggregates_token_usage(self):
         assert isinstance(resp.extra["usage"], dict)
         assert resp.extra["usage"]["total_tokens"] == 15
         assert req.shared_data["_metrics"]["token_usage"]["total_tokens"] == 15
+
+
+# ---------------------------------------------------------------------------
+# model_dump includes new cache fields
+# ---------------------------------------------------------------------------
+
+
+class TestModelDump:
+    def test_includes_cached_input_tokens(self):
+        usage = TokenUsage(
+            input_tokens=100,
+            output_tokens=50,
+            cached_input_tokens=80,
+            cache_creation_input_tokens=10,
+            model_name="gpt-4o",
+        )
+        data = usage.model_dump()
+        assert data["cached_input_tokens"] == 80
+        assert data["cache_creation_input_tokens"] == 10
+
+    def test_defaults_zero(self):
+        usage = TokenUsage(input_tokens=100, output_tokens=50)
+        data = usage.model_dump()
+        assert data["cached_input_tokens"] == 0
+        assert data["cache_creation_input_tokens"] == 0