Remove problematic token usage tracking

Author: satwikkansal

mxgo/agents/email_agent.py

@@ -11,17 +11,22 @@
 from smolagents import Tool, ToolCallingAgent
 from smolagents.monitoring import TokenUsage
 
-# Monkey patch TokenUsage to handle None values
+# Monkey patch TokenUsage to handle None values robustly
 original_post_init = TokenUsage.__post_init__
 
 
 def patched_post_init(self):
-    # Handle None values by setting them to 0
+    # Handle None values by setting them to 0 before any operations
     if self.input_tokens is None:
         self.input_tokens = 0
     if self.output_tokens is None:
         self.output_tokens = 0
-    original_post_init(self)
+    # Call original post_init which will now work with non-None values
+    try:
+        original_post_init(self)
+    except Exception:
+        # If original post_init still fails, set total_tokens manually
+        self.total_tokens = self.input_tokens + self.output_tokens
 
 
 TokenUsage.__post_init__ = patched_post_init

mxgo/routed_litellm_model.py

@@ -48,10 +48,6 @@ def __init__(self, current_handle: ProcessingInstructions | None = None, target_
             )
             raise exceptions.EnvironmentVariableNotFoundError(msg)
 
-        # Initialize token count attributes before calling super().__init__
-        self._last_input_token_count = 0
-        self._last_output_token_count = 0
-
         # Set environment variables for Azure OpenAI models before initializing the parent class
         # This is required because LiteLLM's Azure provider looks for specific environment variables
         self._set_azure_environment_variables(model_list)
@@ -240,15 +236,6 @@ def generate(
 
         response = self.client.completion(**completion_kwargs)
 
-        # Safely handle response usage tracking
-        if response and hasattr(response, "usage") and response.usage:
-            self._last_input_token_count = response.usage.prompt_tokens
-            self._last_output_token_count = response.usage.completion_tokens
-        else:
-            # Set default values if usage information is not available
-            self._last_input_token_count = 0
-            self._last_output_token_count = 0
-
         return ChatMessage.from_dict(
             response.choices[0].message.model_dump(include={"role", "content", "tool_calls"}),
             raw=response,
@@ -305,32 +292,3 @@ def __call__(
             raise RuntimeError(msg) from e
         else:
             return chat_message
-
-    @property
-    def last_input_token_count(self) -> int:
-        """Safely return the last input token count."""
-        return getattr(self, "_last_input_token_count", 0)
-
-    @last_input_token_count.setter
-    def last_input_token_count(self, value: int | None) -> None:
-        """Safely set the last input token count."""
-        self._last_input_token_count = value if value is not None else 0
-
-    @property
-    def last_output_token_count(self) -> int:
-        """Safely return the last output token count."""
-        return getattr(self, "_last_output_token_count", 0)
-
-    @last_output_token_count.setter
-    def last_output_token_count(self, value: int | None) -> None:
-        """Safely set the last output token count."""
-        self._last_output_token_count = value if value is not None else 0
-
-    def __getattr__(self, name: str):
-        """Handle any missing attribute access gracefully, especially for token-related properties."""
-        # Handle various token-related attribute names that might be accessed
-        if name in ("input_tokens", "output_tokens", "total_tokens", "prompt_tokens", "completion_tokens", "usage"):
-            return 0
-        # For other missing attributes, raise AttributeError as normal
-        msg = f"'{self.__class__.__name__}' object has no attribute '{name}'"
-        raise AttributeError(msg)