Update app.py

app.py

@@ -8,7 +8,7 @@ import pandas as pd
 import google.generativeai as genai
 from smolagents import CodeAgent, DuckDuckGoSearchTool
 
-# Define the system prompt
+# System prompt used by the agent
 SYSTEM_PROMPT = """You are a general AI assistant. I will ask you a question.
 Report your thoughts, and finish your answer with just the answer — no prefixes like "FINAL ANSWER:".
 Your answer should be a number OR as few words as possible OR a comma-separated list of numbers and/or strings.
@@ -17,16 +17,17 @@ If you're asked for a string, don’t use articles or abbreviations (e.g. for ci
 
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
-# Wrapper class for model output
+# Generation result wrapper to match smolagents expectations
 class GenerationResult:
-    def __init__(self, content, token_usage=None):
+    def __init__(self, content, token_usage=None, input_tokens=0, output_tokens=0):
         self.content = content
         self.token_usage = token_usage or {}
+        self.input_tokens = input_tokens
+        self.output_tokens = output_tokens
 
 # Gemini model wrapper
 class GeminiFlashModel:
     def __init__(self, model_id="gemini-1.5-flash", api_key=None):
-        import google.generativeai as genai
         genai.configure(api_key=api_key or os.getenv("GEMINI_API_KEY"))
         self.model = genai.GenerativeModel(model_id)
         self.system_prompt = SYSTEM_PROMPT
@@ -46,11 +47,21 @@ class GeminiFlashModel:
 
         try:
             response = self.model.generate_content(prompt)
-            return GenerationResult(content=response.text.strip(), token_usage={})  # Gemini doesn't expose usage
+            return GenerationResult(
+                content=response.text.strip(),
+                token_usage={},
+                input_tokens=0,
+                output_tokens=0
+            )
         except Exception as e:
-            return GenerationResult(content=f"GENERATION ERROR: {e}", token_usage={})
-
-# Agent using Gemini
+            return GenerationResult(
+                content=f"GENERATION ERROR: {e}",
+                token_usage={},
+                input_tokens=0,
+                output_tokens=0
+            )
+
+# Agent wrapper
 class MyAgent:
     def __init__(self):
         self.model = GeminiFlashModel(model_id="gemini-1.5-flash")
@@ -59,6 +70,7 @@ class MyAgent:
     def __call__(self, question: str) -> str:
         return self.agent.run(question)
 
+# Main evaluation function
 def run_and_submit_all(profile: gr.OAuthProfile | None):
     space_id = os.getenv("SPACE_ID")
 
@@ -127,7 +139,7 @@ def run_and_submit_all(profile: gr.OAuthProfile | None):
     except Exception as e:
         return f"Submission failed: {e}", pd.DataFrame(results_log)
 
-# Gradio UI
+# Gradio UI setup
 with gr.Blocks() as demo:
     gr.Markdown("# Basic Agent Evaluation Runner")
     gr.Markdown("""