lmstudio_gradio.py

#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
High-Performance Chat Interface for LM Studio

This script creates a robust and efficient chat interface using Gradio,
facilitating seamless interactions with the LM Studio API. It leverages
GPU capabilities for accelerated processing and adheres to best practices
in modern Python programming. Comprehensive logging and error handling
ensure reliability and ease of maintenance.

Author: Your Name
Date: YYYY-MM-DD
"""

import gradio as gr
import httpx  # Replacing 'requests' with 'httpx' for asynchronous HTTP calls
import logging
import json
import os
import numpy as np
import torch
import asyncio

# ===========================
# Configuration and Constants
# ===========================

# Set up logging for detailed diagnostics
logging.basicConfig(
    level=logging.DEBUG,  # Set to DEBUG for more verbose output
    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s"
)
logger = logging.getLogger(__name__)

# LM Studio REST API Base URL
BASE_URL = os.getenv("LMSTUDIO_API_BASE_URL", "http://localhost:1234/v1")

# GPU Availability and Device Configuration
USE_GPU = torch.cuda.is_available()
DEVICE = torch.device("cuda" if USE_GPU else "cpu")
logger.info(f"Using device: {DEVICE}")

# Constants for Dynamic max_tokens Calculation
MODEL_MAX_TOKENS = 32768          # Model's maximum context length
AVERAGE_CHARS_PER_TOKEN = 4       # Approximate average characters per token
BUFFER_TOKENS = 2000              # Reserved tokens for system prompts and overhead
MIN_OUTPUT_TOKENS = 1000          # Minimum tokens to ensure meaningful responses

# Maximum number of embeddings to store to optimize memory usage
MAX_EMBEDDINGS = 100

# HTTPX Timeout Configuration
HTTPX_TIMEOUT = 300  # seconds, adjust as needed for longer processing times

# ===========================
# Utility Functions
# ===========================

def calculate_max_tokens(message, model_max_tokens=MODEL_MAX_TOKENS,
                        buffer=BUFFER_TOKENS, avg_chars_per_token=AVERAGE_CHARS_PER_TOKEN,
                        min_tokens=MIN_OUTPUT_TOKENS):
    """
    Calculate the maximum number of tokens for the output based on the input message length.

    Args:
        message (str): The input message from the user.
        model_max_tokens (int): The total token capacity of the model.
        buffer (int): Reserved tokens for system prompts and overhead.
        avg_chars_per_token (int): Approximate number of characters per token.
        min_tokens (int): Minimum number of tokens to ensure a meaningful response.

    Returns:
        int: The calculated maximum tokens for the output.
    """
    input_length = len(message)
    input_tokens = input_length / avg_chars_per_token
    max_tokens = model_max_tokens - int(input_tokens) - buffer
    calculated_max = max(max_tokens, min_tokens)
    logger.debug(f"Input length (chars): {input_length}, "
                 f"Estimated input tokens: {input_tokens}, "
                 f"Max tokens for output: {calculated_max}")
    return calculated_max

async def get_embeddings(text):
    """
    Retrieve embeddings for the given text from the LM Studio API.

    Args:
        text (str): The input text to generate embeddings for.

    Returns:
        list or None: The embedding vector as a list if successful, else None.
    """
    url = f"{BASE_URL}/embeddings"
    payload = {"model": "nomad_embed_text_v1_5_Q8_0", "input": text}
    logger.info(f"Requesting embeddings for input: {text[:100]}...")
    async with httpx.AsyncClient(timeout=HTTPX_TIMEOUT) as client:
        try:
            response = await client.post(
                url,
                json=payload,  # Proper JSON serialization
                headers={
                    "Content-Type": "application/json"  # Ensuring correct Content-Type
                }
            )
            logger.info(f"Embeddings response status code: {response.status_code}")
            response.raise_for_status()
            data = response.json()
            logger.debug(f"Embeddings response data: {data}")
            if "data" in data and len(data["data"]) > 0:
                embedding = np.array(data["data"][0]["embedding"])
                if USE_GPU:
                    embedding = torch.tensor(embedding, device=DEVICE).tolist()  # Convert to list for serialization
                return embedding
            else:
                logger.error("Invalid response structure for embeddings.")
                return None
        except httpx.RequestError as e:
            logger.error(f"Failed to retrieve embeddings: {e}")
            return None
        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error while retrieving embeddings: {e}")
            return None
        except json.JSONDecodeError as e:
            logger.error(f"JSON decode error: {e}")
            return None

def calculate_similarity(vec1, vec2):
    """
    Calculate the cosine similarity between two vectors using GPU acceleration.

    Args:
        vec1 (list or torch.Tensor): The first embedding vector.
        vec2 (list or torch.Tensor): The second embedding vector.

    Returns:
        float: The cosine similarity score.
    """
    if vec1 is None or vec2 is None:
        logger.warning("One or both vectors for similarity calculation are None.")
        return 0.0
    logger.debug("Calculating similarity between vectors.")
    vec1_tensor = torch.tensor(vec1, device=DEVICE) if not isinstance(vec1, torch.Tensor) else vec1.to(DEVICE)
    vec2_tensor = torch.tensor(vec2, device=DEVICE) if not isinstance(vec2, torch.Tensor) else vec2.to(DEVICE)
    similarity = torch.nn.functional.cosine_similarity(vec1_tensor.unsqueeze(0), vec2_tensor.unsqueeze(0)).item()
    logger.debug(f"Calculated similarity: {similarity}")
    return similarity

# ===========================
# API Interaction Handling
# ===========================

async def chat_with_lmstudio(messages, max_tokens):
    """
    Handle chat completions with the LM Studio API using streaming.

    Args:
        messages (list): A list of message dictionaries following OpenAI's format.
        max_tokens (int): The maximum number of tokens to generate in the response.

    Yields:
        str: Chunks of the generated response.
    """
    url = f"{BASE_URL}/chat/completions"
    payload = {
        "model": "Qwen2.5-Coder-32B-Instruct",  # Adjusted model name if necessary
        "messages": messages,
        "temperature": 0.7,
        "max_tokens": max_tokens,
        "stream": True,
    }
    logger.info(f"Sending request to chat/completions with max_tokens: {max_tokens}")
    async with httpx.AsyncClient(timeout=HTTPX_TIMEOUT) as client:
        try:
            async with client.stream("POST", url, json=payload, headers={"Content-Type": "application/json"}) as response:
                logger.info(f"chat/completions response status code: {response.status_code}")
                response.raise_for_status()
                async for line in response.aiter_lines():
                    if line:
                        try:
                            decoded_line = line.strip()
                            if decoded_line.startswith("data: "):
                                data = json.loads(decoded_line[6:])
                                logger.debug(f"Received chunk: {data}")
                                content = data.get("choices", [{}])[0].get("delta", {}).get("content", "")
                                yield content
                        except json.JSONDecodeError as e:
                            logger.error(f"JSON decode error: {e}")
        except httpx.RequestError as e:
            logger.error(f"LM Studio chat/completions request failed: {e}")
            yield "An error occurred while generating a response."
        except httpx.HTTPStatusError as e:
            logger.error(f"HTTP error during chat/completions: {e}")
            yield "An HTTP error occurred while generating a response."

# ===========================
# User Interface Implementation
# ===========================

def gradio_chat_interface():
    """
    Create and launch the Gradio Blocks interface for the chat application.
    """
    with gr.Blocks() as interface:
        gr.Markdown("# 🚀 High-Performance Chat Interface for LM Studio")

        # Chatbot component to display the conversation
        chatbot = gr.Chatbot(label="Conversation", type="messages")

        # User input textbox
        user_input = gr.Textbox(
            label="Your Message",
            placeholder="Type your message here...",
            lines=2,
            interactive=True
        )

        # File upload component for context files
        file_input = gr.File(
            label="Upload Context File (.txt)",
            type="binary",  # Correct value as per Gradio's expectations
            interactive=True
        )

        # Display relevant context based on similarity
        context_display = gr.Textbox(
            label="Relevant Context",
            interactive=False
        )

        # State to store embeddings and message history
        embeddings_state = gr.State({"embeddings": [], "messages_history": []})

        async def chat_handler(message, file, state):
            """
            Handle user input, process embeddings, retrieve context, and generate responses.

            Args:
                message (str): The user's input message.
                file (UploadedFile): The uploaded context file.
                state (dict): The current state containing embeddings and message history.

            Yields:
                list: Updated chatbot messages, new state, and context display text.
            """
            embeddings = state.get("embeddings", [])
            messages_history = state.get("messages_history", [])

            # ===========================
            # File Processing
            # ===========================
            if file:
                try:
                    file_content = file.read().decode("utf-8")
                    message += f"\n[File Content]:\n{file_content}"
                    logger.info("Successfully processed uploaded file.")
                except Exception as e:
                    error_msg = f"Error reading file: {e}"
                    logger.error(error_msg)
                    yield [error_msg, state, ""]
                    return  # Terminate the generator after yielding the error

            # ===========================
            # Embeddings Generation
            # ===========================
            user_embedding = await get_embeddings(message)
            if user_embedding is not None:
                embeddings.append(user_embedding)
                messages_history.append({"role": "user", "content": message})
                logger.info("Embeddings generated and appended to state.")
            else:
                error_msg = "Failed to generate embeddings."
                logger.error(error_msg)
                yield [error_msg, state, ""]
                return  # Terminate the generator after yielding the error

            # Limit the number of stored embeddings to optimize memory usage
            if len(embeddings) > MAX_EMBEDDINGS:
                embeddings = embeddings[-MAX_EMBEDDINGS:]
                messages_history = messages_history[-MAX_EMBEDDINGS:]

            # ===========================
            # Similarity Calculation and Context Retrieval
            # ===========================
            history = [{"role": "user", "content": message}]
            context_text = ""
            if len(embeddings) > 1:
                similarities = [
                    (calculate_similarity(user_embedding, emb), idx)
                    for idx, emb in enumerate(embeddings[:-1])
                ]
                similarities.sort(reverse=True, key=lambda x: x[0])
                top_context = similarities[:3]
                for similarity, idx in top_context:
                    context_message = messages_history[idx]
                    history.insert(0, {"role": "system", "content": context_message["content"]})
                    context_text += f"Context: {context_message['content'][:100]}...\n"
                logger.info("Relevant context retrieved based on similarity.")

            # ===========================
            # Dynamic max_tokens Calculation
            # ===========================
            max_tokens = calculate_max_tokens(message)
            logger.info(f"Calculated max_tokens for output: {max_tokens}")

            # ===========================
            # Chat with LM Studio API
            # ===========================
            response = ""
            try:
                async for chunk in chat_with_lmstudio(history, max_tokens):
                    response += chunk
                    # Ensure response is a string
                    if not isinstance(response, str):
                        response = str(response)
                    # Handle empty response
                    if not response.strip():
                        response = "Sorry, I couldn't process your request."

                    # Update chatbot in real-time with partial responses
                    updated_chat = chatbot.value.copy()
                    updated_chat.append({"role": "user", "content": message})
                    updated_chat.append({"role": "assistant", "content": response})
                    logger.debug(f"Updated Chat: {updated_chat}")
                    yield [
                        updated_chat,
                        {"embeddings": embeddings, "messages_history": messages_history},
                        context_text
                    ]
                logger.info("Response generation completed.")
            except Exception as e:
                error_msg = f"An error occurred while generating a response: {e}"
                logger.error(error_msg)
                yield [error_msg, state, ""]
                return  # Terminate the generator after yielding the error

            # ===========================
            # Final State Update
            # ===========================
            messages_history.append({"role": "assistant", "content": response})
            new_state = {"embeddings": embeddings, "messages_history": messages_history}
            updated_chat = chatbot.value.copy()
            updated_chat.append({"role": "user", "content": message})
            updated_chat.append({"role": "assistant", "content": response})

            # Final yield
            try:
                logger.debug(f"Final Updated Chat: {updated_chat}")
                yield [
                    updated_chat,
                    new_state,
                    context_text
                ]
            except Exception as e:
                error_msg = f"Error updating chatbot: {e}"
                logger.error(error_msg)
                yield ["An error occurred while updating the chat.", state, ""]

        # ===========================
        # Send Button Configuration
        # ===========================
        send_button = gr.Button("Send")
        send_button.click(
            chat_handler,
            inputs=[user_input, file_input, embeddings_state],
            outputs=[chatbot, embeddings_state, context_display],
            show_progress=True
        )

    # ===========================
    # Launch the Interface
    # ===========================
    interface.launch(share=True, server_name="0.0.0.0", server_port=7860)

# ===========================
# Main Execution
# ===========================

if __name__ == "__main__":
    asyncio.run(gradio_chat_interface())