app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os
import json
import base64
from PIL import Image
import io
import atexit

# 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")
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:
        print("No image path provided")
        return None
    
    try:
        print(f"Encoding image from path: {image_path}")
        if isinstance(image_path, Image.Image):
            image = image_path
        else:
            image = Image.open(image_path)
        
        if image.mode == 'RGBA':
            image = image.convert('RGB')
        
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")
        print("Image encoded successfully")
        return img_str
    except Exception as e:
        print(f"Error encoding image: {e}")
        return None

# Modified respond function
def respond(
    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,
    temperature,
    top_p,
    frequency_penalty,
    seed,
    provider,
    custom_api_key,
    custom_model,    
    model_search_term, # Not directly used in this function but passed by UI
    selected_model     # From radio
):
    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 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}"}
                    })
    
    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

    # --- 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

        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 = "..."


        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)

# GRADIO UI
with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
    chatbot = gr.Chatbot(
        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
    )
    
    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"]
    )
    
    with gr.Accordion("Settings", open=False):
        system_message_box = gr.Textbox(value="You are a helpful AI assistant.", label="System Prompt")
        with gr.Row():
            # ... (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)")

    # --- 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")
        
        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."

        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):
        return gr.update(choices=[m for m in models_list if search_term.lower() in m.lower()])

    # Function to set custom model from radio (remains the same)
    def set_custom_model_from_radio(selected):
        return selected # Updates custom_model_box with the selected featured model

    # 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.
    
    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", [])

        # 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": []}


    msg_input_box.submit(
        handle_submit,
        [msg_input_box, chatbot],
        [chatbot, msg_input_box] # Output to chatbot and clear msg_input_box
    )
    
    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)

# 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__":
    print("Launching the Serverless TextGen Hub demo application.")
    demo.launch(show_api=False) # show_api can be True if needed for other purposes