Update app.py

app.py

@@ -52,18 +52,21 @@ class MyAgent(CodeAgent):
     def __init__(self):
         model = HfApiModel(
             model="cognitivecomputations/dolphin-2.6-mixtral-8x7b",
-            api_key=os.getenv("HF_API_TOKEN", "").strip(),
+            api_key=os.getenv("HF_API_TOKEN", "").strip()
         )
         super().__init__(model=model, tools=tools)
         self.system_prompt = SYSTEM_PROMPT
 
     def __call__(self, task):
-        # Inject the system prompt manually
-        if isinstance(task, dict) and "question" in task:
-            task["messages"] = [
-                {"role": "system", "content": self.system_prompt},
-                {"role": "user", "content": task["question"]}
-            ]
+        # If task contains messages, prepend system prompt
+        if isinstance(task, dict):
+            if "question" in task:
+                task["messages"] = [
+                    {"role": "system", "content": self.system_prompt},
+                    {"role": "user", "content": task["question"]}
+                ]
+            elif "messages" in task:
+                task["messages"].insert(0, {"role": "system", "content": self.system_prompt})
         return super().__call__(task)