Update app.py

app.py

@@ -1,7 +1,7 @@
 import yaml
 import os
 from smolagents import GradioUI, CodeAgent, LiteLLMModel
-from litellm import token_counter  # fallback when provider doesn't return usage
+# from litellm import token_counter  # fallback when provider doesn't return usage
 
 # Get current directory path
 CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))
@@ -13,40 +13,40 @@ 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
+# 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)
+#     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)
+#         # 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
+#         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
+#         return out
 
-model = CountedLiteLLMModel(
+model = LiteLLMModel(
     model_id='claude-3-5-sonnet-latest',
     api_key=os.getenv('ANTHROPIC_API_KEY'),
 )