agent.py

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)}"