from smolagents import (
    load_tool,
    CodeAgent,
    LiteLLMModel,
    GradioUI,
    MCPClient
)
from mcp import StdioServerParameters
import os
from dotenv import load_dotenv
load_dotenv()

kgb_server_parameters = StdioServerParameters(
    command="npx",
    args=[        
        "mcp-remote",
        "https://agents-mcp-hackathon-kgb-mcp.hf.space/gradio_api/mcp/sse",
        "--transport",
        "sse-only"],
)

# Initialize the model
model = LiteLLMModel(
    model_id="gemini/gemini-2.0-flash-exp",
    api_key=os.getenv("GEMINI_API_TOKEN")
)

# Load tools from all MCP servers using MCPClient
server_parameters = [kgb_server_parameters]

mcp = MCPClient(server_parameters)
with mcp:
    all_tools = mcp.get_tools() 
    agent = CodeAgent(tools=all_tools, model=model, add_base_tools=True)

    GradioUI(agent).launch()