send back proper summaries

Author: jsonbailey

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

@@ -1,14 +1,15 @@
 """ManagedAgentGraph — LaunchDarkly managed wrapper for agent graph execution."""
 
-from typing import Any
+from typing import Any, Dict
 
 from ldai.agent_graph import AgentGraphDefinition
 from ldai.providers import AgentGraphRunner
 from ldai.providers.types import (
     AgentGraphRunnerResult,
-    GraphMetricSummary,
+    LDAIMetrics,
     ManagedGraphResult,
 )
+from ldai.tracker import LDAIMetricSummary
 
 
 class ManagedAgentGraph:
@@ -42,24 +43,20 @@ async def run(self, input: Any) -> ManagedGraphResult:
         Run the agent graph with the given input.
 
         Delegates to the underlying AgentGraphRunner, then drives all
-        LaunchDarkly tracking from ``result.metrics``:
-
-        - Graph-level events (path, duration, success/failure, total tokens) via
-          the graph tracker obtained from the graph definition.
-        - Per-node events (tokens, duration, tool calls, success) via per-node
-          trackers for each key present in ``result.metrics.node_metrics``.
+        LaunchDarkly tracking from ``result.metrics``.
 
         :param input: The input prompt or structured input for the graph
         :return: ManagedGraphResult containing the content, metric summary,
             and raw response.
         """
-        result = await self._runner.run(input)
-
-        summary = self._build_summary_from_runner_result(result)
-
         graph_tracker = self._graph.create_tracker()
-        self._flush_graph_tracking(result, graph_tracker)
-        self._flush_node_tracking(result)
+        result = await graph_tracker.track_graph_metrics_of_async(
+            lambda r: r.metrics,
+            lambda: self._runner.run(input),
+        )
+
+        summary = graph_tracker.get_summary()
+        summary.node_metrics = self._track_node_metrics(result.metrics.node_metrics)
 
         return ManagedGraphResult(
             content=result.content,
@@ -68,51 +65,22 @@ async def run(self, input: Any) -> ManagedGraphResult:
             evaluations=None,
         )
 
-    def _build_summary_from_runner_result(
-        self,
-        result: AgentGraphRunnerResult,
-    ) -> GraphMetricSummary:
-        """Build a GraphMetricSummary from an AgentGraphRunnerResult."""
-        m = result.metrics
-        return GraphMetricSummary(
-            success=m.success,
-            path=list(m.path),
-            duration_ms=m.duration_ms,
-            usage=m.usage,
-            node_metrics=dict(m.node_metrics),
-        )
-
-    def _flush_graph_tracking(self, result: AgentGraphRunnerResult, tracker: Any) -> None:
-        """
-        Drive graph-level LaunchDarkly tracking events from runner result metrics.
+    def _track_node_metrics(
+        self, node_metrics: Dict[str, LDAIMetrics]
+    ) -> Dict[str, LDAIMetricSummary]:
         """
-        m = result.metrics
-        if m.path:
-            tracker.track_path(m.path)
-        if m.duration_ms is not None:
-            tracker.track_duration(m.duration_ms)
-        if m.success:
-            tracker.track_invocation_success()
-        else:
-            tracker.track_invocation_failure()
-        if m.usage is not None:
-            tracker.track_total_tokens(m.usage)
-
-    def _flush_node_tracking(self, result: AgentGraphRunnerResult) -> None:
-        """
-        Drive per-node LaunchDarkly tracking events from ``result.metrics.node_metrics``.
+        Drive per-node LaunchDarkly tracking events and collect node metric summaries.
 
         For each node key present in ``node_metrics``, obtains the node's
-        config tracker via the graph definition and fires token, duration,
-        tool call, and success/error events.
+        config tracker via the graph definition, fires tracking events, and
+        returns a map of node key to the tracker's metric summary.
         """
-        for node_key, node_ldai_metrics in result.metrics.node_metrics.items():
+        node_summaries: Dict[str, LDAIMetricSummary] = {}
+        for node_key, node_ldai_metrics in node_metrics.items():
             node = self._graph.get_node(node_key)
             if node is None:
                 continue
             node_tracker = node.get_config().create_tracker()
-            if node_tracker is None:
-                continue
 
             if node_ldai_metrics.usage is not None:
                 node_tracker.track_tokens(node_ldai_metrics.usage)
@@ -125,6 +93,9 @@ def _flush_node_tracking(self, result: AgentGraphRunnerResult) -> None:
             else:
                 node_tracker.track_error()
 
+            node_summaries[node_key] = node_tracker.get_summary()
+        return node_summaries
+
     def get_agent_graph_runner(self) -> AgentGraphRunner:
         """
         Return the underlying AgentGraphRunner for advanced use.

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

@@ -138,8 +138,8 @@ class GraphMetrics:
 class GraphMetricSummary:
     """Contains a summary of metrics for an agent graph run."""
 
-    success: bool
-    """Whether the graph run succeeded."""
+    success: Optional[bool] = None
+    """Whether the graph run succeeded. Absent if invocation status has not been tracked."""
 
     path: List[str] = field(default_factory=list)
     """Ordered list of node keys visited during the run."""
@@ -150,8 +150,8 @@ class GraphMetricSummary:
     usage: Optional[TokenUsage] = None
     """Optional aggregate token usage information across all nodes in the graph run."""
 
-    node_metrics: Dict[str, LDAIMetrics] = field(default_factory=dict)
-    """Per-node metrics keyed by node key."""
+    node_metrics: Dict[str, LDAIMetricSummary] = field(default_factory=dict)
+    """Per-node metric summaries keyed by node key."""
 
     resumption_token: Optional[str] = None
     """Optional resumption token from the graph tracker for cross-process resumption."""

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

@@ -13,7 +13,7 @@
 from ldai import log
 
 if TYPE_CHECKING:
-    from ldai.providers.types import LDAIMetrics
+    from ldai.providers.types import GraphMetrics, GraphMetricSummary, LDAIMetrics
 
 
 class FeedbackKind(Enum):
@@ -616,7 +616,11 @@ def _openai_to_token_usage(data: dict) -> TokenUsage:
 
 class AIGraphTracker:
     """
-    Tracks graph-level, node-level, and edge-level metrics for AI agent graph operations.
+    Tracks graph-level metrics for AI agent graph operations.
+
+    Maintains an internal :class:`~ldai.providers.types.GraphMetricSummary`
+    that is updated as tracking methods are called. Retrieve it via
+    :meth:`get_summary`.
     """
 
     def __init__(
@@ -642,11 +646,22 @@ def __init__(
         self._version = version
         self._context = context
 
+        from ldai.providers.types import GraphMetricSummary
+        self._summary = GraphMetricSummary()
+
     @property
     def graph_key(self) -> str:
         """Graph configuration key used in tracking payloads."""
         return self._graph_key
 
+    def get_summary(self) -> GraphMetricSummary:
+        """
+        Get the current summary of graph-level metrics.
+
+        :return: Summary of graph metrics tracked so far.
+        """
+        return self._summary
+
     def __get_track_data(self):
         """
         Get tracking data for events.
@@ -664,6 +679,10 @@ def track_invocation_success(self) -> None:
         """
         Track a successful graph invocation.
         """
+        if self._summary.success is not None:
+            log.warning("Invocation status has already been tracked for this graph execution. %s", self.__get_track_data())
+            return
+        self._summary.success = True
         self._ld_client.track(
             "$ld:ai:graph:invocation_success",
             self._context,
@@ -675,6 +694,10 @@ def track_invocation_failure(self) -> None:
         """
         Track an unsuccessful graph invocation.
         """
+        if self._summary.success is not None:
+            log.warning("Invocation status has already been tracked for this graph execution. %s", self.__get_track_data())
+            return
+        self._summary.success = False
         self._ld_client.track(
             "$ld:ai:graph:invocation_failure",
             self._context,
@@ -688,6 +711,10 @@ def track_duration(self, duration: int) -> None:
 
         :param duration: Duration in milliseconds.
         """
+        if self._summary.duration_ms is not None:
+            log.warning("Duration has already been tracked for this graph execution. %s", self.__get_track_data())
+            return
+        self._summary.duration_ms = duration
         self._ld_client.track(
             "$ld:ai:graph:duration:total",
             self._context,
@@ -703,6 +730,10 @@ def track_total_tokens(self, tokens: Optional[TokenUsage] = None) -> None:
         """
         if tokens is None or tokens.total <= 0:
             return
+        if self._summary.usage is not None:
+            log.warning("Token usage has already been tracked for this graph execution. %s", self.__get_track_data())
+            return
+        self._summary.usage = tokens
         self._ld_client.track(
             "$ld:ai:graph:total_tokens",
             self._context,
@@ -716,6 +747,10 @@ def track_path(self, path: List[str]) -> None:
 
         :param path: An array of configuration keys representing the sequence of nodes executed during graph traversal.
         """
+        if self._summary.path:
+            log.warning("Path has already been tracked for this graph execution. %s", self.__get_track_data())
+            return
+        self._summary.path = list(path)
         track_data = {**self.__get_track_data(), "path": path}
         self._ld_client.track(
             "$ld:ai:graph:path",
@@ -780,3 +815,92 @@ def track_handoff_failure(self, source_key: str, target_key: str) -> None:
             track_data,
             1,
         )
+
+    def _track_from_graph_metrics(
+        self,
+        result: Any,
+        metrics_extractor: Callable[[Any], Optional[GraphMetrics]],
+        elapsed_ms: int,
+    ) -> None:
+        metrics: Optional[GraphMetrics] = None
+        try:
+            metrics = metrics_extractor(result)
+        except Exception as exc:
+            log.warning("Failed to extract graph metrics: %s", exc)
+
+        if metrics is None:
+            self.track_duration(elapsed_ms)
+            return
+
+        self.track_duration(metrics.duration_ms if metrics.duration_ms is not None else elapsed_ms)
+        if metrics.success:
+            self.track_invocation_success()
+        else:
+            self.track_invocation_failure()
+        if metrics.path:
+            self.track_path(metrics.path)
+        if metrics.usage is not None:
+            self.track_total_tokens(metrics.usage)
+
+    def track_graph_metrics_of(
+        self,
+        metrics_extractor: Callable[[Any], Optional[GraphMetrics]],
+        func: Callable[[], Any],
+    ) -> Any:
+        """
+        Track graph-level metrics for a synchronous graph operation.
+
+        Times the operation, extracts :class:`~ldai.providers.types.GraphMetrics`
+        via the provided extractor, and fires graph-level tracking events
+        (path, duration, success/failure, total tokens).
+
+        If the extracted ``GraphMetrics`` has a non-``None`` ``duration_ms``,
+        that value is used instead of the wall-clock elapsed time.
+
+        Node-level metrics are not tracked by this method.
+
+        For async operations, use :meth:`track_graph_metrics_of_async`.
+
+        :param metrics_extractor: Function that extracts GraphMetrics from the result
+        :param func: Synchronous callable that runs the graph operation
+        :return: The result of the operation
+        """
+        start_ns = time.perf_counter_ns()
+        try:
+            result = func()
+        except Exception as err:
+            duration = (time.perf_counter_ns() - start_ns) // 1_000_000
+            self.track_duration(duration)
+            self.track_invocation_failure()
+            raise err
+
+        elapsed_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
+        self._track_from_graph_metrics(result, metrics_extractor, elapsed_ms)
+        return result
+
+    async def track_graph_metrics_of_async(
+        self,
+        metrics_extractor: Callable[[Any], Optional[GraphMetrics]],
+        func: Callable[[], Any],
+    ) -> Any:
+        """
+        Track graph-level metrics for an async graph operation (``func`` is awaited).
+
+        Same event semantics as :meth:`track_graph_metrics_of`.
+
+        :param metrics_extractor: Function that extracts GraphMetrics from the result
+        :param func: Async callable that runs the graph operation
+        :return: The result of the operation
+        """
+        start_ns = time.perf_counter_ns()
+        try:
+            result = await func()
+        except Exception as err:
+            duration = (time.perf_counter_ns() - start_ns) // 1_000_000
+            self.track_duration(duration)
+            self.track_invocation_failure()
+            raise err
+
+        elapsed_ms = (time.perf_counter_ns() - start_ns) // 1_000_000
+        self._track_from_graph_metrics(result, metrics_extractor, elapsed_ms)
+        return result