Skip to content

Commit 6e108c1

Browse files
authored
feat(core): zero-out token costs for cache hits (#32437)
1 parent bc4251b commit 6e108c1

File tree

2 files changed

+30
-0
lines changed

2 files changed

+30
-0
lines changed

libs/core/langchain_core/language_models/chat_models.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -666,6 +666,16 @@ def _convert_cached_generations(self, cache_val: list) -> list[ChatGeneration]:
666666
converted_generations.append(chat_gen)
667667
else:
668668
# Already a ChatGeneration or other expected type
669+
if hasattr(gen, "message") and isinstance(gen.message, AIMessage):
670+
# We zero out cost on cache hits
671+
gen.message = gen.message.model_copy(
672+
update={
673+
"usage_metadata": {
674+
**(gen.message.usage_metadata or {}),
675+
"total_cost": 0,
676+
}
677+
}
678+
)
669679
converted_generations.append(gen)
670680
return converted_generations
671681

libs/core/tests/unit_tests/language_models/chat_models/test_cache.py

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -458,3 +458,23 @@ def test_cleanup_serialized() -> None:
458458
"name": "CustomChat",
459459
"type": "constructor",
460460
}
461+
462+
463+
def test_token_costs_are_zeroed_out() -> None:
464+
# We zero-out token costs for cache hits
465+
local_cache = InMemoryCache()
466+
messages = [
467+
AIMessage(
468+
content="Hello, how are you?",
469+
usage_metadata={"input_tokens": 5, "output_tokens": 10, "total_tokens": 15},
470+
),
471+
]
472+
model = GenericFakeChatModel(messages=iter(messages), cache=local_cache)
473+
first_response = model.invoke("Hello")
474+
assert isinstance(first_response, AIMessage)
475+
assert first_response.usage_metadata
476+
477+
second_response = model.invoke("Hello")
478+
assert isinstance(second_response, AIMessage)
479+
assert second_response.usage_metadata
480+
assert second_response.usage_metadata["total_cost"] == 0 # type: ignore[typeddict-item]

0 commit comments

Comments
 (0)