Update agent.py

agent.py

@@ -1,42 +1,36 @@
 import os
 import requests
 
-API_URL = "https://api-inference.huggingface.co/models/google/flan-t5-base"
-HF_TOKEN = os.getenv("HF_TOKEN")
+API_URL = "https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-alpha"
+HEADERS = {"Authorization": f"Bearer {os.getenv('HF_TOKEN')}"}
 
 system_prompt = (
-    "You are a helpful assistant participating in the GAIA benchmark. "
-    "Always return direct, factual answers with no explanation. Output only the final answer."
+    "You are an expert AI assistant taking a benchmark test for reasoning. "
+    "Answer only the main question directly. Do not explain. Do not show work. "
+    "If you see images or documents mentioned, assume you have access and extract answer. "
+    "Use search or tools if needed. Answer format must be plain with no prefix or suffix."
 )
 
-headers = {
-    "Authorization": f"Bearer {HF_TOKEN}",
-    "Content-Type": "application/json"
-}
-
 class BasicAgent:
     def __init__(self):
-        print("Flan-T5 Agent initialized using Hugging Face API")
+        print("✅ Agent initialized with Zephyr 7B.")
 
     def __call__(self, question: str) -> str:
-        prompt = f"{system_prompt}\n\nQuestion:\n{question}\n\nAnswer:"
-        payload = {
-            "inputs": prompt,
-            "parameters": {
-                "temperature": 0.0,
-                "max_new_tokens": 128
-            }
-        }
-
+        prompt = f"{system_prompt}\n\nQuestion: {question}\nAnswer:"
         try:
-            response = requests.post(API_URL, headers=headers, json=payload, timeout=30)
+            response = requests.post(
+                API_URL,
+                headers=HEADERS,
+                json={"inputs": prompt},
+                timeout=40
+            )
             response.raise_for_status()
-            data = response.json()
-            if isinstance(data, list) and "generated_text" in data[0]:
-                return data[0]["generated_text"].strip()
-            elif "generated_text" in data[0]:
-                return data[0]["generated_text"].strip()
+            result = response.json()
+            if isinstance(result, list):
+                return result[0]["generated_text"].split("Answer:")[-1].strip()
+            elif "generated_text" in result:
+                return result["generated_text"].strip()
             else:
-                return str(data)
+                return "ERROR: Unexpected response format"
         except Exception as e:
             return f"AGENT ERROR: {e}"