Skip to content
2 changes: 2 additions & 0 deletions packages/sdk/server-ai/src/ldai/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@
LDAIAgentConfig,
LDAIAgentDefaults,
LDMessage,
LDTool,
ModelConfig,
ProviderConfig,
)
Expand Down Expand Up @@ -64,6 +65,7 @@
'Judge',
'JudgeConfiguration',
'JudgeResult',
'LDTool',
'LDMessage',
'ModelConfig',
'ProviderConfig',
Expand Down
30 changes: 28 additions & 2 deletions packages/sdk/server-ai/src/ldai/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@
Edge,
JudgeConfiguration,
LDMessage,
LDTool,
ModelConfig,
ProviderConfig,
)
Expand All @@ -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."""

Expand Down Expand Up @@ -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),
Expand All @@ -101,6 +123,7 @@ def _completion_config(
provider=provider,
create_tracker=tracker_factory,
judge_configuration=judge_configuration,
tools=tools,
)

return config
Expand Down Expand Up @@ -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),
Expand All @@ -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:
Expand Down
37 changes: 37 additions & 0 deletions packages/sdk/server-ai/src/ldai/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -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']
Expand Down Expand Up @@ -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:
"""
Expand All @@ -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


Expand All @@ -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:
"""
Expand All @@ -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


Expand All @@ -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]:
"""
Expand All @@ -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


Expand All @@ -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]:
"""
Expand All @@ -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


Expand Down
135 changes: 135 additions & 0 deletions packages/sdk/server-ai/tests/test_tools.py
Original file line number Diff line number Diff line change
@@ -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'}
Loading