import gradio as gr
import os

from mcp.client.stdio import StdioServerParameters
from smolagents import InferenceClientModel, CodeAgent, ToolCollection
from smolagents.mcp_client import MCPClient


mcp_client = None
api_key1 = os.getenv('mcp')

# Initialize MCPClient with URL
try:
   
    mcp_client = MCPClient(
        ## Try this working example on the hub:
        # {"url": "https://abidlabs-mcp-tools.hf.space/gradio_api/mcp/sse"}
        {"url":"https://wlchee-mcp-sentiment.hf.space/gradio_api/mcp/sse"},
     )
    # Retrieve tools from the MCP client
    tools = mcp_client.get_tools()
    if not tools:
        raise Exception("No tools retrieved from the MCP client.")
    
    # Create the model and agent
    model = InferenceClientModel(model_id="Qwen/Qwen2.5-Omni-7B")
    agent = CodeAgent(tools=[*tools], model=model)

    # Define the Gradio interface
    demo = gr.ChatInterface(
        fn=lambda message, history: str(agent.run(message)),
        type="messages",
        examples=["Prime factorization of 68"],
        title="Agent with MCP Tools",
        description="This is a simple agent that uses MCP tools to answer questions.",
    )

    # Launch the interface
    demo.launch(share=True)

finally:
    # Ensure MCPClient disconnects on script termination
    if mcp_client:
        mcp_client.disconnect()