add back missing dataclass

remove UnderscoreToken class to match TypeScript

Author: Daniel OBrien

ldai/tracker.py

@@ -1,7 +1,7 @@
 import time
 from typing import Dict, Union
 from ldclient import Context, LDClient
-from ldai.types import BedrockTokenUsage, FeedbackKind, OpenAITokenUsage, TokenUsage, UnderscoreTokenUsage
+from ldai.types import BedrockTokenUsage, FeedbackKind, OpenAITokenUsage, TokenUsage
 
 class LDAIConfigTracker:
     def __init__(self, ld_client: LDClient, version_key: str, config_key: str, context: Context):
@@ -55,7 +55,7 @@ def track_bedrock_converse(self, res: dict) -> dict:
             self.track_tokens(BedrockTokenUsage(res['usage']))
         return res
 
-    def track_tokens(self, tokens: Union[TokenUsage, UnderscoreTokenUsage, BedrockTokenUsage]) -> None:
+    def track_tokens(self, tokens: Union[TokenUsage, BedrockTokenUsage]) -> None:
         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'])

ldai/types.py

@@ -32,6 +32,7 @@ class FeedbackKind(Enum):
     Positive = "positive"
     Negative = "negative"
 
+@dataclass
 class TokenUsage():
     total_tokens: int
     prompt_tokens: int
@@ -58,19 +59,6 @@ def to_metrics(self) -> TokenMetrics:
             'output': self.completion_tokens,
         }
 
-class UnderscoreTokenUsage:
-    def __init__(self, data: dict):
-        self.total_tokens = data.get('total_tokens', 0)
-        self.prompt_tokens = data.get('prompt_tokens', 0)
-        self.completion_tokens = data.get('completion_tokens', 0)
-
-    def to_metrics(self) -> TokenMetrics:
-        return {
-            'total': self.total_tokens,
-            'input': self.prompt_tokens,
-            'output': self.completion_tokens,
-        }
-
 class BedrockTokenUsage:
     def __init__(self, data: dict):
         self.totalTokens = data.get('totalTokens', 0)