import gradio as gr
import random
from smolagents import GradioUI, CodeAgent, HfApiModel

# Import our custom tools from their modules
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()


# Do not allow use some base tools from smolagents

# from smolagents.prompts import CODE_SYSTEM_PROMPT
# modified_system_prompt = CODE_SYSTEM_PROMPT + "\nAlways use tools that you are provided, do not use tools VisitWebpageTool and DuckDuckGoSearchTool from base toolbox." #


# Create Alfred with all the tools
alfred = CodeAgent(
    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool], 
    model=model,
    add_base_tools=False,  # Add any additional base tools
    planning_interval=3,   # Enable planning every 3 steps
    # system_prompt=modified_system_prompt
)

print(alfred.system_prompt)

if __name__ == "__main__":
    GradioUI(alfred).launch()