import gradio as gr
import random
import os
from smolagents import GradioUI, CodeAgent, HfApiModel, LiteLLMModel, ApiModel

# 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(model_id="Qwen/Qwen2.5-Coder-32B-Instruct")
# model = LiteLLMModel(model_id="gpt-4")

# Use Gemini 1.5 Flash - Long context (1M tokens) and cheap
model = LiteLLMModel(
    model_id="gemini/gemini-1.5-flash",
    api_key=os.getenv("GEMINI_API_KEY"),
    max_tokens=2000  # Can be higher due to long context window
)

# 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()

# Create Alfred with full configuration - Gemini can handle more tokens
alfred = CodeAgent(
    tools=[guest_info_tool, weather_info_tool, hub_stats_tool, search_tool], 
    model=model,
    add_base_tools=True,  # Re-enable base tools
    planning_interval=3   # Re-enable planning with shorter interval
)

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