Refactor AI configuration tests to use AICompletionConfigDefault and update async tracking for OpenAI metrics

Author: edwinokonkwo

ldai/providers/ai_provider_factory.py

@@ -113,6 +113,8 @@ async def _try_create_provider(
                     )
                 return None
 
+        # TODO: REL-10773 OpenAI provider
+        # TODO: REL-10776 Vercel provider
         # For future external providers, use dynamic import
         provider_mappings = {
             # 'openai': ('launchdarkly_server_sdk_ai_openai', 'OpenAIProvider'),

ldai/testing/test_model_config.py

@@ -2,7 +2,7 @@
 from ldclient import Config, Context, LDClient
 from ldclient.integrations.test_data import TestData
 
-from ldai import AIConfig, LDAIClient, LDMessage, ModelConfig
+from ldai import AICompletionConfigDefault, LDAIClient, LDMessage, ModelConfig
 
 
 @pytest.fixture
@@ -133,14 +133,14 @@ def test_model_config_handles_custom():
 
 def test_uses_default_on_invalid_flag(ldai_client: LDAIClient):
     context = Context.create('user-key')
-    default_value = AIConfig(
+    default_value = AICompletionConfigDefault(
         enabled=True,
         model=ModelConfig('fakeModel', parameters={'temperature': 0.5, 'maxTokens': 4096}),
         messages=[LDMessage(role='system', content='Hello, {{name}}!')],
     )
     variables = {'name': 'World'}
 
-    config, _ = ldai_client.config('missing-flag', context, default_value, variables)
+    config = ldai_client.config('missing-flag', context, default_value, variables)
 
     assert config.messages is not None
     assert len(config.messages) > 0
@@ -155,14 +155,14 @@ def test_uses_default_on_invalid_flag(ldai_client: LDAIClient):
 
 def test_model_config_interpolation(ldai_client: LDAIClient):
     context = Context.create('user-key')
-    default_value = AIConfig(
+    default_value = AICompletionConfigDefault(
         enabled=True,
         model=ModelConfig('fakeModel'),
         messages=[LDMessage(role='system', content='Hello, {{name}}!')],
     )
     variables = {'name': 'World'}
 
-    config, _ = ldai_client.config('model-config', context, default_value, variables)
+    config = ldai_client.config('model-config', context, default_value, variables)
 
     assert config.messages is not None
     assert len(config.messages) > 0
@@ -177,9 +177,9 @@ def test_model_config_interpolation(ldai_client: LDAIClient):
 
 def test_model_config_no_variables(ldai_client: LDAIClient):
     context = Context.create('user-key')
-    default_value = AIConfig(enabled=True, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=True, model=ModelConfig('fake-model'), messages=[])
 
-    config, _ = ldai_client.config('model-config', context, default_value, {})
+    config = ldai_client.config('model-config', context, default_value, {})
 
     assert config.messages is not None
     assert len(config.messages) > 0
@@ -194,21 +194,21 @@ def test_model_config_no_variables(ldai_client: LDAIClient):
 
 def test_provider_config_handling(ldai_client: LDAIClient):
     context = Context.builder('user-key').name("Sandy").build()
-    default_value = AIConfig(enabled=True, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=True, model=ModelConfig('fake-model'), messages=[])
     variables = {'name': 'World'}
 
-    config, _ = ldai_client.config('model-config', context, default_value, variables)
+    config = ldai_client.config('model-config', context, default_value, variables)
 
     assert config.provider is not None
     assert config.provider.name == 'fakeProvider'
 
 
 def test_context_interpolation(ldai_client: LDAIClient):
     context = Context.builder('user-key').name("Sandy").set('last', 'Beaches').build()
-    default_value = AIConfig(enabled=True, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=True, model=ModelConfig('fake-model'), messages=[])
     variables = {'name': 'World'}
 
-    config, _ = ldai_client.config(
+    config = ldai_client.config(
         'ctx-interpolation', context, default_value, variables
     )
 
@@ -228,10 +228,10 @@ def test_multi_context_interpolation(ldai_client: LDAIClient):
     user_context = Context.builder('user-key').name("Sandy").build()
     org_context = Context.builder('org-key').kind('org').name("LaunchDarkly").set('shortname', 'LD').build()
     context = Context.multi_builder().add(user_context).add(org_context).build()
-    default_value = AIConfig(enabled=True, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=True, model=ModelConfig('fake-model'), messages=[])
     variables = {'name': 'World'}
 
-    config, _ = ldai_client.config(
+    config = ldai_client.config(
         'multi-ctx-interpolation', context, default_value, variables
     )
 
@@ -249,10 +249,10 @@ def test_multi_context_interpolation(ldai_client: LDAIClient):
 
 def test_model_config_multiple(ldai_client: LDAIClient):
     context = Context.create('user-key')
-    default_value = AIConfig(enabled=True, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=True, model=ModelConfig('fake-model'), messages=[])
     variables = {'name': 'World', 'day': 'Monday'}
 
-    config, _ = ldai_client.config(
+    config = ldai_client.config(
         'multiple-messages', context, default_value, variables
     )
 
@@ -270,9 +270,9 @@ def test_model_config_multiple(ldai_client: LDAIClient):
 
 def test_model_config_disabled(ldai_client: LDAIClient):
     context = Context.create('user-key')
-    default_value = AIConfig(enabled=False, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=False, model=ModelConfig('fake-model'), messages=[])
 
-    config, _ = ldai_client.config('off-config', context, default_value, {})
+    config = ldai_client.config('off-config', context, default_value, {})
 
     assert config.model is not None
     assert config.enabled is False
@@ -283,9 +283,9 @@ def test_model_config_disabled(ldai_client: LDAIClient):
 
 def test_model_initial_config_disabled(ldai_client: LDAIClient):
     context = Context.create('user-key')
-    default_value = AIConfig(enabled=False, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=False, model=ModelConfig('fake-model'), messages=[])
 
-    config, _ = ldai_client.config('initial-config-disabled', context, default_value, {})
+    config = ldai_client.config('initial-config-disabled', context, default_value, {})
 
     assert config.enabled is False
     assert config.model is None
@@ -295,9 +295,9 @@ def test_model_initial_config_disabled(ldai_client: LDAIClient):
 
 def test_model_initial_config_enabled(ldai_client: LDAIClient):
     context = Context.create('user-key')
-    default_value = AIConfig(enabled=False, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=False, model=ModelConfig('fake-model'), messages=[])
 
-    config, _ = ldai_client.config('initial-config-enabled', context, default_value, {})
+    config = ldai_client.config('initial-config-enabled', context, default_value, {})
 
     assert config.enabled is True
     assert config.model is None
@@ -318,9 +318,9 @@ def test_config_method_tracking(ldai_client: LDAIClient):
 
     client = LDAIClient(mock_client)
     context = Context.create('user-key')
-    default_value = AIConfig(enabled=False, model=ModelConfig('fake-model'), messages=[])
+    default_value = AICompletionConfigDefault(enabled=False, model=ModelConfig('fake-model'), messages=[])
 
-    config, tracker = client.config('test-config-key', context, default_value)
+    config = client.config('test-config-key', context, default_value)
 
     mock_client.track.assert_called_once_with(
         '$ld:ai:config:function:single',

ldai/testing/test_tracker.py

@@ -276,7 +276,8 @@ def test_tracks_bedrock_metrics_with_error(client: LDClient):
     assert tracker.get_summary().usage == TokenUsage(330, 220, 110)
 
 
-def test_tracks_openai_metrics(client: LDClient):
+@pytest.mark.asyncio
+async def test_tracks_openai_metrics(client: LDClient):
     context = Context.create("user-key")
     tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
 
@@ -292,7 +293,10 @@ def to_dict(self):
                 "completion_tokens": 110,
             }
 
-    tracker.track_openai_metrics(lambda: Result())
+    async def get_result():
+        return Result()
+
+    await tracker.track_openai_metrics(get_result)
 
     calls = [
         call(
@@ -326,15 +330,16 @@ def to_dict(self):
     assert tracker.get_summary().usage == TokenUsage(330, 220, 110)
 
 
-def test_tracks_openai_metrics_with_exception(client: LDClient):
+@pytest.mark.asyncio
+async def test_tracks_openai_metrics_with_exception(client: LDClient):
     context = Context.create("user-key")
     tracker = LDAIConfigTracker(client, "variation-key", "config-key", 3, "fakeModel", "fakeProvider", context)
 
-    def raise_exception():
+    async def raise_exception():
         raise ValueError("Something went wrong")
 
     try:
-        tracker.track_openai_metrics(raise_exception)
+        await tracker.track_openai_metrics(raise_exception)
         assert False, "Should have thrown an exception"
     except ValueError:
         pass