from smolagents import CodeAgent, InferenceClientModel
from smolagents import VisitWebpageTool, WebSearchTool, WikipediaSearchTool, PythonInterpreterTool,FinalAnswerTool


class GaiaAgent:
    def __init__(self):
        self.model = InferenceClientModel(
            max_tokens = 8096,
            temperature = 0.5, 
            model_id = "deepseek-ai/DeepSeek-R1-0528"
            provider = "huggingface"         #'https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud',
            custom_role_conversations=None
        )
#or 
#model=InferenceClientModel(model_id="Qwen/Qwen2.5-Coder-32B-Instruct", provider="together")

#with open("prompts.yaml", 'r') as stream:
 #   prompt_templates = yaml.safe_load(stream)

#Creating the agent
        self.agent = CodeAgent(
            model=self.model,
            tools=[
                VisitWebpageTool(),
                WebSearchTool(),
                WikipediaSearchTool(),
                PythonInterpreterTool(),
                FinalAnswerTool(),
            ], 
            max_steps=10,
            verbosity_level = 1,        #or other parameter
            grammar = None,             
            planning_interval=None,      
            #name="web_agent",            
            #description="Brows the web to find information",      
            #additional_autorized_imports = ["pandas"],
        )
    def __call__(self, question: str) -> str:
        try:
            return self.agent.run(question)
        except Exception as e:
            return f"ERROR: {e}"

#manager_agent = CodeAgent(
 #   model=InferenceClientModel("deepseek-ai/DeepSeek-R1", provider="together", max_tokens=8096),
  #  managed_agent=[web_agent],
   # additional_authorized_imports=[
    #    "datetime",
     #   "pandas",
      #  "numpy",
       # "json"
   # ],
    #planning_intervals=5,
    #verbosity_level=2,
    #max_steps=15,
#)