Update app.py

app.py

@@ -17,7 +17,12 @@ 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"
 
-# Gemini model wrapper (lightweight, no smolagents.model.base)
+# Wrapper class for model output
+class GenerationResult:
+    def __init__(self, content):
+        self.content = content
+
+# Gemini model wrapper
 class GeminiFlashModel:
     def __init__(self, model_id="gemini-1.5-flash", api_key=None):
         import google.generativeai as genai
@@ -26,14 +31,12 @@ class GeminiFlashModel:
         self.system_prompt = SYSTEM_PROMPT
 
     def generate(self, messages, stop_sequences=None, **kwargs):
-        # ✅ 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()
@@ -42,10 +45,9 @@ class GeminiFlashModel:
 
         try:
             response = self.model.generate_content(prompt)
-            return {"content": response.text.strip()}  # ✅ SmolAgent expects this format
+            return GenerationResult(response.text.strip())
         except Exception as e:
-            return {"content": f"GENERATION ERROR: {e}"}
-
+            return GenerationResult(f"GENERATION ERROR: {e}")
 
 # Agent using Gemini
 class MyAgent: