Update app.py

app.py

@@ -12,11 +12,42 @@ from tools.catering_service_tool import SimpleTool as CateringServiceTool
 from tools.superhero_party_theme_generator import SuperheroPartyThemeTool as SuperheroPartyThemeGenerator
 from tools.final_answer import FinalAnswerTool as FinalAnswer
 
+class CountedLiteLLMModel(LiteLLMModel):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        # ensure attrs exist (instrumentor will read them)
+        self.last_input_token_count = 0
+        self.last_output_token_count = 0
 
+    def generate(self, *args, **kwargs):
+        # call as usual
+        out = super().generate(*args, **kwargs)
+
+        # try to read usage from the last response if LiteLLM provided it
+        usage = getattr(self, "last_response_usage", None)
+        # some smolagents versions store usage on the response object instead:
+        if usage is None:
+            usage = getattr(out, "usage", None)
+
+        if usage:
+            # common OpenAI-style keys; adjust if your provider uses different names
+            self.last_input_token_count = \
+                usage.get("prompt_tokens") or usage.get("input_tokens") or 0
+            self.last_output_token_count = \
+                usage.get("completion_tokens") or usage.get("output_tokens") or 0
+        else:
+            # fallback: estimate with LiteLLM's token_counter to avoid zeros
+            prompt = kwargs.get("prompt") or (args[0] if args else "")
+            try:
+                self.last_input_token_count = token_counter(model=self.model_id, text=prompt) or 0
+            except Exception:
+                pass  # leave prior value if estimation fails
+
+        return out
 
 model = CountedLiteLLMModel(
 model_id='claude-3-5-sonnet-latest',
-api_base=None,
+api_key=os.getenv('ANTHROPIC_API_KEY'),
 )
 
 web_search = WebSearch()