fix: Only track usage for entry methods (#94)

Author: jsonbailey

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

@@ -45,24 +45,13 @@ def __init__(self, client: LDClient):
             1,
         )
 
-    def completion_config(
+    def _completion_config(
         self,
         key: str,
         context: Context,
         default_value: AICompletionConfigDefault,
         variables: Optional[Dict[str, Any]] = None,
     ) -> AICompletionConfig:
-        """
-        Get the value of a completion configuration.
-
-        :param key: The key of the completion configuration.
-        :param context: The context to evaluate the completion configuration in.
-        :param default_value: The default value of the completion configuration.
-        :param variables: Additional variables for the completion configuration.
-        :return: The completion configuration with a tracker used for gathering metrics.
-        """
-        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
         )
@@ -79,6 +68,26 @@ def completion_config(
 
         return config
 
+    def completion_config(
+        self,
+        key: str,
+        context: Context,
+        default_value: AICompletionConfigDefault,
+        variables: Optional[Dict[str, Any]] = None,
+    ) -> AICompletionConfig:
+        """
+        Get the value of a completion configuration.
+
+        :param key: The key of the completion configuration.
+        :param context: The context to evaluate the completion configuration in.
+        :param default_value: The default value of the completion configuration.
+        :param variables: Additional variables for the completion configuration.
+        :return: The completion configuration with a tracker used for gathering metrics.
+        """
+        self._client.track(_TRACK_USAGE_COMPLETION_CONFIG, context, key, 1)
+
+        return self._completion_config(key, context, default_value, variables)
+
     def config(
         self,
         key: str,
@@ -99,24 +108,13 @@ def config(
         """
         return self.completion_config(key, context, default_value, variables)
 
-    def judge_config(
+    def _judge_config(
         self,
         key: str,
         context: Context,
         default_value: AIJudgeConfigDefault,
         variables: Optional[Dict[str, Any]] = None,
     ) -> AIJudgeConfig:
-        """
-        Get the value of a judge configuration.
-
-        :param key: The key of the judge configuration.
-        :param context: The context to evaluate the judge configuration in.
-        :param default_value: The default value of the judge configuration.
-        :param variables: Additional variables for the judge configuration.
-        :return: The judge configuration with a tracker used for gathering metrics.
-        """
-        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
         )
@@ -150,6 +148,26 @@ def _extract_evaluation_metric_key(variation: Dict[str, Any]) -> Optional[str]:
 
         return config
 
+    def judge_config(
+        self,
+        key: str,
+        context: Context,
+        default_value: AIJudgeConfigDefault,
+        variables: Optional[Dict[str, Any]] = None,
+    ) -> AIJudgeConfig:
+        """
+        Get the value of a judge configuration.
+
+        :param key: The key of the judge configuration.
+        :param context: The context to evaluate the judge configuration in.
+        :param default_value: The default value of the judge configuration.
+        :param variables: Additional variables for the judge configuration.
+        :return: The judge configuration with a tracker used for gathering metrics.
+        """
+        self._client.track(_TRACK_USAGE_JUDGE_CONFIG, context, key, 1)
+
+        return self._judge_config(key, context, default_value, variables)
+
     async def create_judge(
         self,
         key: str,
@@ -204,7 +222,7 @@ async def create_judge(
             extended_variables['message_history'] = '{{message_history}}'
             extended_variables['response_to_evaluate'] = '{{response_to_evaluate}}'
 
-            judge_config = self.judge_config(key, context, default_value, extended_variables)
+            judge_config = self._judge_config(key, context, default_value, extended_variables)
 
             if not judge_config.enabled or not judge_config.tracker:
                 return None
@@ -304,7 +322,7 @@ async def create_chat(
         """
         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)
+        config = self._completion_config(key, context, default_value, variables)
 
         if not config.enabled or not config.tracker:
             return None