from smolagents import CodeAgent,DuckDuckGoSearchTool, HfApiModel,load_tool,tool
import datetime
import requests
import pytz
import yaml
from tools.final_answer import FinalAnswerTool

from Gradio_UI import GradioUI

import requests

class FinalAnswerTool:
    def __call__(self, result: str):
        return f"\n Final Answer:\n{result}"

class ToolCallingAgent:
    def __init__(self, tools, final_tool, model, max_steps=10):
        self.tools = tools
        self.final_tool = final_tool
        self.model = model
        self.max_steps = max_steps

    def run(self, input_text):
        print(f"\n Agent received: {input_text}")
        for tool in self.tools:
            try:
                result = tool(input_text)
                if result:
                    final_output = self.final_tool(str(result))
                    print(final_output)
                    return final_output  # Stop after first useful tool response
            except Exception as e:
                print(f"\n[{tool.__class__.__name__} Error]: {str(e)}")

# Initialize all tools
weather_tool = WeatherTool(api_key="YOUR_API_KEY_HERE")
visit_webpage = VisitWebpageTool()
duckduckgo_tool = CustomDuckDuckGoSearchTool()
final_answer = FinalAnswerTool()
model = DummyModel()

# Initialize the agent with the final answer tool
web_agent = ToolCallingAgent(
    tools=[duckduckgo_tool, visit_webpage, weather_tool],
    final_tool=final_answer,
    model=model,
    max_steps=10
)

# Run with a query
web_agent.run("Tokyo")  # Works for weather, search, or URL


final_answer = FinalAnswerTool()

# If the agent does not answer, the model is overloaded, please use another model or the following Hugging Face Endpoint that also contains qwen2.5 coder:
# model_id='https://pflgm2locj2t89co.us-east-1.aws.endpoints.huggingface.cloud' 

model = HfApiModel(
max_tokens=2096,
temperature=0.5,
model_id='Qwen/Qwen2.5-Coder-32B-Instruct',# it is possible that this model may be overloaded
custom_role_conversions=None,
)


# Import tool from Hub
image_generation_tool = load_tool("agents-course/text-to-image", trust_remote_code=True)

with open("prompts.yaml", 'r') as stream:
    prompt_templates = yaml.safe_load(stream)
    
agent = CodeAgent(
    model=model,
    tools=[final_answer], ## add your tools here (don't remove final answer)
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name=None,
    description=None,
    prompt_templates=prompt_templates
)


GradioUI(agent).launch()