Update app.py

app.py

@@ -5,47 +5,44 @@ import json
 import base64
 from PIL import Image
 import io
-import requests # Keep for potential future use, though not directly used in core logic now
-from smolagents.mcp_client import MCPClient # Ensure this is correctly installed and importable
+import requests
+from smolagents.mcp_client import MCPClient
+from mcp import ToolResult # For type hinting, good practice
+from mcp.common.content_block import ValueContentBlock # To access the actual tool return value
+import numpy as np # For handling audio array
+import soundfile as sf # For converting audio array to WAV
 
 ACCESS_TOKEN = os.getenv("HF_TOKEN")
-if ACCESS_TOKEN:
-    print("Access token loaded from HF_TOKEN environment variable.")
-else:
-    print("Warning: HF_TOKEN environment variable not set. Some operations might fail.")
+print("Access token loaded.")
 
 # Function to encode image to base64
-def encode_image(image_path_or_pil):
-    if not image_path_or_pil:
-        print("No image path or PIL Image provided")
+def encode_image(image_path):
+    if not image_path:
+        print("No image path provided")
         return None
     
     try:
-        if isinstance(image_path_or_pil, Image.Image):
-            image = image_path_or_pil
-            print(f"Encoding PIL Image object.")
-        elif isinstance(image_path_or_pil, str):
-            print(f"Encoding image from path: {image_path_or_pil}")
-            if not os.path.exists(image_path_or_pil):
-                print(f"Error: Image file not found at {image_path_or_pil}")
-                return None
-            image = Image.open(image_path_or_pil)
+        print(f"Encoding image from path: {image_path}")
+        
+        # If it's already a PIL Image
+        if isinstance(image_path, Image.Image):
+            image = image_path
         else:
-            print(f"Error: Unsupported image input type: {type(image_path_or_pil)}")
-            return None
+            # 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") # Or PNG if preferred, ensure consistency
+        image.save(buffered, format="JPEG")
         img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
-        print("Image encoded successfully to base64.")
+        print("Image encoded successfully")
         return img_str
     except Exception as e:
         print(f"Error encoding image: {e}")
-        import traceback
-        traceback.print_exc()
         return None
 
 # Dictionary to store active MCP connections
@@ -54,606 +51,748 @@ mcp_connections = {}
 def connect_to_mcp_server(server_url, server_name=None):
     """Connect to an MCP server and return available tools"""
     if not server_url:
-        return None, "No server URL provided. Please enter a valid URL."
+        return None, "No server URL provided"
     
     try:
-        print(f"Attempting to connect to MCP server at URL: {server_url}")
-        client = MCPClient({"url": server_url}) # This might block or raise if connection fails
-        tools = client.get_tools() # This should also be a blocking call until tools are fetched
+        # Create an MCP client and connect to the server
+        client = MCPClient({"url": server_url})
+        # Get available tools
+        tools = client.get_tools()
         
-        name = server_name.strip() if server_name and server_name.strip() else f"Server_{len(mcp_connections) + 1}"
+        # Store the connection for later use
+        name = server_name or f"Server_{len(mcp_connections)}"
         mcp_connections[name] = {"client": client, "tools": tools, "url": server_url}
         
-        print(f"Successfully connected to MCP server: {name} with {len(tools)} tools.")
-        return name, f"Successfully connected to '{name}' ({server_url}). Found {len(tools)} tool(s)."
+        return name, f"Successfully connected to {name} with {len(tools)} available tools"
     except Exception as e:
-        print(f"Error connecting to MCP server at {server_url}: {e}")
-        import traceback
-        traceback.print_exc()
-        return None, f"Error connecting to MCP server '{server_url}': {str(e)}"
+        print(f"Error connecting to MCP server: {e}")
+        return None, f"Error connecting to MCP server: {str(e)}"
 
 def list_mcp_tools(server_name):
     """List available tools for a connected MCP server"""
     if server_name not in mcp_connections:
-        return "Server not connected or name not found."
+        return "Server not connected"
     
     tools = mcp_connections[server_name]["tools"]
     tool_info = []
     for tool in tools:
-        tool_info.append(f"- **{tool.name}**: {tool.description}")
+        tool_info.append(f"- {tool.name}: {tool.description}")
     
     if not tool_info:
-        return "No tools available for this server."
+        return "No tools available for this server"
     
     return "\n".join(tool_info)
 
 def call_mcp_tool(server_name, tool_name, **kwargs):
-    """Call a specific tool from an MCP server and process its result."""
+    """Call a specific tool from an MCP server"""
     if server_name not in mcp_connections:
-        return {"type": "error", "message": f"Server '{server_name}' not connected."}
+        return {"error": f"Server '{server_name}' not connected"} # Return dict for consistency
     
-    mcp_client_instance = mcp_connections[server_name]["client"]
+    client_data = mcp_connections[server_name]
+    client = client_data["client"]
+    server_tools = client_data["tools"]
     
+    # Find the requested tool
+    tool = next((t for t in server_tools if t.name == tool_name), None)
+    if not tool:
+        return {"error": f"Tool '{tool_name}' not found on server '{server_name}'"}
+
     try:
-        print(f"Calling MCP tool: {server_name}.{tool_name} with args: {kwargs}")
-        # Assuming mcp_client_instance.call_tool returns an mcp.client.tool.ToolResult object
-        tool_result = mcp_client_instance.call_tool(tool_name, kwargs) 
-        
-        if tool_result and tool_result.content:
-            # Process multiple blocks if present, concatenating text or prioritizing audio
-            audio_block_found = None
-            text_parts = []
-            json_parts = []
-            other_parts = []
-
-            for block in tool_result.content:
-                if hasattr(block, 'uri') and isinstance(block.uri, str) and block.uri.startswith('data:audio/'):
-                    audio_block_found = {
-                        "type": "audio", 
-                        "data_uri": block.uri, 
-                        "name": getattr(block, 'name', 'audio_output.wav') 
-                    }
-                    break # Prioritize first audio block
-                elif hasattr(block, 'text') and block.text is not None:
-                    text_parts.append(str(block.text))
-                elif hasattr(block, 'json_data') and block.json_data is not None:
-                    try:
-                        json_parts.append(json.dumps(block.json_data, indent=2))
-                    except TypeError:
-                        json_parts.append(str(block.json_data)) # Fallback
-                else:
-                    other_parts.append(str(block))
-            
-            if audio_block_found:
-                print(f"MCP tool returned audio: {audio_block_found['name']}")
-                return audio_block_found
-            elif text_parts:
-                full_text = "\n".join(text_parts)
-                print(f"MCP tool returned text: {full_text[:100]}...")
-                return {"type": "text", "value": full_text}
-            elif json_parts:
-                full_json_str = "\n".join(json_parts)
-                print(f"MCP tool returned JSON string.")
-                return {"type": "json_string", "value": full_json_str} # Treat as string for display
-            elif other_parts:
-                print(f"MCP tool returned other content types.")
-                return {"type": "text", "value": "\n".join(other_parts)}
+        # Call the tool with provided arguments
+        mcp_tool_result: ToolResult = client.call_tool(tool_name=tool_name, arguments=kwargs)
+
+        actual_result = None
+        if mcp_tool_result.content:
+            content_block = mcp_tool_result.content[0]
+            if isinstance(content_block, ValueContentBlock):
+                actual_result = content_block.value
+            elif hasattr(content_block, 'text'): # e.g., TextContentBlock
+                actual_result = content_block.text
             else:
-                print("MCP tool executed but returned no interpretable primary content blocks.")
-                return {"type": "text", "value": "Tool executed, but returned no standard content (audio/text/json)."}
+                actual_result = str(content_block) # Fallback
+        else: # No content
+             return {"warning": "Tool returned no content."}
+
+
+        # Special handling for audio result (e.g., from Kokoro TTS)
+        # This checks if the result is a tuple (sample_rate, audio_data_list)
+        # Gradio MCP server serializes numpy arrays to lists.
+        if (server_name == "kokoroTTS" and tool_name == "text_to_audio" and 
+            isinstance(actual_result, tuple) and len(actual_result) == 2 and
+            isinstance(actual_result[0], int) and 
+            (isinstance(actual_result[1], list) or isinstance(actual_result[1], np.ndarray))):
+            
+            print(f"Received audio data from {server_name}.{tool_name}")
+            sample_rate, audio_data_list = actual_result
+            
+            # Convert list to numpy array if necessary
+            audio_data = np.array(audio_data_list)
+
+            # Ensure correct dtype for soundfile (float32 is common, or int16)
+            # Kokoro returns float, likely in [-1, 1] range.
+            if audio_data.dtype != np.float32 and audio_data.dtype != np.int16:
+                # Attempt to normalize if it looks like it's not in [-1, 1] for float
+                if np.issubdtype(audio_data.dtype, np.floating) and (np.min(audio_data) < -1.1 or np.max(audio_data) > 1.1):
+                    print(f"Warning: Audio data for {server_name}.{tool_name} might not be normalized. Min: {np.min(audio_data)}, Max: {np.max(audio_data)}")
+                audio_data = audio_data.astype(np.float32) 
+            
+            wav_io = io.BytesIO()
+            sf.write(wav_io, audio_data, sample_rate, format='WAV')
+            wav_io.seek(0)
+            
+            wav_b64 = base64.b64encode(wav_io.read()).decode('utf-8')
+            
+            return {
+                "type": "audio_b64", 
+                "data": wav_b64, 
+                "message": f"Audio generated by {server_name}.{tool_name}"
+            }
+        
+        # Handle other types of results
+        if isinstance(actual_result, dict):
+            return actual_result
+        elif isinstance(actual_result, str):
+            try: # If string is JSON, parse to dict
+                return json.loads(actual_result)
+            except json.JSONDecodeError:
+                return {"text": actual_result} # Wrap raw string
+        else:
+            return {"value": str(actual_result)} # Fallback for other primitive types
 
-        print("MCP tool executed, but ToolResult or its content was empty.")
-        return {"type": "text", "value": "Tool executed, but returned no content."}
     except Exception as e:
-        print(f"Error calling MCP tool '{tool_name}' or processing its result: {e}")
+        print(f"Error calling MCP tool: {e}")
         import traceback
         traceback.print_exc()
-        return {"type": "error", "message": f"Error during MCP tool '{tool_name}' execution: {str(e)}"}
+        return {"error": f"Error calling MCP tool: {str(e)}"}
 
-def analyze_message_for_tool_call(message, active_mcp_servers, llm_client, llm_model_to_use, base_system_message):
+def analyze_message_for_tool_call(message, active_mcp_servers, client, model_to_use, system_message):
     """Analyze a message to determine if an MCP tool should be called"""
-    if not message or not message.strip() or not active_mcp_servers:
+    if not message or not message.strip():
         return None, None
     
-    tool_info_for_llm = []
-    for server_name_iter in active_mcp_servers:
-        if server_name_iter in mcp_connections:
-            server_tools = mcp_connections[server_name_iter]["tools"]
-            for tool in server_tools:
-                # Provide a concise description for the LLM
-                tool_info_for_llm.append(
-                    f"- Server: '{server_name_iter}', Tool: '{tool.name}', Description: '{tool.description}'"
-                )
+    tool_info = []
+    for server_name in active_mcp_servers:
+        if server_name in mcp_connections:
+            server_tools_raw = list_mcp_tools(server_name) # This returns a string
+            if server_tools_raw != "Server not connected" and server_tools_raw != "No tools available for this server":
+                 # Parse the string from list_mcp_tools
+                for line in server_tools_raw.split("\n"):
+                    if line.startswith("- "):
+                        parts = line[2:].split(":", 1)
+                        if len(parts) == 2:
+                            tool_info.append({
+                                "server_name": server_name,
+                                "tool_name": parts[0].strip(),
+                                "description": parts[1].strip()
+                            })
     
-    if not tool_info_for_llm:
-        print("No active MCP tools found for analysis.")
+    if not tool_info:
         return None, None
     
-    tools_string_for_llm = "\n".join(tool_info_for_llm)
+    tools_desc = []
+    for info in tool_info:
+        tools_desc.append(f"{info['server_name']}.{info['tool_name']}: {info['description']}")
     
-    # More robust system prompt for tool detection
-    analysis_system_prompt = f"""You are an expert assistant that determines if a user's request requires an external tool.
-You have access to the following tools:
-{tools_string_for_llm}
-
-Based on the user's message, decide if any of these tools are appropriate.
-If a tool is needed, respond ONLY with a JSON object containing:
-"server_name": The name of the server providing the tool.
-"tool_name": The name of the tool to be called.
-"parameters": A dictionary of parameters for the tool, inferred from the user's message. Ensure parameter names match what the tool expects (often 'text', 'query', 'speed', etc.).
-
-If NO tool is needed, respond ONLY with the exact string: NO_TOOL_NEEDED
-
-Example 1 (TTS tool):
-User: "Can you say 'hello world' for me at a slightly faster speed?"
-Response: {{"server_name": "kokoroTTS", "tool_name": "text_to_audio", "parameters": {{"text": "hello world", "speed": 1.2}}}}
-
-Example 2 (File tool):
-User: "Read the content of my_document.txt"
-Response: {{"server_name": "FileSystemServer", "tool_name": "readFile", "parameters": {{"path": "my_document.txt"}}}}
-
-Example 3 (No tool):
-User: "What's the weather like today?" (Assuming no weather tool is listed)
-Response: NO_TOOL_NEEDED
-
-User's current message is: "{message}"
-Now, provide your decision:"""
+    tools_string = "\n".join(tools_desc)
+    
+    analysis_system_prompt = f"""You are an assistant that helps determine if a user message requires using an external tool.
+Available tools:
+{tools_string}
+
+Your job is to:
+1. Analyze the user's message.
+2. Determine if they're asking to use one of the tools.
+3. If yes, respond ONLY with a JSON object with "server_name", "tool_name", and "parameters".
+4. If no, respond ONLY with the exact string "NO_TOOL_NEEDED".
+
+Example 1 (User wants TTS):
+User: "Please turn this text into speech: Hello world"
+Response: {{"server_name": "kokoroTTS", "tool_name": "text_to_audio", "parameters": {{"text": "Hello world", "speed": 1.0}}}}
+
+Example 2 (User wants TTS with different server name):
+User: "Use mySpeechTool to say 'good morning'"
+Response: {{"server_name": "mySpeechTool", "tool_name": "text_to_audio", "parameters": {{"text": "good morning"}}}}
+
+Example 3 (User does not want a tool):
+User: "What is the capital of France?"
+Response: NO_TOOL_NEEDED"""
     
     try:
-        print(f"Sending tool analysis request to LLM model: {llm_model_to_use}")
-        response = llm_client.chat_completion(
-            model=llm_model_to_use,
+        response = client.chat_completion(
+            model=model_to_use,
             messages=[
-                # {"role": "system", "content": base_system_message}, # Optional: provide original system message for context
-                {"role": "user", "content": analysis_system_prompt} # The prompt itself is the user message here
+                {"role": "system", "content": analysis_system_prompt},
+                {"role": "user", "content": message}
             ],
-            temperature=0.1, # Low temperature for deterministic tool selection
-            max_tokens=300,
-            stop=["\n\n"] # Stop early if LLM adds extra verbiage
+            temperature=0.1, 
+            max_tokens=300
         )
         
-        analysis_text = response.choices[0].message.content.strip()
-        print(f"LLM tool analysis response: '{analysis_text}'")
+        analysis = response.choices[0].message.content.strip()
+        print(f"Tool analysis LLM response: '{analysis}'")
         
-        if "NO_TOOL_NEEDED" in analysis_text or analysis_text == "NO_TOOL_NEEDED":
-            print("LLM determined no tool needed.")
+        if analysis == "NO_TOOL_NEEDED":
             return None, None
         
-        # Try to extract JSON from the response (handle potential markdown code blocks)
-        if analysis_text.startswith("```json"):
-            analysis_text = analysis_text.replace("```json", "").replace("```", "").strip()
-        elif analysis_text.startswith("```"):
-             analysis_text = analysis_text.replace("```", "").strip()
-
-
-        json_start = analysis_text.find("{")
-        json_end = analysis_text.rfind("}") + 1
-        
-        if json_start == -1 or json_end <= json_start:
-            print(f"Could not find valid JSON object in LLM response: '{analysis_text}'")
-            return None, None
-            
-        json_str = analysis_text[json_start:json_end]
         try:
-            tool_call_data = json.loads(json_str)
-            if "server_name" in tool_call_data and "tool_name" in tool_call_data:
-                print(f"LLM suggested tool call: {tool_call_data}")
-                return tool_call_data.get("server_name"), {
-                    "tool_name": tool_call_data.get("tool_name"),
-                    "parameters": tool_call_data.get("parameters", {})
+            tool_call = json.loads(analysis)
+            if isinstance(tool_call, dict) and "server_name" in tool_call and "tool_name" in tool_call:
+                return tool_call.get("server_name"), {
+                    "tool_name": tool_call.get("tool_name"),
+                    "parameters": tool_call.get("parameters", {})
                 }
             else:
-                print(f"LLM response parsed as JSON but missing server_name or tool_name: {json_str}")
+                print(f"LLM response for tool call was not a valid JSON with required keys: {analysis}")
                 return None, None
-        except json.JSONDecodeError as e:
-            print(f"Failed to parse tool call JSON from LLM response: '{json_str}'. Error: {e}")
+        except json.JSONDecodeError:
+            print(f"Failed to parse tool call JSON from LLM: {analysis}")
             return None, None
             
     except Exception as e:
-        print(f"Error during LLM analysis for tool calls: {str(e)}")
-        import traceback
-        traceback.print_exc()
+        print(f"Error analyzing message for tool calls: {str(e)}")
         return None, None
 
 def respond(
-    message_text_input, # From user function, this is just the text part
-    message_files_input, # From user function, this is the list of file paths
-    history_tuples: list[tuple[tuple[str, list], str]], # History: list of ((user_text, [user_files]), assistant_response)
-    system_message_prompt,
-    max_tokens_val,
-    temperature_val,
-    top_p_val,
-    frequency_penalty_val,
-    seed_val,
-    provider_choice,
-    custom_api_key_val,
-    custom_model_id,    
-    # model_search_term_val, # Not directly used in respond, but kept for signature consistency if UI passes it
-    selected_hf_model_id,
-    mcp_is_enabled,
-    active_mcp_server_names, # List of selected server names
-    mcp_interaction_mode_choice
+    message,
+    image_files,
+    history: list[tuple[str, str]],
+    system_message,
+    max_tokens,
+    temperature,
+    top_p,
+    frequency_penalty,
+    seed,
+    provider,
+    custom_api_key,
+    custom_model,    
+    model_search_term,
+    selected_model,
+    mcp_enabled=False,
+    active_mcp_servers=None,
+    mcp_interaction_mode="Natural Language"
 ):
-    print(f"\n--- RESPOND FUNCTION CALLED ---")
-    print(f"Message Text: '{message_text_input}'")
-    print(f"Message Files: {message_files_input}")
-    # print(f"History (first item type if exists): {type(history_tuples) if history_tuples else 'No history'}")
-    print(f"System Prompt: '{system_message_prompt}'")
-    print(f"Provider: {provider_choice}, MCP Enabled: {mcp_is_enabled}, MCP Mode: {mcp_interaction_mode_choice}")
-    print(f"Active MCP Servers: {active_mcp_server_names}")
-
-    token_to_use_for_llm = custom_api_key_val if custom_api_key_val.strip() else ACCESS_TOKEN
-    if not token_to_use_for_llm and provider_choice != "hf-inference": # Basic check
-        yield "Error: API Key required for non-hf-inference providers."
-        return
+    print(f"Received message: {message}")
+    print(f"Received {len(image_files) if image_files else 0} images")
+    # print(f"History: {history}") # Can be verbose
+    print(f"System message: {system_message}")
+    print(f"Max tokens: {max_tokens}, Temperature: {temperature}, Top-P: {top_p}")
+    print(f"Frequency Penalty: {frequency_penalty}, Seed: {seed}")
+    print(f"Selected provider: {provider}")         
+    print(f"Custom API Key provided: {bool(custom_api_key.strip())}")
+    print(f"Selected model (custom_model): {custom_model}")  
+    print(f"Model search term: {model_search_term}")
+    print(f"Selected model from radio: {selected_model}")
+    print(f"MCP enabled: {mcp_enabled}")
+    print(f"Active MCP servers: {active_mcp_servers}")
+    print(f"MCP interaction mode: {mcp_interaction_mode}")
+
+    token_to_use = custom_api_key if custom_api_key.strip() != "" else ACCESS_TOKEN
     
-    llm_client_instance = InferenceClient(token=token_to_use_for_llm, provider=provider_choice)
+    if custom_api_key.strip() != "":
+        print("USING CUSTOM API KEY: BYOK token provided by user is being used for authentication")
+    else:
+        print("USING DEFAULT API KEY: Environment variable HF_TOKEN is being used for authentication")
     
-    current_seed = None if seed_val == -1 else seed_val
-    model_id_for_llm = custom_model_id.strip() if custom_model_id.strip() else selected_hf_model_id
-    print(f"Using LLM model: {model_id_for_llm} via {provider_choice}")
-
-    # --- MCP Tool Call Logic ---
-    if mcp_is_enabled and (message_text_input or message_files_input) and active_mcp_server_names:
-        tool_call_output_dict = None
-        invoked_tool_display_name = "a tool"
-        invoked_server_display_name = "an MCP server"
-
-        if message_text_input and message_text_input.startswith("/mcp"):
-            print("Processing explicit MCP command...")
-            command_parts = message_text_input.split(" ", 3)
+    client = InferenceClient(token=token_to_use, provider=provider)
+    print(f"Hugging Face Inference Client initialized with {provider} provider.")
+
+    if seed == -1:
+        seed = None
+        
+    model_to_use = custom_model.strip() if custom_model.strip() != "" else selected_model
+    print(f"Model selected for inference: {model_to_use}")
+
+    if mcp_enabled and message:
+        if message.startswith("/mcp"):
+            command_parts = message.split(" ", 3)
             if len(command_parts) < 3:
                 yield "Invalid MCP command. Format: /mcp <server_name> <tool_name> [arguments_json]"
                 return
             
-            _, server_name_cmd, tool_name_cmd = command_parts[:3]
-            invoked_server_display_name = server_name_cmd
-            invoked_tool_display_name = tool_name_cmd
-            args_json_str = "{}" if len(command_parts) < 4 else command_parts
+            _, server_name, tool_name = command_parts[:3]
+            args_json = "{}" if len(command_parts) < 4 else command_parts[3]
             
             try:
-                args_dict_cmd = json.loads(args_json_str)
-                tool_call_output_dict = call_mcp_tool(invoked_server_display_name, invoked_tool_display_name, **args_dict_cmd)
+                args_dict = json.loads(args_json)
+                result = call_mcp_tool(server_name, tool_name, **args_dict)
+
+                if isinstance(result, dict) and result.get("type") == "audio_b64":
+                    yield f"<audio controls src=\"data:audio/wav;base64,{result.get('data')}\"></audio>"
+                elif isinstance(result, dict) and "error" in result:
+                    yield f"Error: {result.get('error')}"
+                elif isinstance(result, dict):
+                    yield json.dumps(result, indent=2)
+                else:
+                    yield str(result) 
+                return
             except json.JSONDecodeError:
-                yield f"Invalid JSON arguments for MCP command: {args_json_str}"
+                yield f"Invalid JSON arguments: {args_json}"
                 return
-            except Exception as e_cmd:
-                yield f"Error preparing MCP command: {str(e_cmd)}"
+            except Exception as e:
+                yield f"Error executing MCP command: {str(e)}"
                 return
-
-        elif mcp_interaction_mode_choice == "Natural Language":
-            print("Analyzing message for natural language tool call...")
-            # For natural language, primarily use message_text_input. Files could be context later.
-            detected_server_nl, tool_info_nl = analyze_message_for_tool_call(
-                message_text_input, 
-                active_mcp_server_names, 
-                llm_client_instance, 
-                model_id_for_llm, 
-                system_message_prompt
+        elif mcp_interaction_mode == "Natural Language" and active_mcp_servers and active_mcp_servers:
+            print("Attempting natural language tool call detection...")
+            server_name, tool_info = analyze_message_for_tool_call(
+                message, active_mcp_servers, client, model_to_use, system_message
             )
             
-            if detected_server_nl and tool_info_nl and tool_info_nl.get("tool_name"):
-                invoked_server_display_name = detected_server_nl
-                invoked_tool_display_name = tool_info_nl['tool_name']
-                tool_params_nl = tool_info_nl.get("parameters", {})
-                tool_call_output_dict = call_mcp_tool(invoked_server_display_name, invoked_tool_display_name, **tool_params_nl)
-
-        # --- Handle MCP Tool Result (if a tool was called) ---
-        if tool_call_output_dict:
-            response_message_parts = [f"I attempted to use the **{invoked_tool_display_name}** tool from **{invoked_server_display_name}**."]
-            
-            if tool_call_output_dict.get("type") == "audio":
-                audio_data_uri = tool_call_output_dict["data_uri"]
-                audio_html_tag = f"<audio controls src='{audio_data_uri}' title='{tool_call_output_dict.get('name', 'Audio Output')}'></audio>"
-                response_message_parts.append(f"Here's the audio output:\n{audio_html_tag}")
-            elif tool_call_output_dict.get("type") == "text":
-                response_message_parts.append(f"\nResult:\n```\n{tool_call_output_dict['value']}\n```")
-            elif tool_call_output_dict.get("type") == "json_string": # Changed from "json" to avoid confusion with dict
-                response_message_parts.append(f"\nResult (JSON):\n```json\n{tool_call_output_dict['value']}\n```")
-            elif tool_call_output_dict.get("type") == "error":
-                response_message_parts.append(f"\nUnfortunately, there was an error: {tool_call_output_dict['message']}")
-            else: # Fallback for unexpected result structure
-                response_message_parts.append(f"\nThe tool returned: {str(tool_call_output_dict)}")
-            
-            yield "\n".join(response_message_parts)
-            return # End here if a tool was called and processed
+            if server_name and tool_info and tool_info.get("tool_name"):
+                try:
+                    print(f"Calling tool via natural language: {server_name}.{tool_info['tool_name']} with parameters: {tool_info['parameters']}")
+                    result = call_mcp_tool(server_name, tool_info['tool_name'], **tool_info.get('parameters', {}))
+                    
+                    response_message = f"I used the **{tool_info['tool_name']}** tool from **{server_name}**."
+                    if isinstance(result, dict) and result.get("message"):
+                         response_message += f" ({result.get('message')})"
+                    response_message += "\n\n"
+
+                    if isinstance(result, dict) and result.get("type") == "audio_b64":
+                        audio_html = f"<audio controls src=\"data:audio/wav;base64,{result.get('data')}\"></audio>"
+                        yield response_message + audio_html
+                    elif isinstance(result, dict) and "error" in result:
+                        result_str = f"Tool Error: {result.get('error')}"
+                        yield response_message + result_str
+                    elif isinstance(result, dict):
+                        result_str = f"Result:\n```json\n{json.dumps(result, indent=2)}\n```"
+                        yield response_message + result_str
+                    else:
+                        result_str = f"Result:\n{str(result)}"
+                        yield response_message + result_str
+                    return 
+                except Exception as e:
+                    print(f"Error executing MCP tool via natural language: {str(e)}")
+                    # yield f"Sorry, I encountered an error trying to use the tool: {str(e)}"
+                    # Fall through to normal LLM response if tool call fails here
+            else:
+                print("No tool call detected by natural language analysis or tool_info incomplete.")
 
-    # --- Regular LLM Response Logic (if no MCP tool was successfully called and returned primary content) ---
-    print("Proceeding with standard LLM response generation.")
+
+    user_content_parts = []
+    if message and message.strip():
+        user_content_parts.append({"type": "text", "text": message})
     
-    # Prepare current user message for LLM (multimodal if files exist)
-    current_user_llm_content = []
-    if message_text_input and message_text_input.strip():
-        current_user_llm_content.append({"type": "text", "text": message_text_input})
+    if image_files and len(image_files) > 0:
+        for img_path in image_files:
+            if img_path:
+                try:
+                    encoded_image = encode_image(img_path)
+                    if encoded_image:
+                        user_content_parts.append({
+                            "type": "image_url",
+                            "image_url": {"url": f"data:image/jpeg;base64,{encoded_image}"}
+                        })
+                except Exception as e:
+                    print(f"Error encoding image {img_path}: {e}")
     
-    if message_files_input:
-        for file_path in message_files_input:
-            if file_path: # file_path is already the actual temp path from gr.File or gr.Image
-                encoded_img_str = encode_image(file_path)
-                if encoded_img_str:
-                    current_user_llm_content.append({
-                        "type": "image_url",
-                        "image_url": {"url": f"data:image/jpeg;base64,{encoded_img_str}"}
-                    })
-                else:
-                    print(f"Warning: Failed to encode image {file_path} for LLM.")
-    
-    if not current_user_llm_content:
-        print("No content (text or valid files) in current user message for LLM.")
-        yield "" # Or some indicator of no action
-        return
-
-    # Augment system message with MCP tool info if enabled
-    augmented_sys_msg = system_message_prompt
-    if mcp_is_enabled and active_mcp_server_names:
-        mcp_tool_descriptions_for_llm = []
-        for server_name_iter in active_mcp_server_names:
+    if not user_content_parts: # If message was only /mcp command and processed
+        print("No further content for LLM after MCP command processing.")
+        # This might happen if an MCP command was fully handled and returned.
+        # If yield was used, the function already exited. If not, we might need to ensure no LLM call.
+        # However, the logic above for MCP commands uses `yield ...; return`, so this path might not be hit often.
+        # If it *is* hit, it means the MCP command didn't yield, and we should not proceed to LLM.
+        if message and message.startswith("/mcp"):
+             return # Ensure we don't fall through after a command that should have yielded.
+
+
+    final_user_content = user_content_parts if len(user_content_parts) > 1 else (user_content_parts[0] if user_content_parts else "")
+
+    augmented_system_message = system_message
+    if mcp_enabled and active_mcp_servers:
+        tool_list_for_prompt = []
+        for server_name_iter in active_mcp_servers:
             if server_name_iter in mcp_connections:
-                # Use the more detailed list_mcp_tools output for the system prompt if desired
-                tools_list_str = list_mcp_tools(server_name_iter) # This returns markdown
-                mcp_tool_descriptions_for_llm.append(f"From server '{server_name_iter}':\n{tools_list_str}")
+                server_tools_str = list_mcp_tools(server_name_iter)
+                if server_tools_str and "not connected" not in server_tools_str and "No tools available" not in server_tools_str:
+                    tool_list_for_prompt.append(f"From server '{server_name_iter}':\n{server_tools_str}")
         
-        if mcp_tool_descriptions_for_llm:
-            full_tools_info_str = "\n\n".join(mcp_tool_descriptions_for_llm)
-            interaction_advice = ""
-            if mcp_interaction_mode_choice == "Command Mode":
-                interaction_advice = "The user can invoke these tools using '/mcp <server_name> <tool_name> <json_args>'."
-            # For Natural Language mode, the LLM doesn't need explicit instruction in system prompt
-            # as `analyze_message_for_tool_call` handles that part.
+        if tool_list_for_prompt:
+            mcp_tools_description = "\n\n".join(tool_list_for_prompt)
             
-            augmented_sys_msg += f"\n\nYou also have access to the following external tools via Model Context Protocol (MCP):\n{full_tools_info_str}\n{interaction_advice}"
-
-    # Prepare messages list for LLM
-    messages_for_llm_api = [{"role": "system", "content": augmented_sys_msg}]
+            if mcp_interaction_mode == "Command Mode":
+                augmented_system_message += f"\n\nYou have access to the following MCP tools. To use them, type a command in the format: /mcp <server_name> <tool_name> <arguments_json>\nTools:\n{mcp_tools_description}"
+            else: # Natural Language
+                augmented_system_message += f"\n\nYou have access to the following MCP tools. You can ask to use them in natural language, and I will try to detect when a tool is needed. If I miss it, you can try being more explicit about the tool name.\nTools:\n{mcp_tools_description}"
     
-    for hist_user_turn, hist_assist_response in history_tuples:
-        hist_user_text, hist_user_files = hist_user_turn # Unpack ((text, [files]))
+    messages_for_api = [{"role": "system", "content": augmented_system_message}]
+    print("Initial messages array constructed.")
+
+    for val in history:
+        past_user_msg, past_assistant_msg = val
         
-        history_user_llm_content = []
-        if hist_user_text and hist_user_text.strip():
-            history_user_llm_content.append({"type": "text", "text": hist_user_text})
-        if hist_user_files:
-            for hist_file_path in hist_user_files:
-                encoded_hist_img = encode_image(hist_file_path)
-                if encoded_hist_img:
-                    history_user_llm_content.append({
-                        "type": "image_url",
-                        "image_url": {"url": f"data:image/jpeg;base64,{encoded_hist_img}"}
-                    })
-        if history_user_llm_content: # Only add if there's actual content
-            messages_for_llm_api.append({"role": "user", "content": history_user_llm_content})
+        # Handle past user messages (could be text or multimodal)
+        if past_user_msg:
+            if isinstance(past_user_msg, list): # Already multimodal
+                 messages_for_api.append({"role": "user", "content": past_user_msg})
+            elif isinstance(past_user_msg, str): # Text only
+                 messages_for_api.append({"role": "user", "content": past_user_msg})
         
-        if hist_assist_response and hist_assist_response.strip():
-            messages_for_llm_api.append({"role": "assistant", "content": hist_assist_response})
+        if past_assistant_msg:
+            messages_for_api.append({"role": "assistant", "content": past_assistant_msg})
 
-    messages_for_llm_api.append({"role": "user", "content": current_user_llm_content})
-    # print(f"Final messages for LLM API: {json.dumps(messages_for_llm_api, indent=2)}")
+    if final_user_content: # Add current user message if it exists
+        messages_for_api.append({"role": "user", "content": final_user_content})
+    
+    print(f"Latest user message appended (content type: {type(final_user_content)})")
+    # print(f"Full messages_for_api: {json.dumps(messages_for_api, indent=2)}") # Can be very verbose
 
+    llm_response_text = ""
+    print(f"Sending request to {provider} provider for model {model_to_use}.")
 
-    llm_parameters = {
-        "max_tokens": max_tokens_val, "temperature": temperature_val, "top_p": top_p_val,
-        "frequency_penalty": frequency_penalty_val,
+    parameters = {
+        "max_tokens": max_tokens,
+        "temperature": temperature,
+        "top_p": top_p,
+        "frequency_penalty": frequency_penalty,
     }
-    if current_seed is not None:
-        llm_parameters["seed"] = current_seed
+    
+    if seed is not None:
+        parameters["seed"] = seed
 
-    print(f"Sending request to LLM: Model={model_id_for_llm}, Params={llm_parameters}")
-    streamed_response_text = ""
     try:
-        llm_stream = llm_client_instance.chat_completion(
-            model=model_id_for_llm,
-            messages=messages_for_llm_api,
+        stream = client.chat_completion(
+            model=model_to_use,
+            messages=messages_for_api,
             stream=True,
-            **llm_parameters
+            **parameters
         )
         
-        # print("Streaming LLM response: ", end="", flush=True)
-        for chunk in llm_stream:
+        # print("Received tokens: ", end="", flush=True) # Can be too noisy
+        
+        for chunk in stream:
             if hasattr(chunk, 'choices') and len(chunk.choices) > 0:
-                delta = chunk.choices.delta
-                if hasattr(delta, 'content') and delta.content:
-                    token = delta.content
-                    # print(token, end="", flush=True)
-                    streamed_response_text += token
-                    yield streamed_response_text
-        # print("\nLLM Stream finished.")
-    except Exception as e_llm:
-        error_msg = f"Error during LLM inference: {str(e_llm)}"
-        print(error_msg)
-        import traceback
-        traceback.print_exc()
-        streamed_response_text += f"\n{error_msg}" # Append error to existing stream if any
-        yield streamed_response_text
+                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) # Can be too noisy
+                        llm_response_text += token_text
+                        yield llm_response_text
+        
+        # print() # Newline after tokens
+    except Exception as e:
+        print(f"Error during LLM inference: {e}")
+        llm_response_text += f"\nLLM Error: {str(e)}"
+        yield llm_response_text
 
-    print(f"--- RESPOND FUNCTION COMPLETED ---")
+    print("Completed LLM response generation.")
 
 
 # GRADIO UI
-with gr.Blocks(theme="Nymbo/Nymbo_Theme", title="Serverless TextGen Hub + MCP") as demo:
-    gr.Markdown("# Serverless TextGen Hub with MCP Client")
+with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
     chatbot = gr.Chatbot(
-        label="Chat",
         height=600, 
         show_copy_button=True, 
-        placeholder="Select a model, connect MCP servers (optional), and start chatting!",
-        bubble_full_width=False,
-        avatar_images=(None, "https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo-square.png") 
+        placeholder="Select a model and begin chatting. Supports multiple inference providers, multimodal inputs, and MCP tools.",
+        layout="panel",
+        show_label=False,
+        render=False # Delay rendering
     )
+    print("Chatbot interface created.")
     
     with gr.Row():
-        msg_textbox = gr.MultimodalTextbox( # Changed from gr.Textbox to gr.MultimodalTextbox
-            placeholder="Type a message or upload images... (Use /mcp for commands)",
+        msg = gr.MultimodalTextbox(
+            placeholder="Type a message or upload images...",
             show_label=False,
             container=False,
             scale=12,
-            file_types=["image"], # Can add more types like "audio", "video" if supported by models
-            file_count="multiple" # Allow multiple image uploads
+            file_types=["image"],
+            file_count="multiple",
+            sources=["upload"],
+            render=False # Delay rendering
         )
-        # submit_button = gr.Button("Send", variant="primary", scale=1, min_width=100) # Optional explicit send button
 
-    with gr.Accordion("LLM Settings", open=False):
-        system_message_prompt_box = gr.Textbox(
-            value="You are a helpful and versatile AI assistant. You can understand text and images. If you have access to MCP tools, you can use them when appropriate or when the user asks.", 
-            label="System Prompt", lines=3
+    chatbot.render()
+    msg.render()
+    
+    with gr.Accordion("Settings", open=False):
+        system_message_box = gr.Textbox(
+            value="You are a helpful AI assistant that can understand images and text. If the user asks you to use a tool, try your best.", 
+            placeholder="You are a helpful assistant.",
+            label="System Prompt"
         )
         
         with gr.Row():
             with gr.Column(scale=1):
-                max_tokens_slider_ui = gr.Slider(minimum=128, maximum=8192, value=1024, step=128, label="Max New Tokens")
-                temperature_slider_ui = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.05, label="Temperature")
-                top_p_slider_ui = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.05, label="Top-P (Nucleus Sampling)")
+                max_tokens_slider = gr.Slider(minimum=1, maximum=8192, value=1024, step=1, label="Max tokens")
+                temperature_slider = gr.Slider(minimum=0.0, maximum=2.0, value=0.7, step=0.01, label="Temperature")
+                top_p_slider = gr.Slider(minimum=0.0, maximum=1.0, value=0.95, step=0.01, label="Top-P")
             with gr.Column(scale=1):
-                frequency_penalty_slider_ui = gr.Slider(minimum=-2.0, maximum=2.0, value=0.0, step=0.1, label="Frequency Penalty")
-                seed_slider_ui = gr.Slider(minimum=-1, maximum=65535, value=-1, step=1, label="Seed (-1 for random)")
+                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)")
         
-        providers_list_ui = [
-            "hf-inference", "cerebras", "together", "sambanova", "novita", 
-            "cohere", "fireworks-ai", "hyperbolic", "nebius",
-        ]
-        provider_radio_ui = gr.Radio(choices=providers_list_ui, value="hf-inference", label="Inference Provider")
-        
-        byok_textbox_ui = gr.Textbox(label="Your Hugging Face API Key (Optional)", placeholder="Enter HF Token if using non-hf-inference providers or private models", type="password")
-        
-        custom_model_id_box = gr.Textbox(label="Custom Model ID (Overrides selection below)", placeholder="e.g., meta-llama/Llama-3-8B-Instruct")
+        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(value="", label="BYOK (Bring Your Own Key)", info="Enter a custom Hugging Face API key here. If empty, only 'hf-inference' provider can be used with the shared token.", placeholder="Enter your Hugging Face API token", type="password")
+        custom_model_box = gr.Textbox(value="", label="Custom Model ID", info="(Optional) Provide a custom Hugging Face model ID. Overrides selected featured model.", placeholder="meta-llama/Llama-3.1-70B-Instruct")
         
-        model_search_box_ui = gr.Textbox(label="Filter Featured Models", placeholder="Search...", lines=1)
+        model_search_box = gr.Textbox(label="Filter Featured Models", placeholder="Search for a featured model...", lines=1)
         
-        # More diverse model list, including some known multimodal ones
-        featured_models_list_data = [
-            "meta-llama/Meta-Llama-3.1-8B-Instruct", # Good default
-            "meta-llama/Meta-Llama-3.1-70B-Instruct",
+        models_list = [
+            "meta-llama/Llama-3.1-405B-Instruct-FP8", # Large model, might be slow/expensive
+            "meta-llama/Llama-3.1-70B-Instruct",
+            "meta-llama/Llama-3.1-8B-Instruct",
             "mistralai/Mistral-Nemo-Instruct-2407",
-            "mistralai/Mixtral-8x22B-Instruct-v0.1",
-            "Qwen/Qwen2-7B-Instruct",
-            "microsoft/Phi-3-medium-128k-instruct",
-            # Multimodal
-            "Salesforce/blip-image-captioning-large", # Example, might not be chat
-            "llava-hf/llava-1.5-7b-hf", # LLaVA example
-            "microsoft/kosmos-2-patch14-224", # Kosmos-2
-            "google/paligemma-3b-mix-448", # PaliGemma
+            "Qwen/Qwen2-72B-Instruct",
+            "Qwen/Qwen2-57B-A14B-Instruct",
+            "CohereForAI/c4ai-command-r-plus",
+            # Multimodal models
+            "Salesforce/LlavaLlama3-8b-hf",
+            "llava-hf/llava-v1.6-mistral-7b-hf",
+            "llava-hf/llava-v1.6-vicuna-13b-hf",
+            "microsoft/Phi-3-vision-128k-instruct",
+            "google/paligemma-3b-mix-448",
+            # Older but still popular
+            "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO",
+            "mistralai/Mixtral-8x7B-Instruct-v0.1",
+            "mistralai/Mistral-7B-Instruct-v0.3",
         ]
-        featured_model_radio_ui = gr.Radio(label="Select a Featured Model", choices=featured_models_list_data, value="meta-llama/Meta-Llama-3.1-8B-Instruct", interactive=True)
-        
-        gr.Markdown("Tip: For multimodal chat, ensure selected model supports image inputs (e.g., LLaVA, PaliGemma, Kosmos-2).")
+        featured_model_radio = gr.Radio(label="Select a Featured Model", choices=models_list, value="meta-llama/Llama-3.1-8B-Instruct", interactive=True)
+        gr.Markdown("[View all Text-to-Text models](https://huggingface.co/models?pipeline_tag=text-generation&sort=trending) | [View all multimodal models](https://huggingface.co/models?pipeline_tag=image-to-text&sort=trending)")
 
-    with gr.Accordion("MCP Client Settings", open=False):
-        mcp_enabled_checkbox_ui = gr.Checkbox(label="Enable MCP Support", value=False, info="Connect to external tools and services via MCP.")
-        
+    with gr.Accordion("MCP Settings", open=False):
+        mcp_enabled_checkbox = gr.Checkbox(label="Enable MCP Support", value=False, info="Enable Model Context Protocol support to connect to external tools and services")
         with gr.Row():
-            mcp_server_url_textbox = gr.Textbox(label="MCP Server URL", placeholder="e.g., https://your-mcp-server.hf.space/gradio_api/mcp/sse")
-            mcp_server_name_textbox = gr.Textbox(label="Friendly Server Name (Optional)", placeholder="MyTTS_Server")
-            mcp_connect_button_ui = gr.Button("Connect", variant="secondary")
+            mcp_server_url = gr.Textbox(label="MCP Server URL", placeholder="https://your-mcp-server.hf.space/gradio_api/mcp/sse", info="URL of the MCP server (usually ends with /gradio_api/mcp/sse for Gradio MCP servers)")
+            mcp_server_name = gr.Textbox(label="Server Name (Optional)", placeholder="e.g., kokoroTTS", info="A friendly name to identify this server")
+            mcp_connect_button = gr.Button("Connect to MCP Server")
         
-        mcp_connection_status_textbox = gr.Textbox(label="MCP Connection Status", placeholder="No MCP servers connected.", interactive=False, lines=2)
+        mcp_status = gr.Textbox(label="MCP Connection Status", placeholder="No MCP servers connected", interactive=False)
+        active_mcp_servers = gr.Dropdown(label="Active MCP Servers for Chat", choices=[], multiselect=True, info="Select which connected MCP servers to make available to the LLM for this chat session")
+        mcp_mode = gr.Radio(label="MCP Interaction Mode", choices=["Natural Language", "Command Mode"], value="Natural Language", info="Choose how to interact with MCP tools")
         
-        active_mcp_servers_dropdown = gr.Dropdown(
-            label="Use Tools From (Select Active MCP Servers)", choices=[], multiselect=True, 
-            info="Choose which connected servers the LLM can use tools from."
-        )
+        gr.Markdown("""
+        ### MCP Interaction Modes & Examples
+        **Natural Language Mode**: Describe what you want.
+        `Please say 'Hello world' using the kokoroTTS server.`
+        `Use my speech tool to read this: "Welcome"`
         
-        mcp_interaction_mode_radio = gr.Radio(
-            label="MCP Interaction Mode", choices=["Natural Language", "Command Mode"], value="Natural Language",
-            info="Natural Language: AI tries to detect tool use. Command Mode: Use '/mcp ...' syntax."
-        )
-        gr.Markdown("Example MCP Command: `/mcp MyTTS text_to_audio {\"text\": \"Hello world!\"}`")
+        **Command Mode**: Use structured commands (server name must match connected server's friendly name).
+        `/mcp <server_name> <tool_name> {"param1": "value1"}`
+        Example: `/mcp kokoroTTS text_to_audio {"text": "Hello world", "speed": 1.0}`
+        """)
+
+    # Chat history state
+    # The chatbot component itself manages history for display.
+    # The `respond` function receives this display history and reconstructs API history.
+
+    def filter_models_ui_update(search_term):
+        print(f"Filtering models with search term: {search_term}")
+        filtered = [m for m in models_list if search_term.lower() in m.lower()]
+        if not filtered: # If search yields no results, show all models
+            filtered = models_list
+        print(f"Filtered models: {filtered}")
+        return gr.Radio(choices=filtered, label="Select a Featured Model", value=featured_model_radio.value if featured_model_radio.value in filtered else (filtered[0] if filtered else None))
+
+    def set_custom_model_from_radio_ui_update(selected_featured_model):
+        print(f"Featured model selected: {selected_featured_model}")
+        return selected_featured_model # This updates the custom_model_box
+
+    def connect_mcp_server_ui_update(url, name_optional):
+        actual_name, status_msg = connect_to_mcp_server(url, name_optional)
+        updated_server_choices = list(mcp_connections.keys())
+        # Keep existing selection if possible
+        current_selection = active_mcp_servers.value if active_mcp_servers.value else []
+        valid_selection = [s for s in current_selection if s in updated_server_choices]
+        if actual_name and actual_name not in valid_selection: # Auto-select newly connected server
+            valid_selection.append(actual_name)
+
+        return status_msg, gr.Dropdown(choices=updated_server_choices, value=valid_selection, label="Active MCP Servers for Chat")
+
+    # This function processes the user's multimodal input and adds it to the chatbot history.
+    # It prepares the history in a way that `bot` can understand.
+    def handle_user_input(multimodal_input, history_list: list):
+        text_content = multimodal_input.get("text", "").strip()
+        files = multimodal_input.get("files", [])
+        
+        # This will be the entry for the user's turn in the history
+        user_turn_for_api = []
+        user_turn_for_display = ""
 
-    # --- Event Handlers ---
-    
-    # Store history as list of tuples: [ ((user_text, [user_files]), assistant_response), ... ]
-    chat_history_state = gr.State([]) 
+        if text_content:
+            user_turn_for_api.append({"type": "text", "text": text_content})
+            user_turn_for_display = text_content
+        
+        if files:
+            display_files_md = ""
+            for file_path in files:
+                if file_path and isinstance(file_path, str): # Gradio provides temp path
+                    encoded_img = encode_image(file_path) # For API
+                    if encoded_img:
+                         user_turn_for_api.append({"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{encoded_img}"}})
+                    # For display, Gradio handles showing the image from MultimodalTextbox output
+                    # We'll just make a note in the display string
+                    display_files_md += f"\n<img src='file={file_path}' style='max-height:150px; display:block;' alt='uploaded image'>" # Gradio can render this!
+            
+            if user_turn_for_display:
+                user_turn_for_display += display_files_md
+            else:
+                user_turn_for_display = display_files_md if display_files_md else "Image(s) uploaded"
+
+
+        if not user_turn_for_display and not user_turn_for_api: # Empty input
+            return history_list, multimodal_input # No change
 
-    def user_interaction(user_multimodal_input, current_chat_history):
-        user_text = user_multimodal_input["text"] if user_multimodal_input and "text" in user_multimodal_input else ""
-        user_files = user_multimodal_input["files"] if user_multimodal_input and "files" in user_multimodal_input else []
+        # The `respond` function expects history as list of [user_api_content, assistant_text_content]
+        # For the current turn, we add [user_api_content, None]
+        # The display history for chatbot is [user_display_content, assistant_text_content]
         
-        # Only add to history if there's text or files
-        if user_text or user_files:
-            current_chat_history.append( ((user_text, user_files), None) ) # Append user turn, assistant response is None initially
-        return current_chat_history, gr.update(value={"text": "", "files": []}) # Clear input textbox
-
-    def bot_response_generator(
-        current_chat_history, system_prompt, max_tokens, temp, top_p_val, freq_penalty, seed_val, 
-        provider_val, api_key_val, custom_model_val, selected_model_val, # Removed search_term as it's not directly used by respond
-        mcp_enabled_val, active_servers_val, mcp_mode_val
-    ):
-        if not current_chat_history or current_chat_history[-1] is not None: # If no user message or last message already has bot response
-            yield current_chat_history # Or simply `return current_chat_history` if not streaming
-            return
+        # We pass the API-formatted user turn to the `message` arg of `respond`
+        # and the existing history to the `history` arg.
+        # The chatbot's display history is updated here.
+        
+        history_list.append([user_turn_for_display, None])
+        return history_list, user_turn_for_api # Return updated history and the API formatted current message
 
-        user_turn_content, _ = current_chat_history[-1] # Get the latest user turn: (text, [files])
-        message_text, message_files = user_turn_content
 
-        # The history passed to `respond` should be all turns *before* the current one
-        history_for_respond = current_chat_history[:-1]
+    # The bot function that calls `respond` generator
+    def call_bot_responder(history_list_for_display, current_user_api_content, sys_msg, max_tok, temp, top_p_val, freq_pen, seed_val, prov, api_key_val, cust_model, _search, sel_model, mcp_on, active_servs, mcp_inter_mode):
+        if not current_user_api_content and not (history_list_for_display and history_list_for_display[-1][0]):
+            print("Bot called with no current message and no history, skipping.")
+            yield history_list_for_display # No change
+            return
 
-        response_stream = respond(
-            message_text, message_files, history_for_respond,
-            system_prompt, max_tokens, temp, top_p_val, freq_penalty, seed_val,
-            provider_val, api_key_val, custom_model_val, selected_model_val,
-            mcp_enabled_val, active_servers_val, mcp_mode_val
-        )
+        # Reconstruct API history from display history
+        # `respond` expects history as list of [user_api_content, assistant_text_content]
+        # The current `history_list_for_display` is [user_display, assistant_text]
+        # This reconstruction is tricky because display != api format.
+        # For simplicity, we'll pass only the text part of history to `respond` for now,
+        # and the full current_user_api_content for the current message.
+        # A more robust solution would store API history separately.
         
-        full_bot_message = ""
-        for chunk in response_stream:
-            full_bot_message = chunk
-            current_chat_history[-1] = (user_turn_content, full_bot_message) # Update last item's assistant part
-            yield current_chat_history
-    
-    # Link UI components to functions
-    msg_textbox.submit(
-        user_interaction,
-        inputs=[msg_textbox, chat_history_state],
-        outputs=[chat_history_state, msg_textbox] # Update history and clear input
-    ).then(
-        bot_response_generator,
-        inputs=[
-            chat_history_state, system_message_prompt_box, max_tokens_slider_ui, temperature_slider_ui, 
-            top_p_slider_ui, frequency_penalty_slider_ui, seed_slider_ui, provider_radio_ui, 
-            byok_textbox_ui, custom_model_id_box, featured_model_radio_ui, 
-            mcp_enabled_checkbox_ui, active_mcp_servers_dropdown, mcp_interaction_mode_radio
-        ],
-        outputs=[chatbot] # Stream to chatbot
-    )
+        # Simplified history for `respond` (text only from past turns)
+        # The `respond` function itself needs to be robust to handle this.
+        # Let's adjust `respond` to take `message` (current API content) and `image_files` (current files)
+        # and `history` (past turns, which we'll simplify here).
 
-    # MCP Connection
-    def handle_mcp_connect(url, name_suggestion):
-        if not url or not url.strip():
-            return "MCP Server URL cannot be empty.", gr.update(choices=list(mcp_connections.keys()))
+        # The `respond` function is already structured to take `message` (text) and `image_files`
+        # The `current_user_api_content` is what we need to pass as `message` (if text) or `image_files`
         
-        _, status_msg = connect_to_mcp_server(url, name_suggestion)
-        # Update dropdown choices with current server names
-        new_choices = list(mcp_connections.keys())
-        # Preserve selected values if they are still valid connections
-        # current_selected = active_mcp_servers_dropdown.value # This might not work directly
-        # new_selected = [s for s in current_selected if s in new_choices]
-        return status_msg, gr.update(choices=new_choices) #, value=new_selected)
-
-    mcp_connect_button_ui.click(
-        handle_mcp_connect,
-        inputs=[mcp_server_url_textbox, mcp_server_name_textbox],
-        outputs=[mcp_connection_status_textbox, active_mcp_servers_dropdown]
-    )
+        current_message_text = ""
+        current_image_paths = []
+
+        if isinstance(current_user_api_content, list): # Multimodal
+            for part in current_user_api_content:
+                if part["type"] == "text":
+                    current_message_text = part["text"]
+                elif part["type"] == "image_url":
+                    # We can't easily get back the path from base64 for `respond`'s current design
+                    # This indicates a slight mismatch. `respond` expects paths for current images.
+                    # For now, let's assume `respond` can handle base64 if passed correctly.
+                    # Or, we modify `handle_user_input` to also pass original paths if needed by `respond`.
+                    # Let's assume `respond`'s `image_files` param can take base64 strings for now.
+                    # This is a simplification.
+                    # The `encode_image` in `respond` expects paths.
+                    # For now, we'll pass None for image_files if it's already in current_user_api_content.
+                    # This part needs careful review of how `respond` handles current images.
+                    # The `respond` function's `image_files` parameter is for new uploads.
+                    # If `current_user_api_content` already has encoded images, `respond` should use that.
+                    # The `respond` function's first two args are `message` (text) and `image_files` (paths).
+                    # We need to extract these from `current_user_api_content`.
+                    pass # Images are part of `current_user_api_content` which is passed to `messages_for_api`
+        elif isinstance(current_user_api_content, str): # Text only
+            current_message_text = current_user_api_content
+        
+        # Simplified history for `respond` (text from display)
+        # `respond` will reconstruct its own API history.
+        simplified_past_history = []
+        if len(history_list_for_display) > 1: # Exclude current turn
+            for user_disp, assistant_text in history_list_for_display[:-1]:
+                # Extract text from user_disp for simplified history
+                user_text_for_hist = user_disp
+                if isinstance(user_disp, str) and "<img src" in user_disp : # Basic check if it was image display
+                     # Try to find text part if any, otherwise empty
+                    lines = user_disp.splitlines()
+                    text_lines = [line for line in lines if not line.strip().startswith("<img")]
+                    user_text_for_hist = "\n".join(text_lines).strip() if text_lines else ""
+
+                simplified_past_history.append([user_text_for_hist, assistant_text])
+
+
+        # The `respond` function's first argument is `message` (current text)
+        # and `image_files` (current image paths).
+        # We need to extract these from `current_user_api_content` if it was prepared by `handle_user_input`.
+        # For now, let's assume `respond` will get the full `current_user_api_content` via `messages_for_api`.
+        # The first two args of `respond` are for the *current* turn's text and image paths.
+
+        # Let's get current text and image paths from `current_user_api_content`
+        # This is slightly redundant as `respond` also reconstructs this, but for clarity:
+        _current_text_for_respond = ""
+        _current_image_paths_for_respond = [] # `respond` expects paths
+
+        if isinstance(current_user_api_content, list):
+            for item in current_user_api_content:
+                if item['type'] == 'text':
+                    _current_text_for_respond = item['text']
+                # We can't get paths back from base64 easily.
+                # This highlights that `respond` needs to be able to take already processed multimodal content.
+                # For now, we'll assume `respond` internally uses the `messages_for_api` which has the full content.
+                # So, we can pass `_current_text_for_respond` and `None` for image_files if images are already in API format.
+
+
+        bot_response_stream = respond(
+            message=_current_text_for_respond, # Current text
+            image_files=None, # Assume images are handled by messages_for_api construction in respond
+            history=simplified_past_history, # Past turns
+            system_message=sys_msg,
+            max_tokens=max_tok,
+            temperature=temp,
+            top_p=top_p_val,
+            frequency_penalty=freq_pen,
+            seed=seed_val,
+            provider=prov,
+            custom_api_key=api_key_val,
+            custom_model=cust_model,
+            model_search_term="", # Not directly used by respond
+            selected_model=sel_model,
+            mcp_enabled=mcp_on,
+            active_mcp_servers=active_servs,
+            mcp_interaction_mode=mcp_inter_mode
+        )
+
+        for response_chunk in bot_response_stream:
+            history_list_for_display[-1][1] = response_chunk
+            yield history_list_for_display
 
-    # Model Filtering
-    def filter_featured_models(search_query):
-        if not search_query:
-            return gr.update(choices=featured_models_list_data)
-        filtered = [m for m in featured_models_list_data if search_query.lower() in m.lower()]
-        return gr.update(choices=filtered if filtered else ["No models match your search"])
 
-    model_search_box_ui.change(filter_featured_models, inputs=model_search_box_ui, outputs=featured_model_radio_ui)
+    # This state will hold the API-formatted content of the current user message
+    current_api_message_state = gr.State(None)
 
-    # Auto-select hf-inference if BYOK is empty and other provider is chosen
-    def validate_api_key_for_provider(api_key_text, current_provider):
-        if not api_key_text.strip() and current_provider != "hf-inference":
-            gr.Warning("API Key needed for non-hf-inference providers. Defaulting to hf-inference.")
-            return gr.update(value="hf-inference")
-        return current_provider # No change if key provided or hf-inference selected
+    msg.submit(
+        handle_user_input,
+        [msg, chatbot], # chatbot here is the history_list
+        [chatbot, current_api_message_state] # Update history display and current_api_message_state
+    ).then(
+        call_bot_responder,
+        [chatbot, current_api_message_state, 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, mcp_enabled_checkbox, active_mcp_servers, mcp_mode],
+        [chatbot] # Update chatbot display with streaming response
+    ).then(
+        lambda: gr.MultimodalTextbox(value={"text": "", "files": []}), # Clear MultimodalTextbox
+        None,
+        [msg]
+    )
+    
+    mcp_connect_button.click(
+        connect_mcp_server_ui_update,
+        [mcp_server_url, mcp_server_name],
+        [mcp_status, active_mcp_servers]
+    )
+    
+    model_search_box.change(fn=filter_models_ui_update, inputs=model_search_box, outputs=featured_model_radio)
+    featured_model_radio.change(fn=set_custom_model_from_radio_ui_update, inputs=featured_model_radio, outputs=custom_model_box)
+    
+    def validate_provider_ui_update(api_key, current_provider):
+        if not api_key.strip() and current_provider != "hf-inference":
+            gr.Info("No API key provided. Defaulting to 'hf-inference' provider.")
+            return gr.Radio(value="hf-inference") # Update provider_radio
+        return gr.Radio(value=current_provider) # No change needed or keep current
 
-    byok_textbox_ui.change(validate_api_key_for_provider, inputs=[byok_textbox_ui, provider_radio_ui], outputs=provider_radio_ui)
-    provider_radio_ui.change(validate_api_key_for_provider, inputs=[byok_textbox_ui, provider_radio_ui], outputs=provider_radio_ui)
+    byok_textbox.change(fn=validate_provider_ui_update, inputs=[byok_textbox, provider_radio], outputs=provider_radio)
+    provider_radio.change(fn=validate_provider_ui_update, inputs=[byok_textbox, provider_radio], outputs=provider_radio)
 
+print("Gradio interface initialized.")
 
 if __name__ == "__main__":
-    print("Launching Gradio demo...")
-    demo.queue().launch(debug=True, show_api=False) # mcp_server=False as this is a client app
+    print("Launching the demo application.")
+    demo.queue().launch(show_api=False, debug=False) # mcp_server=False as this is a client