aggregate usage metadata

Author: hntrl

libs/core/langchain_core/tracers/langchain.py

@@ -21,6 +21,7 @@
 
 from langchain_core.env import get_runtime_environment
 from langchain_core.load import dumpd
+from langchain_core.messages.ai import add_usage
 from langchain_core.tracers.base import BaseTracer
 from langchain_core.tracers.schemas import Run
 
@@ -72,26 +73,27 @@ def _get_executor() -> ThreadPoolExecutor:
 def _get_usage_metadata_from_generations(
     generations: list[list[dict[str, Any]]],
 ) -> dict[str, Any] | None:
-    """Extract usage_metadata from generations.
+    """Extract and aggregate usage_metadata from generations.
 
-    Iterates through generations to find and return the first usage_metadata
-    found in a message. This is typically present in chat model outputs.
+    Iterates through generations to find and aggregate all usage_metadata
+    found in messages. This is typically present in chat model outputs.
 
     Args:
         generations: List of generation batches, where each batch is a list
             of generation dicts that may contain a "message" key with
             "usage_metadata".
 
     Returns:
-        The usage_metadata dict if found, otherwise None.
+        The aggregated usage_metadata dict if found, otherwise None.
     """
+    output: dict[str, Any] | None = None
     for generation_batch in generations:
         for generation in generation_batch:
             if isinstance(generation, dict) and "message" in generation:
                 message = generation["message"]
                 if isinstance(message, dict) and "usage_metadata" in message:
-                    return message["usage_metadata"]
-    return None
+                    output = add_usage(output, message["usage_metadata"])
+    return output
 
 
 class LangChainTracer(BaseTracer):

libs/core/tests/unit_tests/tracers/test_langchain.py

@@ -179,7 +179,7 @@ def mock_create_run(**kwargs: Any) -> Any:
         # Returns None for empty generations
         ([], None),
         ([[]], None),
-        # Returns first usage_metadata from multiple generations
+        # Aggregates usage_metadata across multiple generations
         (
             [
                 [
@@ -207,7 +207,7 @@ def mock_create_run(**kwargs: Any) -> Any:
                     },
                 ]
             ],
-            {"input_tokens": 5, "output_tokens": 10, "total_tokens": 15},
+            {"input_tokens": 55, "output_tokens": 110, "total_tokens": 165},
         ),
         # Finds usage_metadata across multiple batches
         (
@@ -236,7 +236,7 @@ def mock_create_run(**kwargs: Any) -> Any:
         "returns_none_when_no_message",
         "returns_none_for_empty_list",
         "returns_none_for_empty_batch",
-        "returns_first_from_multiple_generations",
+        "aggregates_across_multiple_generations",
         "finds_across_multiple_batches",
     ],
 )