Update agent.py

agent.py

@@ -1,30 +1,44 @@
 import os
+import json
+from typing import Dict, List
 import requests
-import yaml
-from tools import web_search, wikipedia_search, visit_webpage, final_answer
 
-with open("prompts.yaml", "r") as f:
-    prompts = yaml.safe_load(f)
-
-SYSTEM_PROMPT = prompts["system_prompt"]
+from gaia_benchmark.questions import load_questions
+from gaia_benchmark.run import run_and_submit_all
 
 class GaiaAgent:
     def __init__(self):
-        self.system_prompt = SYSTEM_PROMPT
-        self.endpoint_url = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.1"
+        self.api_url = os.environ["HF_MISTRAL_ENDPOINT"]  # Your Mistral endpoint
+        self.api_key = os.environ["HF_TOKEN"]             # Hugging Face token
         self.headers = {
-            "Authorization": f"Bearer {os.environ['HF_TOKEN']}",
-            "Content-Type": "application/json"
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def generate(self, prompt: str, stop: List[str] = []) -> str:
+        payload = {
+            "inputs": prompt,
+            "parameters": {
+                "temperature": 0.0,
+                "max_new_tokens": 1024,
+                "stop": stop,
+            }
         }
+        response = requests.post(self.api_url, headers=self.headers, json=payload)
+        response.raise_for_status()
+        output = response.json()
+        if isinstance(output, dict) and "generated_text" in output:
+            return output["generated_text"]
+        if isinstance(output, list):
+            return output[0]["generated_text"]
+        return str(output)
+
+    def answer_question(self, question: Dict) -> str:
+        question_text = question["question"]
+        prompt = f"""You are a helpful agent answering a science question.
+Question: {question_text}
+Answer:"""
+        return self.generate(prompt).strip()
 
-    def __call__(self, task):
-        prompt = self.system_prompt + "\nTask:\n" + task
-        payload = {"inputs": prompt}
-        try:
-            response = requests.post(self.endpoint_url, headers=self.headers, json=payload)
-            output = response.json()
-            if isinstance(output, list) and "generated_text" in output[0]:
-                return output[0]["generated_text"]
-            return output
-        except Exception as e:
-            return f"AGENT ERROR: {str(e)}"
+    def run(self):
+        run_and_submit_all(self.answer_question)