Update app.py

app.py

@@ -19,28 +19,32 @@ DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # Gemini model wrapper (lightweight, no smolagents.model.base)
 class GeminiFlashModel:
-    def __init__(self, model_name="gemini-1.5-flash", api_key=None):
-        self.model_name = model_name
-        self.api_key = api_key or os.getenv("GEMINI_API_KEY")
-        if not self.api_key:
-            raise ValueError("GEMINI_API_KEY is not set.")
-        genai.configure(api_key=self.api_key)
-        self.model = genai.GenerativeModel(model_name)
+    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("GOOGLE_API_KEY"))
+        self.model = genai.GenerativeModel(model_id)
+        self.system_prompt = SYSTEM_PROMPT
 
     def generate(self, messages, stop_sequences=None, **kwargs):
-        if isinstance(messages, list):
-            if not any(m["role"] == "system" for m in messages):
-                messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
-            else:
-                raise TypeError("Expected 'messages' to be a list of dicts.")
-                
-                prompt = "\n".join(f"{m['role'].capitalize()}: {m['content']}" for m in messages)
-                
-                try:
-                    response = self.model.generate_content(prompt)
-                    return {"content": response.text.strip()}  # ✅ Return a dict with 'content'
-                except Exception as e:
-                    return {"content": f"GENERATION ERROR: {e}"}  # ✅ Error message also in 'content'
+        # ✅ Validate and prepend system prompt
+        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")
+
+        if not any(m.get("role") == "system" for m in messages):
+            messages = [{"role": "system", "content": self.system_prompt}] + messages
+
+        # ✅ Construct prompt string (you may improve formatting here)
+        prompt = ""
+        for m in messages:
+            role = m["role"].capitalize()
+            content = m["content"]
+            prompt += f"{role}: {content}\n"
+
+        try:
+            response = self.model.generate_content(prompt)
+            return {"content": response.text.strip()}  # ✅ SmolAgent expects this format
+        except Exception as e:
+            return {"content": f"GENERATION ERROR: {e}"}
 
 
 # Agent using Gemini