Update agent.py

agent.py

@@ -1,4 +1,6 @@
 import os
+import json
+
 from tools import TOOLS
 from metadata import load_metadata
 from mistral_hf_wrapper import MistralInference
@@ -6,22 +8,26 @@ from mistral_hf_wrapper import MistralInference
 API_URL = os.getenv("HF_MISTRAL_ENDPOINT")
 API_TOKEN = os.getenv("HF_TOKEN")
 
+mistral = MistralInference(api_url=API_URL, api_token=API_TOKEN)
+
 def load_tasks():
     tasks = []
     with open("metadata.jsonl", "r", encoding="utf-8") as f:
         for line in f:
             item = json.loads(line)
-            question_id = item.get("question_id")
+            question_id = item.get("task_id") or item.get("question_id")
             question = item.get("Question") or item.get("question")
             if question_id and question:
                 tasks.append({
-                    "question_id": task_id,
+                    "question_id": question_id,
                     "question": question
                 })
     return tasks
 
 def solve_task(task, tools=TOOLS):
-    system_prompt = "You are a helpful agent. Use reasoning, tools if needed, and return the answer only."
+    system_prompt = (
+        "You are a helpful agent. Use reasoning, tools if needed, and return the answer only."
+    )
     user_prompt = task["question"]
 
     prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n{user_prompt} [/INST]"
@@ -29,5 +35,6 @@ def solve_task(task, tools=TOOLS):
 
     return {
         "question_id": task["question_id"],
-        "submitted_answer": response.strip()
+        "question": task["question"],
+        "submitted_answer": response.strip() if response else "ERROR: Empty model response"
     }