import gradio as gr
import os

from mcp import StdioServerParameters
from smolagents import InferenceClientModel, CodeAgent, ToolCollection, MCPClient

mcp_client = MCPClient(
    {
        "url": "https://abidlabs-mcp-tool-http.hf.space/gradio_api/mcp/sse"
    }
)

tools = mcp_client.get_tools()

model = InferenceClientModel(token=os.getenv("HF_TOKEN"))

agent = CodeAgent(tools=[*tools], model=model)

demo = gr.ChatInterface(
    fn= lambda message, history: str(agent.run(message)),
    type="messages",
    examples=["Prime factorization of 68"],
    title="Agent with MCP tool",
    description="This agent can use the MCP tool to answer questions.",
)

demo.launch()

finally:
mcp_client.disconnect()