From 5daa84567f1970cbf8ed0cb10cc1f01f4dde2080 Mon Sep 17 00:00:00 2001
From: Viliam Gonda <gonda.v@gjh.sk>
Date: Sat, 31 Oct 2026 11:35:58 +0100
Subject: [PATCH] fix(langchain): avoid double counting Anthropic prompt cache
 creation tokens

---
 langfuse/langchain/CallbackHandler.py | 29 ++++++++++++------
 tests/unit/test_parse_usage_model.py  | 43 ++++++++++++++++++++-------
 2 files changed, 52 insertions(+), 20 deletions(-)

diff --git a/langfuse/langchain/CallbackHandler.py b/langfuse/langchain/CallbackHandler.py
index dac76f4fc..596fa07f5 100644
--- a/langfuse/langchain/CallbackHandler.py
+++ b/langfuse/langchain/CallbackHandler.py
@@ -1749,21 +1749,32 @@ def _parse_usage_model(usage: Union[pydantic.BaseModel, dict]) -> Any:
             usage_model["cache_creation"], dict
         ):
             cache_creation = usage_model.pop("cache_creation")
-            total_cache_creation_tokens = 0
+            tier_key_map = {
+                "ephemeral_5m_input_tokens": "input_cache_creation_5m",
+                "ephemeral_1h_input_tokens": "input_cache_creation_1h",
+            }
+            attributed = 0
 
             for key, value in cache_creation.items():
                 if not isinstance(value, int):
                     continue
 
-                usage_model[f"cache_creation_{key}"] = value
-                total_cache_creation_tokens += value
+                langfuse_key = tier_key_map.get(key, f"input_cache_creation_{key}")
+                usage_model[langfuse_key] = value
+                if value > 0:
+                    attributed += value
 
-            # Aggregate mirrors Anthropic's legacy cache_creation_input_tokens field;
-            # keep the API-provided value if it is already present
-            if total_cache_creation_tokens > 0:
-                usage_model.setdefault(
-                    "cache_creation_input_tokens", total_cache_creation_tokens
-                )
+            if attributed > 0:
+                for aggregate_key in (
+                    "cache_creation_input_tokens",
+                    "input_cache_creation",
+                ):
+                    if aggregate_key in usage_model:
+                        remainder = usage_model[aggregate_key] - attributed
+                        if remainder > 0:
+                            usage_model[aggregate_key] = remainder
+                        else:
+                            del usage_model[aggregate_key]
 
     usage_model = {k: v for k, v in usage_model.items() if isinstance(v, int)}
 
diff --git a/tests/unit/test_parse_usage_model.py b/tests/unit/test_parse_usage_model.py
index fd39ad658..2409c285f 100644
--- a/tests/unit/test_parse_usage_model.py
+++ b/tests/unit/test_parse_usage_model.py
@@ -92,28 +92,28 @@ def test_anthropic_cache_creation_nested_dict():
     result = _parse_usage_model(usage)
     assert result["input"] == 9454
     assert result["output"] == 380
-    assert result["cache_creation_ephemeral_1h_input_tokens"] == 0
-    assert result["cache_creation_ephemeral_5m_input_tokens"] == 2048
-    assert result["cache_creation_input_tokens"] == 2048
+    assert result["input_cache_creation_1h"] == 0
+    assert result["input_cache_creation_5m"] == 2048
+    assert "cache_creation_input_tokens" not in result
     # No nested dict may survive into the final usage model
     assert all(isinstance(v, int) for v in result.values())
 
 
-def test_anthropic_cache_creation_keeps_existing_aggregate():
-    """API-provided cache_creation_input_tokens is not overwritten by the computed sum."""
+def test_anthropic_cache_creation_keeps_unattributed_remainder():
+    """Unattributed cache creation tokens remain on cache_creation_input_tokens."""
     usage = {
         "input_tokens": 100,
         "output_tokens": 10,
         "cache_creation_input_tokens": 3000,
         "cache_creation": {
             "ephemeral_1h_input_tokens": 1000,
-            "ephemeral_5m_input_tokens": 2000,
+            "ephemeral_5m_input_tokens": 1500,
         },
     }
     result = _parse_usage_model(usage)
-    assert result["cache_creation_input_tokens"] == 3000
-    assert result["cache_creation_ephemeral_1h_input_tokens"] == 1000
-    assert result["cache_creation_ephemeral_5m_input_tokens"] == 2000
+    assert result["cache_creation_input_tokens"] == 500
+    assert result["input_cache_creation_1h"] == 1000
+    assert result["input_cache_creation_5m"] == 1500
 
 
 def test_anthropic_cache_creation_all_zero():
@@ -127,8 +127,29 @@ def test_anthropic_cache_creation_all_zero():
         },
     }
     result = _parse_usage_model(usage)
-    assert result["cache_creation_ephemeral_1h_input_tokens"] == 0
-    assert result["cache_creation_ephemeral_5m_input_tokens"] == 0
+    assert result["input_cache_creation_1h"] == 0
+    assert result["input_cache_creation_5m"] == 0
+    assert "cache_creation_input_tokens" not in result
+
+
+def test_anthropic_cache_creation_deduplicates_aggregate():
+    """Anthropic Messages API usage de-duplicates aggregate cache_creation_input_tokens."""
+    raw = {
+        "input_tokens": 4,
+        "output_tokens": 153,
+        "cache_read_input_tokens": 0,
+        "cache_creation_input_tokens": 5544,
+        "cache_creation": {
+            "ephemeral_5m_input_tokens": 5544,
+            "ephemeral_1h_input_tokens": 0,
+        },
+    }
+    result = _parse_usage_model(raw)
+    assert result["input"] == 4
+    assert result["output"] == 153
+    assert result["cache_read_input_tokens"] == 0
+    assert result["input_cache_creation_5m"] == 5544
+    assert result["input_cache_creation_1h"] == 0
     assert "cache_creation_input_tokens" not in result