feat(core): automatically count and store meta for tool call count

libs/core/langchain_core/tracers/__init__.py

@@ -15,6 +15,11 @@
     )
     from langchain_core.tracers.schemas import Run
     from langchain_core.tracers.stdout import ConsoleCallbackHandler
+    from langchain_core.tracers.utils import (
+        count_tool_calls_in_run,
+        get_tool_call_count_from_run,
+        store_tool_call_count_in_run,
+    )
 
 __all__ = (
     "BaseTracer",
@@ -25,6 +30,9 @@
     "Run",
     "RunLog",
     "RunLogPatch",
+    "count_tool_calls_in_run",
+    "get_tool_call_count_from_run",
+    "store_tool_call_count_in_run",
 )
 
 _dynamic_imports = {
@@ -36,6 +44,9 @@
     "RunLogPatch": "log_stream",
     "Run": "schemas",
     "ConsoleCallbackHandler": "stdout",
+    "count_tool_calls_in_run": "utils",
+    "get_tool_call_count_from_run": "utils",
+    "store_tool_call_count_in_run": "utils",
 }
 
 

libs/core/langchain_core/tracers/base.py

@@ -37,13 +37,28 @@ class BaseTracer(_TracerCore, BaseCallbackHandler, ABC):
     def _persist_run(self, run: Run) -> None:
         """Persist a run."""
 
+    def _store_tool_call_metadata(self, run: Run) -> None:
+        """Store tool call count in run metadata automatically."""
+        try:
+            # Avoid circular imports
+            from langchain_core.tracers.utils import (  # noqa: PLC0415
+                store_tool_call_count_in_run,
+            )
+
+            store_tool_call_count_in_run(run)
+        except Exception:  # noqa: S110
+            # Avoid breaking existing functionality
+            pass
+
     def _start_trace(self, run: Run) -> None:
         """Start a trace for a run."""
         super()._start_trace(run)
         self._on_run_create(run)
 
     def _end_trace(self, run: Run) -> None:
         """End a trace for a run."""
+        self._store_tool_call_metadata(run)
+
         if not run.parent_run_id:
             self._persist_run(run)
         self.run_map.pop(str(run.id))
@@ -534,6 +549,19 @@ class AsyncBaseTracer(_TracerCore, AsyncCallbackHandler, ABC):
     async def _persist_run(self, run: Run) -> None:
         """Persist a run."""
 
+    async def _store_tool_call_metadata(self, run: Run) -> None:
+        """Store tool call count in run metadata."""
+        try:
+            # Avoid circular imports
+            from langchain_core.tracers.utils import (  # noqa: PLC0415
+                store_tool_call_count_in_run,
+            )
+
+            store_tool_call_count_in_run(run)
+        except Exception:  # noqa: S110
+            # Avoid breaking existing functionality
+            pass
+
     @override
     async def _start_trace(self, run: Run) -> None:
         """Start a trace for a run.
@@ -551,6 +579,8 @@ async def _end_trace(self, run: Run) -> None:
         Ending a trace will run concurrently with each _on_[run_type]_end method.
         No _on_[run_type]_end callback should depend on operations in _end_trace.
         """
+        await self._store_tool_call_metadata(run)
+
         if not run.parent_run_id:
             await self._persist_run(run)
         self.run_map.pop(str(run.id))

libs/core/langchain_core/tracers/utils.py

@@ -0,0 +1,91 @@
+"""Utility functions for working with Run objects and tracers."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from langchain_core.tracers.schemas import Run
+
+
+def count_tool_calls_in_run(run: Run) -> int:
+    """Count tool calls in a `Run` object by examining messages.
+
+    Args:
+        run: The `Run` object to examine.
+
+    Returns:
+        The total number of tool calls found in the run's messages.
+    """
+    tool_call_count = 0
+
+    # Check inputs for messages containing tool calls
+    inputs = getattr(run, "inputs", {}) or {}
+    if isinstance(inputs, dict) and "messages" in inputs:
+        messages = inputs["messages"]
+        if messages:
+            for msg in messages:
+                # Handle both dict and object representations
+                if hasattr(msg, "tool_calls"):
+                    tool_calls = getattr(msg, "tool_calls", [])
+                    if tool_calls:
+                        tool_call_count += len(tool_calls)
+                elif isinstance(msg, dict) and "tool_calls" in msg:
+                    tool_calls = msg.get("tool_calls", [])
+                    if tool_calls:
+                        tool_call_count += len(tool_calls)
+
+    # Also check outputs for completeness
+    outputs = getattr(run, "outputs", {}) or {}
+    if isinstance(outputs, dict) and "messages" in outputs:
+        messages = outputs["messages"]
+        if messages:
+            for msg in messages:
+                if hasattr(msg, "tool_calls"):
+                    tool_calls = getattr(msg, "tool_calls", [])
+                    if tool_calls:
+                        tool_call_count += len(tool_calls)
+                elif isinstance(msg, dict) and "tool_calls" in msg:
+                    tool_calls = msg.get("tool_calls", [])
+                    if tool_calls:
+                        tool_call_count += len(tool_calls)
+
+    return tool_call_count
+
+
+def store_tool_call_count_in_run(run: Run, *, always_store: bool = False) -> int:
+    """Count tool calls in a `Run` and store the count in run metadata.
+
+    Args:
+        run: The `Run` object to analyze and modify.
+        always_store: If `True`, always store the count even if `0`.
+            If `False`, only store when there are tool calls.
+
+    Returns:
+        The number of tool calls found and stored.
+    """
+    tool_call_count = count_tool_calls_in_run(run)
+
+    # Only store if there are tool calls or if explicitly requested
+    if tool_call_count > 0 or always_store:
+        # Store in run.extra for easy access
+        if not hasattr(run, "extra") or run.extra is None:
+            run.extra = {}
+        run.extra["tool_call_count"] = tool_call_count
+
+    return tool_call_count
+
+
+def get_tool_call_count_from_run(run: Run) -> int | None:
+    """Get the tool call count from run metadata if available.
+
+    Args:
+        run: The `Run` object to check.
+
+    Returns:
+        The tool call count if stored in metadata, otherwise `None`.
+    """
+    extra = getattr(run, "extra", {}) or {}
+    if isinstance(extra, dict):
+        return extra.get("tool_call_count")
+    return None

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

@@ -0,0 +1,127 @@
+"""Test automatic tool call count storage in tracers."""
+
+from __future__ import annotations
+
+from unittest.mock import MagicMock, PropertyMock
+
+from langchain_core.messages import AIMessage
+from langchain_core.messages.tool import ToolCall
+from langchain_core.tracers.base import BaseTracer
+from langchain_core.tracers.schemas import Run
+
+
+class MockTracer(BaseTracer):
+    """Mock tracer for testing automatic metadata storage."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        self.persisted_runs: list[Run] = []
+
+    def _persist_run(self, run: Run) -> None:
+        """Store the run for inspection."""
+        self.persisted_runs.append(run)
+
+
+def test_base_tracer_automatically_stores_tool_call_count() -> None:
+    """Test that `BaseTracer` automatically stores tool call count."""
+    tracer = MockTracer()
+
+    # Create a mock run with tool calls
+    run = MagicMock(spec=Run)
+    run.id = "test-run-id"
+    run.parent_run_id = None  # Root run, will be persisted
+    run.extra = {}
+
+    # Set up messages with tool calls
+    tool_calls = [
+        ToolCall(name="search", args={"query": "test"}, id="call_1"),
+        ToolCall(name="calculator", args={"expression": "2+2"}, id="call_2"),
+    ]
+    messages = [AIMessage(content="Test", tool_calls=tool_calls)]
+    run.inputs = {"messages": messages}
+    run.outputs = {}
+
+    # Add run to tracer's run_map to simulate it being tracked
+    tracer.run_map[str(run.id)] = run
+
+    # End the trace (this should trigger automatic metadata storage)
+    tracer._end_trace(run)
+
+    # Verify tool call count was automatically stored
+    assert "tool_call_count" in run.extra
+    assert run.extra["tool_call_count"] == 2
+
+    # Verify the run was persisted
+    assert len(tracer.persisted_runs) == 1
+    assert tracer.persisted_runs[0] == run
+
+
+def test_base_tracer_handles_no_tool_calls() -> None:
+    """Test that `BaseTracer` handles runs with no tool calls gracefully."""
+    tracer = MockTracer()
+
+    # Create a mock run without tool calls
+    run = MagicMock(spec=Run)
+    run.id = "test-run-id-no-tools"
+    run.parent_run_id = None
+    run.extra = {}
+
+    # Set up messages without tool calls
+    messages = [AIMessage(content="No tools here")]
+    run.inputs = {"messages": messages}
+    run.outputs = {}
+
+    # Add run to tracer's run_map
+    tracer.run_map[str(run.id)] = run
+
+    # End the trace
+    tracer._end_trace(run)
+
+    # Verify tool call count is not stored when there are no tool calls
+    assert "tool_call_count" not in run.extra
+
+
+def test_base_tracer_handles_runs_without_messages() -> None:
+    """Test that `BaseTracer` handles runs without messages gracefully."""
+    tracer = MockTracer()
+
+    # Create a mock run without messages
+    run = MagicMock(spec=Run)
+    run.id = "test-run-id-no-messages"
+    run.parent_run_id = None
+    run.extra = {}
+    run.inputs = {}
+    run.outputs = {}
+
+    # Add run to tracer's run_map
+    tracer.run_map[str(run.id)] = run
+
+    # End the trace
+    tracer._end_trace(run)
+
+    # Verify tool call count is not stored when there are no messages
+    assert "tool_call_count" not in run.extra
+
+
+def test_base_tracer_doesnt_break_on_metadata_error() -> None:
+    """Test that `BaseTracer` continues working if metadata storage fails."""
+    tracer = MockTracer()
+
+    # Create a mock run that will cause an error in tool call counting
+    run = MagicMock(spec=Run)
+    run.id = "test-run-id-error"
+    run.parent_run_id = None
+    run.extra = {}
+
+    # Make the run.inputs property raise an error when accessed
+    type(run).inputs = PropertyMock(side_effect=RuntimeError("Simulated error"))
+
+    # Add run to tracer's run_map
+    tracer.run_map[str(run.id)] = run
+
+    # End the trace - this should not raise an exception
+    tracer._end_trace(run)
+
+    # The run should still be persisted despite the metadata error
+    assert len(tracer.persisted_runs) == 1
+    assert tracer.persisted_runs[0] == run

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

@@ -2,13 +2,16 @@
 
 EXPECTED_ALL = [
     "BaseTracer",
+    "ConsoleCallbackHandler",
     "EvaluatorCallbackHandler",
     "LangChainTracer",
-    "ConsoleCallbackHandler",
+    "LogStreamCallbackHandler",
     "Run",
     "RunLog",
     "RunLogPatch",
-    "LogStreamCallbackHandler",
+    "count_tool_calls_in_run",
+    "get_tool_call_count_from_run",
+    "store_tool_call_count_in_run",
 ]