[REL-11697] Simplify types; add enabled directly to objects; cursor feedback

Author: andrewklatzke

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',

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

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,18 +425,30 @@ def agent_graph(
         self,
         key: str,
         context: Context,
-    ) -> AIAgentGraphResponse:
+    ) -> AgentGraphDefinition:
         """`
         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 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(

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
 
 
 # ============================================================================

packages/sdk/server-ai/tests/test_agent_graph.py

@@ -237,32 +237,32 @@ 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):
     graph = ldai_client.agent_graph(
         "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):
     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
+    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 = {}