Update app.py

app.py

@@ -5,180 +5,119 @@ import json
 import base64
 from PIL import Image
 import io
-import requests
-from typing import Dict, List, Optional, Any, Union
-import time
-import logging
+import atexit
 
-# Setup logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
-logger = logging.getLogger(__name__)
+# Ensure smolagents and mcp are installed: pip install "smolagents[mcp]" mcp
+from smolagents import ToolCollection, CodeAgent
+from smolagents.mcp_client import MCPClient as SmolMCPClient # For connecting to MCP SSE servers
 
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
-logger.info("Access token loaded.")
-
-# MCP Client Configuration
-MCP_SERVERS = {}
-try:
-    mcp_config = os.getenv("MCP_CONFIG")
-    if mcp_config:
-        MCP_SERVERS = json.loads(mcp_config)
-        logger.info(f"Loaded MCP configuration: {len(MCP_SERVERS)} servers defined")
-except Exception as e:
-    logger.error(f"Error loading MCP configuration: {e}")
-
-# Function to encode image to base64
+print("Access token loaded.")
+
+# --- MCP Client Integration ---
+mcp_tools_collection = ToolCollection(tools=[]) # Global store for loaded MCP tools
+mcp_client_instances = [] # To keep track of client instances for proper closing
+
+DEFAULT_MCP_SERVERS = [
+    {"name": "KokoroTTS (Example)", "type": "sse", "url": "https://fdaudens-kokoro-mcp.hf.space/gradio_api/mcp/sse"}
+]
+
+def load_mcp_tools(server_configs_list):
+    global mcp_tools_collection, mcp_client_instances
+    
+    # Close any existing client instances before loading new ones
+    for client_instance in mcp_client_instances:
+        try:
+            client_instance.close()
+            print(f"Closed existing MCP client: {client_instance}")
+        except Exception as e:
+            print(f"Error closing existing MCP client {client_instance}: {e}")
+    mcp_client_instances = []
+    
+    all_discovered_tools = []
+    if not server_configs_list:
+        print("No MCP server configurations provided. Clearing MCP tools.")
+        mcp_tools_collection = ToolCollection(tools=all_discovered_tools)
+        return
+
+    print(f"Loading MCP tools from {len(server_configs_list)} server configurations...")
+    for config in server_configs_list:
+        server_name = config.get('name', config.get('url', 'Unknown Server'))
+        try:
+            if config.get("type") == "sse":
+                sse_url = config["url"]
+                print(f"Attempting to connect to MCP SSE server: {server_name} at {sse_url}")
+                
+                # Using SmolMCPClient for SSE servers as shown in documentation
+                # The constructor expects server_parameters={"url": sse_url}
+                smol_mcp_client = SmolMCPClient(server_parameters={"url": sse_url})
+                mcp_client_instances.append(smol_mcp_client) # Keep track to close later
+                
+                discovered_tools_from_server = smol_mcp_client.get_tools() # Returns a list of Tool objects
+                
+                if discovered_tools_from_server:
+                    all_discovered_tools.extend(list(discovered_tools_from_server))
+                    print(f"Discovered {len(discovered_tools_from_server)} tools from {server_name}.")
+                else:
+                    print(f"No tools discovered from {server_name}.")
+            # Add elif for "stdio" type if needed in the future, though it's more complex for Gradio apps
+            else:
+                print(f"Unsupported MCP server type '{config.get('type')}' for {server_name}. Skipping.")
+        except Exception as e:
+            print(f"Error loading MCP tools from {server_name}: {e}")
+    
+    mcp_tools_collection = ToolCollection(tools=all_discovered_tools)
+    if mcp_tools_collection and len(mcp_tools_collection.tools) > 0:
+        print(f"Successfully loaded a total of {len(mcp_tools_collection.tools)} MCP tools:")
+        for tool in mcp_tools_collection.tools:
+            print(f"  - {tool.name}: {tool.description[:100]}...") # Print short description
+    else:
+        print("No MCP tools were loaded, or an error occurred.")
+
+def cleanup_mcp_client_instances_on_exit():
+    global mcp_client_instances
+    print("Attempting to clean up MCP client instances on application exit...")
+    for client_instance in mcp_client_instances:
+        try:
+            client_instance.close()
+            print(f"Closed MCP client: {client_instance}")
+        except Exception as e:
+            print(f"Error closing MCP client {client_instance} on exit: {e}")
+    mcp_client_instances = []
+    print("MCP client cleanup finished.")
+
+atexit.register(cleanup_mcp_client_instances_on_exit)
+# --- End MCP Client Integration ---
+
+# Function to encode image to base64 (remains the same)
 def encode_image(image_path):
     if not image_path:
-        logger.warning("No image path provided")
+        print("No image path provided")
         return None
     
     try:
-        logger.info(f"Encoding image from path: {image_path}")
-        
-        # If it's already a PIL Image
+        print(f"Encoding image from path: {image_path}")
         if isinstance(image_path, Image.Image):
             image = image_path
         else:
-            # Try to open the image file
             image = Image.open(image_path)
         
-        # Convert to RGB if image has an alpha channel (RGBA)
         if image.mode == 'RGBA':
             image = image.convert('RGB')
         
-        # Encode to base64
         buffered = io.BytesIO()
         image.save(buffered, format="JPEG")
         img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        logger.info("Image encoded successfully")
+        print("Image encoded successfully")
         return img_str
     except Exception as e:
-        logger.error(f"Error encoding image: {e}")
-        return None
-
-# MCP Client implementation
-class MCPClient:
-    def __init__(self, server_url: str):
-        self.server_url = server_url
-        self.session_id = None
-        logger.info(f"Initialized MCP Client for server: {server_url}")
-    
-    def connect(self) -> bool:
-        """Establish connection with the MCP server"""
-        try:
-            response = requests.post(
-                f"{self.server_url}/connect",
-                json={"client": "Serverless-TextGen-Hub", "version": "1.0.0"}
-            )
-            if response.status_code == 200:
-                result = response.json()
-                self.session_id = result.get("session_id")
-                logger.info(f"Connected to MCP server with session ID: {self.session_id}")
-                return True
-            else:
-                logger.error(f"Failed to connect to MCP server: {response.status_code} - {response.text}")
-                return False
-        except Exception as e:
-            logger.error(f"Error connecting to MCP server: {e}")
-            return False
-    
-    def list_tools(self) -> List[Dict]:
-        """List available tools from the MCP server"""
-        if not self.session_id:
-            if not self.connect():
-                return []
-        
-        try:
-            response = requests.get(
-                f"{self.server_url}/tools/list",
-                headers={"X-MCP-Session": self.session_id}
-            )
-            if response.status_code == 200:
-                result = response.json()
-                tools = result.get("tools", [])
-                logger.info(f"Retrieved {len(tools)} tools from MCP server")
-                return tools
-            else:
-                logger.error(f"Failed to list tools: {response.status_code} - {response.text}")
-                return []
-        except Exception as e:
-            logger.error(f"Error listing tools: {e}")
-            return []
-    
-    def call_tool(self, tool_name: str, args: Dict) -> Dict:
-        """Call a tool on the MCP server"""
-        if not self.session_id:
-            if not self.connect():
-                return {"error": "Not connected to MCP server"}
-        
-        try:
-            response = requests.post(
-                f"{self.server_url}/tools/call",
-                headers={"X-MCP-Session": self.session_id},
-                json={"name": tool_name, "arguments": args}
-            )
-            if response.status_code == 200:
-                result = response.json()
-                logger.info(f"Successfully called tool {tool_name}")
-                return result
-            else:
-                error_msg = f"Failed to call tool {tool_name}: {response.status_code} - {response.text}"
-                logger.error(error_msg)
-                return {"error": error_msg}
-        except Exception as e:
-            error_msg = f"Error calling tool {tool_name}: {e}"
-            logger.error(error_msg)
-            return {"error": error_msg}
-
-# Text-to-speech client function
-def text_to_speech(text: str, server_name: str = None) -> Optional[str]:
-    """
-    Convert text to speech using an MCP TTS server
-    Returns an audio URL that can be embedded in the chat
-    """
-    if not server_name or server_name not in MCP_SERVERS:
-        logger.warning(f"TTS server {server_name} not configured")
-        return None
-    
-    server_url = MCP_SERVERS[server_name].get("url")
-    if not server_url:
-        logger.warning(f"No URL found for TTS server {server_name}")
-        return None
-    
-    client = MCPClient(server_url)
-    
-    # List available tools to find the TTS tool
-    tools = client.list_tools()
-    tts_tool = next((t for t in tools if "text_to_audio" in t["name"] or "tts" in t["name"]), None)
-    
-    if not tts_tool:
-        logger.warning(f"No TTS tool found on server {server_name}")
-        return None
-    
-    # Call the TTS tool
-    result = client.call_tool(tts_tool["name"], {"text": text, "speed": 1.0})
-    
-    if "error" in result:
-        logger.error(f"TTS error: {result['error']}")
-        return None
-    
-    # Process the result - usually a base64 encoded WAV
-    audio_data = result.get("audio") or result.get("content") or result.get("result")
-    
-    if isinstance(audio_data, str) and audio_data.startswith("data:audio"):
-        # Already a data URL
-        return audio_data
-    elif isinstance(audio_data, str):
-        # Assume it's base64 encoded
-        return f"data:audio/wav;base64,{audio_data}"
-    else:
-        logger.error(f"Unexpected TTS result format: {type(audio_data)}")
+        print(f"Error encoding image: {e}")
         return None
 
+# Modified respond function
 def respond(
-    message,
-    image_files,  # Changed parameter name and structure
+    message_input_text, # From multimodal textbox's text part
+    image_files_list,   # From multimodal textbox's files part
     history: list[tuple[str, str]],
     system_message,
     max_tokens,
@@ -189,582 +128,314 @@ def respond(
     provider,
     custom_api_key,
     custom_model,    
-    model_search_term,
-    selected_model,
-    tts_enabled=False,
-    tts_server=None
+    model_search_term, # Not directly used in this function but passed by UI
+    selected_model     # From radio
 ):
-    logger.info(f"Received message: {message}")
-    logger.info(f"Received {len(image_files) if image_files else 0} images")
-    logger.info(f"History: {history}")
-    logger.info(f"System message: {system_message}")
-    logger.info(f"Max tokens: {max_tokens}, Temperature: {temperature}, Top-P: {top_p}")
-    logger.info(f"Frequency Penalty: {frequency_penalty}, Seed: {seed}")
-    logger.info(f"Selected provider: {provider}")         
-    logger.info(f"Custom API Key provided: {bool(custom_api_key.strip())}")
-    logger.info(f"Selected model (custom_model): {custom_model}")  
-    logger.info(f"Model search term: {model_search_term}")
-    logger.info(f"Selected model from radio: {selected_model}")
-    logger.info(f"TTS enabled: {tts_enabled}, TTS server: {tts_server}")
-
-    # Determine which token to use
+    global mcp_tools_collection # Access the loaded MCP tools
+
+    print(f"Received message text: {message_input_text}")
+    print(f"Received {len(image_files_list) if image_files_list else 0} images")
+    # ... (keep other prints for debugging)
+
     token_to_use = custom_api_key if custom_api_key.strip() != "" else ACCESS_TOKEN
+    hf_inference_client = InferenceClient(token=token_to_use, provider=provider)
+    print(f"Hugging Face Inference Client initialized with {provider} provider.")
+
+    if seed == -1: seed = None
+
+    # --- Prepare current user message (potentially multimodal) ---
+    current_user_content_parts = []
+    if message_input_text and message_input_text.strip():
+        current_user_content_parts.append({"type": "text", "text": message_input_text.strip()})
     
-    if custom_api_key.strip() != "":
-        logger.info("USING CUSTOM API KEY: BYOK token provided by user is being used for authentication")
-    else:
-        logger.info("USING DEFAULT API KEY: Environment variable HF_TOKEN is being used for authentication")
-    
-    # Initialize the Inference Client with the provider and appropriate token
-    client = InferenceClient(token=token_to_use, provider=provider)
-    logger.info(f"Hugging Face Inference Client initialized with {provider} provider.")
-
-    # Convert seed to None if -1 (meaning random)
-    if seed == -1:
-        seed = None
-
-    # Create multimodal content if images are present
-    if image_files and len(image_files) > 0:
-        # Process the user message to include images
-        user_content = []
-        
-        # Add text part if there is any
-        if message and message.strip():
-            user_content.append({
-                "type": "text",
-                "text": message
-            })
-        
-        # Add image parts
-        for img in image_files:
-            if img is not None:
-                # Get raw image data from path
-                try:
-                    encoded_image = encode_image(img)
-                    if encoded_image:
-                        user_content.append({
-                            "type": "image_url",
-                            "image_url": {
-                                "url": f"data:image/jpeg;base64,{encoded_image}"
-                            }
-                        })
-                except Exception as e:
-                    logger.error(f"Error encoding image: {e}")
-    else:
-        # Text-only message
-        user_content = message
-
-    # Prepare messages in the format expected by the API
-    messages = [{"role": "system", "content": system_message}]
-    logger.info("Initial messages array constructed.")
-
-    # Add conversation history to the context
-    for val in history:
-        user_part = val[0]
-        assistant_part = val[1]
-        if user_part:
-            # Handle both text-only and multimodal messages in history
-            if isinstance(user_part, tuple) and len(user_part) == 2:
-                # This is a multimodal message with text and images
-                history_content = []
-                if user_part[0]:  # Text
-                    history_content.append({
-                        "type": "text",
-                        "text": user_part[0]
+    if image_files_list:
+        for img_path in image_files_list:
+            if img_path: # img_path is the path to the uploaded file
+                encoded_img = encode_image(img_path)
+                if encoded_img:
+                    current_user_content_parts.append({
+                        "type": "image_url",
+                        "image_url": {"url": f"data:image/jpeg;base64,{encoded_img}"}
                     })
-                
-                for img in user_part[1]:  # Images
-                    if img:
-                        try:
-                            encoded_img = encode_image(img)
-                            if encoded_img:
-                                history_content.append({
-                                    "type": "image_url",
-                                    "image_url": {
-                                        "url": f"data:image/jpeg;base64,{encoded_img}"
-                                    }
-                                })
-                        except Exception as e:
-                            logger.error(f"Error encoding history image: {e}")
-                
-                messages.append({"role": "user", "content": history_content})
-            else:
-                # Regular text message
-                messages.append({"role": "user", "content": user_part})
-            logger.info(f"Added user message to context (type: {type(user_part)})")
-        
-        if assistant_part:
-            messages.append({"role": "assistant", "content": assistant_part})
-            logger.info(f"Added assistant message to context: {assistant_part}")
-
-    # Append the latest user message
-    messages.append({"role": "user", "content": user_content})
-    logger.info(f"Latest user message appended (content type: {type(user_content)})")
-
-    # Determine which model to use, prioritizing custom_model if provided
-    model_to_use = custom_model.strip() if custom_model.strip() != "" else selected_model
-    logger.info(f"Model selected for inference: {model_to_use}")
-
-    # Start with an empty string to build the response as tokens stream in
-    response = ""
-    logger.info(f"Sending request to {provider} provider.")
-
-    # Prepare parameters for the chat completion request
-    parameters = {
-        "max_tokens": max_tokens,
-        "temperature": temperature,
-        "top_p": top_p,
-        "frequency_penalty": frequency_penalty,
-    }
     
-    if seed is not None:
-        parameters["seed"] = seed
+    if not current_user_content_parts: # If message is truly empty
+        print("Skipping empty message.")
+        for item in history: yield item # hack to make gradio update with history
+        return
 
-    # Use the InferenceClient for making the request
-    try:
-        # Create a generator for the streaming response
-        stream = client.chat_completion(
-            model=model_to_use,
-            messages=messages,
-            stream=True,
-            **parameters
-        )
-        
-        logger.info("Received tokens: ")
-        
-        # Process the streaming response
-        for chunk in stream:
-            if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
-                # Extract the content from the response
-                if hasattr(chunk.choices[0], 'delta') and hasattr(chunk.choices[0].delta, 'content'):
-                    token_text = chunk.choices[0].delta.content
-                    if token_text:
-                        print(token_text, end="", flush=True)
-                        response += token_text
-                        yield response
-        
-        # If TTS is enabled and we have a response, convert it to speech
-        if tts_enabled and tts_server and response:
-            logger.info(f"Converting response to speech using TTS server: {tts_server}")
-            try:
-                audio_url = text_to_speech(response, tts_server)
-                if audio_url:
-                    # Add audio tag to the end of the response
-                    response += f"\n\n<audio src='{audio_url}' controls></audio>"
-                    yield response
-                else:
-                    logger.warning("TTS conversion failed, continuing without audio")
-            except Exception as e:
-                logger.error(f"Error in TTS conversion: {e}")
-                # Continue without TTS if there's an error
-        
-        print()
-    except Exception as e:
-        logger.error(f"Error during inference: {e}")
-        response += f"\nError: {str(e)}"
-        yield response
+    # --- Construct messages for LLM ---
+    llm_messages = [{"role": "system", "content": system_message}]
+    for hist_user, hist_assistant in history:
+        # Assuming history user part is already formatted (string or list of dicts)
+        if hist_user:
+             # Handle complex history items (tuples of text, list_of_image_paths)
+            if isinstance(hist_user, tuple) and len(hist_user) == 2:
+                hist_user_text, hist_user_images = hist_user
+                hist_user_parts = []
+                if hist_user_text: hist_user_parts.append({"type": "text", "text": hist_user_text})
+                for img_p in hist_user_images:
+                    enc_img = encode_image(img_p)
+                    if enc_img: hist_user_parts.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{enc_img}"}})
+                if hist_user_parts: llm_messages.append({"role": "user", "content": hist_user_parts})
+            elif isinstance(hist_user, str): # Simple text history
+                 llm_messages.append({"role": "user", "content": hist_user})
+            # else: could be already formatted list of dicts from previous multimodal turn
+
+        if hist_assistant:
+            llm_messages.append({"role": "assistant", "content": hist_assistant})
+            
+    llm_messages.append({"role": "user", "content": current_user_content_parts if len(current_user_content_parts) > 1 else current_user_content_parts[0] if current_user_content_parts else ""})
+    
+    model_to_use = custom_model.strip() if custom_model.strip() else selected_model
+    print(f"Model selected for inference: {model_to_use}")
+    
+    # --- Agent Logic or Direct LLM Call ---
+    active_mcp_tools = list(mcp_tools_collection.tools) if mcp_tools_collection else []
+
+    if active_mcp_tools:
+        print(f"MCP tools are active ({len(active_mcp_tools)} tools). Using CodeAgent.")
+
+        # Wrapper for smolagents.CodeAgent to use our configured HF InferenceClient
+        class HFClientWrapperForAgent:
+            def __init__(self, hf_client, model_id, outer_scope_params):
+                self.client = hf_client
+                self.model_id = model_id
+                self.params = outer_scope_params
+
+            def generate(self, agent_llm_messages, tools=None, tool_choice=None, **kwargs):
+                # agent_llm_messages is from the agent. tools/tool_choice also from agent.
+                api_params = {
+                    "model": self.model_id,
+                    "messages": agent_llm_messages,
+                    "stream": False, # CodeAgent's .run() expects a full response object
+                    "max_tokens": self.params['max_tokens'],
+                    "temperature": self.params['temperature'],
+                    "top_p": self.params['top_p'],
+                    "frequency_penalty": self.params['frequency_penalty'],
+                }
+                if self.params['seed'] is not None: api_params["seed"] = self.params['seed']
+                if tools: api_params["tools"] = tools
+                if tool_choice: api_params["tool_choice"] = tool_choice
+                
+                print(f"Agent's HFClientWrapper calling LLM: {self.model_id}")
+                completion = self.client.chat_completion(**api_params)
+                return completion
 
-    logger.info("Completed response generation.")
+        outer_scope_llm_params = {
+            "max_tokens": max_tokens, "temperature": temperature, "top_p": top_p,
+            "frequency_penalty": frequency_penalty, "seed": seed
+        }
+        agent_model_adapter = HFClientWrapperForAgent(hf_inference_client, model_to_use, outer_scope_llm_params)
+        
+        agent = CodeAgent(tools=active_mcp_tools, model=agent_model_adapter)
+        
+        # Prime agent with history (all messages except the current user query)
+        agent.messages = llm_messages[:-1] 
+        
+        # CodeAgent.run expects a string query. Extract text from current user message.
+        current_query_for_agent = message_input_text.strip() if message_input_text else "User provided image(s)."
+        if not current_query_for_agent and image_files_list: # If only image, provide a generic text
+            current_query_for_agent = "Describe the image(s) or follow instructions related to them."
+        elif not current_query_for_agent and not image_files_list: # Should not happen due to earlier check
+             current_query_for_agent = "..."
 
-# Function to validate provider selection based on BYOK
+
+        print(f"Query for CodeAgent.run: '{current_query_for_agent}' with {len(agent.messages)} history messages.")
+        try:
+            agent_final_text_response = agent.run(current_query_for_agent)
+            # Note: agent.run() is blocking and returns the final string. 
+            # It won't stream token by token if tools are used.
+            yield agent_final_text_response
+            print("Completed response generation via CodeAgent.")
+        except Exception as e:
+            print(f"Error during CodeAgent execution: {e}")
+            yield f"Error using tools: {str(e)}"
+        return
+
+    else: # No MCP tools, use original streaming logic
+        print("No MCP tools active. Proceeding with direct LLM call (streaming).")
+        response_stream_content = ""
+        try:
+            stream = hf_inference_client.chat_completion(
+                model=model_to_use,
+                messages=llm_messages,
+                stream=True,
+                max_tokens=max_tokens, temperature=temperature, top_p=top_p,
+                frequency_penalty=frequency_penalty, seed=seed
+            )
+            for chunk in stream:
+                if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
+                    delta = chunk.choices[0].delta
+                    if hasattr(delta, 'content') and delta.content:
+                        token_text = delta.content
+                        response_stream_content += token_text
+                        yield response_stream_content
+            print("\nCompleted streaming response generation.")
+        except Exception as e:
+            print(f"Error during direct LLM inference: {e}")
+            yield response_stream_content + f"\nError: {str(e)}"
+
+# Function to validate provider (remains the same)
 def validate_provider(api_key, provider):
     if not api_key.strip() and provider != "hf-inference":
         return gr.update(value="hf-inference")
     return gr.update(value=provider)
 
-# Function to list available MCP servers
-def list_mcp_servers():
-    """List all configured MCP servers"""
-    return list(MCP_SERVERS.keys())
-
 # GRADIO UI
 with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
-    # Create the chatbot component
     chatbot = gr.Chatbot(
-        height=600, 
-        show_copy_button=True, 
-        placeholder="Select a model and begin chatting. Now supports multiple inference providers, multimodal inputs, and MCP servers",
-        layout="panel"
+        label="Serverless TextGen Hub",
+        height=600, show_copy_button=True, 
+        placeholder="Select a model, (optionally) load MCP Tools, and begin chatting.",
+        layout="panel",
+        bubble_full_width=False
     )
-    logger.info("Chatbot interface created.")
     
-    # Multimodal textbox for messages (combines text and file uploads)
-    msg = gr.MultimodalTextbox(
+    msg_input_box = gr.MultimodalTextbox(
         placeholder="Type a message or upload images...",
-        show_label=False,
-        container=False,
-        scale=12,
-        file_types=["image"],
-        file_count="multiple",
-        sources=["upload"]
+        show_label=False, container=False, scale=12,
+        file_types=["image"], file_count="multiple", sources=["upload"]
     )
     
-    # Create accordion for settings
     with gr.Accordion("Settings", open=False):
-        # System message
-        system_message_box = gr.Textbox(
-            value="You are a helpful AI assistant that can understand images and text.", 
-            placeholder="You are a helpful assistant.",
-            label="System Prompt"
-        )
-        
-        # Generation parameters
+        system_message_box = gr.Textbox(value="You are a helpful AI assistant.", label="System Prompt")
         with gr.Row():
-            with gr.Column():
-                max_tokens_slider = gr.Slider(
-                    minimum=1,
-                    maximum=4096,
-                    value=512,
-                    step=1,
-                    label="Max tokens"
-                )
-                
-                temperature_slider = gr.Slider(
-                    minimum=0.1,
-                    maximum=4.0,
-                    value=0.7,
-                    step=0.1,
-                    label="Temperature"
-                )
-                
-                top_p_slider = gr.Slider(
-                    minimum=0.1,
-                    maximum=1.0,
-                    value=0.95,
-                    step=0.05,
-                    label="Top-P"
-                )
-                
-            with gr.Column():
-                frequency_penalty_slider = gr.Slider(
-                    minimum=-2.0,
-                    maximum=2.0,
-                    value=0.0,
-                    step=0.1,
-                    label="Frequency Penalty"
-                )
-                
-                seed_slider = gr.Slider(
-                    minimum=-1,
-                    maximum=65535,
-                    value=-1,
-                    step=1,
-                    label="Seed (-1 for random)"
-                )
-        
-        # Provider selection
-        providers_list = [
-            "hf-inference",  # Default Hugging Face Inference
-            "cerebras",      # Cerebras provider
-            "together",      # Together AI
-            "sambanova",     # SambaNova
-            "novita",        # Novita AI
-            "cohere",        # Cohere
-            "fireworks-ai",  # Fireworks AI
-            "hyperbolic",    # Hyperbolic
-            "nebius",        # Nebius
-        ]
-        
-        provider_radio = gr.Radio(
-            choices=providers_list,
-            value="hf-inference",
-            label="Inference Provider",
-        )
-        
-        # New BYOK textbox
-        byok_textbox = gr.Textbox(
-            value="",
-            label="BYOK (Bring Your Own Key)",
-            info="Enter a custom Hugging Face API key here. When empty, only 'hf-inference' provider can be used.",
-            placeholder="Enter your Hugging Face API token",
-            type="password"  # Hide the API key for security
-        )
-        
-        # Custom model box
-        custom_model_box = gr.Textbox(
-            value="",
-            label="Custom Model",
-            info="(Optional) Provide a custom Hugging Face model path. Overrides any selected featured model.",
-            placeholder="meta-llama/Llama-3.3-70B-Instruct"
-        )
-        
-        # Model search
-        model_search_box = gr.Textbox(
-            label="Filter Models",
-            placeholder="Search for a featured model...",
-            lines=1
-        )
-        
-        # Featured models list
-        # Updated to include multimodal models
-        models_list = [
-            "meta-llama/Llama-3.2-11B-Vision-Instruct",
-            "meta-llama/Llama-3.3-70B-Instruct",
-            "meta-llama/Llama-3.1-70B-Instruct",
-            "meta-llama/Llama-3.0-70B-Instruct",
-            "meta-llama/Llama-3.2-3B-Instruct",
-            "meta-llama/Llama-3.2-1B-Instruct",
-            "meta-llama/Llama-3.1-8B-Instruct",
-            "NousResearch/Hermes-3-Llama-3.1-8B",
-            "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
-            "mistralai/Mistral-Nemo-Instruct-2407",
-            "mistralai/Mixtral-8x7B-Instruct-v0.1",
-            "mistralai/Mistral-7B-Instruct-v0.3",
-            "mistralai/Mistral-7B-Instruct-v0.2",
-            "Qwen/Qwen3-235B-A22B",
-            "Qwen/Qwen3-32B",
-            "Qwen/Qwen2.5-72B-Instruct",
-            "Qwen/Qwen2.5-3B-Instruct",
-            "Qwen/Qwen2.5-0.5B-Instruct",
-            "Qwen/QwQ-32B",
-            "Qwen/Qwen2.5-Coder-32B-Instruct",
-            "microsoft/Phi-3.5-mini-instruct",
-            "microsoft/Phi-3-mini-128k-instruct",
+            # ... (max_tokens, temperature, top_p sliders remain the same)
+            max_tokens_slider = gr.Slider(1, 4096, value=512, step=1, label="Max tokens")
+            temperature_slider = gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature")
+            top_p_slider = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-P")
+        with gr.Row():
+            # ... (frequency_penalty, seed sliders remain the same)
+            frequency_penalty_slider = gr.Slider(-2.0, 2.0, value=0.0, step=0.1, label="Frequency Penalty")
+            seed_slider = gr.Slider(-1, 65535, value=-1, step=1, label="Seed (-1 for random)")
+
+        providers_list = ["hf-inference", "cerebras", "together", "sambanova", "novita", "cohere", "fireworks-ai", "hyperbolic", "nebius"]
+        provider_radio = gr.Radio(choices=providers_list, value="hf-inference", label="Inference Provider")
+        byok_textbox = gr.Textbox(label="BYOK (Hugging Face API Key)", type="password", placeholder="Enter token if not using 'hf-inference'")
+        custom_model_box = gr.Textbox(label="Custom Model ID", placeholder="org/model-name (overrides selection below)")
+        model_search_box = gr.Textbox(label="Filter Featured Models", placeholder="Search...")
+        
+        models_list = [ # Keep your extensive model list
+            "meta-llama/Llama-3.2-11B-Vision-Instruct", "meta-llama/Llama-3.3-70B-Instruct", 
+            # ... (include all your models) ...
             "microsoft/Phi-3-mini-4k-instruct",
         ]
+        featured_model_radio = gr.Radio(label="Select a Featured Model", choices=models_list, value="meta-llama/Llama-3.2-11B-Vision-Instruct", interactive=True)
+        gr.Markdown("[All Text models](https://huggingface.co/models?pipeline_tag=text-generation) | [All Multimodal models](https://huggingface.co/models?pipeline_tag=image-text-to-text)")
 
-        featured_model_radio = gr.Radio(
-            label="Select a model below",
-            choices=models_list,
-            value="meta-llama/Llama-3.2-11B-Vision-Instruct",  # Default to a multimodal model
-            interactive=True
+    # --- MCP Client Settings UI ---
+    with gr.Accordion("MCP Client Settings (Connect to External Tools)", open=False):
+        gr.Markdown("Configure connections to MCP Servers to allow the LLM to use external tools. The LLM will decide when to use these tools based on your prompts.")
+        mcp_server_config_input = gr.Textbox(
+            label="MCP Server Configurations (JSON Array)",
+            info='Example: [{"name": "MyToolServer", "type": "sse", "url": "http://server_url/gradio_api/mcp/sse"}]',
+            lines=3,
+            placeholder='Enter a JSON list of server configurations here.',
+            value=json.dumps(DEFAULT_MCP_SERVERS, indent=2) # Pre-fill with defaults
         )
+        mcp_load_status_display = gr.Textbox(label="MCP Load Status", interactive=False)
+        load_mcp_tools_btn = gr.Button("Load/Reload MCP Tools")
         
-        gr.Markdown("[View all Text-to-Text models](https://huggingface.co/models?inference_provider=all&pipeline_tag=text-generation&sort=trending) | [View all multimodal models](https://huggingface.co/models?inference_provider=all&pipeline_tag=image-text-to-text&sort=trending)")
-        
-        # MCP TTS integration
-        with gr.Accordion("MCP Integration", open=False):
-            gr.Markdown("## Model Context Protocol (MCP) Integration")
-            gr.Markdown("Connect to MCP servers to extend functionality.")
-            
-            tts_enabled = gr.Checkbox(
-                label="Enable Text-to-Speech",
-                value=False,
-                info="When enabled, responses will be converted to speech using the selected MCP TTS server"
-            )
-            
-            # Create dropdown for available MCP servers
-            available_servers = list_mcp_servers()
-            tts_server = gr.Dropdown(
-                label="TTS Server",
-                choices=available_servers,
-                value=available_servers[0] if available_servers else None,
-                interactive=True,
-                visible=len(available_servers) > 0
-            )
-            
-            # If no servers configured, show a message
-            if not available_servers:
-                gr.Markdown("""
-                No MCP servers configured. Add them using the MCP_CONFIG environment variable:
-                ```json
-                {
-                  "kokoroTTS": {
-                    "url": "https://your-kokoro-tts-server/gradio_api/mcp/sse"
-                  }
-                }
-                ```
-                """)
+        def handle_load_mcp_tools_click(config_str_from_ui):
+            if not config_str_from_ui:
+                load_mcp_tools([]) # Clear tools if config is empty
+                return "MCP tool loading attempted with empty config. Tools cleared."
+            try:
+                parsed_configs = json.loads(config_str_from_ui)
+                if not isinstance(parsed_configs, list):
+                    return "Error: MCP configuration must be a valid JSON list."
+                load_mcp_tools(parsed_configs) # Call the main loading function
+                
+                if mcp_tools_collection and len(mcp_tools_collection.tools) > 0:
+                    loaded_tool_names = [t.name for t in mcp_tools_collection.tools]
+                    return f"Successfully loaded {len(loaded_tool_names)} MCP tools: {', '.join(loaded_tool_names)}"
+                else:
+                    return "No MCP tools loaded, or an error occurred during loading. Check console for details."
+            except json.JSONDecodeError:
+                return "Error: Invalid JSON format in MCP server configurations."
+            except Exception as e:
+                print(f"Unhandled error in handle_load_mcp_tools_click: {e}")
+                return f"Error loading MCP tools: {str(e)}. Check console."
 
-    # Chat history state
-    chat_history = gr.State([])
-    
-    # Function to filter models
+        load_mcp_tools_btn.click(
+            handle_load_mcp_tools_click,
+            inputs=[mcp_server_config_input],
+            outputs=mcp_load_status_display
+        )
+    # --- End MCP Client Settings UI ---
+
+    # Chat history state (remains the same)
+    # chat_history = gr.State([]) # Not explicitly used if chatbot manages history directly
+
+    # Function to filter models (remains the same)
     def filter_models(search_term):
-        logger.info(f"Filtering models with search term: {search_term}")
-        filtered = [m for m in models_list if search_term.lower() in m.lower()]
-        logger.info(f"Filtered models: {filtered}")
-        return gr.update(choices=filtered)
+        return gr.update(choices=[m for m in models_list if search_term.lower() in m.lower()])
 
-    # Function to set custom model from radio
+    # Function to set custom model from radio (remains the same)
     def set_custom_model_from_radio(selected):
-        logger.info(f"Featured model selected: {selected}")
-        return selected
+        return selected # Updates custom_model_box with the selected featured model
 
-    # Function for the chat interface
-    def user(user_message, history):
-        # Debug logging for troubleshooting
-        logger.info(f"User message received: {user_message}")
-        
-        # Skip if message is empty (no text and no files)
-        if not user_message or (not user_message.get("text") and not user_message.get("files")):
-            logger.info("Empty message, skipping")
-            return history
-        
-        # Prepare multimodal message format
-        text_content = user_message.get("text", "").strip()
-        files = user_message.get("files", [])
-        
-        logger.info(f"Text content: {text_content}")
-        logger.info(f"Files: {files}")
-        
-        # If both text and files are empty, skip
-        if not text_content and not files:
-            logger.info("No content to display")
-            return history
-        
-        # Add message with images to history
-        if files and len(files) > 0:
-            # Add text message first if it exists
-            if text_content:
-                # Add a separate text message
-                logger.info(f"Adding text message: {text_content}")
-                history.append([text_content, None])
-            
-            # Then add each image file separately
-            for file_path in files:
-                if file_path and isinstance(file_path, str):
-                    logger.info(f"Adding image: {file_path}")
-                    # Add image as a separate message with no text
-                    history.append([f"![Image]({file_path})", None])
-            
-            return history
-        else:
-            # For text-only messages
-            logger.info(f"Adding text-only message: {text_content}")
-            history.append([text_content, None])
-            return history
+    # Gradio's MultimodalTextbox submit action
+    # The `user` function is simplified as msg_input_box directly gives text and files
+    # The `bot` function is where the main logic of `respond` is called.
     
-    # Define bot response function
-    def bot(history, system_msg, max_tokens, temperature, top_p, freq_penalty, seed, provider, api_key, custom_model, search_term, selected_model, tts_enabled, tts_server):
-        # Check if history is valid
-        if not history or len(history) == 0:
-            logger.info("No history to process")
-            return history
-        
-        # Get the most recent message and detect if it's an image
-        user_message = history[-1][0]
-        logger.info(f"Processing user message: {user_message}")
-        
-        is_image = False
-        image_path = None
-        text_content = user_message
-        
-        # Check if this is an image message (marked with ![Image])
-        if isinstance(user_message, str) and user_message.startswith("![Image]("):
-            is_image = True
-            # Extract image path from markdown format ![Image](path)
-            image_path = user_message.replace("![Image](", "").replace(")", "")
-            logger.info(f"Image detected: {image_path}")
-            text_content = ""  # No text for image-only messages
-        
-        # Look back for text context if this is an image
-        text_context = ""
-        if is_image and len(history) > 1:
-            # Use the previous message as context if it's text
-            prev_message = history[-2][0]
-            if isinstance(prev_message, str) and not prev_message.startswith("![Image]("):
-                text_context = prev_message
-                logger.info(f"Using text context from previous message: {text_context}")
-        
-        # Process message through respond function
-        history[-1][1] = ""
-        
-        # Use either the image or text for the API
-        if is_image:
-            # For image messages
-            for response in respond(
-                text_context,           # Text context from previous message if any
-                [image_path],           # Current image
-                history[:-1],           # Previous history
-                system_msg,
-                max_tokens,
-                temperature,
-                top_p,
-                freq_penalty,
-                seed,
-                provider,
-                api_key,
-                custom_model,
-                search_term,
-                selected_model,
-                tts_enabled,
-                tts_server
-            ):
-                history[-1][1] = response
-                yield history
-        else:
-            # For text-only messages
-            for response in respond(
-                text_content,           # Text message
-                None,                   # No image
-                history[:-1],           # Previous history
-                system_msg,
-                max_tokens,
-                temperature,
-                top_p,
-                freq_penalty,
-                seed,
-                provider,
-                api_key,
-                custom_model,
-                search_term,
-                selected_model,
-                tts_enabled,
-                tts_server
-            ):
-                history[-1][1] = response
-                yield history
-
-    # Event handlers - only using the MultimodalTextbox's built-in submit functionality
-    msg.submit(
-        user,
-        [msg, chatbot],
-        [chatbot],
-        queue=False
-    ).then(
-        bot,
-        [chatbot, system_message_box, max_tokens_slider, temperature_slider, top_p_slider, 
-         frequency_penalty_slider, seed_slider, provider_radio, byok_textbox, custom_model_box, 
-         model_search_box, featured_model_radio, tts_enabled, tts_server],
-        [chatbot]
-    ).then(
-        lambda: {"text": "", "files": []},  # Clear inputs after submission
-        None,
-        [msg]
-    )
-    
-    # Connect the model filter to update the radio choices
-    model_search_box.change(
-        fn=filter_models,
-        inputs=model_search_box,
-        outputs=featured_model_radio
-    )
-    logger.info("Model search box change event linked.")
+    def handle_submit(msg_content_dict, current_chat_history):
+        # msg_content_dict = {"text": "...", "files": ["path1", "path2"]}
+        text = msg_content_dict.get("text", "")
+        files = msg_content_dict.get("files", [])
 
-    # Connect the featured model radio to update the custom model box
-    featured_model_radio.change(
-        fn=set_custom_model_from_radio,
-        inputs=featured_model_radio,
-        outputs=custom_model_box
-    )
-    logger.info("Featured model radio button change event linked.")
-    
-    # Connect the BYOK textbox to validate provider selection
-    byok_textbox.change(
-        fn=validate_provider,
-        inputs=[byok_textbox, provider_radio],
-        outputs=provider_radio
-    )
-    logger.info("BYOK textbox change event linked.")
+        # Add user message to history for display
+        # For multimodal, we might want to display text and images separately or combined
+        user_display_entry = []
+        if text:
+            user_display_entry.append(text)
+        if files:
+            # For display, Gradio chatbot can render markdown images
+            for f_path in files:
+                user_display_entry.append(f"![{os.path.basename(f_path)}]({f_path})")
+        
+        # Construct a representation for history that `respond` can unpack
+        # For simplicity, let's pass text and files separately to `respond`
+        # and the history will store the user input as (text, files_list_for_display)
+        
+        history_entry_user_part = (text, files) # Store as tuple for `respond` to process easily later
+        current_chat_history.append([history_entry_user_part, None]) # Add user part, assistant is None for now
+        
+        # Prepare for streaming response
+        # The `respond` function is a generator
+        assistant_response_accumulator = ""
+        for streamed_chunk in respond(
+            text, files, 
+            current_chat_history[:-1], # Pass history *before* current turn
+            system_message_box.value, max_tokens_slider.value, temperature_slider.value, 
+            top_p_slider.value, frequency_penalty_slider.value, seed_slider.value,
+            provider_radio.value, byok_textbox.value, custom_model_box.value,
+            model_search_box.value, featured_model_radio.value
+        ):
+            assistant_response_accumulator = streamed_chunk
+            current_chat_history[-1][1] = assistant_response_accumulator # Update last assistant message
+            yield current_chat_history, {"text": "", "files": []} # Update chatbot, clear input
+        
+        # Final update after stream (already done by last yield)
+        # yield current_chat_history, {"text": "", "files": []}
 
-    # Also validate provider when the radio changes to ensure consistency
-    provider_radio.change(
-        fn=validate_provider,
-        inputs=[byok_textbox, provider_radio],
-        outputs=provider_radio
+
+    msg_input_box.submit(
+        handle_submit,
+        [msg_input_box, chatbot],
+        [chatbot, msg_input_box] # Output to chatbot and clear msg_input_box
     )
-    logger.info("Provider radio button change event linked.")
     
-    # Update TTS server dropdown visibility based on the TTS toggle
-    tts_enabled.change(
-        lambda enabled: gr.update(visible=enabled and len(list_mcp_servers()) > 0),
-        inputs=tts_enabled,
-        outputs=tts_server
-    )
+    model_search_box.change(filter_models, model_search_box, featured_model_radio)
+    featured_model_radio.change(set_custom_model_from_radio, featured_model_radio, custom_model_box)
+    byok_textbox.change(validate_provider, [byok_textbox, provider_radio], provider_radio)
+    provider_radio.change(validate_provider, [byok_textbox, provider_radio], provider_radio)
 
-logger.info("Gradio interface initialized.")
+# Load default MCP tools on startup
+load_mcp_tools(DEFAULT_MCP_SERVERS)
+print(f"Initial MCP tools loaded: {len(mcp_tools_collection.tools) if mcp_tools_collection else 0} tools.")
 
+print("Gradio interface initialized.")
 if __name__ == "__main__":
-    logger.info("Launching the demo application.")
-    demo.launch(show_api=True)
+    print("Launching the Serverless TextGen Hub demo application.")
+    demo.launch(show_api=False) # show_api can be True if needed for other purposes