Update agent.py

agent.py

@@ -1 +1,30 @@
-# agent.py - Core agent logic
+import os
+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"]
+
+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.headers = {
+            "Authorization": f"Bearer {os.environ['HF_TOKEN']}",
+            "Content-Type": "application/json"
+        }
+
+    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)}"