Fix priority tier token detail parsing

Co-authored-by: Hassieb Pakzad <hassiebp@users.noreply.github.com>

Author: cursoragent

langfuse/langchain/CallbackHandler.py

@@ -1102,6 +1102,24 @@ def _flatten_comprehension(matrix: Any) -> Any:
     return [item for row in matrix for item in row]
 
 
+_TOKEN_DETAIL_SUBTRACT_KEYS = {
+    "audio",
+    "cache_read",
+    "cache_creation",
+    "reasoning",
+}
+
+
+def _should_subtract_token_detail(detail_key: str) -> bool:
+    normalized_key = detail_key.lower()
+    for subtract_key in _TOKEN_DETAIL_SUBTRACT_KEYS:
+        if normalized_key == subtract_key or normalized_key.endswith(
+            f"_{subtract_key}"
+        ):
+            return True
+    return False
+
+
 def _parse_usage_model(usage: Union[pydantic.BaseModel, dict]) -> Any:
     # maintains a list of key translations. For each key, the usage model is checked
     # and a new object will be created with the new key if the key exists in the usage model
@@ -1177,7 +1195,7 @@ def _parse_usage_model(usage: Union[pydantic.BaseModel, dict]) -> Any:
             for key, value in input_token_details.items():
                 usage_model[f"input_{key}"] = value
 
-                if "input" in usage_model:
+                if "input" in usage_model and _should_subtract_token_detail(key):
                     usage_model["input"] = max(0, usage_model["input"] - value)
 
         if "output_token_details" in usage_model:
@@ -1186,7 +1204,7 @@ def _parse_usage_model(usage: Union[pydantic.BaseModel, dict]) -> Any:
             for key, value in output_token_details.items():
                 usage_model[f"output_{key}"] = value
 
-                if "output" in usage_model:
+                if "output" in usage_model and _should_subtract_token_detail(key):
                     usage_model["output"] = max(0, usage_model["output"] - value)
 
         # Vertex AI

tests/test_langchain_usage.py

@@ -0,0 +1,49 @@
+from langfuse.langchain.CallbackHandler import _parse_usage_model
+
+
+def test_parse_usage_model_skips_priority_subtraction():
+    usage = {
+        "input": 13,
+        "output": 1,
+        "total": 14,
+        "input_token_details": {
+            "audio": 0,
+            "priority_cache_read": 0,
+            "priority": 13,
+        },
+        "output_token_details": {
+            "audio": 0,
+            "priority_reasoning": 0,
+            "priority": 1,
+        },
+    }
+
+    parsed = _parse_usage_model(usage)
+
+    assert parsed["input"] == 13
+    assert parsed["output"] == 1
+    assert parsed["total"] == 14
+
+
+def test_parse_usage_model_subtracts_known_details():
+    usage = {
+        "input": 100,
+        "output": 50,
+        "total": 150,
+        "input_token_details": {
+            "cache_read": 20,
+            "audio": 5,
+        },
+        "output_token_details": {
+            "reasoning": 10,
+        },
+    }
+
+    parsed = _parse_usage_model(usage)
+
+    assert parsed["input"] == 75
+    assert parsed["output"] == 40
+    assert parsed["input_cache_read"] == 20
+    assert parsed["input_audio"] == 5
+    assert parsed["output_reasoning"] == 10
+    assert parsed["total"] == 150