feat: add root-level tools map with customParameters to AI Config types

Adds AITool dataclass and tools map (keyed by tool name) to completion
and agent config types. The root-level tools map is distinct from
model.parameters.tools[] which remains passable to LLM providers as-is.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>

Author: jsonbailey

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

@@ -20,6 +20,7 @@
     AIConfig,
     AIJudgeConfig,
     AIJudgeConfigDefault,
+    AITool,
     Edge,
     JudgeConfiguration,
     LDAIAgent,

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

@@ -21,6 +21,7 @@
     AICompletionConfigDefault,
     AIJudgeConfig,
     AIJudgeConfigDefault,
+    AITool,
     Edge,
     JudgeConfiguration,
     LDMessage,
@@ -50,6 +51,22 @@
 _DISABLED_JUDGE_DEFAULT = AIJudgeConfigDefault.disabled()
 
 
+def _parse_tools(tools_data: Optional[Dict[str, Any]]) -> Optional[Dict[str, AITool]]:
+    """Parse the root-level tools map from a flag variation dict."""
+    if not tools_data or not isinstance(tools_data, dict):
+        return None
+    result = {}
+    for tool_name, tool_dict in tools_data.items():
+        if isinstance(tool_dict, dict):
+            result[tool_name] = AITool(
+                name=tool_dict.get('name', tool_name),
+                type=tool_dict.get('type'),
+                parameters=tool_dict.get('parameters'),
+                custom_parameters=tool_dict.get('customParameters'),
+            )
+    return result or None
+
+
 class LDAIClient:
     """The LaunchDarkly AI SDK client object."""
 
@@ -89,10 +106,13 @@ def _completion_config(
         variables: Optional[Dict[str, Any]] = None,
     ) -> AICompletionConfig:
         (model, provider, messages, instructions,
-         tracker_factory, enabled, judge_configuration, _) = self.__evaluate(
+         tracker_factory, enabled, judge_configuration, variation) = self.__evaluate(
             key, context, default.to_dict(), variables
         )
 
+        tools_data = variation.get('tools')
+        tools = _parse_tools(tools_data) if tools_data is not None else default.tools
+
         config = AICompletionConfig(
             key=key,
             enabled=bool(enabled),
@@ -101,6 +121,7 @@ def _completion_config(
             provider=provider,
             create_tracker=tracker_factory,
             judge_configuration=judge_configuration,
+            tools=tools,
         )
 
         return config
@@ -891,13 +912,16 @@ def __evaluate_agent(
         :return: Configured AIAgentConfig instance.
         """
         (model, provider, messages, instructions,
-         tracker_factory, enabled, judge_configuration, _) = self.__evaluate(
+         tracker_factory, enabled, judge_configuration, variation) = self.__evaluate(
             key, context, default.to_dict(), variables, graph_key=graph_key
         )
 
         # For agents, prioritize instructions over messages
         final_instructions = instructions if instructions is not None else default.instructions
 
+        tools_data = variation.get('tools')
+        tools = _parse_tools(tools_data) if tools_data is not None else default.tools
+
         return AIAgentConfig(
             key=key,
             enabled=bool(enabled) if enabled is not None else (default.enabled or False),
@@ -906,6 +930,7 @@ def __evaluate_agent(
             instructions=final_instructions,
             create_tracker=tracker_factory,
             judge_configuration=judge_configuration or default.judge_configuration,
+            tools=tools,
         )
 
     def __interpolate_template(self, template: str, variables: Dict[str, Any]) -> str:

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

@@ -5,6 +5,28 @@
 from typing_extensions import Self
 
 
+@dataclass(frozen=True)
+class AITool:
+    """
+    A single tool entry from the root-level tools map in an AI Config flag variation.
+    Distinct from model.parameters.tools[] which is the raw array passed to LLM providers.
+    """
+    name: str
+    type: Optional[str] = None
+    parameters: Optional[Dict[str, Any]] = None
+    custom_parameters: Optional[Dict[str, Any]] = None
+
+    def to_dict(self) -> dict:
+        result: Dict[str, Any] = {'name': self.name}
+        if self.type is not None:
+            result['type'] = self.type
+        if self.parameters is not None:
+            result['parameters'] = self.parameters
+        if self.custom_parameters is not None:
+            result['customParameters'] = self.custom_parameters  # camelCase in wire format
+        return result
+
+
 @dataclass
 class LDMessage:
     role: Literal['system', 'user', 'assistant']
@@ -208,6 +230,7 @@ class AICompletionConfigDefault(AIConfigDefault):
     """
     messages: Optional[List[LDMessage]] = None
     judge_configuration: Optional[JudgeConfiguration] = None
+    tools: Optional[Dict[str, 'AITool']] = None
 
     def to_dict(self) -> dict:
         """
@@ -217,6 +240,8 @@ 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:
+            result['tools'] = {k: v.to_dict() for k, v in self.tools.items()}
         return result
 
 
@@ -227,6 +252,7 @@ class AICompletionConfig(AIConfig):
     """
     messages: Optional[List[LDMessage]] = None
     judge_configuration: Optional[JudgeConfiguration] = None
+    tools: Optional[Dict[str, 'AITool']] = None
 
     def to_dict(self) -> dict:
         """
@@ -236,6 +262,8 @@ 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:
+            result['tools'] = {k: v.to_dict() for k, v in self.tools.items()}
         return result
 
 
@@ -250,6 +278,7 @@ class AIAgentConfigDefault(AIConfigDefault):
     """
     instructions: Optional[str] = None
     judge_configuration: Optional[JudgeConfiguration] = None
+    tools: Optional[Dict[str, 'AITool']] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -260,6 +289,8 @@ 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:
+            result['tools'] = {k: v.to_dict() for k, v in self.tools.items()}
         return result
 
 
@@ -270,6 +301,7 @@ class AIAgentConfig(AIConfig):
     """
     instructions: Optional[str] = None
     judge_configuration: Optional[JudgeConfiguration] = None
+    tools: Optional[Dict[str, 'AITool']] = None
 
     def to_dict(self) -> Dict[str, Any]:
         """
@@ -280,6 +312,8 @@ 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:
+            result['tools'] = {k: v.to_dict() for k, v in self.tools.items()}
         return result
 
 

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

@@ -0,0 +1,137 @@
+import pytest
+from ldclient import Config, Context, LDClient
+from ldclient.integrations.test_data import TestData
+
+from ldai import AITool, LDAIClient
+from ldai.models import AIAgentConfigDefault, AICompletionConfigDefault
+
+
+@pytest.fixture
+def td() -> TestData:
+    td = TestData.data_source()
+    td.update(
+        td.flag('completion-with-tools')
+        .variations(
+            {
+                'model': {'name': 'gpt-5', 'parameters': {'temperature': 0.7}},
+                'messages': [{'role': 'user', 'content': 'Hello'}],
+                'tools': {
+                    'web-search-tool': {
+                        'name': 'web-search-tool',
+                        'type': 'function',
+                        'parameters': {'type': 'object', 'properties': {}, 'required': []},
+                        'customParameters': {'some-custom-parameter': 'some-custom-value'},
+                    }
+                },
+                '_ldMeta': {'enabled': True, 'variationKey': 'v1', 'version': 1},
+            },
+        )
+        .variation_for_all(0)
+    )
+
+    td.update(
+        td.flag('completion-no-tools')
+        .variations(
+            {
+                'model': {'name': 'gpt-5'},
+                'messages': [{'role': 'user', 'content': 'Hello'}],
+                '_ldMeta': {'enabled': True, 'variationKey': 'v1', 'version': 1},
+            },
+        )
+        .variation_for_all(0)
+    )
+
+    td.update(
+        td.flag('agent-with-tools')
+        .variations(
+            {
+                'model': {'name': 'gpt-5'},
+                'instructions': 'You are a helpful agent.',
+                'tools': {
+                    'search-tool': {
+                        'name': 'search-tool',
+                        'type': 'function',
+                        'customParameters': {'maxResults': 10},
+                    }
+                },
+                '_ldMeta': {'enabled': True, 'variationKey': 'v1', 'version': 1, 'mode': 'agent'},
+            },
+        )
+        .variation_for_all(0)
+    )
+
+    return td
+
+
+@pytest.fixture
+def client(td) -> LDAIClient:
+    config = Config('fake-sdk-key', update_processor_class=td, send_events=False)
+    ld_client = LDClient(config=config)
+    return LDAIClient(ld_client)
+
+
+@pytest.fixture
+def context() -> Context:
+    return Context.builder('test-user').name('Test User').build()
+
+
+def test_completion_config_includes_tools_from_variation(client, context):
+    result = client.completion_config('completion-with-tools', context, AICompletionConfigDefault())
+
+    assert result.tools is not None
+    assert 'web-search-tool' in result.tools
+    tool = result.tools['web-search-tool']
+    assert tool.name == 'web-search-tool'
+    assert tool.type == 'function'
+    assert tool.custom_parameters == {'some-custom-parameter': 'some-custom-value'}
+
+
+def test_completion_config_tools_none_when_not_in_variation(client, context):
+    result = client.completion_config('completion-no-tools', context, AICompletionConfigDefault())
+
+    assert result.tools is None
+
+
+def test_completion_config_falls_back_to_default_tools(client, context):
+    default_tool = AITool(name='default-tool', type='function', custom_parameters={'priority': 'high'})
+    default = AICompletionConfigDefault(tools={'default-tool': default_tool})
+
+    result = client.completion_config('completion-no-tools', context, default)
+
+    assert result.tools is not None
+    assert 'default-tool' in result.tools
+    assert result.tools['default-tool'].custom_parameters == {'priority': 'high'}
+
+
+def test_agent_config_includes_tools_from_variation(client, context):
+    result = client.agent_config('agent-with-tools', context, AIAgentConfigDefault())
+
+    assert result.tools is not None
+    assert 'search-tool' in result.tools
+    tool = result.tools['search-tool']
+    assert tool.name == 'search-tool'
+    assert tool.custom_parameters == {'maxResults': 10}
+
+
+def test_aitool_to_dict_serializes_custom_parameters_as_camel_case():
+    tool = AITool(
+        name='my-tool',
+        type='function',
+        parameters={'type': 'object'},
+        custom_parameters={'someKey': 'someValue'},
+    )
+    d = tool.to_dict()
+
+    assert d['name'] == 'my-tool'
+    assert d['type'] == 'function'
+    assert d['parameters'] == {'type': 'object'}
+    assert 'customParameters' in d
+    assert d['customParameters'] == {'someKey': 'someValue'}
+    assert 'custom_parameters' not in d
+
+
+def test_aitool_to_dict_omits_none_fields():
+    tool = AITool(name='bare-tool')
+    d = tool.to_dict()
+
+    assert d == {'name': 'bare-tool'}