Feature: Added generic "usage" stats tracking for API that support "usage" key in thier responses.

Author: amanjaiswal73892

src/agentlab/agents/tool_use_agent/multi_tool_agent.py

@@ -294,6 +294,7 @@ def obs_preprocessor(self, obs):
 
     @cost_tracker_decorator
     def get_action(self, obs: Any) -> float:
+        self.llm.reset_stats()
         if len(self.messages) == 0:
             self.config.goal.apply(self.llm, self.messages, obs)
             self.config.general_hints.apply(self.llm, self.messages)
@@ -309,7 +310,11 @@ def get_action(self, obs: Any) -> float:
         self._responses.append(response)  # may be useful for debugging
         # self.messages.append(response.assistant_message)  # this is tool call
 
-        agent_info = bgym.AgentInfo(think=think, chat_messages=self.messages, stats={})
+        agent_info = bgym.AgentInfo(
+            think=think,
+            chat_messages=self.messages,
+            stats=self.llm.stats.stats_dict,
+        )
         return action, agent_info
 
 

src/agentlab/llm/tracking.py

@@ -1,13 +1,16 @@
+import logging
 import os
 import re
 import threading
+from collections import defaultdict
 from contextlib import contextmanager
+from dataclasses import dataclass, field
 from functools import cache
+from typing import Optional
 
 import requests
 from langchain_community.callbacks import bedrock_anthropic_callback, openai_info
-from typing import Optional
-import logging
+
 
 TRACKER = threading.local()
 
@@ -136,15 +139,25 @@ class TrackAPIPricingMixin:
     Usage: provide the pricing_api to use in the constructor.
     """
 
+    def reset_stats(self):
+        self.stats = Stats()
+
     def __init__(self, *args, **kwargs):
         pricing_api = kwargs.pop("pricing_api", None)
         self._pricing_api = pricing_api
         super().__init__(*args, **kwargs)
         self.set_pricing_attributes()
+        self.reset_stats()
 
     def __call__(self, *args, **kwargs):
         """Call the API and update the pricing tracker."""
         response = self._call_api(*args, **kwargs)
+
+        usage = dict(getattr(response, "usage", {}))
+        usage = {f"usage_{k}": v for k, v in usage.items() if isinstance(v, (int, float))}
+        usage |= {"n_api_calls": 1}
+        self.stats.increment_stats_dict(usage)
+
         self.update_pricing_tracker(response)
         return self._parse_response(response)
 
@@ -215,3 +228,15 @@ def get_tokens_counts_from_response(self, response) -> tuple:
             "Unable to extract input and output tokens from the response. Defaulting to 0."
         )
         return 0, 0
+
+
+@dataclass
+class Stats:
+    stats_dict: dict = field(default_factory=lambda: defaultdict(float))
+
+    def increment_stats_dict(self, stats_dict: dict):
+        """increment the stats_dict with the given values."""
+        for k, v in stats_dict.items():
+            self.stats_dict[k] += v
+
+