app.py

#!/usr/bin/env python3
"""
Gradio Interface for Multimodal Chat with SSH Tunnel Keepalive and API Fallback

This application provides a Gradio web interface for multimodal chat with a
local vLLM model. It establishes an SSH tunnel to a local vLLM server and
provides fallback to Hyperbolic API if that server is unavailable.
"""

import os
import time
import threading
import logging
import base64
import json
from io import BytesIO
import gradio as gr
from openai import OpenAI
from ssh_tunneler import SSHTunnel

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('app')

# Get environment variables
SSH_HOST = os.environ.get('SSH_HOST')
SSH_PORT = int(os.environ.get('SSH_PORT', 22))
SSH_USERNAME = os.environ.get('SSH_USERNAME')
SSH_PASSWORD = os.environ.get('SSH_PASSWORD')
REMOTE_PORT = int(os.environ.get('REMOTE_PORT', 8000))  # vLLM API port on remote machine
LOCAL_PORT = int(os.environ.get('LOCAL_PORT', 8020))  # Local forwarded port
VLLM_MODEL = os.environ.get('MODEL_NAME', 'google/gemma-3-27b-it')
HYPERBOLIC_KEY = os.environ.get('HYPERBOLIC_XYZ_KEY')
FALLBACK_MODEL = 'Qwen/Qwen2.5-VL-72B-Instruct'  # Fallback model at Hyperbolic

# Set the maximum number of concurrent API calls before queuing
MAX_CONCURRENT = int(os.environ.get('MAX_CONCURRENT', 3))  # Default to 3 concurrent calls

# API endpoints
VLLM_ENDPOINT = "http://localhost:" + str(LOCAL_PORT) + "/v1"
HYPERBOLIC_ENDPOINT = "https://api.hyperbolic.xyz/v1"

# Global variables
tunnel = None
use_fallback = False  # Whether to use fallback API instead of local vLLM
tunnel_status = {"is_running": False, "message": "Initializing tunnel..."}

def start_ssh_tunnel():
    """
    Start the SSH tunnel and monitor its status.
    """
    global tunnel, use_fallback, tunnel_status
    
    if not all([SSH_HOST, SSH_USERNAME, SSH_PASSWORD]):
        logger.error("Missing SSH connection details. Falling back to Hyperbolic API.")
        use_fallback = True
        tunnel_status = {"is_running": False, "message": "Missing SSH credentials"}
        return
    
    try:
        logger.info("Starting SSH tunnel...")
        tunnel = SSHTunnel(
            ssh_host=SSH_HOST,
            ssh_port=SSH_PORT,
            username=SSH_USERNAME,
            password=SSH_PASSWORD,
            remote_port=REMOTE_PORT,
            local_port=LOCAL_PORT,
            reconnect_interval=30,
            keep_alive_interval=15
        )
        
        if tunnel.start():
            logger.info("SSH tunnel started successfully")
            use_fallback = False
            tunnel_status = {"is_running": True, "message": "Connected"}
        else:
            logger.warning("Failed to start SSH tunnel. Falling back to Hyperbolic API.")
            use_fallback = True
            tunnel_status = {"is_running": False, "message": "Connection failed"}
    
    except Exception as e:
        logger.error(f"Error starting SSH tunnel: {str(e)}")
        use_fallback = True
        tunnel_status = {"is_running": False, "message": "Connection error"}

def check_vllm_api_health():
    """
    Check if the vLLM API is actually responding by querying the /v1/models endpoint.
    
    Returns:
        tuple: (is_healthy, message)
    """
    try:
        import requests
        response = requests.get(f"{VLLM_ENDPOINT}/models", timeout=5)
        if response.status_code == 200:
            try:
                data = response.json()
                if 'data' in data and len(data['data']) > 0:
                    model_id = data['data'][0].get('id', 'Unknown model')
                    return True, f"API is healthy. Available model: {model_id}"
                else:
                    return True, "API is healthy but no models found"
            except Exception as e:
                return False, f"API returned 200 but invalid JSON: {str(e)}"
        else:
            return False, f"API returned status code: {response.status_code}"
    except Exception as e:
        return False, f"API request failed: {str(e)}"

def monitor_tunnel():
    """
    Monitor the SSH tunnel status and update the global variables.
    """
    global tunnel, use_fallback, tunnel_status
    
    logger.info("Starting tunnel monitoring thread")
    
    while True:
        try:
            if tunnel is not None:
                ssh_status = tunnel.check_status()
                
                # Check if the tunnel is running
                if ssh_status["is_running"]:
                    # Check if vLLM API is actually responding
                    is_healthy, message = check_vllm_api_health()
                    
                    if is_healthy:
                        use_fallback = False
                        tunnel_status = {
                            "is_running": True,
                            "message": f"Connected and healthy. {message}"
                        }
                    else:
                        use_fallback = True
                        tunnel_status = {
                            "is_running": False,
                            "message": "Tunnel connected but vLLM API unhealthy"
                        }
                else:
                    # Log the actual error for troubleshooting but don't expose it in the UI
                    logger.error(f"SSH tunnel disconnected: {ssh_status['error'] or 'Unknown error'}")
                    use_fallback = True
                    tunnel_status = {
                        "is_running": False,
                        "message": "Disconnected - Check server status"
                    }
            else:
                use_fallback = True
                tunnel_status = {"is_running": False, "message": "Tunnel not initialized"}
        
        except Exception as e:
            logger.error(f"Error monitoring tunnel: {str(e)}")
            use_fallback = True
            tunnel_status = {"is_running": False, "message": "Monitoring error"}
        
        time.sleep(5)  # Check every 5 seconds

def get_openai_client(use_fallback_api=None):
    """
    Create and return an OpenAI client configured for the appropriate endpoint.
    
    Args:
        use_fallback_api (bool): If True, use Hyperbolic API. If False, use local vLLM.
                                 If None, use the global use_fallback setting.
    
    Returns:
        OpenAI: Configured OpenAI client
    """
    global use_fallback
    
    # Determine which API to use
    if use_fallback_api is None:
        use_fallback_api = use_fallback
    
    if use_fallback_api:
        logger.info("Using Hyperbolic API")
        return OpenAI(
            api_key=HYPERBOLIC_KEY,
            base_url=HYPERBOLIC_ENDPOINT
        )
    else:
        logger.info("Using local vLLM API")
        return OpenAI(
            api_key="EMPTY",  # vLLM doesn't require an actual API key
            base_url=VLLM_ENDPOINT
        )

def get_model_name(use_fallback_api=None):
    """
    Return the appropriate model name based on the API being used.
    
    Args:
        use_fallback_api (bool): If True, use fallback model. If None, use the global setting.
    
    Returns:
        str: Model name
    """
    global use_fallback
    
    if use_fallback_api is None:
        use_fallback_api = use_fallback
    
    return FALLBACK_MODEL if use_fallback_api else VLLM_MODEL

def convert_files_to_base64(files):
    """
    Convert uploaded files to base64 strings.
    
    Args:
        files (list): List of file paths
    
    Returns:
        list: List of base64-encoded strings
    """
    base64_images = []
    for file in files:
        with open(file, "rb") as image_file:
            # Read image data and encode to base64
            base64_data = base64.b64encode(image_file.read()).decode("utf-8")
            base64_images.append(base64_data)
    return base64_images

def process_chat(message_dict, history):
    """
    Process user message and send to the appropriate API.
    
    Args:
        message_dict (dict): User message containing text and files
        history (list): Chat history
    
    Returns:
        list: Updated chat history
    """
    global use_fallback
    
    text = message_dict.get("text", "")
    files = message_dict.get("files", [])
    
    # Add user message to history first
    if not history:
        history = []
    
    # Add user message to chat history
    if files:
        # For each file, add a separate user message
        for file in files:
            history.append({"role": "user", "content": (file,)})
    
    # Add text message if not empty
    if text.strip():
        history.append({"role": "user", "content": text})
    else:
        # If no text but files exist, don't add an empty message
        if not files:
            history.append({"role": "user", "content": ""})
    
    # Convert all files to base64
    base64_images = convert_files_to_base64(files)
    
    # Prepare conversation history in OpenAI format
    openai_messages = []
    
    # Convert history to OpenAI format
    for h in history:
        if h["role"] == "user":
            # Handle user messages
            if isinstance(h["content"], tuple):
                # This is a file-only message, skip for now
                continue
            else:
                # Text message
                openai_messages.append({
                    "role": "user",
                    "content": h["content"]
                })
        elif h["role"] == "assistant":
            openai_messages.append({
                "role": "assistant",
                "content": h["content"]
            })
    
    # Handle images for the last user message if needed
    if base64_images:
        # Update the last user message to include image content
        if openai_messages and openai_messages[-1]["role"] == "user":
            # Get the last message
            last_msg = openai_messages[-1]
            
            # Format for OpenAI multimodal content structure
            content_list = []
            
            # Add text if there is any
            if last_msg["content"]:
                content_list.append({"type": "text", "text": last_msg["content"]})
            
            # Add images
            for img_b64 in base64_images:
                content_list.append({
                    "type": "image_url",
                    "image_url": {
                        "url": f"data:image/jpeg;base64,{img_b64}"
                    }
                })
            
            # Replace the content with the multimodal content list
            last_msg["content"] = content_list
    
    # Try primary API first, fall back if needed
    try:
        # First try with the currently selected API (vLLM or fallback)
        client = get_openai_client()
        model = get_model_name()
        
        response = client.chat.completions.create(
            model=model,
            messages=openai_messages,
            stream=True  # Use streaming for better UX
        )
        
        # Stream the response
        assistant_message = ""
        for chunk in response:
            if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content is not None:
                assistant_message += chunk.choices[0].delta.content
                # Update in real-time
                history_with_stream = history.copy()
                history_with_stream.append({"role": "assistant", "content": assistant_message})
                yield history_with_stream
        
        # Ensure we have the final message added
        if not assistant_message:
            assistant_message = "No response received from the model."
        
        # Add assistant response to history if not already added
        if not history or history[-1]["role"] != "assistant":
            history.append({"role": "assistant", "content": assistant_message})
        
        return history
    
    except Exception as primary_error:
        logger.error(f"Primary API error: {str(primary_error)}")
        
        # If we're not already using fallback, try that
        if not use_fallback:
            try:
                logger.info("Falling back to Hyperbolic API")
                client = get_openai_client(use_fallback_api=True)
                model = get_model_name(use_fallback_api=True)
                
                response = client.chat.completions.create(
                    model=model,
                    messages=openai_messages,
                    stream=True
                )
                
                # Stream the response
                assistant_message = ""
                for chunk in response:
                    if hasattr(chunk.choices[0].delta, 'content') and chunk.choices[0].delta.content is not None:
                        assistant_message += chunk.choices[0].delta.content
                        # Update in real-time
                        history_with_stream = history.copy()
                        history_with_stream.append({"role": "assistant", "content": assistant_message})
                        yield history_with_stream
                
                # Ensure we have the final message added
                if not assistant_message:
                    assistant_message = "No response received from the fallback model."
                
                # Add assistant response to history if not already added
                if not history or history[-1]["role"] != "assistant":
                    history.append({"role": "assistant", "content": assistant_message})
                
                # Update fallback status (global already declared at function start)
                use_fallback = True
                
                return history
            
            except Exception as fallback_error:
                logger.error(f"Fallback API error: {str(fallback_error)}")
                error_msg = "Error connecting to both primary and fallback APIs."
                history.append({"role": "assistant", "content": error_msg})
                return history
        else:
            # Already using fallback, just report the error
            error_msg = "An error occurred with the model service."
            history.append({"role": "assistant", "content": error_msg})
            return history

def get_tunnel_status_message():
    """
    Return a formatted status message for display in the UI.
    """
    global tunnel_status, use_fallback, MAX_CONCURRENT
    
    api_mode = "Hyperbolic API" if use_fallback else "Local vLLM API"
    model = get_model_name()
    
    status_color = "🟢" if (tunnel_status["is_running"] and not use_fallback) else "🔴"
    status_text = tunnel_status["message"]
    
    return f"{status_color} Tunnel Status: {status_text}\nCurrent API: {api_mode}\nCurrent Model: {model}\nConcurrent Requests: {MAX_CONCURRENT}"

def toggle_api():
    """
    Toggle between local vLLM and Hyperbolic API.
    """
    global use_fallback
    use_fallback = not use_fallback
    
    api_mode = "Hyperbolic API" if use_fallback else "Local vLLM API"
    model = get_model_name()
    
    return f"Switched to {api_mode} using {model}"

def update_concurrency(new_value):
    """
    Update the MAX_CONCURRENT value.
    
    Args:
        new_value (str): New concurrency value as string
    
    Returns:
        str: Status message
    """
    global MAX_CONCURRENT
    try:
        value = int(new_value)
        if value < 1:
            return f"Error: Concurrency must be at least 1. Keeping current value: {MAX_CONCURRENT}"
        
        MAX_CONCURRENT = value
        # Note: This only updates the value for future event handlers
        # Existing event handlers keep their original concurrency_limit
        # A page refresh is needed for this to fully take effect
        return f"Concurrency updated to {MAX_CONCURRENT}. You may need to refresh the page for all changes to take effect."
    except ValueError:
        return f"Error: Invalid number. Keeping current value: {MAX_CONCURRENT}"

# Start the SSH tunnel in a background thread
if __name__ == "__main__":
    # Start the SSH tunnel
    start_ssh_tunnel()
    
    # Start the monitoring thread
    monitor_thread = threading.Thread(target=monitor_tunnel, daemon=True)
    monitor_thread.start()
    
    # Create Gradio application with Blocks for more control
    with gr.Blocks(theme="soft") as demo:
        gr.Markdown("# Multimodal Chat Interface")
        
        # Create chatbot component with message type
        chatbot = gr.Chatbot(
            label="Conversation",
            type="messages",
            show_copy_button=True,
            avatar_images=("👤", "🗣️"),
            height=400
        )
        
        # Create multimodal textbox for input
        with gr.Row():
            textbox = gr.MultimodalTextbox(
                file_types=["image", "video"],
                file_count="multiple",
                placeholder="Type your message here and/or upload images...",
                label="Message",
                show_label=False,
                scale=9
            )
            submit_btn = gr.Button("Send", size="sm", scale=1)
        
        # Clear button
        clear_btn = gr.Button("Clear Chat")
        
        # Set up submit event chain with concurrency limit
        submit_event = textbox.submit(
            fn=process_chat,
            inputs=[textbox, chatbot],
            outputs=chatbot,
            concurrency_limit=MAX_CONCURRENT  # Set concurrency limit for this event
        ).then(
            fn=lambda: {"text": "", "files": []},
            inputs=None,
            outputs=textbox
        )
        
        # Connect the submit button to the same functions with same concurrency limit
        submit_btn.click(
            fn=process_chat,
            inputs=[textbox, chatbot],
            outputs=chatbot,
            concurrency_limit=MAX_CONCURRENT  # Set concurrency limit for this event
        ).then(
            fn=lambda: {"text": "", "files": []},
            inputs=None,
            outputs=textbox
        )
        
        # Set up clear button
        clear_btn.click(lambda: [], None, chatbot)

        # Load example images if they exist
        examples = []
        
        # Define example images with paths
        example_images = {
            "dog_pic.jpg": "What breed is this?",
            "ghostimg.png": "What's in this image?",
            "newspaper.png": "Provide a python list of dicts about everything on this page."
        }
        
        # Check each image and add to examples if it exists
        for img_name, prompt_text in example_images.items():
            img_path = os.path.join(os.path.dirname(__file__), img_name)
            if os.path.exists(img_path):
                examples.append([{"text": prompt_text, "files": [img_path]}])
        
        # Add examples if we have any
        if examples:
            gr.Examples(
                examples=examples,
                inputs=textbox
            )
        
        # Add status display
        status_text = gr.Textbox(
            label="Tunnel and API Status",
            value=get_tunnel_status_message(),
            interactive=False
        )
        
        # Refresh status button and toggle API button
        with gr.Row():
            refresh_btn = gr.Button("Refresh Status")

        # Set up refresh status button
        refresh_btn.click(
            fn=get_tunnel_status_message,
            inputs=None,
            outputs=status_text
        )
        
        # Just load the initial status without auto-refresh
        demo.load(
            fn=get_tunnel_status_message,
            inputs=None,
            outputs=status_text
        )
    
    # Launch the interface with the specified concurrency setting
    demo.queue(default_concurrency_limit=MAX_CONCURRENT)
    demo.launch()