from smolagents import CodeAgent, DuckDuckGoSearchTool, HfApiModel, load_tool, tool
import re
import yaml
from tools.final_answer import FinalAnswerTool
from Gradio_UI import GradioUI

    @tool
    def get_distance_to_planet(planet: str) -> str:
        """A tool that fetches the distance (in kilometers) from Earth to a given planet using an online lookup.
        Args:
            planet: Name of the planet in the solar system.
        """
        try:
            # Build a query to search for the distance information.
            query = f"distance from Earth to {planet} in km"
            # Use the DuckDuckGoSearchTool to perform the search.
            search_results = DuckDuckGoSearchTool.run(query=query)
            # Use a regex to extract a number (which may include commas or decimals) followed by 'km'
            match = re.search(r'([\d,\.]+)\s*km', search_results)
            if match:
                distance = match.group(1)
                return f"The distance from Earth to {planet} is approximately {distance} km."
            else:
                return f"Unable to extract the distance for {planet} from the search results."
        except Exception as e:
            return f"Error fetching the distance to {planet}: {str(e)}"

    final_answer = FinalAnswerTool()

model = HfApiModel(
    max_tokens=2096,
    temperature=0.5,
    model_id='Qwen/Qwen2.5-Coder-32B-Instruct',
    custom_role_conversions=None,
    )

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

agent = CodeAgent(
    model=model,
    tools=[final_answer, get_distance_to_planet],
    max_steps=6,
    verbosity_level=1,
    grammar=None,
    planning_interval=None,
    name="Planet_Distance_Agent",
    description="Agent that provides the distance from Earth to any planet in the solar system using a web search.",
    prompt_templates=prompt_templates
    )

GradioUI(agent).launch()