Update app.py

app.py

@@ -5,10 +5,9 @@ import gradio as gr
 import requests
 import pandas as pd
 
-import google.generativeai as genai
-from smolagents import CodeAgent, DuckDuckGoSearchTool
+from smolagents import LiteLLMModel, CodeAgent, DuckDuckGoSearchTool
 
-# System prompt used by the agent
+# System prompt for 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,66 +16,29 @@ 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"
 
-# Generation result wrapper to match smolagents expectations
-class GenerationResult:
-    def __init__(self, content, input_tokens=0, output_tokens=0, token_usage=None):
-        self.content = content
-        self.input_tokens = input_tokens
-        self.output_tokens = output_tokens
-        self.token_usage = token_usage or {}
-
-# Gemini model wrapper
-class GeminiFlashModel:
-    def __init__(self, model_id="gemini-1.5-flash", api_key=None):
-        genai.configure(api_key=api_key or os.getenv("GEMINI_API_KEY"))
-        self.model = genai.GenerativeModel(model_id)
-        self.system_prompt = SYSTEM_PROMPT
-
-    def generate(self, messages, **kwargs):
-        if not isinstance(messages, list) or not all(isinstance(m, dict) for m in messages):
-            raise TypeError("Expected 'messages' to be a list of dicts")
-
-        # Ensure system prompt is first message
-        if not any(m.get("role") == "system" for m in messages):
-            messages = [{"role": "system", "content": self.system_prompt}] + messages
-
-        # Build prompt text by concatenating messages with roles
-        prompt = ""
-        for m in messages:
-            role = m["role"].capitalize()
-            content = m["content"]
-            prompt += f"{role}: {content}\n"
-
-        try:
-            response = self.model.generate_content(prompt)
-            # Always wrap the result in GenerationResult
-            return GenerationResult(
-                content=response.text.strip(),
-                input_tokens=0,   # Could add token counts here if available
-                output_tokens=0,
-            )
-        except Exception as e:
-            # Wrap errors too, so agent doesn't fail
-            return GenerationResult(
-                content=f"GENERATION ERROR: {e}",
-                input_tokens=0,
-                output_tokens=0,
-            )
-
-# Agent wrapper
+# Agent wrapper using LiteLLMModel
 class MyAgent:
     def __init__(self):
-        self.model = GeminiFlashModel(model_id="gemini-1.5-flash")
-        self.agent = CodeAgent(tools=[DuckDuckGoSearchTool()], model=self.model)
+        gemini_api_key = os.getenv("GEMINI_API_KEY")
+        if not gemini_api_key:
+            raise ValueError("GEMINI_API_KEY not set in environment variables.")
+        
+        # Instantiate LiteLLMModel with Gemini API key and model id
+        self.model = LiteLLMModel(
+            model_id="gemini/gemini-2.0-flash-lite",
+            api_key=gemini_api_key,
+            system_prompt=SYSTEM_PROMPT
+        )
+        
+        # Create the CodeAgent with optional base tools and DuckDuckGo search
+        self.agent = CodeAgent(
+            tools=[DuckDuckGoSearchTool()],
+            model=self.model,
+            add_base_tools=True,
+        )
 
     def __call__(self, question: str) -> str:
-        result = self.agent.run(question)
-        # result can be GenerationResult or maybe dict or str - normalize:
-        if hasattr(result, "content"):
-            return result.content
-        if isinstance(result, dict):
-            return result.get("content", str(result))
-        return str(result)
+        return self.agent.run(question)
 
 # Main evaluation function
 def run_and_submit_all(profile: gr.OAuthProfile | None):