Update app.py

app.py

@@ -16,11 +16,28 @@ of numbers and/or strings. If you are asked for a number, don't use comma to wri
 DEFAULT_API_URL = "https://agents-course-unit4-scoring.hf.space"
 
 # Custom model wrapper that injects system message
+#class PatchedOpenAIServerModel(OpenAIServerModel):
+ #   def generate(self, messages, *args, **kwargs):
+ #       messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
+ #       return super().generate(messages, *args, **kwargs)
+
 class PatchedOpenAIServerModel(OpenAIServerModel):
     def generate(self, messages, *args, **kwargs):
-        messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
+        # Handle string inputs (raw prompts)
+        if isinstance(messages, str):
+            messages = [
+                {"role": "system", "content": SYSTEM_PROMPT},
+                {"role": "user", "content": messages}
+            ]
+        elif isinstance(messages, list):
+            # Assume already in proper format
+            messages = [{"role": "system", "content": SYSTEM_PROMPT}] + messages
+        else:
+            raise TypeError("Expected messages to be a string or list of dicts")
+        
         return super().generate(messages, *args, **kwargs)
 
+
 class MyAgent:
     def __init__(self):
         self.model = PatchedOpenAIServerModel(