@@ -120,14 +120,17 @@ def ldai_client(client: LDClient) -> LDAIClient:
120120def test_single_agent_method (ldai_client : LDAIClient ):
121121 """Test the single agent() method functionality."""
122122 context = Context .builder ('user-key' ).set ('expertise' , 'advanced' ).build ()
123- defaults = LDAIAgentDefaults (
124- enabled = False ,
125- model = ModelConfig ('fallback-model' ),
126- instructions = "Default instructions"
123+ config = LDAIAgentConfig (
124+ key = 'research-agent' ,
125+ default_value = LDAIAgentDefaults (
126+ enabled = False ,
127+ model = ModelConfig ('fallback-model' ),
128+ instructions = "Default instructions"
129+ ),
130+ variables = {'topic' : 'quantum computing' }
127131 )
128- variables = {'topic' : 'quantum computing' }
129132
130- agent = ldai_client .agent ('research-agent' , context , defaults , variables )
133+ agent = ldai_client .agent (config , context )
131134
132135 assert agent .enabled is True
133136 assert agent .model is not None
@@ -143,15 +146,18 @@ def test_single_agent_method(ldai_client: LDAIClient):
143146def test_single_agent_with_defaults (ldai_client : LDAIClient ):
144147 """Test single agent method with non-existent flag using defaults."""
145148 context = Context .create ('user-key' )
146- defaults = LDAIAgentDefaults (
147- enabled = True ,
148- model = ModelConfig ('default-model' , parameters = {'temp' : 0.8 }),
149- provider = ProviderConfig ('default-provider' ),
150- instructions = "You are a default assistant for {{task}}."
149+ config = LDAIAgentConfig (
150+ key = 'non-existent-agent' ,
151+ default_value = LDAIAgentDefaults (
152+ enabled = True ,
153+ model = ModelConfig ('default-model' , parameters = {'temp' : 0.8 }),
154+ provider = ProviderConfig ('default-provider' ),
155+ instructions = "You are a default assistant for {{task}}."
156+ ),
157+ variables = {'task' : 'general assistance' }
151158 )
152- variables = {'task' : 'general assistance' }
153159
154- agent = ldai_client .agent ('non-existent-agent' , context , defaults , variables )
160+ agent = ldai_client .agent (config , context )
155161
156162 assert agent .enabled is True
157163 assert agent .model is not None and agent .model .name == 'default-model'
@@ -236,7 +242,7 @@ def test_agents_method_different_variables_per_agent(ldai_client: LDAIClient):
236242
237243def test_agents_with_multi_context_interpolation (ldai_client : LDAIClient ):
238244 """Test agents method with multi-context interpolation."""
239- user_context = Context .builder ('user-key' ).name ('Bob ' ).build ()
245+ user_context = Context .builder ('user-key' ).name ('Alice ' ).build ()
240246 org_context = Context .builder ('org-key' ).kind ('org' ).name ('LaunchDarkly' ).set ('tier' , 'Enterprise' ).build ()
241247 context = Context .multi_builder ().add (user_context ).add (org_context ).build ()
242248
@@ -252,21 +258,24 @@ def test_agents_with_multi_context_interpolation(ldai_client: LDAIClient):
252258 ]
253259
254260 agents = ldai_client .agents (agent_configs , context )
255- agent = agents ['multi-context-agent' ]
256261
257- expected_instructions = 'Welcome Bob from LaunchDarkly! Your organization tier is Enterprise.'
258- assert agent .instructions == expected_instructions
262+ agent = agents [ 'multi-context-agent' ]
263+ assert agent .instructions == 'Welcome Alice from LaunchDarkly! Your organization tier is Enterprise.'
259264
260265
261266def test_disabled_agent_single_method (ldai_client : LDAIClient ):
262267 """Test that disabled agents are properly handled in single agent method."""
263268 context = Context .create ('user-key' )
264- defaults = LDAIAgentDefaults (enabled = True , instructions = "Default" )
269+ config = LDAIAgentConfig (
270+ key = 'disabled-agent' ,
271+ default_value = LDAIAgentDefaults (enabled = False ),
272+ variables = {}
273+ )
265274
266- agent = ldai_client .agent ('disabled-agent' , context , defaults )
275+ agent = ldai_client .agent (config , context )
267276
268277 assert agent .enabled is False
269- assert agent .instructions == 'This agent is disabled.'
278+ assert agent .tracker is not None
270279
271280
272281def test_disabled_agent_multiple_method (ldai_client : LDAIClient ):
@@ -276,116 +285,35 @@ def test_disabled_agent_multiple_method(ldai_client: LDAIClient):
276285 agent_configs = [
277286 LDAIAgentConfig (
278287 key = 'disabled-agent' ,
279- default_value = LDAIAgentDefaults (enabled = True , instructions = "Default" ),
288+ default_value = LDAIAgentDefaults (enabled = False ),
280289 variables = {}
281290 )
282291 ]
283292
284293 agents = ldai_client .agents (agent_configs , context )
285- agent = agents ['disabled-agent' ]
286294
287- assert agent . enabled is False
288- assert agent . instructions == 'This agent is disabled.'
295+ assert len ( agents ) == 1
296+ assert agents [ 'disabled- agent' ]. enabled is False
289297
290298
291299def test_agent_with_missing_metadata (ldai_client : LDAIClient ):
292300 """Test agent handling when metadata is minimal or missing."""
293301 context = Context .create ('user-key' )
294- defaults = LDAIAgentDefaults (
295- enabled = False ,
296- model = ModelConfig ('default-model' ),
297- instructions = "Default instructions"
302+ config = LDAIAgentConfig (
303+ key = 'minimal-agent' ,
304+ default_value = LDAIAgentDefaults (
305+ enabled = False ,
306+ model = ModelConfig ('default-model' ),
307+ instructions = "Default instructions"
308+ )
298309 )
299310
300- agent = ldai_client .agent ('minimal-agent' , context , defaults )
311+ agent = ldai_client .agent (config , context )
301312
302313 assert agent .enabled is True # From flag
303314 assert agent .instructions == 'Minimal agent configuration.'
304- assert agent .model == defaults .model # Falls back to default
305- assert agent .tracker is not None
306-
307-
308- def test_empty_agents_list (ldai_client : LDAIClient ):
309- """Test agents method with empty agent configs list."""
310- context = Context .create ('user-key' )
311-
312- agents = ldai_client .agents ([], context )
313-
314- assert len (agents ) == 0
315- assert agents == {}
316-
317-
318- def test_agent_tracker_functionality (ldai_client : LDAIClient ):
319- """Test that agent tracker works correctly."""
320- context = Context .create ('user-key' )
321- defaults = LDAIAgentDefaults (enabled = True , instructions = "Default" )
322-
323- agent = ldai_client .agent ('customer-support-agent' , context , defaults )
324-
315+ assert agent .model == config .default_value .model # Falls back to default
325316 assert agent .tracker is not None
326- assert hasattr (agent .tracker , 'track_success' )
327- assert hasattr (agent .tracker , 'track_duration' )
328- assert hasattr (agent .tracker , 'track_tokens' )
329-
330-
331- def test_agent_tracking_calls (ldai_client : LDAIClient ):
332- """Test that tracking calls are made for agent usage."""
333- from unittest .mock import MagicMock , patch
334-
335- context = Context .create ('user-key' )
336- defaults = LDAIAgentDefaults (enabled = True , instructions = "Default" )
337-
338- # Test single agent tracking
339- with patch .object (ldai_client ._client , 'track' ) as mock_track :
340- ldai_client .agent ('customer-support-agent' , context , defaults )
341- mock_track .assert_called_with (
342- "$ld:ai:agent:function:single" ,
343- context ,
344- 'customer-support-agent' ,
345- 1
346- )
347-
348- # Test multiple agents tracking
349- agent_configs = [
350- LDAIAgentConfig (
351- key = 'customer-support-agent' ,
352- default_value = defaults ,
353- variables = {}
354- ),
355- LDAIAgentConfig (
356- key = 'sales-assistant' ,
357- default_value = defaults ,
358- variables = {}
359- )
360- ]
361-
362- with patch .object (ldai_client ._client , 'track' ) as mock_track :
363- ldai_client .agents (agent_configs , context )
364- mock_track .assert_called_with (
365- "$ld:ai:agent:function:multiple" ,
366- context ,
367- 2 ,
368- 2
369- )
370-
371-
372- def test_backwards_compatibility_with_config (ldai_client : LDAIClient ):
373- """Test that the existing config method still works after agent additions."""
374- from ldai .client import AIConfig , LDMessage
375-
376- context = Context .create ('user-key' )
377- default_value = AIConfig (
378- enabled = True ,
379- model = ModelConfig ('test-model' ),
380- messages = [LDMessage (role = 'system' , content = 'Test message' )]
381- )
382-
383- # This should still work as before
384- config , tracker = ldai_client .config ('customer-support-agent' , context , default_value )
385-
386- assert config .enabled is True
387- assert config .model is not None
388- assert tracker is not None
389317
390318
391319def test_agent_config_dataclass ():
@@ -412,49 +340,3 @@ def test_agent_config_dataclass():
412340
413341 assert config_no_vars .key == 'test-agent-2'
414342 assert config_no_vars .variables is None
415-
416-
417- def test_agent_config_optional_default_value ():
418- """Test that LDAIAgentConfig defaults to {enabled: False} when default_value is not provided."""
419- config = LDAIAgentConfig (key = 'test-agent' )
420-
421- assert config .key == 'test-agent'
422- assert config .default_value is not None
423- assert config .default_value .enabled is False
424- assert config .variables is None
425-
426-
427- def test_single_agent_optional_default_value (ldai_client : LDAIClient ):
428- """Test the single agent() method with optional default_value."""
429- context = Context .create ('user-key' )
430-
431- # Should work with no default_value provided (defaults to {enabled: False})
432- agent = ldai_client .agent ('non-existent-agent' , context )
433-
434- assert agent .enabled is False # Should default to False
435- assert agent .tracker is not None
436-
437-
438- def test_agents_method_with_optional_defaults (ldai_client : LDAIClient ):
439- """Test agents method with optional default_value configurations."""
440- context = Context .create ('user-key' )
441-
442- agent_configs = [
443- LDAIAgentConfig (key = 'customer-support-agent' ), # No default_value
444- LDAIAgentConfig (
445- key = 'sales-assistant' ,
446- default_value = LDAIAgentDefaults (enabled = True , instructions = "Custom sales assistant" )
447- )
448- ]
449-
450- agents = ldai_client .agents (agent_configs , context )
451-
452- assert len (agents ) == 2
453-
454- # First agent should use default {enabled: False} from auto-generated default_value
455- support_agent = agents ['customer-support-agent' ]
456- assert support_agent .enabled is True # From flag configuration
457-
458- # Second agent should use custom default
459- sales_agent = agents ['sales-assistant' ]
460- assert sales_agent .enabled is True
0 commit comments