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

from smolagents import ToolCollection, CodeAgent
from smolagents.mcp_client import MCPClient as SmolMCPClient

ACCESS_TOKEN = os.getenv("HF_TOKEN")
print("Access token loaded.")

mcp_tools_collection = ToolCollection(tools=[])
mcp_client_instances = []

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
    
    # No explicit close for SmolMCPClient instances as it's not available directly
    # Rely on script termination or GC for now.
    # If you were using ToolCollection per server: tc.close() would be the way.
    print(f"Clearing {len(mcp_client_instances)} previous MCP client instance references.")
    mcp_client_instances = [] # Clear references; old objects will be GC'd if not referenced elsewhere
    
    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}")
                smol_mcp_client = SmolMCPClient(server_parameters={"url": sse_url})
                mcp_client_instances.append(smol_mcp_client)
                discovered_tools_from_server = smol_mcp_client.get_tools()
                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}.")
            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]}...")
    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 clear MCP client instance references on application exit...")
    # No explicit close called here as per previous fix
    mcp_client_instances = []
    print("MCP client instance reference cleanup finished.")

atexit.register(cleanup_mcp_client_instances_on_exit)

def encode_image(image_path):
    if not image_path: return None
    try:
        image = Image.open(image_path) if not isinstance(image_path, Image.Image) else image_path
        if image.mode == 'RGBA': image = image.convert('RGB')
        buffered = io.BytesIO()
        image.save(buffered, format="JPEG")
        return base64.b64encode(buffered.getvalue()).decode("utf-8")
    except Exception as e:
        print(f"Error encoding image {image_path}: {e}")
        return None

def respond(
    message_input_text,
    image_files_list,
    history: list[tuple[str, str]], # history will be list of (user_str_display, assistant_str_display)
    system_message,
    max_tokens,
    temperature,
    top_p,
    frequency_penalty,
    seed,
    provider,
    custom_api_key,
    custom_model,    
    model_search_term,
    selected_model
):
    global mcp_tools_collection
    print(f"Respond: Text='{message_input_text}', Images={len(image_files_list) if image_files_list else 0}")

    token_to_use = custom_api_key if custom_api_key.strip() else ACCESS_TOKEN
    hf_inference_client = InferenceClient(token=token_to_use, provider=provider)
    if seed == -1: seed = None

    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:
            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:
        for item in history: yield item # Should not happen if handle_submit filters empty
        return

    llm_messages = [{"role": "system", "content": system_message}]
    for hist_user_str, hist_assistant in history: # hist_user_str is display string
        # For LLM context, we only care about the text part of history if it was multimodal.
        # Current image handling is only for the *current* turn.
        # If you need to re-process history for multimodal context for LLM, this part needs more logic.
        # For now, assuming hist_user_str is sufficient as text context from past turns.
        if hist_user_str:
             llm_messages.append({"role": "user", "content": hist_user_str})
        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 "")})
    
    # FIX for Issue 1: 'NoneType' object has no attribute 'strip'
    model_to_use = (custom_model.strip() if custom_model else "") or selected_model
    print(f"Model selected for inference: {model_to_use}")
    
    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.")
        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):
                api_params = {
                    "model": self.model_id, "messages": agent_llm_messages, "stream": False,
                    "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} with params: {api_params}")
                completion = self.client.chat_completion(**api_params)
                
                # FIX for Issue 2 (Potential): Ensure content is not None for text responses
                if completion.choices and completion.choices[0].message and \
                   completion.choices[0].message.content is None and \
                   (not completion.choices[0].message.tool_calls or not completion.choices[0].message.tool_calls):
                    print("Warning (HFClientWrapperForAgent): Model returned None content. Setting to empty string.")
                    completion.choices[0].message.content = ""
                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, messages_constructor=lambda: llm_messages[:-1].copy()) # Prime with history

        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:
            current_query_for_agent = "Process the provided image(s) or follow related instructions."
        elif not current_query_for_agent and not image_files_list:
             current_query_for_agent = "..." # Should be caught by earlier check

        print(f"Query for CodeAgent.run: '{current_query_for_agent}' with {len(llm_messages)-1} history messages for priming.")
        try:
            agent_final_text_response = agent.run(current_query_for_agent)
            yield agent_final_text_response
            print("Completed response generation via CodeAgent.")
        except Exception as e:
            print(f"Error during CodeAgent execution: {e}") # This will now print the actual underlying error
            yield f"Error using tools: {str(e)}" # The str(e) might be the user-facing error
        return
    else:
        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)}"

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)

with gr.Blocks(theme="Nymbo/Nymbo_Theme") as demo:
    # UserWarning for type='tuples' is known. Consider changing to type='messages' later for robustness.
    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_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_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 = [
            "meta-llama/Llama-3.2-11B-Vision-Instruct", "meta-llama/Llama-3.3-70B-Instruct",
            "meta-llama/Llama-3.1-70B-Instruct", "meta-llama/Llama-3.0-70B-Instruct",
            "meta-llama/Llama-3.2-3B-Instruct", "meta-llama/Llama-3.2-1B-Instruct",
            "meta-llama/Llama-3.1-8B-Instruct", "NousResearch/Hermes-3-Llama-3.1-8B",
            "NousResearch/Nous-Hermes-2-Mixtral-8x7B-DPO", "mistralai/Mistral-Nemo-Instruct-2407",
            "mistralai/Mixtral-8x7B-Instruct-v0.1", "mistralai/Mistral-7B-Instruct-v0.3",
            "mistralai/Mistral-7B-Instruct-v0.2", "Qwen/Qwen3-235B-A22B", "Qwen/Qwen3-32B",
            "Qwen/Qwen2.5-72B-Instruct", "Qwen/Qwen2.5-3B-Instruct", "Qwen/Qwen2.5-0.5B-Instruct",
            "Qwen/QwQ-32B", "Qwen/Qwen2.5-Coder-32B-Instruct", "microsoft/Phi-3.5-mini-instruct",
            "microsoft/Phi-3-mini-128k-instruct", "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)")

    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)
        )
        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([])
                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)
                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. 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)

    def filter_models(search_term):
        return gr.update(choices=[m for m in models_list if search_term.lower() in m.lower()])
    def set_custom_model_from_radio(selected):
        return selected

    def handle_submit(msg_content_dict, current_chat_history):
        text = msg_content_dict.get("text", "").strip()
        files = msg_content_dict.get("files", []) # list of file paths

        if not text and not files: # Skip if both are empty
            print("Skipping empty submission from multimodal textbox.")
            # Yield current history to prevent Gradio from complaining about no output
            yield current_chat_history, {"text": "", "files": []} # Clear input
            return

        # FIX for Issue 4: Pydantic FileMessage error by ensuring user part of history is a string
        user_display_parts = []
        if text:
            user_display_parts.append(text)
        if files:
            for f_path in files:
                base_name = os.path.basename(f_path) if f_path else "file"
                f_path_str = f_path if f_path else ""
                user_display_parts.append(f"\n![{base_name}]({f_path_str})")
        user_display_message_for_chatbot = " ".join(user_display_parts).strip()
        
        current_chat_history.append([user_display_message_for_chatbot, None])
        
        # Prepare history for respond function (ensure user part is string)
        history_for_respond = []
        for user_h, assistant_h in current_chat_history[:-1]: # History before current turn
             history_for_respond.append((str(user_h) if user_h is not None else "", assistant_h))


        assistant_response_accumulator = ""
        for streamed_chunk in respond(
            text, files, 
            history_for_respond,
            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
            yield current_chat_history, {"text": "", "files": []}
    
    msg_input_box.submit(
        handle_submit,
        [msg_input_box, chatbot],
        [chatbot, 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_mcp_tools(DEFAULT_MCP_SERVERS) # Load defaults on startup
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)