feat!: Split track_metrics_of into sync and async variants (#112)

feat: Add optional graph_key to all LDAIConfigTracker track_* methods for graph correlation
feat: Add track_tool_call/track_tool_calls to LDAIConfigTracker
fix: make AIGraphTracker.track_total_tokens accept Optional[TokenUsage], skip when None or total <= 0
feat: Add get_tool_calls_from_response and sum_token_usage_from_messages to langchain_helper
feat: Add get_ai_usage_from_response to openai_helper
fix!: Remove node-scoped methods from AIGraphTracker (track_node_invocation, track_tool_call, track_node_judge_response), use related AIConfigTracker methods instead
fix: use time.perf_counter_ns() for sub-millisecond precision in duration calculations

Author: jsonbailey

packages/ai-providers/server-ai-langchain/README.md

@@ -138,7 +138,7 @@ provider = await LangChainProvider.create(config)
 async def invoke():
     return await provider.invoke_model(messages)
 
-response = await config.tracker.track_metrics_of(
+response = await config.tracker.track_metrics_of_async(
     invoke,
     lambda r: r.metrics
 )

packages/ai-providers/server-ai-langchain/src/ldai_langchain/__init__.py

@@ -3,7 +3,9 @@
     create_langchain_model,
     get_ai_metrics_from_response,
     get_ai_usage_from_response,
+    get_tool_calls_from_response,
     map_provider,
+    sum_token_usage_from_messages,
 )
 from ldai_langchain.langchain_model_runner import LangChainModelRunner
 from ldai_langchain.langchain_runner_factory import LangChainRunnerFactory
@@ -18,5 +20,7 @@
     'create_langchain_model',
     'get_ai_metrics_from_response',
     'get_ai_usage_from_response',
+    'get_tool_calls_from_response',
     'map_provider',
+    'sum_token_usage_from_messages',
 ]

packages/ai-providers/server-ai-langchain/src/ldai_langchain/langchain_helper.py

@@ -115,3 +115,41 @@ def get_ai_metrics_from_response(response: Any) -> LDAIMetrics:
     :return: LDAIMetrics with success status and token usage
     """
     return LDAIMetrics(success=True, usage=get_ai_usage_from_response(response))
+
+
+def get_tool_calls_from_response(response: Any) -> List[str]:
+    """
+    Get tool call names from a LangChain provider response.
+
+    :param response: The response from the LangChain model
+    :return: List of tool names in order, or empty list if none
+    """
+    names: List[str] = []
+    if hasattr(response, 'tool_calls') and isinstance(response.tool_calls, list):
+        for tc in response.tool_calls:
+            n = tc.get('name')
+            if n:
+                names.append(str(n))
+    return names
+
+
+def sum_token_usage_from_messages(messages: List[Any]) -> Optional[TokenUsage]:
+    """
+    Sum token usage across LangChain messages using get_ai_usage_from_response per message.
+
+    :param messages: List of message objects (e.g. from a graph state)
+    :return: Aggregated TokenUsage, or None if no usage on any message
+    """
+    in_sum = 0
+    out_sum = 0
+    total_sum = 0
+    for m in messages:
+        u = get_ai_usage_from_response(m)
+        if u is None:
+            continue
+        in_sum += u.input
+        out_sum += u.output
+        total_sum += u.total
+    if in_sum == 0 and out_sum == 0 and total_sum == 0:
+        return None
+    return TokenUsage(total=total_sum, input=in_sum, output=out_sum)

packages/ai-providers/server-ai-langchain/tests/test_langchain_provider.py

@@ -7,7 +7,15 @@
 
 from ldai import LDMessage
 
-from ldai_langchain import LangChainModelRunner, LangChainRunnerFactory, convert_messages_to_langchain, get_ai_metrics_from_response, map_provider
+from ldai_langchain import (
+    LangChainModelRunner,
+    LangChainRunnerFactory,
+    convert_messages_to_langchain,
+    get_ai_metrics_from_response,
+    get_tool_calls_from_response,
+    map_provider,
+    sum_token_usage_from_messages,
+)
 
 
 class TestConvertMessages:
@@ -237,6 +245,82 @@ async def test_returns_success_false_when_structured_model_invocation_throws_err
         assert result.metrics.usage is None
 
 
+class TestGetToolCallsFromResponse:
+    """Tests for get_tool_calls_from_response."""
+
+    def test_returns_tool_call_names_in_order(self):
+        """Should return tool call names from response.tool_calls."""
+        mock_response = MagicMock()
+        mock_response.tool_calls = [
+            {'name': 'search', 'args': {}},
+            {'name': 'calculator', 'args': {}},
+        ]
+        assert get_tool_calls_from_response(mock_response) == ['search', 'calculator']
+
+    def test_returns_empty_list_when_tool_calls_is_empty(self):
+        """Should return empty list when tool_calls is an empty list."""
+        mock_response = MagicMock()
+        mock_response.tool_calls = []
+        assert get_tool_calls_from_response(mock_response) == []
+
+    def test_returns_empty_list_when_no_tool_calls_attribute(self):
+        """Should return empty list when response has no tool_calls attribute."""
+        mock_response = MagicMock(spec=[])
+        assert get_tool_calls_from_response(mock_response) == []
+
+    def test_returns_empty_list_when_tool_calls_is_not_a_list(self):
+        """Should return empty list when tool_calls is not a list."""
+        mock_response = MagicMock()
+        mock_response.tool_calls = 'not-a-list'
+        assert get_tool_calls_from_response(mock_response) == []
+
+    def test_skips_tool_calls_without_name(self):
+        """Should skip tool calls that have no name."""
+        mock_response = MagicMock()
+        mock_response.tool_calls = [{'args': {}}, {'name': 'search', 'args': {}}]
+        assert get_tool_calls_from_response(mock_response) == ['search']
+
+
+class TestSumTokenUsageFromMessages:
+    """Tests for sum_token_usage_from_messages."""
+
+    def test_sums_usage_across_messages(self):
+        """Should sum token usage from all messages."""
+        msg1 = AIMessage(content='a')
+        msg1.usage_metadata = {'total_tokens': 10, 'input_tokens': 6, 'output_tokens': 4}
+        msg2 = AIMessage(content='b')
+        msg2.usage_metadata = {'total_tokens': 20, 'input_tokens': 12, 'output_tokens': 8}
+
+        result = sum_token_usage_from_messages([msg1, msg2])
+
+        assert result is not None
+        assert result.total == 30
+        assert result.input == 18
+        assert result.output == 12
+
+    def test_returns_none_when_no_usage_on_any_message(self):
+        """Should return None when no message has usage metadata."""
+        msg = AIMessage(content='hello')
+        assert sum_token_usage_from_messages([msg]) is None
+
+    def test_returns_none_for_empty_list(self):
+        """Should return None for an empty message list."""
+        assert sum_token_usage_from_messages([]) is None
+
+    def test_skips_messages_without_usage(self):
+        """Should skip messages that have no usage and sum the rest."""
+        msg1 = AIMessage(content='a')
+        msg2 = AIMessage(content='b')
+        msg2.usage_metadata = {'total_tokens': 5, 'input_tokens': 3, 'output_tokens': 2}
+
+        result = sum_token_usage_from_messages([msg1, msg2])
+
+        assert result is not None
+        assert result.total == 5
+        assert result.input == 3
+        assert result.output == 2
+
+
 class TestGetLlm:
     """Tests for LangChainModelRunner.get_llm."""
 

packages/sdk/server-ai/README.md

@@ -150,8 +150,8 @@ async def main():
     # Create LangChain model from configuration
     llm = await LangChainProvider.create_langchain_model(ai_config)
 
-    # Use with tracking
-    response = await ai_config.tracker.track_metrics_of(
+    # Use with tracking (sync invoke)
+    response = ai_config.tracker.track_metrics_of(
         lambda: llm.invoke(messages),
         lambda result: LangChainProvider.get_ai_metrics_from_response(result)
     )
@@ -190,7 +190,7 @@ async def main():
             temperature=ai_config.model.get_parameter('temperature') if ai_config.model else 0.5,
         )
 
-    result = await ai_config.tracker.track_metrics_of(
+    result = await ai_config.tracker.track_metrics_of_async(
         call_custom_provider,
         map_custom_provider_metrics
     )

packages/sdk/server-ai/src/ldai/judge/__init__.py

@@ -71,7 +71,7 @@ async def evaluate(
             messages = self._construct_evaluation_messages(input_text, output_text)
             assert self._evaluation_response_structure is not None
 
-            response = await self._ai_config_tracker.track_metrics_of(
+            response = await self._ai_config_tracker.track_metrics_of_async(
                 lambda: self._model_runner.invoke_structured_model(messages, self._evaluation_response_structure),
                 lambda result: result.metrics,
             )

packages/sdk/server-ai/src/ldai/managed_model.py

@@ -48,7 +48,7 @@ async def invoke(self, prompt: str) -> ModelResponse:
         config_messages = self._ai_config.messages or []
         all_messages = config_messages + self._messages
 
-        response = await self._tracker.track_metrics_of(
+        response = await self._tracker.track_metrics_of_async(
             lambda: self._model_runner.invoke_model(all_messages),
             lambda result: result.metrics,
         )

packages/sdk/server-ai/src/ldai/providers/runner_factory.py

@@ -77,7 +77,7 @@ def _with_fallback(
                     continue
                 result = fn(provider_factory)
                 if result is not None:
-                    log.debug(f"Successfully created capability using provider '{provider_type}'")
+                    log.debug(f"Successfully invoked create function with provider '{provider_type}'")
                     return result
             except Exception as exc:
                 log.warning(f"Provider '{provider_type}' failed: {exc}")