From ffa0931671aaf6c499159bf913e880075645ebb5 Mon Sep 17 00:00:00 2001 From: chelsealong Date: Wed, 28 Oct 2026 23:56:17 +0000 Subject: [PATCH] fix(langchain): avoid double counting Anthropic cache writes --- langfuse/langchain/CallbackHandler.py | 40 +++++++++++++---------- tests/unit/test_parse_usage_model.py | 47 +++++++++++++++++---------- 2 files changed, 53 insertions(+), 34 deletions(-) diff --git a/langfuse/langchain/CallbackHandler.py b/langfuse/langchain/CallbackHandler.py index dac76f4fc..f593df099 100644 --- a/langfuse/langchain/CallbackHandler.py +++ b/langfuse/langchain/CallbackHandler.py @@ -1743,27 +1743,33 @@ def _parse_usage_model(usage: Union[pydantic.BaseModel, dict]) -> Any: 0, usage_model[f"input_modality_{item['modality']}"] - value ) - # Anthropic extended prompt caching — cache_creation is a nested dict - # keyed by cache tier, e.g. {"ephemeral_5m_input_tokens": 2048, "ephemeral_1h_input_tokens": 0} + # Anthropic cache_creation tiers are parts of cache_creation_input_tokens. + # Keep only the unattributed remainder to avoid counting writes twice. if "cache_creation" in usage_model and isinstance( usage_model["cache_creation"], dict ): cache_creation = usage_model.pop("cache_creation") - total_cache_creation_tokens = 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 - - # 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 - ) + attributed = 0 + has_tier = False + for tier, output_key in ( + ("ephemeral_5m_input_tokens", "input_cache_creation_5m"), + ("ephemeral_1h_input_tokens", "input_cache_creation_1h"), + ): + value = cache_creation.get(tier) + if isinstance(value, int) and value >= 0: + has_tier = True + usage_model[output_key] = value + attributed += value + + if has_tier: + aggregate = usage_model.get("cache_creation_input_tokens", 0) + remainder = ( + aggregate if isinstance(aggregate, int) else 0 + ) - attributed + if remainder > 0: + usage_model["cache_creation_input_tokens"] = remainder + else: + usage_model.pop("cache_creation_input_tokens", None) 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..abe76277b 100644 --- a/tests/unit/test_parse_usage_model.py +++ b/tests/unit/test_parse_usage_model.py @@ -79,45 +79,58 @@ def test_prompt_tokens_details_dict_empty(): def test_anthropic_cache_creation_nested_dict(): - """Anthropic extended prompt caching: cache_creation nested dict is flattened.""" + """Tiered cache writes contribute to usage exactly once.""" usage = { - "input_tokens": 9454, - "output_tokens": 380, + "input_tokens": 4, + "output_tokens": 153, "cache_read_input_tokens": 0, + "cache_creation_input_tokens": 5544, "cache_creation": { "ephemeral_1h_input_tokens": 0, - "ephemeral_5m_input_tokens": 2048, + "ephemeral_5m_input_tokens": 5544, }, } 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"] == 4 + assert result["output"] == 153 + assert result["input_cache_creation_1h"] == 0 + assert result["input_cache_creation_5m"] == 5544 + assert "cache_creation_input_tokens" not in result + assert sum(result.values()) == 5701 # 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.""" + """Only the unattributed cache writes remain in the aggregate bucket.""" usage = { "input_tokens": 100, "output_tokens": 10, - "cache_creation_input_tokens": 3000, + "cache_creation_input_tokens": 3500, "cache_creation": { "ephemeral_1h_input_tokens": 1000, "ephemeral_5m_input_tokens": 2000, }, } 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"] == 2000 + assert sum(result.values()) == 3610 + + +def test_anthropic_cache_creation_without_aggregate(): + usage = { + "input_tokens": 7, + "output_tokens": 3, + "cache_creation": {"ephemeral_5m_input_tokens": 20}, + } + result = _parse_usage_model(usage) + assert result == {"input": 7, "output": 3, "input_cache_creation_5m": 20} def test_anthropic_cache_creation_all_zero(): - """All-zero cache_creation tiers: flattened keys kept, no aggregate invented.""" + """All-zero cache_creation tiers do not invent an aggregate.""" usage = { "input_tokens": 50, "output_tokens": 5, @@ -127,8 +140,8 @@ 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