[REL-11697] resolve linting errors

andrewklatzke · andrewklatzke · commit a4c735aa0450 · 2026-01-14T16:35:10.000-09:00
diff --git 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
@@ -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
@@ -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
@@ -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
