From 68bc11e33fb63b16caf403f1cab3e4840939a2cb Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 9 Feb 2026 13:29:34 +0000 Subject: [PATCH 1/2] Fix priority tier token detail parsing Co-authored-by: Hassieb Pakzad --- langfuse/langchain/CallbackHandler.py | 22 ++++++++++-- tests/test_langchain_usage.py | 49 +++++++++++++++++++++++++++ 2 files changed, 69 insertions(+), 2 deletions(-) create mode 100644 tests/test_langchain_usage.py diff --git a/langfuse/langchain/CallbackHandler.py b/langfuse/langchain/CallbackHandler.py index a2e9816da..b853eddf7 100644 --- a/langfuse/langchain/CallbackHandler.py +++ b/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 diff --git a/tests/test_langchain_usage.py b/tests/test_langchain_usage.py new file mode 100644 index 000000000..e174600fd --- /dev/null +++ b/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 From 04fca344c56bf58446f33855be841356872381ba Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 9 Feb 2026 13:33:40 +0000 Subject: [PATCH 2/2] Skip priority token details when subtracting Co-authored-by: Hassieb Pakzad --- langfuse/langchain/CallbackHandler.py | 15 +-------------- tests/test_langchain_usage.py | 8 ++++++-- 2 files changed, 7 insertions(+), 16 deletions(-) diff --git a/langfuse/langchain/CallbackHandler.py b/langfuse/langchain/CallbackHandler.py index b853eddf7..dbff2ce6b 100644 --- a/langfuse/langchain/CallbackHandler.py +++ b/langfuse/langchain/CallbackHandler.py @@ -1102,22 +1102,9 @@ 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 + return not normalized_key.startswith("priority") def _parse_usage_model(usage: Union[pydantic.BaseModel, dict]) -> Any: diff --git a/tests/test_langchain_usage.py b/tests/test_langchain_usage.py index e174600fd..91f87c1cd 100644 --- a/tests/test_langchain_usage.py +++ b/tests/test_langchain_usage.py @@ -33,17 +33,21 @@ def test_parse_usage_model_subtracts_known_details(): "input_token_details": { "cache_read": 20, "audio": 5, + "custom_detail": 3, }, "output_token_details": { "reasoning": 10, + "custom_output": 2, }, } parsed = _parse_usage_model(usage) - assert parsed["input"] == 75 - assert parsed["output"] == 40 + assert parsed["input"] == 72 + assert parsed["output"] == 38 assert parsed["input_cache_read"] == 20 assert parsed["input_audio"] == 5 + assert parsed["input_custom_detail"] == 3 assert parsed["output_reasoning"] == 10 + assert parsed["output_custom_output"] == 2 assert parsed["total"] == 150