Skip to content
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
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 @@ -28,6 +28,7 @@
LDMessage,
ModelConfig,
ProviderConfig,
ToolCustomParametersMap,
)
from ldai.providers import (
AgentGraphResult,
Expand Down Expand Up @@ -68,6 +69,7 @@
'LDMessage',
'ModelConfig',
'ProviderConfig',
'ToolCustomParametersMap',
'log',
# Deprecated exports
'AIConfig',
Expand Down
43 changes: 38 additions & 5 deletions packages/sdk/server-ai/src/ldai/client.py
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@
LDMessage,
ModelConfig,
ProviderConfig,
ToolCustomParametersMap,
)
from ldai.providers import ToolRegistry
from ldai.providers.runner_factory import RunnerFactory
Expand Down Expand Up @@ -68,7 +69,9 @@ def _completion_config(
default: AICompletionConfigDefault,
variables: Optional[Dict[str, Any]] = None,
) -> AICompletionConfig:
model, provider, messages, instructions, tracker, enabled, judge_configuration, _ = self.__evaluate(
(model, provider, messages, instructions,
tracker, enabled, judge_configuration,
_, tool_custom_parameters) = self.__evaluate(
key, context, default.to_dict(), variables
)

Expand All @@ -80,6 +83,7 @@ def _completion_config(
provider=provider,
tracker=tracker,
judge_configuration=judge_configuration,
tool_custom_parameters=tool_custom_parameters,
)

return config
Expand Down Expand Up @@ -134,7 +138,9 @@ def _judge_config(
default: AIJudgeConfigDefault,
variables: Optional[Dict[str, Any]] = None,
) -> AIJudgeConfig:
model, provider, messages, instructions, tracker, enabled, judge_configuration, variation = self.__evaluate(
(model, provider, messages, instructions,
tracker, enabled, judge_configuration,
variation, _) = self.__evaluate(
key, context, default.to_dict(), variables
)

Expand Down Expand Up @@ -750,7 +756,8 @@ def __evaluate(
variables: Optional[Dict[str, Any]] = None,
) -> Tuple[
Optional[ModelConfig], Optional[ProviderConfig], Optional[List[LDMessage]],
Optional[str], LDAIConfigTracker, bool, Optional[Any], Dict[str, Any]
Optional[str], LDAIConfigTracker, bool, Optional[Any], Dict[str, Any],
Optional[ToolCustomParametersMap]
]:
"""
Internal method to evaluate a configuration and extract components.
Expand Down Expand Up @@ -828,7 +835,30 @@ def __evaluate(
if judges:
judge_configuration = JudgeConfiguration(judges=judges)

return model, provider_config, messages, instructions, tracker, enabled, judge_configuration, variation
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):
parsed: ToolCustomParametersMap = {}
for t in tool_defs_raw:
if not isinstance(t, dict) or not t.get('name'):
continue
cp = t.get('customParameters')
if isinstance(cp, dict) and cp:
parsed[t['name']] = cp
if parsed:
tool_custom_parameters = parsed

return (model, provider_config, messages, instructions,
tracker, enabled, judge_configuration,
variation, tool_custom_parameters)

def __evaluate_agent(
self,
Expand All @@ -846,7 +876,9 @@ def __evaluate_agent(
:param variables: Variables for interpolation.
:return: Configured AIAgentConfig instance.
"""
model, provider, messages, instructions, tracker, enabled, judge_configuration, _ = self.__evaluate(
(model, provider, messages, instructions,
tracker, enabled, judge_configuration,
_, tool_custom_parameters) = self.__evaluate(
key, context, default.to_dict(), variables
)

Expand All @@ -861,6 +893,7 @@ def __evaluate_agent(
instructions=final_instructions,
tracker=tracker,
judge_configuration=judge_configuration or default.judge_configuration,
tool_custom_parameters=tool_custom_parameters or default.tool_custom_parameters,
)

def __interpolate_template(self, template: str, variables: Dict[str, Any]) -> str:
Expand Down
104 changes: 102 additions & 2 deletions packages/sdk/server-ai/src/ldai/models.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,58 @@

from ldai.tracker import LDAIConfigTracker

# Type alias for tool custom parameters: maps tool name -> custom params dict
ToolCustomParametersMap = Dict[str, Dict[str, Any]]


def _get_tool_custom_parameter(
tool_custom_parameters: Optional['ToolCustomParametersMap'],
tool_name: str,
key: str,
) -> Any:
"""Retrieve a custom parameter for a specific tool.

:param tool_custom_parameters: The tool custom parameters map.
: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 tool_custom_parameters is None:
return None
tool_params = tool_custom_parameters.get(tool_name)
if tool_params is None:
return None
return tool_params.get(key)


def _serialize_tool_custom_parameters(
result: Dict[str, Any],
tool_custom_parameters: Optional['ToolCustomParametersMap'],
) -> None:
"""Serialize tool_custom_parameters into the result dict.

Injects tools into ``result['model']['parameters']['tools']``.
Iteration order follows dict insertion order (Python 3.7+); if the
original flag payload order matters, the caller must ensure it is
preserved at parse time.

:param result: The mutable dict being built by to_dict().
:param tool_custom_parameters: The tool custom parameters map.
"""
if tool_custom_parameters is None:
return
model = result.get('model') or {}
params = model.get('parameters') or {}
tools_list = []
for name, custom_params in 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


@dataclass
class LDMessage:
Expand Down Expand Up @@ -208,15 +260,33 @@ class AICompletionConfigDefault(AIConfigDefault):
"""
messages: Optional[List[LDMessage]] = None
judge_configuration: Optional[JudgeConfiguration] = None
tool_custom_parameters: Optional[ToolCustomParametersMap] = None

def get_tool_custom_parameter(
self, tool_name: str, key: str,
) -> Any:
"""Retrieve a custom parameter for a specific tool."""
return _get_tool_custom_parameter(
self.tool_custom_parameters, tool_name, key,
)

def to_dict(self) -> dict:
"""
Render the given default values as an AICompletionConfigDefault-compatible dictionary object.
"""Render the given default values as an AICompletionConfigDefault-compatible dictionary object.

Note: tool_custom_parameters are serialized into
``model.parameters.tools`` so that they are carried inside the
variation dict used as the evaluation fallback. This is how
completion-config defaults reach ``_completion_config`` — there
is no separate ``or default.tool_custom_parameters`` fallback
like agent configs have.
"""
result = self._base_to_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()
_serialize_tool_custom_parameters(
result, self.tool_custom_parameters,
)
return result


Expand All @@ -227,6 +297,15 @@ class AICompletionConfig(AIConfig):
"""
messages: Optional[List[LDMessage]] = None
judge_configuration: Optional[JudgeConfiguration] = None
tool_custom_parameters: Optional[ToolCustomParametersMap] = None

def get_tool_custom_parameter(
self, tool_name: str, key: str,
) -> Any:
"""Retrieve a custom parameter for a specific tool."""
return _get_tool_custom_parameter(
self.tool_custom_parameters, tool_name, key,
)

def to_dict(self) -> dict:
"""
Expand All @@ -250,6 +329,15 @@ class AIAgentConfigDefault(AIConfigDefault):
"""
instructions: Optional[str] = None
judge_configuration: Optional[JudgeConfiguration] = None
tool_custom_parameters: Optional[ToolCustomParametersMap] = None

def get_tool_custom_parameter(
self, tool_name: str, key: str,
) -> Any:
"""Retrieve a custom parameter for a specific tool."""
return _get_tool_custom_parameter(
self.tool_custom_parameters, tool_name, key,
)

def to_dict(self) -> Dict[str, Any]:
"""
Expand All @@ -260,6 +348,9 @@ 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()
_serialize_tool_custom_parameters(
result, self.tool_custom_parameters,
)
return result


Expand All @@ -270,6 +361,15 @@ class AIAgentConfig(AIConfig):
"""
instructions: Optional[str] = None
judge_configuration: Optional[JudgeConfiguration] = None
tool_custom_parameters: Optional[ToolCustomParametersMap] = None

def get_tool_custom_parameter(
self, tool_name: str, key: str,
) -> Any:
"""Retrieve a custom parameter for a specific tool."""
return _get_tool_custom_parameter(
self.tool_custom_parameters, tool_name, key,
)

def to_dict(self) -> Dict[str, Any]:
"""
Expand Down
72 changes: 72 additions & 0 deletions packages/sdk/server-ai/tests/test_agents.py
Original file line number Diff line number Diff line change
Expand Up @@ -103,6 +103,30 @@ def td() -> TestData:
.variation_for_all(0)
)

# Agent with tools and custom parameters
td.update(
td.flag('agent-with-tools')
.variations(
{
'model': {
'name': 'gpt-4',
'parameters': {
'temperature': 0.3,
'tools': [
{'name': 'get-order', 'customParameters': {'includeHistory': True, 'maxItems': 5}},
{'name': 'search-products', 'customParameters': {'category': 'electronics'}},
{'name': 'send-email'},
],
},
},
'provider': {'name': 'openai'},
'instructions': 'You are a support agent with tools.',
'_ldMeta': {'enabled': True, 'variationKey': 'tools-v1', 'version': 1, 'mode': 'agent'},
}
)
.variation_for_all(0)
)

return td


Expand Down Expand Up @@ -363,3 +387,51 @@ def test_agents_request_without_default_uses_disabled(ldai_client: LDAIClient):

assert 'missing-agent' in agents
assert agents['missing-agent'].enabled is False


def test_agent_config_has_tools(ldai_client: LDAIClient):
"""Test that agent configs parse tools with custom parameters from flag variations."""
context = Context.create('user-key')

agent = ldai_client.agent_config('agent-with-tools', context)

assert agent.enabled is True
assert agent.tool_custom_parameters is not None
# Only tools with non-empty customParameters are included
assert len(agent.tool_custom_parameters) == 2

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'
# 'send-email' has no customParameters, so it is not in the map
assert 'send-email' not in agent.tool_custom_parameters
assert agent.get_tool_custom_parameter('send-email', 'anything') is None


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 = LDAIAgentDefaults(
enabled=False,
model=ModelConfig('fallback-model'),
instructions='Default instructions',
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.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 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.tool_custom_parameters is None
Loading
Loading