Update agent.py

agent.py

@@ -3,26 +3,27 @@ import json
 
 from tools import TOOLS
 from metadata import load_metadata  
-from mistral_inference import query_mistral  
+from mistral_hf_wrapper import MistralInference
 
+# Load environment variables
 API_URL = os.getenv("HF_MISTRAL_URL")
 API_TOKEN = os.getenv("HF_TOKEN")
 
-HEADERS = {
-    "Authorization": f"Bearer {API_TOKEN}",
-    "Content-Type": "application/json"
-}
-
 # Load all tasks from metadata.jsonl
 def load_tasks():
     return load_metadata("metadata.jsonl")
 
+# Initialize Mistral
+mistral = MistralInference()
+
 # Solve a single task
 def solve_task(task, tools=TOOLS):
     system_prompt = "You are a helpful agent. Use reasoning, tools if needed, and return the answer only."
     user_prompt = task["question"]
 
-    response = query_mistral(API_URL, HEADERS, system_prompt, user_prompt)
+    prompt = f"<s>[INST] <<SYS>>\n{system_prompt}\n<</SYS>>\n{user_prompt} [/INST]"
+    response = mistral.run(prompt)
+
     return {
         "task_id": task["question_id"],
         "submitted_answer": response.strip()