1"""Callback Handler that tracks `AIMessage.usage_metadata`."""23import threading4from collections.abc import Generator5from contextlib import contextmanager6from contextvars import ContextVar7from typing import Any89from typing_extensions import override1011from langchain_core.callbacks import BaseCallbackHandler12from langchain_core.messages import AIMessage13from langchain_core.messages.ai import UsageMetadata, add_usage14from langchain_core.outputs import ChatGeneration, LLMResult15from langchain_core.tracers.context import register_configure_hook161718class UsageMetadataCallbackHandler(BaseCallbackHandler):19    """Callback Handler that tracks `AIMessage.usage_metadata`.2021    Example:22        ```python23        from langchain.chat_models import init_chat_model24        from langchain_core.callbacks import UsageMetadataCallbackHandler2526        llm_1 = init_chat_model(model="openai:gpt-5.5")27        llm_2 = init_chat_model(model="anthropic:claude-haiku-4-5-20251001")2829        callback = UsageMetadataCallbackHandler()30        result_1 = llm_1.invoke("Hello", config={"callbacks": [callback]})31        result_2 = llm_2.invoke("Hello", config={"callbacks": [callback]})32        callback.usage_metadata33        ```3435    !!! version-added "Added in `langchain-core` 0.3.49"3637    """3839    def __init__(self) -> None:40        """Initialize the `UsageMetadataCallbackHandler`."""41        super().__init__()42        self._lock = threading.Lock()43        self.usage_metadata: dict[str, UsageMetadata] = {}4445    @override46    def __repr__(self) -> str:47        return str(self.usage_metadata)4849    @override50    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:51        """Collect token usage."""52        # Check for usage_metadata (langchain-core >= 0.2.2)53        try:54            generation = response.generations[0][0]55        except IndexError:56            generation = None5758        usage_metadata = None59        model_name = None60        if isinstance(generation, ChatGeneration):61            try:62                message = generation.message63                if isinstance(message, AIMessage):64                    usage_metadata = message.usage_metadata65                    model_name = message.response_metadata.get("model_name")66            except AttributeError:67                pass6869        # update shared state behind lock70        if usage_metadata and model_name:71            with self._lock:72                if model_name not in self.usage_metadata:73                    self.usage_metadata[model_name] = usage_metadata74                else:75                    self.usage_metadata[model_name] = add_usage(76                        self.usage_metadata[model_name], usage_metadata77                    )787980@contextmanager81def get_usage_metadata_callback(82    name: str = "usage_metadata_callback",83) -> Generator[UsageMetadataCallbackHandler, None, None]:84    """Get usage metadata callback.8586    Get context manager for tracking usage metadata across chat model calls using87    [`AIMessage.usage_metadata`][langchain.messages.AIMessage.usage_metadata].8889    Args:90        name: The name of the context variable.9192    Yields:93        The usage metadata callback.9495    Example:96        ```python97        from langchain.chat_models import init_chat_model98        from langchain_core.callbacks import get_usage_metadata_callback99100        llm_1 = init_chat_model(model="openai:gpt-5.5")101        llm_2 = init_chat_model(model="anthropic:claude-haiku-4-5-20251001")102103        with get_usage_metadata_callback() as cb:104            llm_1.invoke("Hello")105            llm_2.invoke("Hello")106            print(cb.usage_metadata)107        ```108109    !!! version-added "Added in `langchain-core` 0.3.49"110111    """112    usage_metadata_callback_var: ContextVar[UsageMetadataCallbackHandler | None] = (113        ContextVar(name, default=None)114    )115    register_configure_hook(usage_metadata_callback_var, inheritable=True)116    cb = UsageMetadataCallbackHandler()117    usage_metadata_callback_var.set(cb)118    yield cb119    usage_metadata_callback_var.set(None)