From 8f30b0a11ea53244b4f531afbbcedc6fdf2d3595 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Wed, 4 Feb 2026 15:09:00 -0900 Subject: [PATCH 1/2] [REL-12055] feat: add agent graph tracker --- packages/sdk/server-ai/src/ldai/__init__.py | 2 + .../src/ldai/agent_graph/__init__.py | 15 +- packages/sdk/server-ai/src/ldai/client.py | 19 +- packages/sdk/server-ai/src/ldai/tracker.py | 244 +++++++++++++++++- 4 files changed, 277 insertions(+), 3 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 72c94f5..c509856 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -13,6 +13,7 @@ Edge, JudgeConfiguration, LDAIAgent, LDAIAgentConfig, LDAIAgentDefaults, LDMessage, ModelConfig, ProviderConfig) from ldai.providers.types import EvalScore, JudgeResponse +from ldai.tracker import AIGraphTracker __all__ = [ 'LDAIClient', @@ -21,6 +22,7 @@ 'AIAgentConfigRequest', 'AIAgents', 'AIAgentGraphConfig', + 'AIGraphTracker', 'Edge', 'AICompletionConfig', 'AICompletionConfigDefault', diff --git a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py index a10ab23..03030e6 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -1,12 +1,15 @@ """Graph implementation for managing AI agent graphs.""" from dataclasses import dataclass -from typing import Any, Callable, Dict, List, Optional, Set +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set from ldclient import Context from ldai.models import AIAgentConfig, AIAgentGraphConfig, Edge +if TYPE_CHECKING: + from ldai.tracker import AIGraphTracker + DEFAULT_FALSE = AIAgentConfig(key="", enabled=False) @@ -54,11 +57,21 @@ def __init__( nodes: Dict[str, AgentGraphNode], context: Context, enabled: bool, + tracker: Optional["AIGraphTracker"] = None, ): self._agent_graph = agent_graph self._context = context self._nodes = nodes self.enabled = enabled + self._tracker = tracker + + def get_tracker(self) -> Optional["AIGraphTracker"]: + """ + Get the graph tracker for this graph definition. + + :return: The AIGraphTracker instance, or None if not available. + """ + return self._tracker def is_enabled(self) -> bool: """Check if the graph is enabled.""" diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index a4605e3..82754a9 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -15,7 +15,7 @@ JudgeConfiguration, LDMessage, ModelConfig, ProviderConfig) from ldai.providers.ai_provider_factory import AIProviderFactory -from ldai.tracker import LDAIConfigTracker +from ldai.tracker import AIGraphTracker, LDAIConfigTracker class LDAIClient: @@ -435,6 +435,19 @@ def agent_graph( """ variation = self._client.variation(key, context, {}) + # Extract variation metadata for tracker + variation_key = variation.get("_ldMeta", {}).get("variationKey", "") + version = int(variation.get("_ldMeta", {}).get("version", 1)) + + # Create graph tracker + tracker = AIGraphTracker( + self._client, + variation_key, + key, + version, + context, + ) + if not variation.get("root"): log.debug(f"Agent graph {key} is disabled, no root config key found") return AgentGraphDefinition( @@ -447,6 +460,7 @@ def agent_graph( nodes={}, context=context, enabled=False, + tracker=tracker, ) edge_keys = list[str](variation.get("edges", {}).keys()) @@ -474,6 +488,7 @@ def agent_graph( nodes={}, context=context, enabled=False, + tracker=tracker, ) try: @@ -504,6 +519,7 @@ def agent_graph( nodes={}, context=context, enabled=False, + tracker=tracker, ) nodes = AgentGraphDefinition.build_nodes( @@ -516,6 +532,7 @@ def agent_graph( nodes=nodes, context=context, enabled=agent_graph_config.enabled, + tracker=tracker, ) def agents( diff --git a/packages/sdk/server-ai/src/ldai/tracker.py b/packages/sdk/server-ai/src/ldai/tracker.py index e5d7ed2..fa6e635 100644 --- a/packages/sdk/server-ai/src/ldai/tracker.py +++ b/packages/sdk/server-ai/src/ldai/tracker.py @@ -1,7 +1,7 @@ import time from dataclasses import dataclass from enum import Enum -from typing import Any, Dict, Optional +from typing import Any, Dict, List, Optional from ldclient import Context, LDClient @@ -407,3 +407,245 @@ def _openai_to_token_usage(data: dict) -> TokenUsage: input=data.get("prompt_tokens", 0), output=data.get("completion_tokens", 0), ) + + +class AIGraphTracker: + """ + Tracks graph-level, node-level, and edge-level metrics for AI agent graph operations. + """ + + def __init__( + self, + ld_client: LDClient, + variation_key: str, + graph_key: str, + version: int, + context: Context, + ): + """ + Initialize an AI Graph tracker. + + :param ld_client: LaunchDarkly client instance. + :param variation_key: Variation key for tracking. + :param graph_key: Graph configuration key for tracking. + :param version: Version of the variation. + :param context: Context for evaluation. + """ + self._ld_client = ld_client + self._variation_key = variation_key + self._graph_key = graph_key + self._version = version + self._context = context + + def __get_track_data(self): + """ + Get tracking data for events. + + :return: Dictionary containing variation, graph key, and version. + """ + track_data = { + "variationKey": self._variation_key, + "graphKey": self._graph_key, + "version": self._version, + } + # Note: aiSdkName and aiSdkVersion are optional and not included for now + return track_data + + def track_invocation_success(self) -> None: + """ + Track a successful graph invocation. + """ + self._ld_client.track( + "$ld:ai:graph:invocation_success", + self._context, + self.__get_track_data(), + 1, + ) + + def track_invocation_failure(self) -> None: + """ + Track an unsuccessful graph invocation. + """ + self._ld_client.track( + "$ld:ai:graph:invocation_failure", + self._context, + self.__get_track_data(), + 1, + ) + + def track_latency(self, duration: int) -> None: + """ + Track the total latency of graph execution. + + :param duration: Duration in milliseconds. + """ + self._ld_client.track( + "$ld:ai:graph:latency", + self._context, + self.__get_track_data(), + duration, + ) + + def track_total_tokens(self, tokens: TokenUsage) -> None: + """ + Track aggregated token usage across the entire graph invocation. + + :param tokens: Token usage data. + """ + self._ld_client.track( + "$ld:ai:graph:total_tokens", + self._context, + self.__get_track_data(), + tokens.total, + ) + + def track_path(self, path: List[str]) -> None: + """ + Track the execution path through the graph. + + :param path: An array of configuration keys representing the sequence of nodes executed during graph traversal. + """ + track_data = {**self.__get_track_data(), "path": path} + self._ld_client.track( + "$ld:ai:graph:path", + self._context, + track_data, + 1, + ) + + def track_judge_response(self, response: Any) -> None: + """ + Track judge responses for the final graph output. + + :param response: JudgeResponse object containing evals and success status. + """ + from ldai.providers.types import EvalScore, JudgeResponse + + if isinstance(response, JudgeResponse): + if response.evals: + track_data = self.__get_track_data() + if response.judge_config_key: + track_data = {**track_data, "judgeConfigKey": response.judge_config_key} + + for metric_key, eval_score in response.evals.items(): + if isinstance(eval_score, EvalScore): + self._ld_client.track( + metric_key, + self._context, + track_data, + eval_score.score, + ) + + def track_node_invocation(self, config_key: str) -> None: + """ + Track when a node is invoked during graph execution. + + :param config_key: The configuration key of the node being invoked. + """ + track_data = {**self.__get_track_data(), "configKey": config_key} + self._ld_client.track( + "$ld:ai:graph:node_invocation", + self._context, + track_data, + 1, + ) + + def track_tool_call(self, config_key: str, tool_key: str) -> None: + """ + Track tool calls made by nodes during graph execution. + + :param config_key: The configuration key of the node making the tool call. + :param tool_key: The key of the tool being called. + """ + track_data = { + **self.__get_track_data(), + "configKey": config_key, + "toolKey": tool_key, + } + self._ld_client.track( + "$ld:ai:graph:tool_call", + self._context, + track_data, + 1, + ) + + def track_node_judge_response(self, config_key: str, response: Any) -> None: + """ + Track judge responses for a specific node. + + :param config_key: The configuration key of the node being evaluated. + :param response: JudgeResponse object containing evals and success status. + """ + from ldai.providers.types import EvalScore, JudgeResponse + + if isinstance(response, JudgeResponse): + if response.evals: + track_data = {**self.__get_track_data(), "configKey": config_key} + if response.judge_config_key: + track_data = {**track_data, "judgeConfigKey": response.judge_config_key} + + for metric_key, eval_score in response.evals.items(): + if isinstance(eval_score, EvalScore): + self._ld_client.track( + metric_key, + self._context, + track_data, + eval_score.score, + ) + + def track_redirect(self, source_key: str, redirected_target: str) -> None: + """ + Track when a node redirects to a different target than originally specified. + + :param source_key: The configuration key of the source node. + :param redirected_target: The configuration key of the target node that was redirected to. + """ + track_data = { + **self.__get_track_data(), + "sourceKey": source_key, + "redirectedTarget": redirected_target, + } + self._ld_client.track( + "$ld:ai:graph:redirect", + self._context, + track_data, + 1, + ) + + def track_handoff_success(self, source_key: str, target_key: str) -> None: + """ + Track successful handoffs between nodes. + + :param source_key: The configuration key of the source node. + :param target_key: The configuration key of the target node. + """ + track_data = { + **self.__get_track_data(), + "sourceKey": source_key, + "targetKey": target_key, + } + self._ld_client.track( + "$ld:ai:graph:handoff_success", + self._context, + track_data, + 1, + ) + + def track_handoff_failure(self, source_key: str, target_key: str) -> None: + """ + Track failed handoffs between nodes. + + :param source_key: The configuration key of the source node. + :param target_key: The configuration key of the target node. + """ + track_data = { + **self.__get_track_data(), + "sourceKey": source_key, + "targetKey": target_key, + } + self._ld_client.track( + "$ld:ai:graph:handoff_failure", + self._context, + track_data, + 1, + ) From f07f60e2ab5f7660873b699b6bb7628e5f82dc63 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Wed, 4 Feb 2026 15:49:48 -0900 Subject: [PATCH 2/2] [REL-12055] remove unnecessary check on typing --- .../sdk/server-ai/src/ldai/agent_graph/__init__.py | 11 ++++------- 1 file changed, 4 insertions(+), 7 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py index 03030e6..c19552d 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -1,14 +1,11 @@ """Graph implementation for managing AI agent graphs.""" -from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional, Set +from typing import Any, Callable, Dict, List, Optional, Set from ldclient import Context from ldai.models import AIAgentConfig, AIAgentGraphConfig, Edge - -if TYPE_CHECKING: - from ldai.tracker import AIGraphTracker +from ldai.tracker import AIGraphTracker DEFAULT_FALSE = AIAgentConfig(key="", enabled=False) @@ -57,7 +54,7 @@ def __init__( nodes: Dict[str, AgentGraphNode], context: Context, enabled: bool, - tracker: Optional["AIGraphTracker"] = None, + tracker: Optional[AIGraphTracker] = None, ): self._agent_graph = agent_graph self._context = context @@ -65,7 +62,7 @@ def __init__( self.enabled = enabled self._tracker = tracker - def get_tracker(self) -> Optional["AIGraphTracker"]: + def get_tracker(self) -> Optional[AIGraphTracker]: """ Get the graph tracker for this graph definition.