refactor: remove ToolDefinition, merge custom parameters into config types

Co-Authored-By: traci@launchdarkly.com <traci@launchdarkly.com>

Author: devin-ai-integration[bot]

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

@@ -28,7 +28,6 @@
     LDMessage,
     ModelConfig,
     ProviderConfig,
-    ToolDefinition,
 )
 from ldai.providers import (
     AgentGraphResult,
@@ -69,7 +68,6 @@
     'LDMessage',
     'ModelConfig',
     'ProviderConfig',
-    'ToolDefinition',
     'log',
     # Deprecated exports
     'AIConfig',

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

@@ -25,7 +25,6 @@
     LDMessage,
     ModelConfig,
     ProviderConfig,
-    ToolDefinition,
 )
 from ldai.providers import ToolRegistry
 from ldai.providers.runner_factory import RunnerFactory
@@ -69,7 +68,9 @@ def _completion_config(
         default: AICompletionConfigDefault,
         variables: Optional[Dict[str, Any]] = None,
     ) -> AICompletionConfig:
-        model, provider, messages, instructions, tracker, enabled, judge_configuration, _, tools = self.__evaluate(
+        (model, provider, messages, instructions,
+         tracker, enabled, judge_configuration,
+         _, tool_custom_parameters) = self.__evaluate(
             key, context, default.to_dict(), variables
         )
 
@@ -81,7 +82,7 @@ def _completion_config(
             provider=provider,
             tracker=tracker,
             judge_configuration=judge_configuration,
-            tools=tools,
+            tool_custom_parameters=tool_custom_parameters,
         )
 
         return config
@@ -138,7 +139,7 @@ def _judge_config(
     ) -> AIJudgeConfig:
         (model, provider, messages, instructions,
          tracker, enabled, judge_configuration,
-         variation, tools) = self.__evaluate(
+         variation, _) = self.__evaluate(
             key, context, default.to_dict(), variables
         )
 
@@ -755,7 +756,7 @@ def __evaluate(
     ) -> Tuple[
         Optional[ModelConfig], Optional[ProviderConfig], Optional[List[LDMessage]],
         Optional[str], LDAIConfigTracker, bool, Optional[Any], Dict[str, Any],
-        Optional[List[ToolDefinition]]
+        Optional[Dict[str, Dict[str, Any]]]
     ]:
         """
         Internal method to evaluate a configuration and extract components.
@@ -833,23 +834,21 @@ def __evaluate(
                 if judges:
                     judge_configuration = JudgeConfiguration(judges=judges)
 
-        tools = None
+        tool_custom_parameters = None
         model_raw = variation.get('model')
         params_raw = model_raw.get('parameters') if isinstance(model_raw, dict) else None
         tool_defs_raw = params_raw.get('tools') if isinstance(params_raw, dict) else None
         if isinstance(tool_defs_raw, list):
-            tools = [
-                ToolDefinition(
-                    name=t.get('name', ''),
-                    custom_parameters=t.get('customParameters', None)
-                )
-                for t in tool_defs_raw
-                if isinstance(t, dict) and t.get('name')
-            ]
-            if not tools:
-                tools = None
+            parsed: Dict[str, Dict[str, Any]] = {}
+            for t in tool_defs_raw:
+                if isinstance(t, dict) and t.get('name'):
+                    parsed[t['name']] = t.get('customParameters') or {}
+            if parsed:
+                tool_custom_parameters = parsed
 
-        return model, provider_config, messages, instructions, tracker, enabled, judge_configuration, variation, tools
+        return (model, provider_config, messages, instructions,
+                tracker, enabled, judge_configuration,
+                variation, tool_custom_parameters)
 
     def __evaluate_agent(
         self,
@@ -867,7 +866,9 @@ def __evaluate_agent(
         :param variables: Variables for interpolation.
         :return: Configured AIAgentConfig instance.
         """
-        model, provider, messages, instructions, tracker, enabled, judge_configuration, _, tools = self.__evaluate(
+        (model, provider, messages, instructions,
+         tracker, enabled, judge_configuration,
+         _, tool_custom_parameters) = self.__evaluate(
             key, context, default.to_dict(), variables
         )
 
@@ -882,7 +883,7 @@ def __evaluate_agent(
             instructions=final_instructions,
             tracker=tracker,
             judge_configuration=judge_configuration or default.judge_configuration,
-            tools=tools or default.tools,
+            tool_custom_parameters=tool_custom_parameters or default.tool_custom_parameters,
         )
 
     def __interpolate_template(self, template: str, variables: Dict[str, Any]) -> str:

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

@@ -77,57 +77,6 @@ def to_dict(self) -> dict:
         }
 
 
-class ToolDefinition:
-    """
-    Definition of a tool available to an AI configuration.
-
-    Each tool has a name used to match against the tool registry,
-    and optional custom parameters that can be configured via the
-    LaunchDarkly dashboard.
-    """
-
-    def __init__(
-        self,
-        name: str,
-        custom_parameters: Optional[Dict[str, Any]] = None,
-    ):
-        """
-        :param name: The name of the tool.
-        :param custom_parameters: Optional custom parameters for
-            the tool, configured via the LaunchDarkly dashboard.
-        """
-        self._name = name
-        self._custom_parameters = custom_parameters
-
-    @property
-    def name(self) -> str:
-        """
-        The name of the tool.
-        """
-        return self._name
-
-    def get_custom_parameter(self, key: str) -> Any:
-        """
-        Retrieve a custom parameter by key.
-
-        :param key: The custom parameter key to look up.
-        :return: The parameter value, or None if not found.
-        """
-        if self._custom_parameters is None:
-            return None
-
-        return self._custom_parameters.get(key)
-
-    def to_dict(self) -> dict:
-        """
-        Render the tool definition as a dictionary object.
-        """
-        result: Dict[str, Any] = {'name': self._name}
-        if self._custom_parameters is not None:
-            result['customParameters'] = self._custom_parameters
-        return result
-
-
 class ProviderConfig:
     """
     Configuration related to the provider.
@@ -259,7 +208,22 @@ class AICompletionConfigDefault(AIConfigDefault):
     """
     messages: Optional[List[LDMessage]] = None
     judge_configuration: Optional[JudgeConfiguration] = None
-    tools: Optional[List[ToolDefinition]] = None
+    tool_custom_parameters: Optional[Dict[str, Dict[str, Any]]] = None
+
+    def get_tool_custom_parameter(self, tool_name: str, key: str) -> Any:
+        """
+        Retrieve a custom parameter for a specific tool.
+
+        :param tool_name: The name of the tool.
+        :param key: The custom parameter key to look up.
+        :return: The parameter value, or None if not found.
+        """
+        if self.tool_custom_parameters is None:
+            return None
+        tool_params = self.tool_custom_parameters.get(tool_name)
+        if tool_params is None:
+            return None
+        return tool_params.get(key)
 
     def to_dict(self) -> dict:
         """
@@ -269,10 +233,16 @@ def to_dict(self) -> dict:
         result['messages'] = [message.to_dict() for message in self.messages] if self.messages else None
         if self.judge_configuration is not None:
             result['judgeConfiguration'] = self.judge_configuration.to_dict()
-        if self.tools is not None:
+        if self.tool_custom_parameters is not None:
             model = result.get('model') or {}
             params = model.get('parameters') or {}
-            params['tools'] = [tool.to_dict() for tool in self.tools]
+            tools_list = []
+            for name, custom_params in self.tool_custom_parameters.items():
+                tool_entry: Dict[str, Any] = {'name': name}
+                if custom_params:
+                    tool_entry['customParameters'] = custom_params
+                tools_list.append(tool_entry)
+            params['tools'] = tools_list
             model['parameters'] = params
             result['model'] = model
         return result
@@ -285,7 +255,22 @@ class AICompletionConfig(AIConfig):
     """
     messages: Optional[List[LDMessage]] = None
     judge_configuration: Optional[JudgeConfiguration] = None
-    tools: Optional[List[ToolDefinition]] = None
+    tool_custom_parameters: Optional[Dict[str, Dict[str, Any]]] = None
+
+    def get_tool_custom_parameter(self, tool_name: str, key: str) -> Any:
+        """
+        Retrieve a custom parameter for a specific tool.
+
+        :param tool_name: The name of the tool.
+        :param key: The custom parameter key to look up.
+        :return: The parameter value, or None if not found.
+        """
+        if self.tool_custom_parameters is None:
+            return None
+        tool_params = self.tool_custom_parameters.get(tool_name)
+        if tool_params is None:
+            return None
+        return tool_params.get(key)
 
     def to_dict(self) -> dict:
         """
@@ -309,7 +294,22 @@ class AIAgentConfigDefault(AIConfigDefault):
     """
     instructions: Optional[str] = None
     judge_configuration: Optional[JudgeConfiguration] = None
-    tools: Optional[List[ToolDefinition]] = None
+    tool_custom_parameters: Optional[Dict[str, Dict[str, Any]]] = None
+
+    def get_tool_custom_parameter(self, tool_name: str, key: str) -> Any:
+        """
+        Retrieve a custom parameter for a specific tool.
+
+        :param tool_name: The name of the tool.
+        :param key: The custom parameter key to look up.
+        :return: The parameter value, or None if not found.
+        """
+        if self.tool_custom_parameters is None:
+            return None
+        tool_params = self.tool_custom_parameters.get(tool_name)
+        if tool_params is None:
+            return None
+        return tool_params.get(key)
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -320,10 +320,16 @@ def to_dict(self) -> Dict[str, Any]:
             result['instructions'] = self.instructions
         if self.judge_configuration is not None:
             result['judgeConfiguration'] = self.judge_configuration.to_dict()
-        if self.tools is not None:
+        if self.tool_custom_parameters is not None:
             model = result.get('model') or {}
             params = model.get('parameters') or {}
-            params['tools'] = [tool.to_dict() for tool in self.tools]
+            tools_list = []
+            for name, custom_params in self.tool_custom_parameters.items():
+                tool_entry: Dict[str, Any] = {'name': name}
+                if custom_params:
+                    tool_entry['customParameters'] = custom_params
+                tools_list.append(tool_entry)
+            params['tools'] = tools_list
             model['parameters'] = params
             result['model'] = model
         return result
@@ -336,7 +342,22 @@ class AIAgentConfig(AIConfig):
     """
     instructions: Optional[str] = None
     judge_configuration: Optional[JudgeConfiguration] = None
-    tools: Optional[List[ToolDefinition]] = None
+    tool_custom_parameters: Optional[Dict[str, Dict[str, Any]]] = None
+
+    def get_tool_custom_parameter(self, tool_name: str, key: str) -> Any:
+        """
+        Retrieve a custom parameter for a specific tool.
+
+        :param tool_name: The name of the tool.
+        :param key: The custom parameter key to look up.
+        :return: The parameter value, or None if not found.
+        """
+        if self.tool_custom_parameters is None:
+            return None
+        tool_params = self.tool_custom_parameters.get(tool_name)
+        if tool_params is None:
+            return None
+        return tool_params.get(key)
 
     def to_dict(self) -> Dict[str, Any]:
         """

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

@@ -3,7 +3,7 @@
 from ldclient.integrations.test_data import TestData
 
 from ldai import (LDAIAgentConfig, LDAIAgentDefaults, LDAIClient, ModelConfig,
-                  ProviderConfig, ToolDefinition)
+                  ProviderConfig)
 
 
 @pytest.fixture
@@ -396,49 +396,40 @@ def test_agent_config_has_tools(ldai_client: LDAIClient):
     agent = ldai_client.agent_config('agent-with-tools', context)
 
     assert agent.enabled is True
-    assert agent.tools is not None
-    assert len(agent.tools) == 3
+    assert agent.tool_custom_parameters is not None
+    assert len(agent.tool_custom_parameters) == 3
 
-    get_order = agent.tools[0]
-    assert get_order.name == 'get-order'
-    assert get_order.get_custom_parameter('includeHistory') is True
-    assert get_order.get_custom_parameter('maxItems') == 5
-
-    search = agent.tools[1]
-    assert search.name == 'search-products'
-    assert search.get_custom_parameter('category') == 'electronics'
-
-    send_email = agent.tools[2]
-    assert send_email.name == 'send-email'
-    assert send_email.get_custom_parameter('anything') is None
+    assert agent.get_tool_custom_parameter('get-order', 'includeHistory') is True
+    assert agent.get_tool_custom_parameter('get-order', 'maxItems') == 5
+    assert agent.get_tool_custom_parameter('search-products', 'category') == 'electronics'
+    assert agent.get_tool_custom_parameter('send-email', 'anything') is None
+    assert 'send-email' in agent.tool_custom_parameters
 
 
 def test_agent_config_tools_fallback_to_default(ldai_client: LDAIClient):
     """Test that agent config falls back to default tools when flag has no tools."""
     context = Context.create('user-key')
-    default_tools = [ToolDefinition('default-tool', custom_parameters={'timeout': 30})]
     default = LDAIAgentDefaults(
         enabled=False,
         model=ModelConfig('fallback-model'),
         instructions='Default instructions',
-        tools=default_tools,
+        tool_custom_parameters={'default-tool': {'timeout': 30}},
     )
 
     agent = ldai_client.agent_config('customer-support-agent', context, default)
 
     assert agent.enabled is True
     # customer-support-agent has no tools in the flag, so falls back to default
-    assert agent.tools is not None
-    assert len(agent.tools) == 1
-    assert agent.tools[0].name == 'default-tool'
-    assert agent.tools[0].get_custom_parameter('timeout') == 30
+    assert agent.tool_custom_parameters is not None
+    assert len(agent.tool_custom_parameters) == 1
+    assert agent.get_tool_custom_parameter('default-tool', 'timeout') == 30
 
 
 def test_agent_config_no_tools(ldai_client: LDAIClient):
-    """Test that agent tools is None when neither flag nor default has tools."""
+    """Test that tool_custom_parameters is None when neither flag nor default has tools."""
     context = Context.create('user-key')
 
     agent = ldai_client.agent_config('customer-support-agent', context)
 
     assert agent.enabled is True
-    assert agent.tools is None
+    assert agent.tool_custom_parameters is None