Skip to content

Commit 97aa9f4

Browse files
committed
feat: Added agent support to SDK
1 parent 7226d56 commit 97aa9f4

File tree

2 files changed

+520
-7
lines changed

2 files changed

+520
-7
lines changed

ldai/client.py

Lines changed: 175 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -125,6 +125,69 @@ def to_dict(self) -> dict:
125125
}
126126

127127

128+
@dataclass(frozen=True)
129+
class LDAIAgent:
130+
"""
131+
Represents an AI agent configuration with instructions and model settings.
132+
133+
An agent is similar to an AIConfig but focuses on instructions rather than messages,
134+
making it suitable for AI assistant/agent use cases.
135+
"""
136+
enabled: Optional[bool] = None
137+
model: Optional[ModelConfig] = None
138+
provider: Optional[ProviderConfig] = None
139+
instructions: Optional[str] = None
140+
tracker: Optional[LDAIConfigTracker] = None
141+
142+
def to_dict(self) -> dict:
143+
"""
144+
Render the given agent as a dictionary object.
145+
"""
146+
result = {
147+
'_ldMeta': {
148+
'enabled': self.enabled or False,
149+
},
150+
'model': self.model.to_dict() if self.model else None,
151+
'provider': self.provider.to_dict() if self.provider else None,
152+
}
153+
if self.instructions is not None:
154+
result['instructions'] = self.instructions
155+
return result
156+
157+
158+
@dataclass(frozen=True)
159+
class LDAIAgentDefaults:
160+
"""
161+
Default values for AI agent configurations.
162+
163+
Similar to LDAIAgent but without tracker and with optional enabled field,
164+
used as fallback values when agent configurations are not available.
165+
"""
166+
enabled: Optional[bool] = None
167+
model: Optional[ModelConfig] = None
168+
provider: Optional[ProviderConfig] = None
169+
instructions: Optional[str] = None
170+
171+
def to_dict(self) -> dict:
172+
"""
173+
Render the given agent defaults as a dictionary object.
174+
"""
175+
result = {
176+
'_ldMeta': {
177+
'enabled': self.enabled or False,
178+
},
179+
'model': self.model.to_dict() if self.model else None,
180+
'provider': self.provider.to_dict() if self.provider else None,
181+
}
182+
if self.instructions is not None:
183+
result['instructions'] = self.instructions
184+
return result
185+
186+
187+
# Type alias for multiple agents
188+
LDAIAgents = Dict[str, LDAIAgent]
189+
190+
128191
class LDAIClient:
129192
"""The LaunchDarkly AI SDK client object."""
130193

@@ -147,13 +210,88 @@ def config(
147210
:param variables: Additional variables for the model configuration.
148211
:return: The value of the model configuration along with a tracker used for gathering metrics.
149212
"""
150-
variation = self._client.variation(key, context, default_value.to_dict())
213+
model, provider, messages, tracker, enabled = self.__evaluate(key, context, default_value.to_dict(), variables)
214+
215+
config = AIConfig(
216+
enabled=bool(enabled),
217+
model=model,
218+
messages=messages,
219+
provider=provider,
220+
)
221+
222+
return config, tracker
223+
224+
def agents(
225+
self,
226+
keys: List[str],
227+
context: Context,
228+
default_value: LDAIAgentDefaults,
229+
variables: Optional[Dict[str, Any]] = None,
230+
) -> LDAIAgents:
231+
"""
232+
Get multiple AI agent configurations.
233+
234+
This method allows you to retrieve multiple agent configurations in a single call,
235+
with each agent having its instructions dynamically interpolated with the provided
236+
variables and context data.
237+
238+
Example:
239+
```python
240+
agents = client.agents(
241+
['customer-support', 'sales-assistant'],
242+
context,
243+
LDAIAgentDefaults(
244+
enabled=True,
245+
model=ModelConfig('gpt-4'),
246+
instructions="You are a helpful assistant."
247+
),
248+
{'company_name': 'Acme Corp'}
249+
)
250+
251+
support_agent = agents['customer-support']
252+
if support_agent.enabled:
253+
print(support_agent.instructions) # Instructions with interpolated variables
254+
# Use support_agent.tracker for metrics tracking
255+
```
256+
257+
:param keys: List of agent configuration keys to retrieve.
258+
:param context: The context to evaluate the agent configurations in.
259+
:param default_value: Default agent configuration values to use as fallback.
260+
:param variables: Additional variables for template interpolation in instructions.
261+
:return: Dictionary mapping agent keys to their LDAIAgent configurations.
262+
"""
263+
result: LDAIAgents = {}
264+
265+
for key in keys:
266+
agent = self.__evaluate_agent(key, context, default_value, variables)
267+
result[key] = agent
268+
269+
return result
270+
271+
def __evaluate(
272+
self,
273+
key: str,
274+
context: Context,
275+
default_dict: Dict[str, Any],
276+
variables: Optional[Dict[str, Any]] = None,
277+
) -> Tuple[Optional[ModelConfig], Optional[ProviderConfig], Optional[List[LDMessage]], Optional[str], LDAIConfigTracker]:
278+
"""
279+
Internal method to evaluate a configuration and extract components.
280+
281+
:param key: The configuration key.
282+
:param context: The evaluation context.
283+
:param default_dict: Default configuration as dictionary.
284+
:param variables: Variables for interpolation.
285+
:return: Tuple of (model, provider, messages, instructions, tracker, enabled).
286+
"""
287+
variation = self._client.variation(key, context, default_dict)
151288

152289
all_variables = {}
153290
if variables:
154291
all_variables.update(variables)
155292
all_variables['ldctx'] = context.to_dict()
156293

294+
# Extract messages
157295
messages = None
158296
if 'messages' in variation and isinstance(variation['messages'], list) and all(
159297
isinstance(entry, dict) for entry in variation['messages']
@@ -168,11 +306,18 @@ def config(
168306
for entry in variation['messages']
169307
]
170308

309+
# Extract instructions
310+
instructions = None
311+
if 'instructions' in variation and isinstance(variation['instructions'], str):
312+
instructions = self.__interpolate_template(variation['instructions'], all_variables)
313+
314+
# Extract provider config
171315
provider_config = None
172316
if 'provider' in variation and isinstance(variation['provider'], dict):
173317
provider = variation['provider']
174318
provider_config = ProviderConfig(provider.get('name', ''))
175319

320+
# Extract model config
176321
model = None
177322
if 'model' in variation and isinstance(variation['model'], dict):
178323
parameters = variation['model'].get('parameters', None)
@@ -183,6 +328,7 @@ def config(
183328
custom=custom
184329
)
185330

331+
# Create tracker
186332
tracker = LDAIConfigTracker(
187333
self._client,
188334
variation.get('_ldMeta', {}).get('variationKey', ''),
@@ -192,14 +338,36 @@ def config(
192338
)
193339

194340
enabled = variation.get('_ldMeta', {}).get('enabled', False)
195-
config = AIConfig(
196-
enabled=bool(enabled),
197-
model=model,
198-
messages=messages,
199-
provider=provider_config,
341+
342+
return model, provider_config, messages, instructions, tracker, enabled
343+
344+
def __evaluate_agent(
345+
self,
346+
key: str,
347+
context: Context,
348+
default_value: LDAIAgentDefaults,
349+
variables: Optional[Dict[str, Any]] = None,
350+
) -> LDAIAgent:
351+
"""
352+
Internal method to evaluate an agent configuration.
353+
354+
:param key: The agent configuration key.
355+
:param context: The evaluation context.
356+
:param default_value: Default agent values.
357+
:param variables: Variables for interpolation.
358+
:return: Configured LDAIAgent instance.
359+
"""
360+
model, provider, instructions, tracker, enabled = self.__evaluate(
361+
key, context, default_value.to_dict(), variables
200362
)
201363

202-
return config, tracker
364+
return LDAIAgent(
365+
enabled=bool(enabled) if enabled is not None else None,
366+
model=model or default_value.model,
367+
provider=provider or default_value.provider,
368+
instructions=instructions,
369+
tracker=tracker,
370+
)
203371

204372
def __interpolate_template(self, template: str, variables: Dict[str, Any]) -> str:
205373
"""

0 commit comments

Comments
 (0)