codekingpro/portable-devtools

1"""Callback Handler that tracks `AIMessage.usage_metadata`."""2 3import threading4from collections.abc import Generator5from contextlib import contextmanager6from contextvars import ContextVar7from typing import Any8 9from typing_extensions import override10 11from 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_hook16 17 18class UsageMetadataCallbackHandler(BaseCallbackHandler):19    """Callback Handler that tracks `AIMessage.usage_metadata`.20 21    Example:22        ```python23        from langchain.chat_models import init_chat_model24        from langchain_core.callbacks import UsageMetadataCallbackHandler25 26        llm_1 = init_chat_model(model="openai:gpt-4o-mini")27        llm_2 = init_chat_model(model="anthropic:claude-haiku-4-5-20251001")28 29        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        ```34 35        ```txt36        {'gpt-4o-mini-2024-07-18': {'input_tokens': 8,37          'output_tokens': 10,38          'total_tokens': 18,39          'input_token_details': {'audio': 0, 'cache_read': 0},40          'output_token_details': {'audio': 0, 'reasoning': 0}},41         'claude-haiku-4-5-20251001': {'input_tokens': 8,42          'output_tokens': 21,43          'total_tokens': 29,44          'input_token_details': {'cache_read': 0, 'cache_creation': 0}}}45        ```46 47    !!! version-added "Added in `langchain-core` 0.3.49"48 49    """50 51    def __init__(self) -> None:52        """Initialize the `UsageMetadataCallbackHandler`."""53        super().__init__()54        self._lock = threading.Lock()55        self.usage_metadata: dict[str, UsageMetadata] = {}56 57    @override58    def __repr__(self) -> str:59        return str(self.usage_metadata)60 61    @override62    def on_llm_end(self, response: LLMResult, **kwargs: Any) -> None:63        """Collect token usage."""64        # Check for usage_metadata (langchain-core >= 0.2.2)65        try:66            generation = response.generations[0][0]67        except IndexError:68            generation = None69 70        usage_metadata = None71        model_name = None72        if isinstance(generation, ChatGeneration):73            try:74                message = generation.message75                if isinstance(message, AIMessage):76                    usage_metadata = message.usage_metadata77                    model_name = message.response_metadata.get("model_name")78            except AttributeError:79                pass80 81        # update shared state behind lock82        if usage_metadata and model_name:83            with self._lock:84                if model_name not in self.usage_metadata:85                    self.usage_metadata[model_name] = usage_metadata86                else:87                    self.usage_metadata[model_name] = add_usage(88                        self.usage_metadata[model_name], usage_metadata89                    )90 91 92@contextmanager93def get_usage_metadata_callback(94    name: str = "usage_metadata_callback",95) -> Generator[UsageMetadataCallbackHandler, None, None]:96    """Get usage metadata callback.97 98    Get context manager for tracking usage metadata across chat model calls using99    [`AIMessage.usage_metadata`][langchain.messages.AIMessage.usage_metadata].100 101    Args:102        name: The name of the context variable.103 104    Yields:105        The usage metadata callback.106 107    Example:108        ```python109        from langchain.chat_models import init_chat_model110        from langchain_core.callbacks import get_usage_metadata_callback111 112        llm_1 = init_chat_model(model="openai:gpt-4o-mini")113        llm_2 = init_chat_model(model="anthropic:claude-haiku-4-5-20251001")114 115        with get_usage_metadata_callback() as cb:116            llm_1.invoke("Hello")117            llm_2.invoke("Hello")118            print(cb.usage_metadata)119        ```120 121        ```txt122        {123            "gpt-4o-mini-2024-07-18": {124                "input_tokens": 8,125                "output_tokens": 10,126                "total_tokens": 18,127                "input_token_details": {"audio": 0, "cache_read": 0},128                "output_token_details": {"audio": 0, "reasoning": 0},129            },130            "claude-haiku-4-5-20251001": {131                "input_tokens": 8,132                "output_tokens": 21,133                "total_tokens": 29,134                "input_token_details": {"cache_read": 0, "cache_creation": 0},135            },136        }137        ```138 139    !!! version-added "Added in `langchain-core` 0.3.49"140 141    """142    usage_metadata_callback_var: ContextVar[UsageMetadataCallbackHandler | None] = (143        ContextVar(name, default=None)144    )145    register_configure_hook(usage_metadata_callback_var, inheritable=True)146    cb = UsageMetadataCallbackHandler()147    usage_metadata_callback_var.set(cb)148    yield cb149    usage_metadata_callback_var.set(None)150