import gradio as gr
from smolagents import GradioUI, CodeAgent, HfApiModel
from tools import DuckDuckGoSearchTool, WeatherInfoTool, HubStatsTool
from retriever import load_guest_dataset

# Initialize the Hugging Face model
model = HfApiModel()

# Initialize the web search tool
search_tool = DuckDuckGoSearchTool()

# Initialize the weather tool
weather_info_tool = WeatherInfoTool()

# Initialize the Hub stats tool
hub_stats_tool = HubStatsTool()

# Load the guest dataset and initialize the guest info tool
guest_info_tool = load_guest_dataset()

# Custom agent class to format responses
class FormattedCodeAgent(CodeAgent):
    def run(self, prompt, reset=True, **kwargs):
        # Call the parent class's run method
        response = super().run(prompt, reset=reset, **kwargs)
        
        # Convert response to a string if it’s a dictionary
        if isinstance(response, dict):
            response_str = "\n".join([f"{key}: {value}" for key, value in response.items()])
        else:
            response_str = str(response)
        
        return response_str

# Create Alfred with formatted responses
alfred = FormattedCodeAgent(
    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool],
    model=model,
    add_base_tools=True,  # Add any additional base tools
    planning_interval=3   # Enable planning every 3 steps
)

# Launch the GradioUI interface
GradioUI(alfred).launch()