from smolagents import LiteLLMModel, CodeAgent, WebSearchTool, GradioUI, Tool, ToolCollection, tool, load_tool

model = LiteLLMModel(model_id="groq/llama-3.3-70b-versatile")


weather_tool = load_tool("ngandugilbert/weather_tool", trust_remote_code=True)

image_generation_tool = Tool.from_space(
    "black-forest-labs/FLUX.1-schnell",
    name="image_generator",
    description="Generate an image from a prompt"
)
# with ToolCollection.from_mcp({"url": "http://127.0.0.1:9000/mcp", "transport": "streamable-http"}, trust_remote_code=True) as tool_collection:
agent = CodeAgent(model=model, tools=[WebSearchTool(),image_generation_tool, weather_tool])

GradioUI(agent).launch()