import yaml
import spaces
import os
from smolagents import GradioUI, CodeAgent, TransformersModel

# Get current directory path
CURRENT_DIR = os.path.dirname(os.path.abspath(__file__))

from tools.get_coordinates import SimpleTool as GetCoordinates
from tools.get_current_weather import SimpleTool as GetCurrentWeather
from tools.final_answer import FinalAnswerTool as FinalAnswer


get_coordinates = GetCoordinates()
get_current_weather = GetCurrentWeather()
final_answer = FinalAnswer()


with open(os.path.join(CURRENT_DIR, "prompts.yaml"), 'r') as stream:
    prompt_templates = yaml.safe_load(stream)

@spaces.GPU
def call_agent():
    model = TransformersModel(
        max_new_tokens=1000,
        model_id='Qwen/Qwen2.5-Coder-3B-Instruct',
    )
    
    agent_get_weather_agent = CodeAgent(
        model=model,
        tools=[get_coordinates, get_current_weather],
        managed_agents=[],
        max_steps=20,
        verbosity_level=2,
        grammar=None,
        planning_interval=None,
        name='get_weather_agent',
        description=None,
        executor_type='local',
        executor_kwargs={},
        max_print_outputs_length=None,
        prompt_templates=prompt_templates
    )
    GradioUI(agent_get_weather_agent).launch()


if __name__ == "__main__":
    call_agent()