[AIC-1636] fix: Improve usage reporting

Co-Authored-By: jbailey@launchdarkly.com <accounts@sidewaysgravity.com>

Author: devin-ai-integration[bot]

packages/sdk/server-ai/src/ldai/client.py

@@ -4,6 +4,7 @@
 from ldclient import Context
 from ldclient.client import LDClient
 
+from ldai import __version__ as ai_sdk_version
 from ldai import log
 from ldai.agent_graph import AgentGraphDefinition
 from ldai.chat import Chat
@@ -17,12 +18,32 @@
 from ldai.providers.ai_provider_factory import AIProviderFactory
 from ldai.tracker import LDAIConfigTracker
 
+_TRACK_SDK_INFO = '$ld:ai:sdk-info'
+_TRACK_USAGE_COMPLETION_CONFIG = '$ld:ai:usage:completion-config'
+_TRACK_USAGE_CREATE_CHAT = '$ld:ai:usage:create-chat'
+_TRACK_USAGE_JUDGE_CONFIG = '$ld:ai:usage:judge-config'
+_TRACK_USAGE_CREATE_JUDGE = '$ld:ai:usage:create-judge'
+_TRACK_USAGE_AGENT_CONFIG = '$ld:ai:usage:agent-config'
+_TRACK_USAGE_AGENT_CONFIGS = '$ld:ai:usage:agent-configs'
+
+_INIT_TRACK_CONTEXT = Context.builder('ld-internal-tracking').anonymous(True).build()
+
 
 class LDAIClient:
     """The LaunchDarkly AI SDK client object."""
 
     def __init__(self, client: LDClient):
         self._client = client
+        self._client.track(
+            _TRACK_SDK_INFO,
+            _INIT_TRACK_CONTEXT,
+            {
+                'aiSdkName': 'launchdarkly-server-sdk-ai',
+                'aiSdkVersion': ai_sdk_version,
+                'aiSdkLanguage': 'python',
+            },
+            1,
+        )
 
     def completion_config(
         self,
@@ -40,7 +61,7 @@ def completion_config(
         :param variables: Additional variables for the completion configuration.
         :return: The completion configuration with a tracker used for gathering metrics.
         """
-        self._client.track('$ld:ai:config:function:single', context, key, 1)
+        self._client.track(_TRACK_USAGE_COMPLETION_CONFIG, context, key, 1)
 
         model, provider, messages, instructions, tracker, enabled, judge_configuration, _ = self.__evaluate(
             key, context, default_value.to_dict(), variables
@@ -94,7 +115,7 @@ def judge_config(
         :param variables: Additional variables for the judge configuration.
         :return: The judge configuration with a tracker used for gathering metrics.
         """
-        self._client.track('$ld:ai:judge:function:single', context, key, 1)
+        self._client.track(_TRACK_USAGE_JUDGE_CONFIG, context, key, 1)
 
         model, provider, messages, instructions, tracker, enabled, judge_configuration, variation = self.__evaluate(
             key, context, default_value.to_dict(), variables
@@ -170,7 +191,7 @@ async def create_judge(
                     if relevance_eval:
                         print('Relevance score:', relevance_eval.score)
         """
-        self._client.track('$ld:ai:judge:function:createJudge', context, key, 1)
+        self._client.track(_TRACK_USAGE_CREATE_JUDGE, context, key, 1)
 
         try:
             if variables:
@@ -281,7 +302,7 @@ async def create_chat(
                 messages = chat.get_messages()
                 print(f"Conversation has {len(messages)} messages")
         """
-        self._client.track('$ld:ai:config:function:createChat', context, key, 1)
+        self._client.track(_TRACK_USAGE_CREATE_CHAT, context, key, 1)
         log.debug(f"Creating chat for key: {key}")
         config = self.completion_config(key, context, default_value, variables)
 
@@ -340,7 +361,7 @@ def agent_config(
         :return: Configured AIAgentConfig instance.
         """
         self._client.track(
-            "$ld:ai:agent:function:single",
+            _TRACK_USAGE_AGENT_CONFIG,
             context,
             key,
             1
@@ -406,7 +427,7 @@ def agent_configs(
         """
         agent_count = len(agent_configs)
         self._client.track(
-            "$ld:ai:agent:function:multiple",
+            _TRACK_USAGE_AGENT_CONFIGS,
             context,
             agent_count,
             agent_count

packages/sdk/server-ai/tests/test_model_config.py

@@ -307,7 +307,7 @@ def test_model_initial_config_enabled(ldai_client: LDAIClient):
 
 
 def test_config_method_tracking(ldai_client: LDAIClient):
-    from unittest.mock import Mock
+    from unittest.mock import Mock, call
 
     mock_client = Mock()
     mock_client.variation.return_value = {
@@ -323,9 +323,30 @@ def test_config_method_tracking(ldai_client: LDAIClient):
 
     config = client.config('test-config-key', context, default_value)
 
-    mock_client.track.assert_called_once_with(
-        '$ld:ai:config:function:single',
+    mock_client.track.assert_any_call(
+        '$ld:ai:usage:completion-config',
         context,
         'test-config-key',
         1
     )
+
+
+def test_sdk_info_tracked_on_init():
+    from unittest.mock import Mock, call
+
+    from ldai.client import _INIT_TRACK_CONTEXT
+
+    mock_client = Mock()
+
+    client = LDAIClient(mock_client)
+
+    mock_client.track.assert_called_once_with(
+        '$ld:ai:sdk-info',
+        _INIT_TRACK_CONTEXT,
+        {
+            'aiSdkName': 'launchdarkly-server-sdk-ai',
+            'aiSdkVersion': '0.14.0',
+            'aiSdkLanguage': 'python',
+        },
+        1,
+    )