From 2d9e84234883aa98f87764dc3c32bb8fe273cdf6 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 8 Jan 2026 15:58:26 -0900 Subject: [PATCH 01/11] local commit --- packages/sdk/server-ai/src/ldai/__init__.py | 6 +- .../src/ldai/agent_graph/__init__.py | 287 ++++++++++++++++++ packages/sdk/server-ai/src/ldai/client.py | 89 +++++- packages/sdk/server-ai/src/ldai/models.py | 26 ++ 4 files changed, 406 insertions(+), 2 deletions(-) create mode 100644 packages/sdk/server-ai/src/ldai/agent_graph/__init__.py diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 0bc0e76..104a926 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -4,12 +4,13 @@ from ldai.chat import Chat from ldai.client import LDAIClient +from ldai.agent_graph import AgentGraph from ldai.judge import Judge from ldai.models import ( # Deprecated aliases for backward compatibility AIAgentConfig, AIAgentConfigDefault, AIAgentConfigRequest, AIAgents, AICompletionConfig, AICompletionConfigDefault, AIConfig, AIJudgeConfig, AIJudgeConfigDefault, JudgeConfiguration, LDAIAgent, LDAIAgentConfig, - LDAIAgentDefaults, LDMessage, ModelConfig, ProviderConfig) + LDAIAgentDefaults, LDMessage, ModelConfig, ProviderConfig, AIAgentGraph, AIAgentGraphEdge) from ldai.providers.types import EvalScore, JudgeResponse __all__ = [ @@ -18,12 +19,15 @@ 'AIAgentConfigDefault', 'AIAgentConfigRequest', 'AIAgents', + 'AIAgentGraph', + 'AIAgentGraphEdge', 'AICompletionConfig', 'AICompletionConfigDefault', 'AIJudgeConfig', 'AIJudgeConfigDefault', 'Chat', 'EvalScore', + 'AgentGraph', 'Judge', 'JudgeConfiguration', 'JudgeResponse', diff --git a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py new file mode 100644 index 0000000..83b3755 --- /dev/null +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -0,0 +1,287 @@ +"""Graph implementation for managing AI agent graphs.""" + +from typing import Any, Callable, Dict, List, Optional, Set +from ldai.models import AIAgentGraph, AIAgentConfig, AIAgentGraphEdge +from ldclient import Context + + +class AgentGraphNode: + """ + Node in an agent graph. + """ + + default_false = AIAgentConfig(key="", enabled=False) + + def __init__( + self, + key: str, + config: AIAgentConfig, + children: List[AIAgentGraphEdge], + parent_graph: "AgentGraph", + ): + self._key = key + self._config = config + self._children = children + self._parent_graph = parent_graph + + def get_key(self) -> str: + """Get the key of the node.""" + return self._key + + def get_config(self) -> AIAgentConfig: + """Get the config of the node.""" + return self._config + + def get_edges(self) -> List[AIAgentGraphEdge]: + """Get the edges of the node.""" + return self._children + + def get_child_nodes(self) -> List["AgentGraphNode"]: + """Get the child nodes of the node as AgentGraphNode objects.""" + return [ + self._parent_graph.get_node(edge.targetConfig) for edge in self._children + ] + + def is_terminal(self) -> bool: + """Check if the node is a terminal node.""" + return len(self._children) == 0 + + def get_parent_nodes(self) -> List["AgentGraphNode"]: + """Get the parent nodes of the node as AgentGraphNode objects.""" + return [ + self._parent_graph.get_node(edge.sourceConfig) + for edge in self._parent_graph._get_parent_edges(self._key) + ] + + def traverse( + self, fn: Callable[["AgentGraphNode", Dict[str, Any]], None], execution_context: Dict[str, Any] = {}, visited: Optional[Set[str]] = None + ) -> None: + """Traverse the graph downwardly from this node, calling fn on each node.""" + if visited is None: + visited = set() + + # Avoid cycles by tracking visited nodes + if self._key in visited: + return + + visited.add(self._key) + fn(self, execution_context) + + for child in self._children: + node = self._parent_graph.get_node(child.targetConfig) + if node is not None: + node.traverse(fn, execution_context, visited) + + def reverse_traverse( + self, + fn: Callable[["AgentGraphNode", Dict[str, Any]], None], + execution_context: Dict[str, Any] = {}, + visited: Optional[Set[str]] = None, + ) -> None: + """Reverse traverse the graph upwardly from this node, calling fn on each node.""" + if visited is None: + visited = set() + + # Avoid cycles by tracking visited nodes + if self._key in visited: + return + + visited.add(self._key) + fn(self, execution_context) + + for parent in self._parent_graph._get_parent_edges(self._key): + node = self._parent_graph.get_node(parent.sourceConfig) + if node is not None: + node.reverse_traverse(fn, execution_context, visited) + + +class AgentGraph: + """ + Graph implementation for managing AI agent graphs. + """ + + default_false = AIAgentConfig(key="", enabled=False) + + def __init__( + self, + agent_graph: AIAgentGraph, + context: Context, + get_agent: Callable[[str, Context, dict], AIAgentConfig], + ): + self._agent_graph = agent_graph + self._context = context + self._get_agent = get_agent + self._nodes = self._build_nodes() + + def _build_nodes(self) -> Dict[str, AgentGraphNode]: + """Build the nodes of the graph into AgentGraphNode objects.""" + nodes = { + self._agent_graph.rootConfigKey: AgentGraphNode( + self._agent_graph.rootConfigKey, + self._get_agent( + self._agent_graph.rootConfigKey, self._context, self.default_false + ), + self._get_child_edges(self._agent_graph.rootConfigKey), + self, + ), + } + + for edge in self._agent_graph.edges: + nodes[edge.targetConfig] = AgentGraphNode( + edge.targetConfig, + self._get_agent(edge.targetConfig, self._context, self.default_false), + self._get_child_edges(edge.targetConfig), + self, + ) + + return nodes + + def _get_child_edges(self, config_key: str) -> List[AIAgentGraphEdge]: + """Get the child edges of the given config.""" + return [ + edge + for edge in self._agent_graph.edges + if edge.sourceConfig == config_key + ] + + def _get_parent_edges(self, config_key: str) -> List[AIAgentGraphEdge]: + """Get the parent edges of the given config.""" + return [ + edge + for edge in self._agent_graph.edges + if edge.targetConfig == config_key + ] + + def _collect_nodes( + self, + node: AgentGraphNode, + node_depths: Dict[str, int], + nodes_by_depth: Dict[int, List[AgentGraphNode]], + visited: Set[str], + ) -> None: + """Collect all reachable nodes from the given node and group them by depth.""" + node_key = node.get_key() + if node_key in visited: + return + visited.add(node_key) + + node_depth = node_depths.get(node_key, 0) + if node_depth not in nodes_by_depth: + nodes_by_depth[node_depth] = [] + nodes_by_depth[node_depth].append(node) + + for child in node.get_child_nodes(): + self._collect_nodes(child, node_depths, nodes_by_depth, visited) + + def terminal_nodes(self) -> List[AgentGraphNode]: + """Get the terminal nodes of the graph, meaning any nodes without children.""" + return [ + node for node in self._nodes.values() if len(node.get_child_nodes()) == 0 + ] + + def root(self) -> AgentGraphNode | None: + """Get the root node of the graph.""" + config = self._get_agent( + self._agent_graph.rootConfigKey, self._context, self.default_false + ) + + if config.enabled is False: + return None + + children = [ + edge + for edge in self._agent_graph.edges + if edge.sourceConfig == self._agent_graph.rootConfigKey + ] + + node = AgentGraphNode(self._agent_graph.rootConfigKey, config, children, self) + + return node + + def traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], None], execution_context: Dict[str, Any] = {}) -> None: + """Traverse from the root down to terminal nodes, visiting nodes in order of depth. + Nodes with the longest paths from the root (deepest nodes) will always be visited last.""" + root_node = self.root() + if root_node is None: + return + + node_depths: Dict[str, int] = {root_node.get_key(): 0} + current_level: List[AgentGraphNode] = [root_node] + depth = 0 + max_depth_limit = 10 # Infinite loop protection limit + + while current_level and depth < max_depth_limit: + next_level: List[AgentGraphNode] = [] + depth += 1 + + for node in current_level: + for child in node.get_child_nodes(): + child_key = child.get_key() + # Defer this child to the next level if it's at a longer path + if child_key not in node_depths or ( + depth > node_depths[child_key] and depth < max_depth_limit + ): + node_depths[child_key] = depth + next_level.append(child) + + current_level = next_level + + # Group all nodes by depth + nodes_by_depth: Dict[int, List[AgentGraphNode]] = {} + visited: Set[str] = set() + + self._collect_nodes(root_node, node_depths, nodes_by_depth, visited) + # Execute the lambda at this level for the nodes at this depth + for depth_level in sorted(nodes_by_depth.keys()): + for node in nodes_by_depth[depth_level]: + execution_context[node.get_key()] = fn(node, execution_context) + + return execution_context[self._agent_graph.rootConfigKey] + + def reverse_traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], execution_context: Dict[str, Any] = {}) -> None: + """Traverse from terminal nodes up to the root, visiting nodes level by level. + The root node will always be visited last, even if multiple paths converge at it.""" + terminal_nodes = self.terminal_nodes() + if not terminal_nodes: + return + + visited: Set[str] = set() + current_level: List[AgentGraphNode] = terminal_nodes + root_key = self._agent_graph.rootConfigKey + root_node_seen = False + + while current_level: + next_level: List[AgentGraphNode] = [] + + for node in current_level: + node_key = node.get_key() + if node_key in visited: + continue + + visited.add(node_key) + # Skip the root node if we reach a terminus, it will be visited last + if node_key == root_key: + root_node_seen = True + continue + + execution_context[node_key] = fn(node, execution_context) + + for parent in node.get_parent_nodes(): + parent_key = parent.get_key() + if parent_key not in visited: + next_level.append(parent) + + current_level = next_level + + # If we saw the root node, append it at the end as it'll always be the last node in a + # reverse traversal (this should always happen, non-contiguous graphs are invalid) + if root_node_seen: + root_node = self.root() + if root_node is not None: + execution_context[root_node.get_key()] = fn(root_node, execution_context) + + return execution_context[self._agent_graph.rootConfigKey] + + def get_node(self, key: str) -> AgentGraphNode | None: + """Get a node by its key.""" + return self._nodes.get(key) diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 47465ef..76b9c4a 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -6,12 +6,13 @@ from ldai import log from ldai.chat import Chat +from ldai.agent_graph import AgentGraph from ldai.judge import Judge from ldai.models import (AIAgentConfig, AIAgentConfigDefault, AIAgentConfigRequest, AIAgents, AICompletionConfig, AICompletionConfigDefault, AIJudgeConfig, AIJudgeConfigDefault, JudgeConfiguration, LDMessage, - ModelConfig, ProviderConfig) + ModelConfig, ProviderConfig, AIAgentGraph, AIAgentGraphEdge) from ldai.providers.ai_provider_factory import AIProviderFactory from ldai.tracker import LDAIConfigTracker @@ -419,6 +420,92 @@ def agent_configs( return result + def agent_graph( + self, + key: str, + context: Context, + ) -> AIAgentGraph: + """ + Retrieve an AI agent graph. + """ + variation = self._client.variation(key, context, {}) + mock_variation = { + "key": "test-agent-graph", + "name": "Test Agent Graph", + "rootConfigKey": "cruise-ship-information-agent", + "description": "Test Agent Graph Description", + "edges": [ + { + "key": "edge-cruise-ship-information-agent-get-weather-for-location", + "sourceConfig": "cruise-ship-information-agent", + "targetConfig": "get-weather-for-location", + "handoff": {}, + }, + { + "key": "edge-cruise-ship-information-agent-ships-in-port-agent", + "sourceConfig": "cruise-ship-information-agent", + "targetConfig": "ships-in-port-agent", + "handoff": {}, + }, + { + "key": "edge-ships-in-port-agent-vessel-details-agent", + "sourceConfig": "ships-in-port-agent", + "targetConfig": "vessel-details-agent", + "handoff": {}, + }, + { + "key": "edge-vessel-details-agent-cruise-information-synthesizer", + "sourceConfig": "vessel-details-agent", + "targetConfig": "cruise-information-synthesizer", + "handoff": {}, + }, + { + "key": "edge-get-weather-for-location-cruise-information-synthesizer", + "sourceConfig": "get-weather-for-location", + "targetConfig": "cruise-information-synthesizer", + "handoff": {}, + }, + ], + } + return AgentGraph( + agent_graph=AIAgentGraph( + key=mock_variation['key'], + name=mock_variation['name'], + rootConfigKey=mock_variation['rootConfigKey'], + edges=[ + AIAgentGraphEdge( + key=edge.get('key', ''), + sourceConfig=edge.get('sourceConfig', ''), + targetConfig=edge.get('targetConfig', ''), + handoff=edge.get('handoff', {}), + ) + for edge in mock_variation['edges'] + ], + description=mock_variation['description'], + ), + context=context, + get_agent=self.agent_config, + ) + + return AgentGraph( + agent_graph=AIAgentGraph( + key=variation.key, + name=variation.name, + rootConfigKey=variation.rootConfigKey, + edges=[ + AIAgentGraphEdge( + key=edge.key, + sourceConfig=edge.sourceConfig, + targetConfig=edge.targetConfig, + handoff=edge.handoff + ) + for edge in variation.edges + ] + ), + context=context, + get_variation=self._client.variation, + ) + def agents( self, agent_configs: List[AIAgentConfigRequest], diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index 988d97d..eadd072 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -338,6 +338,32 @@ class AIAgentConfigRequest: # Type alias for all AI Config variants AIConfigKind = Union[AIAgentConfig, AICompletionConfig, AIJudgeConfig] +# ============================================================================ +# AI Config Agent Graph Edge Type +# ============================================================================ +@dataclass +class AIAgentGraphEdge: + """ + Edge configuration for an agent graph. + """ + key: str + sourceConfig: str + targetConfig: str + handoff: Optional[dict] = None + +# ============================================================================ +# AI Config Agent Graph +# ============================================================================ +@dataclass +class AIAgentGraph: + """ + Agent graph configuration. + """ + key: str + name: str + rootConfigKey: str + edges: List[AIAgentGraphEdge] + description: Optional[str] = '' # ============================================================================ # Deprecated Type Aliases for Backward Compatibility From 3ff3b8b39e44ad3354bbc4a79268a803e8d0722e Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Wed, 14 Jan 2026 14:35:10 -0900 Subject: [PATCH 02/11] [REL-11697] Update PoC per spec and implement tests --- packages/sdk/server-ai/src/ldai/__init__.py | 10 +- .../src/ldai/agent_graph/__init__.py | 183 +++----- packages/sdk/server-ai/src/ldai/client.py | 123 ++--- packages/sdk/server-ai/src/ldai/models.py | 12 +- .../sdk/server-ai/tests/test_agent_graph.py | 441 ++++++++++++++++++ 5 files changed, 564 insertions(+), 205 deletions(-) create mode 100644 packages/sdk/server-ai/tests/test_agent_graph.py diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 104a926..2566b2a 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -4,13 +4,13 @@ from ldai.chat import Chat from ldai.client import LDAIClient -from ldai.agent_graph import AgentGraph +from ldai.agent_graph import AgentGraphDefinition from ldai.judge import Judge from ldai.models import ( # Deprecated aliases for backward compatibility AIAgentConfig, AIAgentConfigDefault, AIAgentConfigRequest, AIAgents, AICompletionConfig, AICompletionConfigDefault, AIConfig, AIJudgeConfig, AIJudgeConfigDefault, JudgeConfiguration, LDAIAgent, LDAIAgentConfig, - LDAIAgentDefaults, LDMessage, ModelConfig, ProviderConfig, AIAgentGraph, AIAgentGraphEdge) + LDAIAgentDefaults, LDMessage, ModelConfig, ProviderConfig, AIAgentGraphConfig, Edge) from ldai.providers.types import EvalScore, JudgeResponse __all__ = [ @@ -19,15 +19,15 @@ 'AIAgentConfigDefault', 'AIAgentConfigRequest', 'AIAgents', - 'AIAgentGraph', - 'AIAgentGraphEdge', + 'AIAgentGraphConfig', + 'Edge', 'AICompletionConfig', 'AICompletionConfigDefault', 'AIJudgeConfig', 'AIJudgeConfigDefault', 'Chat', 'EvalScore', - 'AgentGraph', + 'AgentGraphDefinition', 'Judge', 'JudgeConfiguration', 'JudgeResponse', 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 83b3755..86127a2 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -1,28 +1,24 @@ """Graph implementation for managing AI agent graphs.""" -from typing import Any, Callable, Dict, List, Optional, Set -from ldai.models import AIAgentGraph, AIAgentConfig, AIAgentGraphEdge +from typing import Any, Callable, Dict, List, Set +from ldai.models import AIAgentGraphConfig, AIAgentConfig, Edge from ldclient import Context - +DEFAULT_FALSE = AIAgentConfig(key="", enabled=False) class AgentGraphNode: """ Node in an agent graph. """ - default_false = AIAgentConfig(key="", enabled=False) - def __init__( self, key: str, config: AIAgentConfig, - children: List[AIAgentGraphEdge], - parent_graph: "AgentGraph", + children: List[Edge], ): self._key = key self._config = config self._children = children - self._parent_graph = parent_graph def get_key(self) -> str: """Get the key of the node.""" @@ -32,124 +28,85 @@ def get_config(self) -> AIAgentConfig: """Get the config of the node.""" return self._config - def get_edges(self) -> List[AIAgentGraphEdge]: - """Get the edges of the node.""" - return self._children - - def get_child_nodes(self) -> List["AgentGraphNode"]: - """Get the child nodes of the node as AgentGraphNode objects.""" - return [ - self._parent_graph.get_node(edge.targetConfig) for edge in self._children - ] - def is_terminal(self) -> bool: """Check if the node is a terminal node.""" - return len(self._children) == 0 - - def get_parent_nodes(self) -> List["AgentGraphNode"]: - """Get the parent nodes of the node as AgentGraphNode objects.""" - return [ - self._parent_graph.get_node(edge.sourceConfig) - for edge in self._parent_graph._get_parent_edges(self._key) - ] - - def traverse( - self, fn: Callable[["AgentGraphNode", Dict[str, Any]], None], execution_context: Dict[str, Any] = {}, visited: Optional[Set[str]] = None - ) -> None: - """Traverse the graph downwardly from this node, calling fn on each node.""" - if visited is None: - visited = set() - - # Avoid cycles by tracking visited nodes - if self._key in visited: - return - - visited.add(self._key) - fn(self, execution_context) - - for child in self._children: - node = self._parent_graph.get_node(child.targetConfig) - if node is not None: - node.traverse(fn, execution_context, visited) - - def reverse_traverse( - self, - fn: Callable[["AgentGraphNode", Dict[str, Any]], None], - execution_context: Dict[str, Any] = {}, - visited: Optional[Set[str]] = None, - ) -> None: - """Reverse traverse the graph upwardly from this node, calling fn on each node.""" - if visited is None: - visited = set() - - # Avoid cycles by tracking visited nodes - if self._key in visited: - return - - visited.add(self._key) - fn(self, execution_context) - - for parent in self._parent_graph._get_parent_edges(self._key): - node = self._parent_graph.get_node(parent.sourceConfig) - if node is not None: - node.reverse_traverse(fn, execution_context, visited) + return len(self._children) == 0 + def get_edges(self) -> List[Edge]: + """Get the edges of the node.""" + return self._children -class AgentGraph: +class AgentGraphDefinition: """ Graph implementation for managing AI agent graphs. """ - - default_false = AIAgentConfig(key="", enabled=False) - def __init__( self, - agent_graph: AIAgentGraph, + agent_graph: AIAgentGraphConfig, + nodes: Dict[str, AgentGraphNode], context: Context, - get_agent: Callable[[str, Context, dict], AIAgentConfig], ): self._agent_graph = agent_graph self._context = context - self._get_agent = get_agent - self._nodes = self._build_nodes() + self._nodes = nodes - def _build_nodes(self) -> Dict[str, AgentGraphNode]: + @staticmethod + def build_nodes( + agent_graph: AIAgentGraphConfig, + graph_nodes: Dict[str, AIAgentConfig], + ) -> Dict[str, "AgentGraphNode"]: """Build the nodes of the graph into AgentGraphNode objects.""" nodes = { - self._agent_graph.rootConfigKey: AgentGraphNode( - self._agent_graph.rootConfigKey, - self._get_agent( - self._agent_graph.rootConfigKey, self._context, self.default_false - ), - self._get_child_edges(self._agent_graph.rootConfigKey), - self, + agent_graph.root_config_key: AgentGraphNode( + agent_graph.root_config_key, + graph_nodes[agent_graph.root_config_key], + [ + edge + for edge in agent_graph.edges + if edge.source_config == agent_graph.root_config_key + ], ), } - for edge in self._agent_graph.edges: - nodes[edge.targetConfig] = AgentGraphNode( - edge.targetConfig, - self._get_agent(edge.targetConfig, self._context, self.default_false), - self._get_child_edges(edge.targetConfig), - self, + for edge in agent_graph.edges: + nodes[edge.target_config] = AgentGraphNode( + edge.target_config, + graph_nodes[edge.target_config], + [ + e + for e in agent_graph.edges + if e.source_config == edge.target_config + ], ) return nodes - def _get_child_edges(self, config_key: str) -> List[AIAgentGraphEdge]: + def get_node(self, key: str) -> AgentGraphNode | None: + """Get a node by its key.""" + return self._nodes.get(key) + + def _get_child_edges(self, config_key: str) -> List[Edge]: """Get the child edges of the given config.""" return [ edge for edge in self._agent_graph.edges - if edge.sourceConfig == config_key + if edge.source_config == config_key ] - def _get_parent_edges(self, config_key: str) -> List[AIAgentGraphEdge]: - """Get the parent edges of the given config.""" + def get_child_nodes(self, node_key: str) -> List[AgentGraphNode]: + """Get the child nodes of the given node key as AgentGraphNode objects.""" return [ - edge + self.get_node(edge.target_config) + for edge in self._agent_graph.edges + if edge.source_config == node_key and self.get_node(edge.target_config) is not None + ] + + def get_parent_nodes(self, node_key: str) -> List[AgentGraphNode]: + """Get the parent nodes of the given node key as AgentGraphNode objects.""" + return [ + self.get_node(edge.source_config) for edge in self._agent_graph.edges - if edge.targetConfig == config_key + if edge.target_config == node_key and self.get_node(edge.source_config) is not None ] def _collect_nodes( @@ -170,33 +127,18 @@ def _collect_nodes( nodes_by_depth[node_depth] = [] nodes_by_depth[node_depth].append(node) - for child in node.get_child_nodes(): + for child in self.get_child_nodes(node_key): self._collect_nodes(child, node_depths, nodes_by_depth, visited) def terminal_nodes(self) -> List[AgentGraphNode]: """Get the terminal nodes of the graph, meaning any nodes without children.""" return [ - node for node in self._nodes.values() if len(node.get_child_nodes()) == 0 + node for node in self._nodes.values() if len(self.get_child_nodes(node.get_key())) == 0 ] def root(self) -> AgentGraphNode | None: """Get the root node of the graph.""" - config = self._get_agent( - self._agent_graph.rootConfigKey, self._context, self.default_false - ) - - if config.enabled is False: - return None - - children = [ - edge - for edge in self._agent_graph.edges - if edge.sourceConfig == self._agent_graph.rootConfigKey - ] - - node = AgentGraphNode(self._agent_graph.rootConfigKey, config, children, self) - - return node + return self._nodes[self._agent_graph.root_config_key] def traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], None], execution_context: Dict[str, Any] = {}) -> None: """Traverse from the root down to terminal nodes, visiting nodes in order of depth. @@ -215,7 +157,8 @@ def traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], None], execu depth += 1 for node in current_level: - for child in node.get_child_nodes(): + node_key = node.get_key() + for child in self.get_child_nodes(node_key): child_key = child.get_key() # Defer this child to the next level if it's at a longer path if child_key not in node_depths or ( @@ -236,7 +179,7 @@ def traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], None], execu for node in nodes_by_depth[depth_level]: execution_context[node.get_key()] = fn(node, execution_context) - return execution_context[self._agent_graph.rootConfigKey] + return execution_context[self._agent_graph.root_config_key] def reverse_traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], execution_context: Dict[str, Any] = {}) -> None: """Traverse from terminal nodes up to the root, visiting nodes level by level. @@ -247,7 +190,7 @@ def reverse_traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any] visited: Set[str] = set() current_level: List[AgentGraphNode] = terminal_nodes - root_key = self._agent_graph.rootConfigKey + root_key = self._agent_graph.root_config_key root_node_seen = False while current_level: @@ -266,7 +209,7 @@ def reverse_traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any] execution_context[node_key] = fn(node, execution_context) - for parent in node.get_parent_nodes(): + for parent in self.get_parent_nodes(node_key): parent_key = parent.get_key() if parent_key not in visited: next_level.append(parent) @@ -280,8 +223,6 @@ def reverse_traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any] if root_node is not None: execution_context[root_node.get_key()] = fn(root_node, execution_context) - return execution_context[self._agent_graph.rootConfigKey] + return execution_context[self._agent_graph.root_config_key] + - def get_node(self, key: str) -> AgentGraphNode | None: - """Get a node by its key.""" - return self._nodes.get(key) diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 76b9c4a..083a961 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -6,13 +6,13 @@ from ldai import log from ldai.chat import Chat -from ldai.agent_graph import AgentGraph +from ldai.agent_graph import AgentGraphDefinition from ldai.judge import Judge from ldai.models import (AIAgentConfig, AIAgentConfigDefault, - AIAgentConfigRequest, AIAgents, AICompletionConfig, + AIAgentConfigRequest, AIAgentGraphConfig, AIAgents, AICompletionConfig, AICompletionConfigDefault, AIJudgeConfig, AIJudgeConfigDefault, JudgeConfiguration, LDMessage, - ModelConfig, ProviderConfig, AIAgentGraph, AIAgentGraphEdge) + ModelConfig, ProviderConfig, Edge) from ldai.providers.ai_provider_factory import AIProviderFactory from ldai.tracker import LDAIConfigTracker @@ -424,87 +424,64 @@ def agent_graph( self, key: str, context: Context, - ) -> AIAgentGraph: + ) -> AgentGraphDefinition: """ Retrieve an AI agent graph. """ variation = self._client.variation(key, context, {}) - mock_variation = { - "key": "test-agent-graph", - "name": "Test Agent Graph", - "rootConfigKey": "cruise-ship-information-agent", - "description": "Test Agent Graph Description", - "edges": [ - { - "key": "edge-cruise-ship-information-agent-get-weather-for-location", - "sourceConfig": "cruise-ship-information-agent", - "targetConfig": "get-weather-for-location", - "handoff": {}, - }, - { - "key": "edge-cruise-ship-information-agent-ships-in-port-agent", - "sourceConfig": "cruise-ship-information-agent", - "targetConfig": "ships-in-port-agent", - "handoff": {}, - }, - { - "key": "edge-ships-in-port-agent-vessel-details-agent", - "sourceConfig": "ships-in-port-agent", - "targetConfig": "vessel-details-agent", - "handoff": {}, - }, - { - "key": "edge-vessel-details-agent-cruise-information-synthesizer", - "sourceConfig": "vessel-details-agent", - "targetConfig": "cruise-information-synthesizer", - "handoff": {}, - }, - { - "key": "edge-get-weather-for-location-cruise-information-synthesizer", - "sourceConfig": "get-weather-for-location", - "targetConfig": "cruise-information-synthesizer", - "handoff": {}, - }, - ], + + if not variation.get("rootConfigKey"): + log.debug(f"Agent graph {key} is disabled, no root config key found") + return { "enabled": False, "graph": None } + + all_agent_keys = [variation["rootConfigKey"]] + [edge["targetConfig"] for edge in variation["edges"]] + agent_configs = { + key: self.agent_config(key, context, AIAgentConfigDefault(enabled=False)) + for key in all_agent_keys } - return AgentGraph( - agent_graph=AIAgentGraph( - key=mock_variation['key'], - name=mock_variation['name'], - rootConfigKey=mock_variation['rootConfigKey'], + + + if not all(config.enabled for config in agent_configs.values()): + log.debug(f"Agent graph {key} is disabled, not all agent configs are enabled") + return { + "enabled": False, + "graph": None, + } + + try: + agent_graph_config = AIAgentGraphConfig( + key=variation["key"], + name=variation["name"], + root_config_key=variation["rootConfigKey"], edges=[ - AIAgentGraphEdge( - key=edge.get('key', ''), - sourceConfig=edge.get('sourceConfig', ''), - targetConfig=edge.get('targetConfig', ''), - handoff=edge.get('handoff', {}), + Edge( + key=edge.get("key", ""), + source_config=edge.get("sourceConfig", ""), + target_config=edge.get("targetConfig", ""), + handoff=edge.get("handoff", {}), ) - for edge in mock_variation['edges'] + for edge in variation["edges"] ], - description=mock_variation['description'], - ), - context=context, - get_agent=self.agent_config, + description=variation["description"], + ) + except Exception as e: + log.debug(f"Agent graph {key} is disabled, invalid agent graph config") + return { "enabled": False, "graph": None } + + + nodes = AgentGraphDefinition.build_nodes( + agent_graph_config, + agent_configs, ) - return AgentGraph( - agent_graph=AIAgentGraph( - key=variation.key, - name=variation.name, - rootConfigKey=variation.rootConfigKey, - edges=[ - AIAgentGraphEdge( - key=edge.key, - sourceConfig=edge.sourceConfig, - targetConfig=edge.targetConfig, - handoff=edge.handoff - ) - for edge in variation.edges - ] + return { + "enabled": True, + "graph": AgentGraphDefinition( + agent_graph=agent_graph_config, + nodes=nodes, + context=context, ), - context=context, - get_variation=self._client.variation, - ) + } def agents( self, diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index eadd072..a449989 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -342,27 +342,27 @@ class AIAgentConfigRequest: # AI Config Agent Graph Edge Type # ============================================================================ @dataclass -class AIAgentGraphEdge: +class Edge: """ Edge configuration for an agent graph. """ key: str - sourceConfig: str - targetConfig: str + source_config: str + target_config: str handoff: Optional[dict] = None # ============================================================================ # AI Config Agent Graph # ============================================================================ @dataclass -class AIAgentGraph: +class AIAgentGraphConfig: """ Agent graph configuration. """ key: str name: str - rootConfigKey: str - edges: List[AIAgentGraphEdge] + root_config_key: str + edges: List[Edge] description: Optional[str] = '' # ============================================================================ diff --git a/packages/sdk/server-ai/tests/test_agent_graph.py b/packages/sdk/server-ai/tests/test_agent_graph.py new file mode 100644 index 0000000..2e82e51 --- /dev/null +++ b/packages/sdk/server-ai/tests/test_agent_graph.py @@ -0,0 +1,441 @@ +import pytest +from ldclient import Config, Context, LDClient +from ldclient.integrations.test_data import TestData + +from ldai import ( + LDAIClient, + AIAgentGraphConfig, + AgentGraphDefinition, + AIAgentConfig, + Edge, +) + + +@pytest.fixture +def td() -> TestData: + td = TestData.data_source() + # Agent graph with depth of 1 + td.update( + td.flag("test-agent-graph") + .variations( + { + "key": "test-agent-graph", + "name": "Test Agent Graph", + "rootConfigKey": "customer-support-agent", + "edges": [ + { + "key": "edge-customer-support-agent-personalized-agent", + "sourceConfig": "customer-support-agent", + "targetConfig": "personalized-agent", + }, + { + "key": "edge-customer-support-agent-multi-context-agent", + "sourceConfig": "customer-support-agent", + "targetConfig": "multi-context-agent", + }, + { + "key": "edge-customer-support-agent-minimal-agent", + "sourceConfig": "customer-support-agent", + "targetConfig": "minimal-agent", + }, + ], + "description": "Test agent graph", + "_ldMeta": { + "enabled": True, + "variationKey": "test-agent-graph", + "version": 1, + }, + } + ) + .variation_for_all(0) + ) + # Agent graph with depth of 3 + td.update( + td.flag("test-agent-graph-depth-3") + .variations( + { + "key": "test-agent-graph-depth-3", + "name": "Test Agent Graph with Depth of 3", + "rootConfigKey": "customer-support-agent", + "edges": [ + { + "key": "edge-customer-support-agent-personalized-agent", + "sourceConfig": "customer-support-agent", + "targetConfig": "personalized-agent", + "handoff": { "state": "from-root-to-personalized" } + }, + { + "key": "edge-personalized-agent-multi-context-agent", + "sourceConfig": "personalized-agent", + "targetConfig": "multi-context-agent", + }, + { + "key": "edge-multi-context-agent-minimal-agent", + "sourceConfig": "multi-context-agent", + "targetConfig": "minimal-agent", + "handoff": {"state": "from-multi-context-to-minimal"}, + }, + { + "key": "edge-customer-support-agent-minimal-agent", + "sourceConfig": "customer-support-agent", + "targetConfig": "minimal-agent", + "handoff": { "state": "from-root-to-minimal" } + }, + ], + "description": "Test agent graph with depth of 3", + "_ldMeta": { + "enabled": True, + "variationKey": "test-agent-graph-depth-3", + "version": 1, + }, + } + ) + .variation_for_all(0) + ) + + # Agent graph with disabled agent included - invalid + td.update( + td.flag("test-agent-graph-disabled-agent") + .variations( + { + "key": "test-agent-graph-disabled-agent", + "name": "Test Agent Graph with Disabled Agent", + "rootConfigKey": "customer-support-agent", + "edges": [ + { + "key": "edge-customer-support-agent-personalized-agent", + "sourceConfig": "customer-support-agent", + "targetConfig": "disabled-agent", + }, + ], + "description": "Test agent graph with disabled agent", + "_ldMeta": { + "enabled": True, + "variationKey": "test-agent-graph-disabled-agent", + "version": 1, + }, + } + ) + .variation_for_all(0) + ) + + # Agent graph with no root key - invalid + td.update( + td.flag("test-agent-graph-no-root-key") + .variations( + { + "name": "Test Agent Graph with No Root Key", + "key": "test-agent-graph-no-root-key", + "edges": [], + } + ) + .variation_for_all(0) + ) + + # Single agent with instructions + td.update( + td.flag("customer-support-agent") + .variations( + { + "model": { + "name": "gpt-4", + "parameters": {"temperature": 0.3, "maxTokens": 2048}, + }, + "provider": {"name": "openai"}, + "instructions": "You are a helpful customer support agent for {{company_name}}. Always be polite and professional.", + "_ldMeta": { + "enabled": True, + "variationKey": "agent-v1", + "version": 1, + "mode": "agent", + }, + } + ) + .variation_for_all(0) + ) + + # Agent with context interpolation + td.update( + td.flag("personalized-agent") + .variations( + { + "model": {"name": "claude-3", "parameters": {"temperature": 0.5}}, + "instructions": "Hello {{ldctx.name}}! I am your personal assistant. Your user key is {{ldctx.key}}.", + "_ldMeta": { + "enabled": True, + "variationKey": "personal-v1", + "version": 2, + "mode": "agent", + }, + } + ) + .variation_for_all(0) + ) + + # Agent with multi-context interpolation + td.update( + td.flag("multi-context-agent") + .variations( + { + "model": {"name": "gpt-3.5-turbo"}, + "instructions": "Welcome {{ldctx.user.name}} from {{ldctx.org.name}}! Your organization tier is {{ldctx.org.tier}}.", + "_ldMeta": { + "enabled": True, + "variationKey": "multi-v1", + "version": 1, + "mode": "agent", + }, + } + ) + .variation_for_all(0) + ) + + # Disabled agent + td.update( + td.flag("disabled-agent") + .variations( + { + "model": {"name": "gpt-4"}, + "instructions": "This agent is disabled.", + "_ldMeta": { + "enabled": False, + "variationKey": "disabled-v1", + "version": 1, + "mode": "agent", + }, + } + ) + .variation_for_all(0) + ) + + # Agent with minimal metadata + td.update( + td.flag("minimal-agent") + .variations( + { + "instructions": "Minimal agent configuration.", + "_ldMeta": {"enabled": True}, + } + ) + .variation_for_all(0) + ) + + return td + + +@pytest.fixture +def client(td: TestData) -> LDClient: + config = Config("sdk-key", update_processor_class=td, send_events=False) + return LDClient(config=config) + + +@pytest.fixture +def ldai_client(client: LDClient) -> LDAIClient: + return LDAIClient(client) + + +def test_agent_graph_method(ldai_client: LDAIClient): + graph = ldai_client.agent_graph("test-agent-graph", Context.create("user-key")) + + assert graph["enabled"] is True + assert graph["graph"] is not None + assert graph["graph"].root() is not None + assert graph["graph"].root().get_key() == "customer-support-agent" + assert len(graph["graph"].get_child_nodes("customer-support-agent")) == 3 + assert len(graph["graph"].get_child_nodes("personalized-agent")) == 0 + assert len(graph["graph"].get_child_nodes("multi-context-agent")) == 0 + assert len(graph["graph"].get_child_nodes("minimal-agent")) == 0 + + +def test_agent_graph_method_disabled_agent(ldai_client: LDAIClient): + graph = ldai_client.agent_graph( + "test-agent-graph-disabled-agent", Context.create("user-key") + ) + + assert graph["enabled"] is False + assert graph["graph"] is None + + +def test_agent_graph_method_no_root_key(ldai_client: LDAIClient): + graph = ldai_client.agent_graph( + "test-agent-graph-no-root-key", Context.create("user-key") + ) + + assert graph["enabled"] is False + assert graph["graph"] is None + + +def test_agent_graph_build_nodes(ldai_client: LDAIClient): + graph_config = ldai_client._client.variation( + "test-agent-graph", Context.create("user-key"), {} + ) + + ai_graph_config = AIAgentGraphConfig( + key=graph_config["key"], + name=graph_config["name"], + root_config_key=graph_config["rootConfigKey"], + edges=[ + Edge( + key=edge.get("key", ""), + source_config=edge.get("sourceConfig", ""), + target_config=edge.get("targetConfig", ""), + handoff=edge.get("handoff", {}), + ) + for edge in graph_config["edges"] + ], + description=graph_config["description"], + ) + + nodes = AgentGraphDefinition.build_nodes( + ai_graph_config, + { + "customer-support-agent": AIAgentConfig( + key="customer-support-agent", enabled=True + ), + "personalized-agent": AIAgentConfig(key="personalized-agent", enabled=True), + "multi-context-agent": AIAgentConfig( + key="multi-context-agent", enabled=True + ), + "minimal-agent": AIAgentConfig(key="minimal-agent", enabled=True), + }, + ) + + assert nodes["customer-support-agent"] is not None + assert nodes["personalized-agent"] is not None + assert nodes["multi-context-agent"] is not None + assert nodes["minimal-agent"] is not None + + assert len(nodes["customer-support-agent"].get_edges()) == 3 + assert len(nodes["personalized-agent"].get_edges()) == 0 + assert len(nodes["multi-context-agent"].get_edges()) == 0 + assert len(nodes["minimal-agent"].get_edges()) == 0 + + assert type(nodes["customer-support-agent"].get_config()) is AIAgentConfig + assert type(nodes["personalized-agent"].get_config()) is AIAgentConfig + assert type(nodes["multi-context-agent"].get_config()) is AIAgentConfig + assert type(nodes["minimal-agent"].get_config()) is AIAgentConfig + + assert type(nodes["customer-support-agent"].get_edges()[0]) is Edge + + +def test_agent_graph_get_methods(ldai_client: LDAIClient): + graph = ldai_client.agent_graph("test-agent-graph", Context.create("user-key"))[ + "graph" + ] + + assert graph.root() is not None + assert graph.root().get_key() == "customer-support-agent" + assert graph.get_node("customer-support-agent") is not None + assert graph.get_node("personalized-agent") is not None + assert graph.get_node("multi-context-agent") is not None + + children = graph.get_child_nodes("customer-support-agent") + assert len(children) == 3 + assert children[0].get_key() == "personalized-agent" + assert children[1].get_key() == "multi-context-agent" + assert children[2].get_key() == "minimal-agent" + + parents = graph.get_parent_nodes("personalized-agent") + assert len(parents) == 1 + assert parents[0].get_key() == "customer-support-agent" + + parents = graph.get_parent_nodes("multi-context-agent") + assert len(parents) == 1 + assert parents[0].get_key() == "customer-support-agent" + + terminal = graph.terminal_nodes() + assert len(terminal) == 3 + assert terminal[0].get_key() == "personalized-agent" + assert terminal[1].get_key() == "multi-context-agent" + assert terminal[2].get_key() == "minimal-agent" + + assert graph.root().is_terminal() is False + assert graph.get_node("customer-support-agent").is_terminal() is False + assert graph.get_node("personalized-agent").is_terminal() is True + assert graph.get_node("multi-context-agent").is_terminal() is True + assert graph.get_node("minimal-agent").is_terminal() is True + + +def test_agent_graph_traverse(ldai_client: LDAIClient): + graph = ldai_client.agent_graph( + "test-agent-graph-depth-3", Context.create("user-key") + )["graph"] + + context = {} + order = [] + + def handle_traverse(node, context): + # Asserting that returned values are included in the context + for previousKey in order: + assert previousKey in context + assert context[previousKey] == previousKey + "-test" + order.append(node.get_key()) + return node.get_key() + "-test" + + graph.traverse(handle_traverse, context) + # Asserting that we traverse in the expected order + # This config specifically has nodes connecting from depth 2->3 and root->3 to ensure the root node is visited first + # and minimal-agent is visited last + assert order == [ + "customer-support-agent", + "personalized-agent", + "multi-context-agent", + "minimal-agent", + ] + + +def test_agent_graph_reverse_traverse(ldai_client: LDAIClient): + graph = ldai_client.agent_graph( + "test-agent-graph-depth-3", Context.create("user-key") + )["graph"] + + context = {} + order = [] + + def handle_reverse_traverse(node, context): + # Asserting that returned values are included in the context + for previousKey in order: + assert previousKey in context + assert context[previousKey] == previousKey + "-test" + order.append(node.get_key()) + return node.get_key() + "-test" + + graph.reverse_traverse(handle_reverse_traverse, context) + # Asserting that we traverse in the expected order + # This config specifically has nodes connecting from depth 2->3 and root->3 to ensure the root node is visited last + assert order == [ + "minimal-agent", + "multi-context-agent", + "personalized-agent", + "customer-support-agent", + ] + + +def test_agent_graph_handoff(ldai_client: LDAIClient): + graph = ldai_client.agent_graph( + "test-agent-graph-depth-3", Context.create("user-key") + )["graph"] + + context = {} + + def handle_traverse(node, context): + if node.get_key() == "multi-context-agent": + first_edge = node.get_edges()[0] + assert first_edge.handoff == {"state": "from-multi-context-to-minimal"} + assert first_edge.source_config == "multi-context-agent" + assert first_edge.target_config == "minimal-agent" + assert first_edge.key == "edge-multi-context-agent-minimal-agent" + if node.get_key() == "customer-support-agent": + first_edge = node.get_edges()[0] + second_edge = node.get_edges()[1] + assert first_edge.handoff == {"state": "from-root-to-personalized"} + assert second_edge.handoff == {"state": "from-root-to-minimal"} + assert first_edge.source_config == "customer-support-agent" + assert first_edge.target_config == "personalized-agent" + assert first_edge.key == "edge-customer-support-agent-personalized-agent" + assert second_edge.source_config == "customer-support-agent" + assert second_edge.target_config == "minimal-agent" + assert second_edge.key == "edge-customer-support-agent-minimal-agent" + return None + + graph.traverse(handle_traverse, context) From a4c735aa045012819cd43f84d3f8c65805920032 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Wed, 14 Jan 2026 16:35:10 -0900 Subject: [PATCH 03/11] [REL-11697] resolve linting errors --- packages/sdk/server-ai/src/ldai/__init__.py | 12 ++-- .../src/ldai/agent_graph/__init__.py | 72 +++++++++++-------- packages/sdk/server-ai/src/ldai/client.py | 40 +++++------ packages/sdk/server-ai/src/ldai/models.py | 22 +++++- 4 files changed, 91 insertions(+), 55 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 2566b2a..2a93086 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -2,15 +2,16 @@ from ldclient import log +from ldai.agent_graph import AgentGraphDefinition from ldai.chat import Chat from ldai.client import LDAIClient -from ldai.agent_graph import AgentGraphDefinition from ldai.judge import Judge from ldai.models import ( # Deprecated aliases for backward compatibility - AIAgentConfig, AIAgentConfigDefault, AIAgentConfigRequest, AIAgents, - AICompletionConfig, AICompletionConfigDefault, AIConfig, AIJudgeConfig, - AIJudgeConfigDefault, JudgeConfiguration, LDAIAgent, LDAIAgentConfig, - LDAIAgentDefaults, LDMessage, ModelConfig, ProviderConfig, AIAgentGraphConfig, Edge) + AIAgentConfig, AIAgentConfigDefault, AIAgentConfigRequest, + AIAgentGraphConfig, AIAgentGraphResponse, AIAgents, AICompletionConfig, + AICompletionConfigDefault, AIConfig, AIJudgeConfig, AIJudgeConfigDefault, + Edge, JudgeConfiguration, LDAIAgent, LDAIAgentConfig, LDAIAgentDefaults, + LDMessage, ModelConfig, ProviderConfig) from ldai.providers.types import EvalScore, JudgeResponse __all__ = [ @@ -20,6 +21,7 @@ 'AIAgentConfigRequest', 'AIAgents', 'AIAgentGraphConfig', + 'AIAgentGraphResponse', '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 86127a2..1a2b471 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -1,10 +1,14 @@ """Graph implementation for managing AI agent graphs.""" -from typing import Any, Callable, Dict, List, Set -from ldai.models import AIAgentGraphConfig, AIAgentConfig, Edge +from typing import Any, Callable, Dict, List, Optional, Set + from ldclient import Context +from ldai.models import AIAgentConfig, AIAgentGraphConfig, Edge + DEFAULT_FALSE = AIAgentConfig(key="", enabled=False) + + class AgentGraphNode: """ Node in an agent graph. @@ -36,10 +40,12 @@ def get_edges(self) -> List[Edge]: """Get the edges of the node.""" return self._children + class AgentGraphDefinition: """ Graph implementation for managing AI agent graphs. """ + def __init__( self, agent_graph: AIAgentGraphConfig, @@ -72,42 +78,40 @@ def build_nodes( nodes[edge.target_config] = AgentGraphNode( edge.target_config, graph_nodes[edge.target_config], - [ - e - for e in agent_graph.edges - if e.source_config == edge.target_config - ], + [e for e in agent_graph.edges if e.source_config == edge.target_config], ) return nodes - def get_node(self, key: str) -> AgentGraphNode | None: + def get_node(self, key: str) -> Optional[AgentGraphNode]: """Get a node by its key.""" return self._nodes.get(key) def _get_child_edges(self, config_key: str) -> List[Edge]: """Get the child edges of the given config.""" return [ - edge - for edge in self._agent_graph.edges - if edge.source_config == config_key + edge for edge in self._agent_graph.edges if edge.source_config == config_key ] def get_child_nodes(self, node_key: str) -> List[AgentGraphNode]: """Get the child nodes of the given node key as AgentGraphNode objects.""" - return [ - self.get_node(edge.target_config) - for edge in self._agent_graph.edges - if edge.source_config == node_key and self.get_node(edge.target_config) is not None - ] + nodes: List[AgentGraphNode] = [] + for edge in self._agent_graph.edges: + if edge.source_config == node_key: + node = self.get_node(edge.target_config) + if node is not None: + nodes.append(node) + return nodes def get_parent_nodes(self, node_key: str) -> List[AgentGraphNode]: """Get the parent nodes of the given node key as AgentGraphNode objects.""" - return [ - self.get_node(edge.source_config) - for edge in self._agent_graph.edges - if edge.target_config == node_key and self.get_node(edge.source_config) is not None - ] + nodes: List[AgentGraphNode] = [] + for edge in self._agent_graph.edges: + if edge.target_config == node_key: + node = self.get_node(edge.source_config) + if node is not None: + nodes.append(node) + return nodes def _collect_nodes( self, @@ -133,14 +137,20 @@ def _collect_nodes( def terminal_nodes(self) -> List[AgentGraphNode]: """Get the terminal nodes of the graph, meaning any nodes without children.""" return [ - node for node in self._nodes.values() if len(self.get_child_nodes(node.get_key())) == 0 + node + for node in self._nodes.values() + if len(self.get_child_nodes(node.get_key())) == 0 ] - def root(self) -> AgentGraphNode | None: + def root(self) -> Optional[AgentGraphNode]: """Get the root node of the graph.""" return self._nodes[self._agent_graph.root_config_key] - def traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], None], execution_context: Dict[str, Any] = {}) -> None: + def traverse( + self, + fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], + execution_context: Dict[str, Any] = {}, + ) -> None: """Traverse from the root down to terminal nodes, visiting nodes in order of depth. Nodes with the longest paths from the root (deepest nodes) will always be visited last.""" root_node = self.root() @@ -181,7 +191,11 @@ def traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], None], execu return execution_context[self._agent_graph.root_config_key] - def reverse_traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], execution_context: Dict[str, Any] = {}) -> None: + def reverse_traverse( + self, + fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], + execution_context: Dict[str, Any] = {}, + ) -> None: """Traverse from terminal nodes up to the root, visiting nodes level by level. The root node will always be visited last, even if multiple paths converge at it.""" terminal_nodes = self.terminal_nodes() @@ -208,7 +222,7 @@ def reverse_traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any] continue execution_context[node_key] = fn(node, execution_context) - + for parent in self.get_parent_nodes(node_key): parent_key = parent.get_key() if parent_key not in visited: @@ -221,8 +235,8 @@ def reverse_traverse(self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any] if root_node_seen: root_node = self.root() if root_node is not None: - execution_context[root_node.get_key()] = fn(root_node, execution_context) + execution_context[root_node.get_key()] = fn( + root_node, execution_context + ) return execution_context[self._agent_graph.root_config_key] - - diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 083a961..ac0d5de 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -5,14 +5,15 @@ from ldclient.client import LDClient from ldai import log -from ldai.chat import Chat from ldai.agent_graph import AgentGraphDefinition +from ldai.chat import Chat from ldai.judge import Judge from ldai.models import (AIAgentConfig, AIAgentConfigDefault, - AIAgentConfigRequest, AIAgentGraphConfig, AIAgents, AICompletionConfig, + AIAgentConfigRequest, AIAgentGraphConfig, + AIAgentGraphResponse, AIAgents, AICompletionConfig, AICompletionConfigDefault, AIJudgeConfig, - AIJudgeConfigDefault, JudgeConfiguration, LDMessage, - ModelConfig, ProviderConfig, Edge) + AIJudgeConfigDefault, Edge, JudgeConfiguration, + LDMessage, ModelConfig, ProviderConfig) from ldai.providers.ai_provider_factory import AIProviderFactory from ldai.tracker import LDAIConfigTracker @@ -424,29 +425,29 @@ def agent_graph( self, key: str, context: Context, - ) -> AgentGraphDefinition: - """ + ) -> AIAgentGraphResponse: + """` Retrieve an AI agent graph. """ variation = self._client.variation(key, context, {}) if not variation.get("rootConfigKey"): log.debug(f"Agent graph {key} is disabled, no root config key found") - return { "enabled": False, "graph": None } + return AIAgentGraphResponse(enabled=False, graph=None) - all_agent_keys = [variation["rootConfigKey"]] + [edge["targetConfig"] for edge in variation["edges"]] + all_agent_keys = [variation["rootConfigKey"]] + [ + edge["targetConfig"] for edge in variation["edges"] + ] agent_configs = { key: self.agent_config(key, context, AIAgentConfigDefault(enabled=False)) for key in all_agent_keys } - if not all(config.enabled for config in agent_configs.values()): - log.debug(f"Agent graph {key} is disabled, not all agent configs are enabled") - return { - "enabled": False, - "graph": None, - } + log.debug( + f"Agent graph {key} is disabled, not all agent configs are enabled" + ) + return AIAgentGraphResponse(enabled=False, graph=None) try: agent_graph_config = AIAgentGraphConfig( @@ -466,22 +467,21 @@ def agent_graph( ) except Exception as e: log.debug(f"Agent graph {key} is disabled, invalid agent graph config") - return { "enabled": False, "graph": None } - + return AIAgentGraphResponse(enabled=False, graph=None) nodes = AgentGraphDefinition.build_nodes( agent_graph_config, agent_configs, ) - return { - "enabled": True, - "graph": AgentGraphDefinition( + return AIAgentGraphResponse( + enabled=True, + graph=AgentGraphDefinition( agent_graph=agent_graph_config, nodes=nodes, context=context, ), - } + ) def agents( self, diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index a449989..f9c736c 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -2,6 +2,7 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional, Union +from ldai.agent_graph import AgentGraphDefinition from ldai.tracker import LDAIConfigTracker @@ -341,16 +342,20 @@ class AIAgentConfigRequest: # ============================================================================ # AI Config Agent Graph Edge Type # ============================================================================ + + @dataclass class Edge: """ Edge configuration for an agent graph. """ + key: str source_config: str target_config: str handoff: Optional[dict] = None + # ============================================================================ # AI Config Agent Graph # ============================================================================ @@ -359,11 +364,26 @@ class AIAgentGraphConfig: """ Agent graph configuration. """ + key: str name: str root_config_key: str edges: List[Edge] - description: Optional[str] = '' + description: Optional[str] = "" + + +# ============================================================================ +# AI Config Agent Graph Response +# ============================================================================ +@dataclass +class AIAgentGraphResponse: + """ + Agent graph response. + """ + + enabled: bool + graph: Optional[AgentGraphDefinition] = None + # ============================================================================ # Deprecated Type Aliases for Backward Compatibility From c543ce6bb69bfce7a693954a0d0f6a0b91178f31 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 15 Jan 2026 09:37:56 -0900 Subject: [PATCH 04/11] [REL-11697] move response data class to avoid circular import --- packages/sdk/server-ai/src/ldai/__init__.py | 4 ++-- .../sdk/server-ai/src/ldai/agent_graph/__init__.py | 14 ++++++++++++++ packages/sdk/server-ai/src/ldai/client.py | 12 ++++++------ packages/sdk/server-ai/src/ldai/models.py | 13 ------------- 4 files changed, 22 insertions(+), 21 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index 2a93086..c606226 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -2,13 +2,13 @@ from ldclient import log -from ldai.agent_graph import AgentGraphDefinition +from ldai.agent_graph import AgentGraphDefinition, AIAgentGraphResponse from ldai.chat import Chat from ldai.client import LDAIClient from ldai.judge import Judge from ldai.models import ( # Deprecated aliases for backward compatibility AIAgentConfig, AIAgentConfigDefault, AIAgentConfigRequest, - AIAgentGraphConfig, AIAgentGraphResponse, AIAgents, AICompletionConfig, + AIAgentGraphConfig, AIAgents, AICompletionConfig, AICompletionConfigDefault, AIConfig, AIJudgeConfig, AIJudgeConfigDefault, Edge, JudgeConfiguration, LDAIAgent, LDAIAgentConfig, LDAIAgentDefaults, LDMessage, ModelConfig, ProviderConfig) 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 1a2b471..d124830 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -1,5 +1,6 @@ """Graph implementation for managing AI agent graphs.""" +from dataclasses import dataclass from typing import Any, Callable, Dict, List, Optional, Set from ldclient import Context @@ -240,3 +241,16 @@ def reverse_traverse( ) return execution_context[self._agent_graph.root_config_key] + + +# ============================================================================ +# AI Config Agent Graph Response +# ============================================================================ +@dataclass +class AIAgentGraphResponse: + """ + Agent graph response. + """ + + enabled: bool + graph: Optional[AgentGraphDefinition] = None diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index ac0d5de..57b735a 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -5,15 +5,15 @@ from ldclient.client import LDClient from ldai import log -from ldai.agent_graph import AgentGraphDefinition +from ldai.agent_graph import AgentGraphDefinition, AIAgentGraphResponse from ldai.chat import Chat from ldai.judge import Judge from ldai.models import (AIAgentConfig, AIAgentConfigDefault, - AIAgentConfigRequest, AIAgentGraphConfig, - AIAgentGraphResponse, AIAgents, AICompletionConfig, - AICompletionConfigDefault, AIJudgeConfig, - AIJudgeConfigDefault, Edge, JudgeConfiguration, - LDMessage, ModelConfig, ProviderConfig) + AIAgentConfigRequest, AIAgentGraphConfig, AIAgents, + AICompletionConfig, AICompletionConfigDefault, + AIJudgeConfig, AIJudgeConfigDefault, Edge, + JudgeConfiguration, LDMessage, ModelConfig, + ProviderConfig) from ldai.providers.ai_provider_factory import AIProviderFactory from ldai.tracker import LDAIConfigTracker diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index f9c736c..33d032b 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -372,19 +372,6 @@ class AIAgentGraphConfig: description: Optional[str] = "" -# ============================================================================ -# AI Config Agent Graph Response -# ============================================================================ -@dataclass -class AIAgentGraphResponse: - """ - Agent graph response. - """ - - enabled: bool - graph: Optional[AgentGraphDefinition] = None - - # ============================================================================ # Deprecated Type Aliases for Backward Compatibility # ============================================================================ From 889bd6c7bd2c7baf61cc38ce4b5b9afc1decd153 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 15 Jan 2026 10:14:03 -0900 Subject: [PATCH 05/11] [REL-11697] Simplify types; add enabled directly to objects; cursor feedback --- packages/sdk/server-ai/src/ldai/__init__.py | 3 +- .../src/ldai/agent_graph/__init__.py | 38 ++++++------ packages/sdk/server-ai/src/ldai/client.py | 60 +++++++++++++++---- packages/sdk/server-ai/src/ldai/models.py | 2 +- .../sdk/server-ai/tests/test_agent_graph.py | 34 +++++------ 5 files changed, 85 insertions(+), 52 deletions(-) diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index c606226..fe16129 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -2,7 +2,7 @@ from ldclient import log -from ldai.agent_graph import AgentGraphDefinition, AIAgentGraphResponse +from ldai.agent_graph import AgentGraphDefinition from ldai.chat import Chat from ldai.client import LDAIClient from ldai.judge import Judge @@ -21,7 +21,6 @@ 'AIAgentConfigRequest', 'AIAgents', 'AIAgentGraphConfig', - 'AIAgentGraphResponse', '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 d124830..8ea5378 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -46,16 +46,23 @@ class AgentGraphDefinition: """ Graph implementation for managing AI agent graphs. """ + enabled: bool def __init__( self, agent_graph: AIAgentGraphConfig, nodes: Dict[str, AgentGraphNode], context: Context, + enabled: bool, ): self._agent_graph = agent_graph self._context = context self._nodes = nodes + self.enabled = enabled + + def is_enabled(self) -> bool: + """Check if the graph is enabled.""" + return self.enabled @staticmethod def build_nodes( @@ -143,17 +150,20 @@ def terminal_nodes(self) -> List[AgentGraphNode]: if len(self.get_child_nodes(node.get_key())) == 0 ] - def root(self) -> Optional[AgentGraphNode]: + def root(self) -> AgentGraphNode: """Get the root node of the graph.""" - return self._nodes[self._agent_graph.root_config_key] + return self._nodes.get(self._agent_graph.root_config_key) def traverse( self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], - execution_context: Dict[str, Any] = {}, - ) -> None: + execution_context: Dict[str, Any] = None, + ) -> Any: """Traverse from the root down to terminal nodes, visiting nodes in order of depth. Nodes with the longest paths from the root (deepest nodes) will always be visited last.""" + if execution_context is None: + execution_context = {} + root_node = self.root() if root_node is None: return @@ -195,10 +205,14 @@ def traverse( def reverse_traverse( self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], - execution_context: Dict[str, Any] = {}, - ) -> None: + execution_context: Dict[str, Any] = None, + ) -> Any: + """Traverse from terminal nodes up to the root, visiting nodes level by level. The root node will always be visited last, even if multiple paths converge at it.""" + if execution_context is None: + execution_context = {} + terminal_nodes = self.terminal_nodes() if not terminal_nodes: return @@ -242,15 +256,3 @@ def reverse_traverse( return execution_context[self._agent_graph.root_config_key] - -# ============================================================================ -# AI Config Agent Graph Response -# ============================================================================ -@dataclass -class AIAgentGraphResponse: - """ - Agent graph response. - """ - - enabled: bool - graph: Optional[AgentGraphDefinition] = None diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 57b735a..57d809a 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -5,7 +5,7 @@ from ldclient.client import LDClient from ldai import log -from ldai.agent_graph import AgentGraphDefinition, AIAgentGraphResponse +from ldai.agent_graph import AgentGraphDefinition from ldai.chat import Chat from ldai.judge import Judge from ldai.models import (AIAgentConfig, AIAgentConfigDefault, @@ -425,7 +425,7 @@ def agent_graph( self, key: str, context: Context, - ) -> AIAgentGraphResponse: + ) -> AgentGraphDefinition: """` Retrieve an AI agent graph. """ @@ -433,10 +433,22 @@ def agent_graph( if not variation.get("rootConfigKey"): log.debug(f"Agent graph {key} is disabled, no root config key found") - return AIAgentGraphResponse(enabled=False, graph=None) + return AgentGraphDefinition( + AIAgentGraphConfig( + key=key, + name="", + root_config_key="", + edges=[], + description="", + enabled=False, + ), + nodes={}, + context=context, + enabled=False, + ) all_agent_keys = [variation["rootConfigKey"]] + [ - edge["targetConfig"] for edge in variation["edges"] + edge["targetConfig"] for edge in variation.get("edges", []) ] agent_configs = { key: self.agent_config(key, context, AIAgentConfigDefault(enabled=False)) @@ -447,7 +459,19 @@ def agent_graph( log.debug( f"Agent graph {key} is disabled, not all agent configs are enabled" ) - return AIAgentGraphResponse(enabled=False, graph=None) + return AgentGraphDefinition( + AIAgentGraphConfig( + key=key, + name="", + root_config_key="", + edges=[], + description="", + enabled=False, + ), + nodes={}, + context=context, + enabled=False, + ) try: agent_graph_config = AIAgentGraphConfig( @@ -467,20 +491,30 @@ def agent_graph( ) except Exception as e: log.debug(f"Agent graph {key} is disabled, invalid agent graph config") - return AIAgentGraphResponse(enabled=False, graph=None) + return AgentGraphDefinition( + AIAgentGraphConfig( + key=key, + name="", + root_config_key="", + edges=[], + description="", + enabled=False, + ), + nodes={}, + context=context, + enabled=False, + ) nodes = AgentGraphDefinition.build_nodes( agent_graph_config, agent_configs, ) - return AIAgentGraphResponse( - enabled=True, - graph=AgentGraphDefinition( - agent_graph=agent_graph_config, - nodes=nodes, - context=context, - ), + return AgentGraphDefinition( + agent_graph=agent_graph_config, + nodes=nodes, + context=context, + enabled=agent_graph_config.enabled, ) def agents( diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index 33d032b..f43c0a8 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -2,7 +2,6 @@ from dataclasses import dataclass, field from typing import Any, Dict, List, Literal, Optional, Union -from ldai.agent_graph import AgentGraphDefinition from ldai.tracker import LDAIConfigTracker @@ -370,6 +369,7 @@ class AIAgentGraphConfig: root_config_key: str edges: List[Edge] description: Optional[str] = "" + enabled: bool = True # ============================================================================ diff --git a/packages/sdk/server-ai/tests/test_agent_graph.py b/packages/sdk/server-ai/tests/test_agent_graph.py index 2e82e51..de584de 100644 --- a/packages/sdk/server-ai/tests/test_agent_graph.py +++ b/packages/sdk/server-ai/tests/test_agent_graph.py @@ -237,14 +237,14 @@ def ldai_client(client: LDClient) -> LDAIClient: def test_agent_graph_method(ldai_client: LDAIClient): graph = ldai_client.agent_graph("test-agent-graph", Context.create("user-key")) - assert graph["enabled"] is True - assert graph["graph"] is not None - assert graph["graph"].root() is not None - assert graph["graph"].root().get_key() == "customer-support-agent" - assert len(graph["graph"].get_child_nodes("customer-support-agent")) == 3 - assert len(graph["graph"].get_child_nodes("personalized-agent")) == 0 - assert len(graph["graph"].get_child_nodes("multi-context-agent")) == 0 - assert len(graph["graph"].get_child_nodes("minimal-agent")) == 0 + assert graph.enabled is True + assert graph is not None + assert graph.root() is not None + assert graph.root().get_key() == "customer-support-agent" + assert len(graph.get_child_nodes("customer-support-agent")) == 3 + assert len(graph.get_child_nodes("personalized-agent")) == 0 + assert len(graph.get_child_nodes("multi-context-agent")) == 0 + assert len(graph.get_child_nodes("minimal-agent")) == 0 def test_agent_graph_method_disabled_agent(ldai_client: LDAIClient): @@ -252,8 +252,8 @@ def test_agent_graph_method_disabled_agent(ldai_client: LDAIClient): "test-agent-graph-disabled-agent", Context.create("user-key") ) - assert graph["enabled"] is False - assert graph["graph"] is None + assert graph.enabled is False + assert graph.root() is None def test_agent_graph_method_no_root_key(ldai_client: LDAIClient): @@ -261,8 +261,8 @@ def test_agent_graph_method_no_root_key(ldai_client: LDAIClient): "test-agent-graph-no-root-key", Context.create("user-key") ) - assert graph["enabled"] is False - assert graph["graph"] is None + assert graph.enabled is False + assert graph.root() is None def test_agent_graph_build_nodes(ldai_client: LDAIClient): @@ -319,9 +319,7 @@ def test_agent_graph_build_nodes(ldai_client: LDAIClient): def test_agent_graph_get_methods(ldai_client: LDAIClient): - graph = ldai_client.agent_graph("test-agent-graph", Context.create("user-key"))[ - "graph" - ] + graph = ldai_client.agent_graph("test-agent-graph", Context.create("user-key")) assert graph.root() is not None assert graph.root().get_key() == "customer-support-agent" @@ -359,7 +357,7 @@ def test_agent_graph_get_methods(ldai_client: LDAIClient): def test_agent_graph_traverse(ldai_client: LDAIClient): graph = ldai_client.agent_graph( "test-agent-graph-depth-3", Context.create("user-key") - )["graph"] + ) context = {} order = [] @@ -387,7 +385,7 @@ def handle_traverse(node, context): def test_agent_graph_reverse_traverse(ldai_client: LDAIClient): graph = ldai_client.agent_graph( "test-agent-graph-depth-3", Context.create("user-key") - )["graph"] + ) context = {} order = [] @@ -414,7 +412,7 @@ def handle_reverse_traverse(node, context): def test_agent_graph_handoff(ldai_client: LDAIClient): graph = ldai_client.agent_graph( "test-agent-graph-depth-3", Context.create("user-key") - )["graph"] + ) context = {} From b6a690474ee4eb0af156a9a135535ad77a76a42d Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 15 Jan 2026 10:19:14 -0900 Subject: [PATCH 06/11] [REL-11697] lint --- packages/sdk/server-ai/src/ldai/agent_graph/__init__.py | 8 +++----- 1 file changed, 3 insertions(+), 5 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 8ea5378..6ff3de2 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -150,14 +150,14 @@ def terminal_nodes(self) -> List[AgentGraphNode]: if len(self.get_child_nodes(node.get_key())) == 0 ] - def root(self) -> AgentGraphNode: + def root(self) -> Optional[AgentGraphNode]: """Get the root node of the graph.""" return self._nodes.get(self._agent_graph.root_config_key) def traverse( self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], - execution_context: Dict[str, Any] = None, + execution_context: Optional[Dict[str, Any]] = None, ) -> Any: """Traverse from the root down to terminal nodes, visiting nodes in order of depth. Nodes with the longest paths from the root (deepest nodes) will always be visited last.""" @@ -205,9 +205,8 @@ def traverse( def reverse_traverse( self, fn: Callable[["AgentGraphNode", Dict[str, Any]], Any], - execution_context: Dict[str, Any] = None, + execution_context: Optional[Dict[str, Any]] = None, ) -> Any: - """Traverse from terminal nodes up to the root, visiting nodes level by level. The root node will always be visited last, even if multiple paths converge at it.""" if execution_context is None: @@ -255,4 +254,3 @@ def reverse_traverse( ) return execution_context[self._agent_graph.root_config_key] - From 20bf05ca2b40cf63974893760b86542e12473f07 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 15 Jan 2026 10:44:13 -0900 Subject: [PATCH 07/11] [REL-11697] updated handoff default type --- packages/sdk/server-ai/src/ldai/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index f43c0a8..008064b 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -352,7 +352,7 @@ class Edge: key: str source_config: str target_config: str - handoff: Optional[dict] = None + handoff: Optional[dict] = {} # ============================================================================ From c1aefd10096acbb8fb781a5d6b5bb86bccbd489b Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 15 Jan 2026 10:45:14 -0900 Subject: [PATCH 08/11] [REL-11697] change to field(default_factory) --- packages/sdk/server-ai/src/ldai/models.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index 008064b..6c058a7 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -352,7 +352,7 @@ class Edge: key: str source_config: str target_config: str - handoff: Optional[dict] = {} + handoff: Optional[dict] = field(default_factory=dict) # ============================================================================ From ee3bc90c9dbeb520a1968807874ce4d9aeb89fd7 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 15 Jan 2026 11:07:48 -0900 Subject: [PATCH 09/11] [REL-11697] add protection if graph exhausts max_depth_limit setting --- .../src/ldai/agent_graph/__init__.py | 37 ++++++++++++++----- packages/sdk/server-ai/src/ldai/client.py | 2 +- 2 files changed, 28 insertions(+), 11 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 6ff3de2..9e7742a 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -127,6 +127,7 @@ def _collect_nodes( node_depths: Dict[str, int], nodes_by_depth: Dict[int, List[AgentGraphNode]], visited: Set[str], + max_depth: int, ) -> None: """Collect all reachable nodes from the given node and group them by depth.""" node_key = node.get_key() @@ -134,13 +135,14 @@ def _collect_nodes( return visited.add(node_key) - node_depth = node_depths.get(node_key, 0) + # Use max_depth for nodes not in node_depths to ensure they execute last + node_depth = node_depths.get(node_key, max_depth) if node_depth not in nodes_by_depth: nodes_by_depth[node_depth] = [] nodes_by_depth[node_depth].append(node) for child in self.get_child_nodes(node_key): - self._collect_nodes(child, node_depths, nodes_by_depth, visited) + self._collect_nodes(child, node_depths, nodes_by_depth, visited, max_depth) def terminal_nodes(self) -> List[AgentGraphNode]: """Get the terminal nodes of the graph, meaning any nodes without children.""" @@ -172,8 +174,11 @@ def traverse( current_level: List[AgentGraphNode] = [root_node] depth = 0 max_depth_limit = 10 # Infinite loop protection limit + max_depth_encountered = 0 + visited: Set[str] = {root_node.get_key()} # Track visited nodes in BFS to prevent cycles - while current_level and depth < max_depth_limit: + # Continue BFS to discover all nodes, but stop recording depths after max_depth_limit + while current_level: next_level: List[AgentGraphNode] = [] depth += 1 @@ -181,20 +186,32 @@ def traverse( node_key = node.get_key() for child in self.get_child_nodes(node_key): child_key = child.get_key() - # Defer this child to the next level if it's at a longer path - if child_key not in node_depths or ( - depth > node_depths[child_key] and depth < max_depth_limit - ): - node_depths[child_key] = depth - next_level.append(child) + if depth <= max_depth_limit: + # Defer this child to the next level if it's at a longer path + if child_key not in node_depths or depth > node_depths[child_key]: + node_depths[child_key] = depth + max_depth_encountered = max(max_depth_encountered, depth) + # Add to next level if not already visited (prevents cycles) + if child_key not in visited: + visited.add(child_key) + next_level.append(child) + else: + max_depth_encountered = max(max_depth_encountered, depth) + if child_key not in visited: + # Push this to the next level to be visited + visited.add(child_key) + next_level.append(child) current_level = next_level + # Use max_depth_limit + 1 to ensure they execute after all recorded nodes + max_depth = max(max_depth_limit + 1, max_depth_encountered + 1) + # Group all nodes by depth nodes_by_depth: Dict[int, List[AgentGraphNode]] = {} visited: Set[str] = set() - self._collect_nodes(root_node, node_depths, nodes_by_depth, visited) + self._collect_nodes(root_node, node_depths, nodes_by_depth, visited, max_depth) # Execute the lambda at this level for the nodes at this depth for depth_level in sorted(nodes_by_depth.keys()): for node in nodes_by_depth[depth_level]: diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index 57d809a..a139901 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -448,7 +448,7 @@ def agent_graph( ) all_agent_keys = [variation["rootConfigKey"]] + [ - edge["targetConfig"] for edge in variation.get("edges", []) + edge.get("targetConfig", "") for edge in variation.get("edges", []) if edge.get("targetConfig") ] agent_configs = { key: self.agent_config(key, context, AIAgentConfigDefault(enabled=False)) From 1a3ba51e5af7694734048296a14ad662da111ff7 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 15 Jan 2026 11:10:54 -0900 Subject: [PATCH 10/11] [REL-11697] lint --- .../sdk/server-ai/src/ldai/agent_graph/__init__.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 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 9e7742a..433199d 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -175,9 +175,8 @@ def traverse( depth = 0 max_depth_limit = 10 # Infinite loop protection limit max_depth_encountered = 0 - visited: Set[str] = {root_node.get_key()} # Track visited nodes in BFS to prevent cycles + seen_nodes: Set[str] = {root_node.get_key()} - # Continue BFS to discover all nodes, but stop recording depths after max_depth_limit while current_level: next_level: List[AgentGraphNode] = [] depth += 1 @@ -192,14 +191,14 @@ def traverse( node_depths[child_key] = depth max_depth_encountered = max(max_depth_encountered, depth) # Add to next level if not already visited (prevents cycles) - if child_key not in visited: - visited.add(child_key) + if child_key not in seen_nodes: + seen_nodes.add(child_key) next_level.append(child) else: max_depth_encountered = max(max_depth_encountered, depth) - if child_key not in visited: + if child_key not in seen_nodes: # Push this to the next level to be visited - visited.add(child_key) + seen_nodes.add(child_key) next_level.append(child) current_level = next_level @@ -209,6 +208,7 @@ def traverse( # Group all nodes by depth nodes_by_depth: Dict[int, List[AgentGraphNode]] = {} + # New visited for children nodes visited: Set[str] = set() self._collect_nodes(root_node, node_depths, nodes_by_depth, visited, max_depth) From 3041fad095697ccc4f2dacc7bc819dc0d8498793 Mon Sep 17 00:00:00 2001 From: Andrew Klatzke Date: Thu, 15 Jan 2026 11:12:03 -0900 Subject: [PATCH 11/11] [REL-11697] remove trailing whitespace --- packages/sdk/server-ai/src/ldai/agent_graph/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 433199d..a10ab23 100644 --- a/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py +++ b/packages/sdk/server-ai/src/ldai/agent_graph/__init__.py @@ -175,7 +175,7 @@ def traverse( depth = 0 max_depth_limit = 10 # Infinite loop protection limit max_depth_encountered = 0 - seen_nodes: Set[str] = {root_node.get_key()} + seen_nodes: Set[str] = {root_node.get_key()} while current_level: next_level: List[AgentGraphNode] = []