PR Feedback

Remove tracking utils

return AI Config

Remove errant print

Author: Daniel OBrien

ldai/client.py

@@ -1,4 +1,3 @@
-import json
 from typing import Any, Dict, Optional
 from ldclient import Context
 from ldclient.client import LDClient
@@ -31,7 +30,6 @@ def model_config(self, key: str, context: Context, default_value: str, variables
         if variables:
             all_variables.update(variables)
 
-        print(variation)
         variation['prompt'] = [
             {
                 **entry,
@@ -40,10 +38,7 @@ def model_config(self, key: str, context: Context, default_value: str, variables
             for entry in variation['prompt']
         ]
 
-        return {
-            'config': variation,
-            'tracker': LDAIConfigTracker(self.client, variation['_ldMeta']['variationId'], key, context)
-        }
+        return AIConfig(config=variation, tracker=LDAIConfigTracker(self.client, variation['_ldMeta']['variationId'], key, context))
 
     def interpolate_template(self, template: str, variables: Dict[str, Any]) -> str:
         """Interpolate the template with the given variables.

ldai/tracker.py

@@ -1,6 +1,5 @@
 from typing import Dict, Union
 from ldclient import Context, LDClient
-from ldai.tracking_utils import usage_to_token_metrics
 from ldai.types import BedrockTokenUsage, FeedbackKind, TokenUsage, UnderscoreTokenUsage
 
 class LDAIConfigTracker:
@@ -20,7 +19,7 @@ def track_duration(self, duration: int) -> None:
         self.ld_client.track('$ld:ai:duration:total', self.context, self.get_track_data(), duration)
 
     def track_tokens(self, tokens: Union[TokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> None:
-        token_metrics = usage_to_token_metrics(tokens)
+        token_metrics = tokens.to_metrics()
         if token_metrics['total'] > 0:
             self.ld_client.track('$ld:ai:tokens:total', self.context, self.get_track_data(), token_metrics['total'])
         if token_metrics['input'] > 0:

ldai/tracking_utils.py

@@ -22,11 +22,8 @@ class AITracker(TypedDict):
 
 class AIConfig():
     def __init__(self, config: AIConfigData, tracker: AITracker):
-        self._config = config
-        self._tracker = tracker
-
-    config: AIConfigData
-    tracker: AITracker
+        self.config = config
+        self.tracker = tracker
 
 class FeedbackKind(Enum):
     Positive = "positive"