AI & Agent Dev Bug Sandbox logo
AI & Agent Dev Bug Sandbox
Back to Radar

UsageMetadataCallbackHandler Silently Drops Usage_metadata When Model_name Is Missing

The UsageMetadataCallbackHandler in langchain_core only records token usage when the response's response_metadata contains a 'model_name' key. If model_name is absent, usage_metadata is silently discarded, defeating the purpose of the callback and causing silent data loss in usage tracking for custom or third-party chat models.

mediumConfidence 95%Langchain-CoreAffected V1.6.1

Origin Analysis

In langchain_core/callbacks/usage.py, the on_llm_end method checks `if usage_metadata and model_name:` before recording usage. Since model_name is not guaranteed to be present in response_metadata (especially for custom BaseChatModel subclasses or some integrations), usage_metadata is dropped entirely when model_name is missing.
Run the following Python code: ```python from langchain_core.callbacks import get_usage_metadata_callback from langchain_core.language_models import GenericFakeChatModel from langchain_core.messages import AIMessage from langchain_core.messages.ai import UsageMetadata usage = UsageMetadata(input_tokens=1, output_tokens=2, total_tokens=3) llm = GenericFakeChatModel(messages=iter([AIMessage("hi", usage_metadata=usage)])) with get_usage_metadata_callback() as cb: llm.invoke("hello") print(cb.usage_metadata) # Output is {} instead of containing usage info ```

Fixing Code Block

def on_llm_end(self, response, **kwargs): usage_metadata = getattr(response, "usage_metadata", None) model_name = getattr(response, "response_metadata", {}).get("model_name") or "unknown" if usage_metadata: self.usage_metadata[model_name] = { "input_tokens": usage_metadata.get("input_tokens", 0), "output_tokens": usage_metadata.get("output_tokens", 0), "total_tokens": usage_metadata.get("total_tokens", 0), }
The condition is changed to only require `usage_metadata` to be truthy. The `model_name` is determined using `get("model_name") or "unknown"`, so missing or None model_name falls back to the string "unknown". This ensures usage is always recorded, even when model_name is not provided, without affecting models that do provide it.

Edge Case Audit

This change may introduce an 'unknown' aggregation bucket for models that previously went untracked, which is the intended behavior but could surprise users who expected only known model names. It is backward compatible for existing cases but users should review any downstream analytics that may be sensitive to the presence of an 'unknown' key. Rollback is straightforward by reverting to the original conditional. No concurrency or threading issues are introduced because the callback instance is typically used within a single context and not shared across threads.

Ecosystem Topology