Update agent.py

agent.py

@@ -1,15 +1,13 @@
 import os
-import json
-from typing import Dict, List
 import requests
+from typing import Dict, List
 
-from gaia_benchmark.questions import load_questions
 from gaia_benchmark.run import run_and_submit_all
 
 class GaiaAgent:
     def __init__(self):
-        self.api_url = os.environ["HF_MISTRAL_ENDPOINT"]  # Your Mistral endpoint
-        self.api_key = os.environ["HF_TOKEN"]             # Hugging Face token
+        self.api_url = os.environ["HF_MISTRAL_ENDPOINT"]
+        self.api_key = os.environ["HF_TOKEN"]
         self.headers = {
             "Authorization": f"Bearer {self.api_key}",
             "Content-Type": "application/json",
@@ -27,16 +25,17 @@ class GaiaAgent:
         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):
+        elif isinstance(output, list) and "generated_text" in output[0]:
             return output[0]["generated_text"]
         return str(output)
 
     def answer_question(self, question: Dict) -> str:
-        question_text = question["question"]
+        q = question["question"]
         prompt = f"""You are a helpful agent answering a science question.
-Question: {question_text}
+Question: {q}
 Answer:"""
         return self.generate(prompt).strip()