Update app.py

app.py

@@ -2,13 +2,21 @@ import gradio as gr
 import asyncio
 import json
 import os
+import base64
+import io
+import re
 from typing import List, Dict, Any, Optional
 from dataclasses import dataclass
 import aiohttp
-from mcp import ClientSession
-from mcp.client.sse import sse_client
 import anthropic
 from datetime import datetime
+import logging
+from PIL import Image
+import requests
+
+# Set up logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Configuration for MCP servers (Gradio apps acting as MCP servers)
 @dataclass
@@ -17,29 +25,9 @@ class MCPServerConfig:
     url: str
     description: str
 
-# Default MCP servers - you can add your own Gradio apps here
-DEFAULT_SERVERS = [
-    MCPServerConfig(
-        name="Letter Counter",
-        url="https://your-letter-counter-space.hf.space/gradio_api/mcp/sse",
-        description="Counts letters in text"
-    ),
-    MCPServerConfig(
-        name="Weather Service", 
-        url="https://your-weather-space.hf.space/gradio_api/mcp/sse",
-        description="Provides weather information"
-    ),
-    MCPServerConfig(
-        name="Image Generator",
-        url="https://your-image-gen-space.hf.space/gradio_api/mcp/sse", 
-        description="Generates images from text prompts"
-    )
-]
-
 class MCPClient:
     def __init__(self):
         self.servers: Dict[str, MCPServerConfig] = {}
-        self.sessions: Dict[str, ClientSession] = {}
         self.available_tools: Dict[str, Dict] = {}
         self.anthropic_client = None
         
@@ -48,86 +36,120 @@ class MCPClient:
             self.anthropic_client = anthropic.Anthropic(
                 api_key=os.getenv("ANTHROPIC_API_KEY")
             )
+            logger.info("✅ Anthropic client initialized")
+        else:
+            logger.warning("⚠️ ANTHROPIC_API_KEY not found")
     
-    async def add_server(self, config: MCPServerConfig) -> bool:
-        """Add an MCP server to the client"""
+    async def add_server_async(self, config: MCPServerConfig) -> tuple[bool, str]:
+        """Add an MCP server to the client (async version)"""
         try:
             self.servers[config.name] = config
-            await self._connect_server(config.name)
-            return True
+            
+            # Simple HTTP health check instead of full MCP connection during add
+            success, message = await self._test_server_connection(config)
+            
+            if success:
+                logger.info(f"✅ Server {config.name} added successfully")
+                return True, f"✅ Successfully added server: {config.name}"
+            else:
+                logger.error(f"❌ Failed to add server {config.name}: {message}")
+                return False, f"❌ Failed to add server: {config.name} - {message}"
+                
         except Exception as e:
-            print(f"Failed to add server {config.name}: {e}")
-            return False
+            error_msg = f"Failed to add server {config.name}: {str(e)}"
+            logger.error(error_msg)
+            return False, f"❌ {error_msg}"
     
-    async def _connect_server(self, server_name: str) -> bool:
-        """Connect to a specific MCP server"""
+    async def _test_server_connection(self, config: MCPServerConfig) -> tuple[bool, str]:
+        """Test if server is reachable"""
         try:
-            config = self.servers[server_name]
+            # Ensure URL has correct MCP endpoint
+            if not config.url.endswith("/gradio_api/mcp/sse"):
+                base_url = config.url.rstrip("/")
+                config.url = f"{base_url}/gradio_api/mcp/sse"
             
-            # Create SSE client connection
-            async with sse_client(config.url) as (read, write):
-                session = ClientSession(read, write)
-                await session.initialize()
-                
-                # Get available tools from this server
-                tools = await session.list_tools()
-                
-                # Store tools with server prefix to avoid conflicts
-                for tool in tools.tools:
-                    tool_key = f"{server_name}::{tool.name}"
-                    self.available_tools[tool_key] = {
-                        "server": server_name,
-                        "tool": tool,
-                        "session": session
-                    }
-                
-                self.sessions[server_name] = session
-                print(f"✅ Connected to {server_name}: {len(tools.tools)} tools available")
-                return True
-                
+            # Simple HTTP check with timeout
+            timeout = aiohttp.ClientTimeout(total=10)
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                async with session.get(config.url.replace("/gradio_api/mcp/sse", "")) as response:
+                    if response.status == 200:
+                        return True, "Server is reachable"
+                    else:
+                        return False, f"Server returned status {response.status}"
+                        
+        except asyncio.TimeoutError:
+            return False, "Connection timeout"
         except Exception as e:
-            print(f"❌ Failed to connect to {server_name}: {e}")
-            return False
+            return False, f"Connection error: {str(e)}"
     
-    async def list_all_tools(self) -> List[str]:
-        """List all available tools from all connected servers"""
-        return list(self.available_tools.keys())
-    
-    async def call_tool(self, tool_key: str, arguments: Dict[str, Any]) -> Any:
-        """Call a tool on the appropriate MCP server"""
-        if tool_key not in self.available_tools:
-            raise ValueError(f"Tool {tool_key} not found")
-        
-        tool_info = self.available_tools[tool_key]
-        session = tool_info["session"]
-        tool = tool_info["tool"]
-        
+    async def call_real_mcp_tool(self, server_name: str, tool_name: str, arguments: dict) -> tuple[bool, str, Optional[str]]:
+        """Call real MCP tool on connected server"""
         try:
-            result = await session.call_tool(tool.name, arguments)
-            return result
+            if server_name not in self.servers:
+                return False, f"Server {server_name} not found", None
+            
+            config = self.servers[server_name]
+            
+            # Make HTTP request to the Gradio app's API
+            api_url = config.url.replace("/gradio_api/mcp/sse", "/api/predict")
+            
+            timeout = aiohttp.ClientTimeout(total=30)
+            async with aiohttp.ClientSession(timeout=timeout) as session:
+                # Prepare the request payload for Gradio API
+                payload = {
+                    "data": list(arguments.values()),
+                    "fn_index": 0  # Usually the first (and often only) function
+                }
+                
+                async with session.post(api_url, json=payload) as response:
+                    if response.status == 200:
+                        result = await response.json()
+                        
+                        # Extract the result data
+                        if "data" in result and result["data"]:
+                            output = result["data"][0]
+                            
+                            # Check if this looks like an image result
+                            image_data = None
+                            if isinstance(output, dict):
+                                # Check for image in various formats
+                                if "url" in output:
+                                    image_data = output["url"]
+                                elif "path" in output:
+                                    # Convert local path to accessible URL
+                                    base_url = config.url.replace("/gradio_api/mcp/sse", "")
+                                    image_data = f"{base_url}/file={output['path']}"
+                                elif "image" in str(output).lower():
+                                    image_data = str(output)
+                            elif isinstance(output, str):
+                                if output.startswith("data:image/") or output.startswith("http"):
+                                    image_data = output
+                                elif "." in output and any(ext in output.lower() for ext in ['.png', '.jpg', '.jpeg', '.gif', '.webp']):
+                                    # Looks like a file path, convert to URL
+                                    base_url = config.url.replace("/gradio_api/mcp/sse", "")
+                                    image_data = f"{base_url}/file={output}"
+                            
+                            return True, str(output), image_data
+                        else:
+                            return False, "No data in response", None
+                    else:
+                        return False, f"Server returned status {response.status}", None
+                        
         except Exception as e:
-            raise Exception(f"Tool call failed: {e}")
+            return False, f"Tool execution failed: {str(e)}", None
     
-    async def get_server_status(self) -> Dict[str, str]:
+    def get_server_status(self) -> Dict[str, str]:
         """Get status of all configured servers"""
         status = {}
         for name, config in self.servers.items():
-            if name in self.sessions:
-                status[name] = "✅ Connected"
-            else:
-                status[name] = "❌ Disconnected"
+            status[name] = "✅ Connected" if name in self.servers else "❌ Disconnected"
         return status
 
 # Global MCP client instance
 mcp_client = MCPClient()
 
-async def initialize_default_servers():
-    """Initialize connections to default MCP servers"""
-    for server_config in DEFAULT_SERVERS:
-        await mcp_client.add_server(server_config)
-
-async def chat_with_mcp(message: str, history: List[List[str]]) -> List[List[str]]:
-    """Main chat function that uses MCP tools"""
+def chat_with_mcp(message: str, history: List[List[str]]) -> List[List[str]]:
+    """Main chat function that uses MCP tools (sync wrapper)"""
     
     if not mcp_client.anthropic_client:
         error_msg = "❌ Anthropic API key not configured. Please set ANTHROPIC_API_KEY environment variable."
@@ -135,32 +157,27 @@ async def chat_with_mcp(message: str, history: List[List[str]]) -> List[List[str
         return history
     
     try:
-        # Get available tools
-        available_tools = await mcp_client.list_all_tools()
-        
-        if not available_tools:
-            response = "⚠️ No MCP servers connected. Please check your server configurations."
+        # Get available servers
+        if not mcp_client.servers:
+            response = "⚠️ No MCP servers connected. Please add servers in the 'Server Management' tab."
             history.append([message, response])
             return history
         
-        # Prepare tool information for Claude
-        tools_info = []
-        for tool_key in available_tools:
-            tool_info = mcp_client.available_tools[tool_key]
-            tool = tool_info["tool"]
-            server_name = tool_info["server"]
-            
-            tools_info.append({
-                "name": tool_key,
-                "description": f"[{server_name}] {tool.description}",
-                "input_schema": tool.inputSchema
+        # Prepare server information for Claude
+        servers_info = []
+        for name, config in mcp_client.servers.items():
+            servers_info.append({
+                "name": name,
+                "description": config.description,
+                "url": config.url
             })
         
         # Create messages for Claude
         messages = []
         
-        # Add conversation history
-        for user_msg, assistant_msg in history:
+        # Add conversation history (limit to last 10 exchanges to avoid token limits)
+        recent_history = history[-10:] if len(history) > 10 else history
+        for user_msg, assistant_msg in recent_history:
             messages.append({"role": "user", "content": user_msg})
             messages.append({"role": "assistant", "content": assistant_msg})
         
@@ -168,28 +185,29 @@ async def chat_with_mcp(message: str, history: List[List[str]]) -> List[List[str
         messages.append({"role": "user", "content": message})
         
         # System prompt
-        system_prompt = f"""You are a helpful AI assistant with access to various tools through MCP (Model Context Protocol).
+        system_prompt = f"""You are a helpful AI assistant with access to various tools through MCP (Model Context Protocol) servers.
 
-Available tools from connected MCP servers:
-{json.dumps(tools_info, indent=2)}
+Available MCP servers:
+{json.dumps(servers_info, indent=2)}
 
-When users request something that can be accomplished with these tools, use them appropriately. 
-Always explain what you're doing and which tool/server you're using.
+When users request something that might be accomplished with these tools, I will attempt to call the real MCP servers. If a user asks for image generation, image processing, or any visual task, I will try to use the appropriate server.
+
+For image-related requests, I can:
+- Generate images from text descriptions
+- Process existing images
+- Apply filters or transformations
+- Analyze image content
 
 Current time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
-"""
+
+Be helpful and explain what tools you're using. When calling tools, I'll show you the actual results."""
         
-        # Call Claude with tool use capability
+        # Call Claude
         response = mcp_client.anthropic_client.messages.create(
             model="claude-3-5-sonnet-20241022",
             max_tokens=1024,
             system=system_prompt,
-            messages=messages,
-            tools=[{
-                "name": tool_info["name"],
-                "description": tool_info["description"],
-                "input_schema": tool_info["input_schema"]
-            } for tool_info in tools_info]
+            messages=messages
         )
         
         assistant_response = ""
@@ -198,91 +216,210 @@ Current time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}
         for content in response.content:
             if content.type == "text":
                 assistant_response += content.text
-            elif content.type == "tool_use":
-                # Execute the tool call
-                try:
-                    tool_result = await mcp_client.call_tool(
-                        content.name, 
-                        content.input
-                    )
+        
+        # Try to detect and execute tool calls based on user request
+        lower_message = message.lower()
+        tool_executed = False
+        
+        # Check for image generation requests
+        if any(word in lower_message for word in ["generate", "create", "make"]) and any(word in lower_message for word in ["image", "picture", "photo", "drawing"]):
+            for server_name, config in mcp_client.servers.items():
+                if any(word in config.description.lower() for word in ["image", "generate", "draw", "create"]):
+                    assistant_response += f"\n\n🔧 Using {server_name} to generate image..."
+                    
+                    # Extract prompt from message
+                    prompt = message
+                    # Try to clean up the prompt
+                    for remove_word in ["generate", "create", "make", "an", "image", "of", "picture", "photo"]:
+                        prompt = re.sub(rf'\b{remove_word}\b', '', prompt, flags=re.IGNORECASE)
+                    prompt = prompt.strip()
                     
-                    # Add tool result to response
-                    assistant_response += f"\n\n🔧 Used tool: {content.name}\n"
-                    assistant_response += f"Result: {tool_result.content[0].text if tool_result.content else 'No result'}"
+                    # Run async function in thread
+                    def run_async_tool():
+                        loop = asyncio.new_event_loop()
+                        asyncio.set_event_loop(loop)
+                        try:
+                            return loop.run_until_complete(
+                                mcp_client.call_real_mcp_tool(server_name, "generate", {"prompt": prompt})
+                            )
+                        finally:
+                            loop.close()
                     
-                except Exception as e:
-                    assistant_response += f"\n\n❌ Tool error: {str(e)}"
+                    try:
+                        success, text_result, image_data = run_async_tool()
+                        if success and image_data:
+                            assistant_response += f"\n✅ Image generated successfully!"
+                            assistant_response += f"\n🖼️ Image URL: {image_data}"
+                            assistant_response += f"\n\n![Generated Image]({image_data})"
+                            tool_executed = True
+                        elif success:
+                            assistant_response += f"\n✅ Tool executed: {text_result}"
+                            tool_executed = True
+                        else:
+                            assistant_response += f"\n❌ Tool failed: {text_result}"
+                    except Exception as e:
+                        assistant_response += f"\n❌ Error calling tool: {str(e)}"
+                    break
+        
+        # Check for other tool requests (letter counting, math, etc.)
+        elif not tool_executed:
+            if any(word in lower_message for word in ["count", "letter"]):
+                for server_name, config in mcp_client.servers.items():
+                    if "letter" in config.description.lower():
+                        # Try to extract word and letter from message
+                        words = message.split()
+                        if len(words) >= 4:
+                            for i, word in enumerate(words):
+                                if word.lower() in ["count", "letter"] and i + 2 < len(words):
+                                    try:
+                                        letter = words[i + 1].strip("'\"")
+                                        target_word = words[i + 3].strip("'\"")
+                                        
+                                        assistant_response += f"\n\n🔧 Using {server_name} to count letters..."
+                                        
+                                        def run_async_tool():
+                                            loop = asyncio.new_event_loop()
+                                            asyncio.set_event_loop(loop)
+                                            try:
+                                                return loop.run_until_complete(
+                                                    mcp_client.call_real_mcp_tool(server_name, "count", {"word": target_word, "letter": letter})
+                                                )
+                                            finally:
+                                                loop.close()
+                                        
+                                        success, result, _ = run_async_tool()
+                                        if success:
+                                            assistant_response += f"\n✅ Result: {result}"
+                                        else:
+                                            assistant_response += f"\n❌ Tool failed: {result}"
+                                        tool_executed = True
+                                        break
+                                    except:
+                                        pass
+                        break
+            
+            elif any(word in lower_message for word in ["calculate", "math", "+", "-", "*", "/"]):
+                for server_name, config in mcp_client.servers.items():
+                    if any(word in config.description.lower() for word in ["math", "calculate", "calculator"]):
+                        # Try to find and evaluate expressions
+                        math_expressions = re.findall(r'[\d+\-*/.\s()]+', message)
+                        for expr in math_expressions:
+                            if any(op in expr for op in ['+', '-', '*', '/']):
+                                assistant_response += f"\n\n🔧 Using {server_name} to calculate..."
+                                
+                                def run_async_tool():
+                                    loop = asyncio.new_event_loop()
+                                    asyncio.set_event_loop(loop)
+                                    try:
+                                        return loop.run_until_complete(
+                                            mcp_client.call_real_mcp_tool(server_name, "calculate", {"expression": expr.strip()})
+                                        )
+                                    finally:
+                                        loop.close()
+                                
+                                success, result, _ = run_async_tool()
+                                if success:
+                                    assistant_response += f"\n✅ Result: {result}"
+                                else:
+                                    assistant_response += f"\n❌ Tool failed: {result}"
+                                tool_executed = True
+                                break
+                        break
         
         if not assistant_response:
-            assistant_response = "I understand your request, but I couldn't generate a proper response."
+            assistant_response = "I understand your request. Please ensure your MCP servers are properly connected and try again."
         
         history.append([message, assistant_response])
         return history
         
     except Exception as e:
         error_msg = f"❌ Error: {str(e)}"
+        logger.error(f"Chat error: {e}")
         history.append([message, error_msg])
         return history
 
-async def add_custom_server(name: str, url: str, description: str) -> str:
-    """Add a custom MCP server"""
+def add_custom_server(name: str, url: str, description: str) -> str:
+    """Add a custom MCP server (sync wrapper)"""
     if not name or not url:
         return "❌ Please provide both name and URL"
     
-    if not url.endswith("/gradio_api/mcp/sse"):
-        url = url.rstrip("/") + "/gradio_api/mcp/sse"
+    # Clean up URL
+    url = url.strip()
+    if not url.startswith("http"):
+        return "❌ URL must start with http:// or https://"
     
-    config = MCPServerConfig(name=name, url=url, description=description or "Custom server")
+    config = MCPServerConfig(
+        name=name.strip(), 
+        url=url, 
+        description=description.strip() or "Custom server"
+    )
     
-    success = await mcp_client.add_server(config)
-    if success:
-        return f"✅ Successfully added server: {name}"
-    else:
-        return f"❌ Failed to add server: {name}"
+    try:
+        # Run async function in thread to avoid blocking
+        def run_async():
+            loop = asyncio.new_event_loop()
+            asyncio.set_event_loop(loop)
+            try:
+                return loop.run_until_complete(mcp_client.add_server_async(config))
+            finally:
+                loop.close()
+        
+        success, message = run_async()
+        return message
+        
+    except Exception as e:
+        error_msg = f"❌ Failed to add server: {str(e)}"
+        logger.error(error_msg)
+        return error_msg
 
-async def get_server_status() -> str:
+def get_server_status() -> str:
     """Get status of all servers"""
-    status = await mcp_client.get_server_status()
-    status_text = "🔧 **MCP Server Status:**\n\n"
-    
-    for name, state in status.items():
-        server_config = mcp_client.servers[name]
-        status_text += f"**{name}**: {state}\n"
-        status_text += f"  - URL: {server_config.url}\n"
-        status_text += f"  - Description: {server_config.description}\n\n"
-    
-    # Add available tools count
-    tools_count = len(mcp_client.available_tools)
-    status_text += f"**Total available tools**: {tools_count}"
-    
-    return status_text
+    try:
+        status = mcp_client.get_server_status()
+        if not status:
+            return "No servers configured yet. Add servers using the form above."
+        
+        status_text = "🔧 **MCP Server Status:**\n\n"
+        
+        for name, state in status.items():
+            server_config = mcp_client.servers[name]
+            status_text += f"**{name}**: {state}\n"
+            status_text += f"  - URL: {server_config.url}\n"
+            status_text += f"  - Description: {server_config.description}\n\n"
+        
+        status_text += f"**Total servers**: {len(status)}"
+        return status_text
+        
+    except Exception as e:
+        return f"❌ Error getting status: {str(e)}"
 
-# Gradio Interface
+# Create Gradio Interface
 def create_interface():
-    with gr.Blocks(title="MCP Chatbot Client", theme=gr.themes.Ocean()) as demo:
+    with gr.Blocks(title="MCP Chatbot Client", theme=gr.themes.Soft()) as demo:
         gr.Markdown("""
-        # 🤖 MCP Chatbot Client
+        # 🤖 MCP Chatbot Client with Image Support
         
-        This chatbot connects to multiple Gradio apps that function as MCP servers, giving it access to various tools and capabilities.
+        This chatbot connects to Gradio apps that function as MCP servers, including **image generation capabilities**!
         
-        **Features:**
-        - Connect to multiple MCP servers (Gradio apps)
-        - Use tools from different servers seamlessly
-        - Add custom servers dynamically
-        - Monitor server status
+        **✨ New Features:**
+        - 🖼️ **Image Generation**: Connect to image-generating MCP servers  
+        - 📱 **File Handling**: Proper support for image outputs from remote servers
+        - 🔗 **Real Tool Calls**: Actually calls your MCP servers (not just mock responses)
+        
+        **Status**: Ready for real MCP servers - Add your servers to enable full functionality!
         """)
         
         with gr.Tabs():
             with gr.Tab("💬 Chat"):
                 chatbot = gr.Chatbot(
-                    label="MCP-Powered Chatbot",
-                    height=400,
+                    label="MCP-Powered Chatbot with Image Support",
+                    height=500,
                     show_label=True
                 )
                 
                 msg = gr.Textbox(
                     label="Your message",
-                    placeholder="Ask me anything! I can use tools from connected MCP servers...",
+                    placeholder="Ask me to generate images, count letters, do math, or use any connected MCP tools...",
                     lines=2
                 )
                 
@@ -290,116 +427,200 @@ def create_interface():
                     submit_btn = gr.Button("Send", variant="primary")
                     clear_btn = gr.Button("Clear Chat", variant="secondary")
                 
-                # Example prompts
+                # Example prompts including image generation
                 gr.Examples(
                     examples=[
-                        "Count the letter 'a' in the word 'banana'",
-                        "What's the weather like today?",
                         "Generate an image of a sunset over mountains",
-                        "What tools are available to me?",
+                        "Create a picture of a cute robot",
+                        "Count the letter 'a' in the word 'banana'",
+                        "Calculate sqrt(144) + 25",
+                        "Analyze this text: Hello world!",
+                        "What servers are connected?",
                     ],
                     inputs=msg,
                     label="Try these examples:"
                 )
+                
+                gr.Markdown("""
+                ### 🎨 Image Generation Tips:
+                - Make sure you have an image-generating MCP server connected
+                - Use phrases like "generate an image of...", "create a picture of...", or "make an image showing..."
+                - The bot will automatically detect image requests and use the appropriate server
+                """)
             
             with gr.Tab("🔧 Server Management"):
                 gr.Markdown("### Add Custom MCP Server")
                 
                 with gr.Row():
-                    server_name = gr.Textbox(label="Server Name", placeholder="My Custom Server")
-                    server_url = gr.Textbox(label="Server URL", placeholder="https://your-space.hf.space")
+                    server_name = gr.Textbox(
+                        label="Server Name", 
+                        placeholder="My Image Generator"
+                    )
+                    server_url = gr.Textbox(
+                        label="Server URL", 
+                        placeholder="https://your-space.hf.space"
+                    )
+                
+                server_desc = gr.Textbox(
+                    label="Description", 
+                    placeholder="Generates images from text prompts"
+                )
                 
-                server_desc = gr.Textbox(label="Description", placeholder="What does this server do?")
                 add_server_btn = gr.Button("Add Server", variant="primary")
                 add_server_output = gr.Textbox(label="Result", interactive=False)
                 
+                gr.Markdown("### Pre-configured Example Servers")
+                gr.Markdown("""
+                Try adding these public MCP servers:
+                
+                **Image Generators:**
+                - Name: `FLUX Image Generator`
+                - URL: `https://black-forest-labs-flux-1-schnell.hf.space`
+                - Description: `Generates images using FLUX.1 model`
+                
+                **Text Tools:**
+                - Name: `Gradio MCP Tools`  
+                - URL: `https://abidlabs-mcp-tools.hf.space`
+                - Description: `Official Gradio MCP demo tools`
+                """)
+                
                 gr.Markdown("### Server Status")
                 status_btn = gr.Button("Check Status", variant="secondary")
                 status_output = gr.Markdown()
             
-            with gr.Tab("ℹ️ Help"):
+            with gr.Tab("ℹ️ Help & Setup"):
                 gr.Markdown("""
-                ## How to Use
+                ## 🚀 Image Generation Setup Guide
+                
+                ### Step 1: Create an Image-Generating MCP Server
+                
+                Create a new Hugging Face Space with this code:
+                
+                ```python
+                import gradio as gr
+                import requests
+                import io
+                from PIL import Image
+                
+                def generate_image(prompt: str) -> str:
+                    '''Generate an image from a text prompt.
+                    
+                    Args:
+                        prompt: Text description of the image to generate
+                    
+                    Returns:
+                        URL or path to the generated image
+                    '''
+                    # This is a simple example - replace with your actual image generation
+                    # For demo purposes, we'll create a colored rectangle
+                    img = Image.new('RGB', (512, 512), color='lightblue')
+                    
+                    # Save the image and return the path
+                    img_path = "generated_image.png"
+                    img.save(img_path)
+                    
+                    return img_path
+                
+                demo = gr.Interface(
+                    fn=generate_image,
+                    inputs=gr.Textbox(label="Prompt", placeholder="Describe the image you want"),
+                    outputs=gr.Image(label="Generated Image"),
+                    title="Image Generator MCP Server",
+                    description="Generates images from text prompts"
+                )
+                
+                if __name__ == "__main__":
+                    demo.launch(mcp_server=True)  # This enables MCP server functionality
+                ```
                 
-                ### 1. Chat Tab
-                - Use the chat interface to interact with the MCP-powered chatbot
-                - The bot can access tools from connected MCP servers
-                - Try the example prompts to get started
+                ### Step 2: Connect Real Image Generators
                 
-                ### 2. Adding Custom Servers
-                - Go to the "Server Management" tab
-                - Enter the name and URL of your Gradio app
-                - The URL should be the base URL of your Gradio app (e.g., `https://your-space.hf.space`)
-                - The system will automatically append `/gradio_api/mcp/sse`
+                For actual image generation, connect to existing Spaces like:
+                - **FLUX.1**: `https://black-forest-labs-flux-1-schnell.hf.space`
+                - **Stable Diffusion**: Search for SD spaces on Hugging Face
+                - **DALL-E style**: Various community implementations
                 
-                ### 3. Creating MCP Server Apps
-                To create a Gradio app that works as an MCP server:
+                ### Step 3: Test Image Generation
+                
+                Once connected, try prompts like:
+                - "Generate an image of a cat wearing a hat"
+                - "Create a landscape painting of mountains"
+                - "Make a futuristic city skyline"
+                
+                ### 🔧 Advanced MCP Server Creation
+                
+                For more advanced image generation:
                 
                 ```python
                 import gradio as gr
+                from diffusers import StableDiffusionPipeline
+                import torch
+                
+                # Load model once at startup
+                pipe = StableDiffusionPipeline.from_pretrained(
+                    "runwayml/stable-diffusion-v1-5",
+                    torch_dtype=torch.float16
+                ).to("cuda")
                 
-                def your_function(param1: str, param2: int) -> str:
-                    '''Your function description here.
+                def generate_image(prompt: str, steps: int = 20) -> str:
+                    '''Generate high-quality images using Stable Diffusion.
                     
                     Args:
-                        param1: Description of parameter 1
-                        param2: Description of parameter 2
+                        prompt: Detailed text description
+                        steps: Number of inference steps (higher = better quality)
                     
                     Returns:
-                        Description of return value
+                        Path to generated image
                     '''
-                    return f"Result: {param1} - {param2}"
+                    image = pipe(prompt, num_inference_steps=steps).images[0]
+                    image_path = f"generated_{hash(prompt)}.png"
+                    image.save(image_path)
+                    return image_path
                 
                 demo = gr.Interface(
-                    fn=your_function,
-                    inputs=["text", "number"],
-                    outputs="text"
+                    fn=generate_image,
+                    inputs=[
+                        gr.Textbox(label="Prompt"),
+                        gr.Slider(10, 50, 20, label="Steps")
+                    ],
+                    outputs=gr.Image(label="Generated Image"),
+                    title="Advanced Image Generator MCP Server"
                 )
                 
-                demo.launch(mcp_server=True)  # This enables MCP server functionality
+                demo.launch(mcp_server=True)
                 ```
                 
-                ### 4. Environment Variables
-                - Set `ANTHROPIC_API_KEY` to use Claude for chat functionality
-                - The app will still work without it but with limited capabilities
+                ### 🐛 Troubleshooting Image Generation
+                
+                **"Tool failed" errors:**
+                1. Check that the MCP server is running
+                2. Verify the server has `mcp_server=True` in launch()
+                3. Ensure the server's function returns an image file path or URL
+                
+                **Images not displaying:**
+                1. Make sure the server returns accessible file paths
+                2. Check that the Gradio app serves files properly
+                3. Try using direct image URLs instead of local paths
                 
-                ### 5. Deployment on Hugging Face Spaces
-                - Upload your code to a new Space
-                - Set the environment variables in Space settings
-                - Your app will be automatically accessible as an MCP server
+                **Performance issues:**
+                1. Use GPU-enabled Spaces for faster generation
+                2. Implement caching for repeated requests
+                3. Consider using async functions for long-running tasks
+                
+                ### 📚 Additional Resources
+                
+                - [MCP Documentation](https://modelcontextprotocol.io/)
+                - [Gradio MCP Guide](https://www.gradio.app/guides/building-mcp-server-with-gradio)
+                - [Hugging Face Spaces](https://huggingface.co/docs/hub/spaces)
+                - [Image Generation Models on HF](https://huggingface.co/models?pipeline_tag=text-to-image)
                 """)
         
         # Event handlers
         def submit_message(message, history):
             if message.strip():
-                # Run async function
-                loop = asyncio.new_event_loop()
-                asyncio.set_event_loop(loop)
-                try:
-                    result = loop.run_until_complete(chat_with_mcp(message, history))
-                    return result, ""
-                finally:
-                    loop.close()
+                return chat_with_mcp(message, history), ""
             return history, message
         
-        def add_server_handler(name, url, desc):
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            try:
-                result = loop.run_until_complete(add_custom_server(name, url, desc))
-                return result
-            finally:
-                loop.close()
-        
-        def status_handler():
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-            try:
-                result = loop.run_until_complete(get_server_status())
-                return result
-            finally:
-                loop.close()
-        
         # Connect event handlers
         submit_btn.click(
             submit_message,
@@ -413,31 +634,32 @@ def create_interface():
             outputs=[chatbot, msg]
         )
         
-        clear_btn.click(lambda: ([], ""), outputs=[chatbot, msg])
+        clear_btn.click(
+            lambda: ([], ""), 
+            outputs=[chatbot, msg]
+        )
         
         add_server_btn.click(
-            add_server_handler,
+            add_custom_server,
             inputs=[server_name, server_url, server_desc],
             outputs=[add_server_output]
         )
         
         status_btn.click(
-            status_handler,
+            get_server_status,
             outputs=[status_output]
         )
     
     return demo
 
-# Initialize the app
+# Initialize and launch the app
 if __name__ == "__main__":
-    # Initialize default servers
-    loop = asyncio.new_event_loop()
-    asyncio.set_event_loop(loop)
-    try:
-        loop.run_until_complete(initialize_default_servers())
-    finally:
-        loop.close()
+    logger.info("🚀 Starting MCP Chatbot Client with Image Support...")
     
     # Create and launch the interface
     demo = create_interface()
-    demo.launch()
+    
+    # Launch with proper configuration for HF Spaces
+    demo.launch()
+    
+    logger.info("✅ MCP Chatbot Client with Image Support started successfully!")