diff --git a/packages/sdk/server-ai/src/ldai/__init__.py b/packages/sdk/server-ai/src/ldai/__init__.py index a3557d8..b7d8752 100644 --- a/packages/sdk/server-ai/src/ldai/__init__.py +++ b/packages/sdk/server-ai/src/ldai/__init__.py @@ -26,6 +26,7 @@ LDAIAgentConfig, LDAIAgentDefaults, LDMessage, + LDTool, ModelConfig, ProviderConfig, ) @@ -64,6 +65,7 @@ 'Judge', 'JudgeConfiguration', 'JudgeResult', + 'LDTool', 'LDMessage', 'ModelConfig', 'ProviderConfig', diff --git a/packages/sdk/server-ai/src/ldai/client.py b/packages/sdk/server-ai/src/ldai/client.py index ac13469..4023574 100644 --- a/packages/sdk/server-ai/src/ldai/client.py +++ b/packages/sdk/server-ai/src/ldai/client.py @@ -24,6 +24,7 @@ Edge, JudgeConfiguration, LDMessage, + LDTool, ModelConfig, ProviderConfig, ) @@ -50,6 +51,25 @@ _DISABLED_JUDGE_DEFAULT = AIJudgeConfigDefault.disabled() +def _parse_tools(tools_data: Optional[Dict[str, Any]]) -> Optional[Dict[str, LDTool]]: + """Parse the root-level tools map from a flag variation dict.""" + if not isinstance(tools_data, dict): + return None + result = {} + for tool_name, tool_dict in tools_data.items(): + if not isinstance(tool_dict, dict): + log.warning('Skipping tool "%s": expected a dict, got %s', tool_name, type(tool_dict).__name__) + continue + result[tool_name] = LDTool( + name=tool_dict.get('name', tool_name), + description=tool_dict.get('description'), + 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 +109,12 @@ 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 = _parse_tools(variation.get('tools')) + config = AICompletionConfig( key=key, enabled=bool(enabled), @@ -101,6 +123,7 @@ def _completion_config( provider=provider, create_tracker=tracker_factory, judge_configuration=judge_configuration, + tools=tools, ) return config @@ -891,13 +914,15 @@ 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 = _parse_tools(variation.get('tools')) + return AIAgentConfig( key=key, enabled=bool(enabled) if enabled is not None else (default.enabled or False), @@ -906,6 +931,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: diff --git a/packages/sdk/server-ai/src/ldai/models.py b/packages/sdk/server-ai/src/ldai/models.py index cf5da1f..28d8cef 100644 --- a/packages/sdk/server-ai/src/ldai/models.py +++ b/packages/sdk/server-ai/src/ldai/models.py @@ -5,6 +5,31 @@ from typing_extensions import Self +@dataclass(frozen=True) +class LDTool: + """ + 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 + description: Optional[str] = None + 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.description is not None: + result['description'] = self.description + 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 +233,7 @@ class AICompletionConfigDefault(AIConfigDefault): """ messages: Optional[List[LDMessage]] = None judge_configuration: Optional[JudgeConfiguration] = None + tools: Optional[Dict[str, 'LDTool']] = None def to_dict(self) -> dict: """ @@ -217,6 +243,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 +255,7 @@ class AICompletionConfig(AIConfig): """ messages: Optional[List[LDMessage]] = None judge_configuration: Optional[JudgeConfiguration] = None + tools: Optional[Dict[str, 'LDTool']] = None def to_dict(self) -> dict: """ @@ -236,6 +265,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 +281,7 @@ class AIAgentConfigDefault(AIConfigDefault): """ instructions: Optional[str] = None judge_configuration: Optional[JudgeConfiguration] = None + tools: Optional[Dict[str, 'LDTool']] = None def to_dict(self) -> Dict[str, Any]: """ @@ -260,6 +292,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 +304,7 @@ class AIAgentConfig(AIConfig): """ instructions: Optional[str] = None judge_configuration: Optional[JudgeConfiguration] = None + tools: Optional[Dict[str, 'LDTool']] = None def to_dict(self) -> Dict[str, Any]: """ @@ -280,6 +315,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 diff --git a/packages/sdk/server-ai/tests/test_tools.py b/packages/sdk/server-ai/tests/test_tools.py new file mode 100644 index 0000000..1c74215 --- /dev/null +++ b/packages/sdk/server-ai/tests/test_tools.py @@ -0,0 +1,135 @@ +import pytest +from ldclient import Config, Context, LDClient +from ldclient.integrations.test_data import TestData + +from ldai import LDTool, 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_tools_none_when_variation_has_no_tools(client, context): + default_tool = LDTool(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 None + + +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 = LDTool( + 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 = LDTool(name='bare-tool') + d = tool.to_dict() + + assert d == {'name': 'bare-tool'}