app.py

import os
import base64
import gradio as gr
import requests
import json
from io import BytesIO
from PIL import Image
import time

# Get API key from environment variable for security
OPENROUTER_API_KEY = os.environ.get("OPENROUTER_API_KEY", "")


# Model information
free_models = [
    ("Google: Gemini Pro 2.0 Experimental (free)", "google/gemini-2.0-pro-exp-02-05:free", 0, 0, 2000000),
    ("Google: Gemini 2.0 Flash Thinking Experimental 01-21 (free)", "google/gemini-2.0-flash-thinking-exp:free", 0, 0, 1048576),
    ("Google: Gemini Flash 2.0 Experimental (free)", "google/gemini-2.0-flash-exp:free", 0, 0, 1048576),
    ("Google: Gemini Pro 2.5 Experimental (free)", "google/gemini-2.5-pro-exp-03-25:free", 0, 0, 1000000),
    ("Google: Gemini Flash 1.5 8B Experimental", "google/gemini-flash-1.5-8b-exp", 0, 0, 1000000),
    ("DeepSeek: DeepSeek R1 Zero (free)", "deepseek/deepseek-r1-zero:free", 0, 0, 163840),
    ("DeepSeek: R1 (free)", "deepseek/deepseek-r1:free", 0, 0, 163840),
    ("DeepSeek: DeepSeek V3 Base (free)", "deepseek/deepseek-v3-base:free", 0, 0, 131072),
    ("DeepSeek: DeepSeek V3 0324 (free)", "deepseek/deepseek-chat-v3-0324:free", 0, 0, 131072),
    ("Google: Gemma 3 4B (free)", "google/gemma-3-4b-it:free", 0, 0, 131072),
    ("Google: Gemma 3 12B (free)", "google/gemma-3-12b-it:free", 0, 0, 131072),
    ("Nous: DeepHermes 3 Llama 3 8B Preview (free)", "nousresearch/deephermes-3-llama-3-8b-preview:free", 0, 0, 131072),
    ("Qwen: Qwen2.5 VL 72B Instruct (free)", "qwen/qwen2.5-vl-72b-instruct:free", 0, 0, 131072),
    ("DeepSeek: DeepSeek V3 (free)", "deepseek/deepseek-chat:free", 0, 0, 131072),
    ("NVIDIA: Llama 3.1 Nemotron 70B Instruct (free)", "nvidia/llama-3.1-nemotron-70b-instruct:free", 0, 0, 131072),
    ("Meta: Llama 3.2 1B Instruct (free)", "meta-llama/llama-3.2-1b-instruct:free", 0, 0, 131072),
    ("Meta: Llama 3.2 11B Vision Instruct (free)", "meta-llama/llama-3.2-11b-vision-instruct:free", 0, 0, 131072),
    ("Meta: Llama 3.1 8B Instruct (free)", "meta-llama/llama-3.1-8b-instruct:free", 0, 0, 131072),
    ("Mistral: Mistral Nemo (free)", "mistralai/mistral-nemo:free", 0, 0, 128000),
    ("Mistral: Mistral Small 3.1 24B (free)", "mistralai/mistral-small-3.1-24b-instruct:free", 0, 0, 96000),
    ("Google: Gemma 3 27B (free)", "google/gemma-3-27b-it:free", 0, 0, 96000),
    ("Qwen: Qwen2.5 VL 3B Instruct (free)", "qwen/qwen2.5-vl-3b-instruct:free", 0, 0, 64000),
    ("DeepSeek: R1 Distill Qwen 14B (free)", "deepseek/deepseek-r1-distill-qwen-14b:free", 0, 0, 64000),
    ("Qwen: Qwen2.5-VL 7B Instruct (free)", "qwen/qwen-2.5-vl-7b-instruct:free", 0, 0, 64000),
    ("Google: LearnLM 1.5 Pro Experimental (free)", "google/learnlm-1.5-pro-experimental:free", 0, 0, 40960),
    ("Qwen: QwQ 32B (free)", "qwen/qwq-32b:free", 0, 0, 40000),
    ("Google: Gemini 2.0 Flash Thinking Experimental (free)", "google/gemini-2.0-flash-thinking-exp-1219:free", 0, 0, 40000),
    ("Bytedance: UI-TARS 72B (free)", "bytedance-research/ui-tars-72b:free", 0, 0, 32768),
    ("Qwerky 72b (free)", "featherless/qwerky-72b:free", 0, 0, 32768),
    ("OlympicCoder 7B (free)", "open-r1/olympiccoder-7b:free", 0, 0, 32768),
    ("OlympicCoder 32B (free)", "open-r1/olympiccoder-32b:free", 0, 0, 32768),
    ("Google: Gemma 3 1B (free)", "google/gemma-3-1b-it:free", 0, 0, 32768),
    ("Reka: Flash 3 (free)", "rekaai/reka-flash-3:free", 0, 0, 32768),
    ("Dolphin3.0 R1 Mistral 24B (free)", "cognitivecomputations/dolphin3.0-r1-mistral-24b:free", 0, 0, 32768),
    ("Dolphin3.0 Mistral 24B (free)", "cognitivecomputations/dolphin3.0-mistral-24b:free", 0, 0, 32768),
    ("Mistral: Mistral Small 3 (free)", "mistralai/mistral-small-24b-instruct-2501:free", 0, 0, 32768),
    ("Qwen2.5 Coder 32B Instruct (free)", "qwen/qwen-2.5-coder-32b-instruct:free", 0, 0, 32768),
    ("Qwen2.5 72B Instruct (free)", "qwen/qwen-2.5-72b-instruct:free", 0, 0, 32768),
    ("Meta: Llama 3.2 3B Instruct (free)", "meta-llama/llama-3.2-3b-instruct:free", 0, 0, 20000),
    ("Qwen: QwQ 32B Preview (free)", "qwen/qwq-32b-preview:free", 0, 0, 16384),
    ("DeepSeek: R1 Distill Qwen 32B (free)", "deepseek/deepseek-r1-distill-qwen-32b:free", 0, 0, 16000),
    ("Qwen: Qwen2.5 VL 32B Instruct (free)", "qwen/qwen2.5-vl-32b-instruct:free", 0, 0, 8192),
    ("Moonshot AI: Moonlight 16B A3B Instruct (free)", "moonshotai/moonlight-16b-a3b-instruct:free", 0, 0, 8192),
    ("DeepSeek: R1 Distill Llama 70B (free)", "deepseek/deepseek-r1-distill-llama-70b:free", 0, 0, 8192),
    ("Qwen 2 7B Instruct (free)", "qwen/qwen-2-7b-instruct:free", 0, 0, 8192),
    ("Google: Gemma 2 9B (free)", "google/gemma-2-9b-it:free", 0, 0, 8192),
    ("Mistral: Mistral 7B Instruct (free)", "mistralai/mistral-7b-instruct:free", 0, 0, 8192),
    ("Microsoft: Phi-3 Mini 128K Instruct (free)", "microsoft/phi-3-mini-128k-instruct:free", 0, 0, 8192),
    ("Microsoft: Phi-3 Medium 128K Instruct (free)", "microsoft/phi-3-medium-128k-instruct:free", 0, 0, 8192),
    ("Meta: Llama 3 8B Instruct (free)", "meta-llama/llama-3-8b-instruct:free", 0, 0, 8192),
    ("OpenChat 3.5 7B (free)", "openchat/openchat-7b:free", 0, 0, 8192),
    ("Meta: Llama 3.3 70B Instruct (free)", "meta-llama/llama-3.3-70b-instruct:free", 0, 0, 8000),
    ("AllenAI: Molmo 7B D (free)", "allenai/molmo-7b-d:free", 0, 0, 4096),
    ("Rogue Rose 103B v0.2 (free)", "sophosympatheia/rogue-rose-103b-v0.2:free", 0, 0, 4096),
    ("Toppy M 7B (free)", "undi95/toppy-m-7b:free", 0, 0, 4096),
    ("Hugging Face: Zephyr 7B (free)", "huggingfaceh4/zephyr-7b-beta:free", 0, 0, 4096),
    ("MythoMax 13B (free)", "gryphe/mythomax-l2-13b:free", 0, 0, 4096),
]

# Helper functions
def encode_image(image):
    """Convert PIL Image to base64 string"""
    buffered = BytesIO()
    image.save(buffered, format="JPEG")
    return base64.b64encode(buffered.getvalue()).decode("utf-8")

def encode_file(file_path):
    """Convert text file to string"""
    try:
        with open(file_path, 'r', encoding='utf-8') as file:
            return file.read()
    except Exception as e:
        return f"Error reading file: {str(e)}"

def process_api_call(messages, model_id, temperature=0.7, top_p=1.0, max_tokens=1000, stream=False):
    """Make API call to OpenRouter"""
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {OPENROUTER_API_KEY}",
        "HTTP-Referer": "https://huggingface.co/spaces",
    }
    
    url = "https://openrouter.ai/api/v1/chat/completions"
    
    data = {
        "model": model_id,
        "messages": messages,
        "stream": stream,
        "temperature": temperature,
        "top_p": top_p,
        "max_tokens": max_tokens
    }
    
    return requests.post(url, headers=headers, json=data, stream=stream)

def update_conversation(message, chat_history, model_choice, uploaded_image=None, uploaded_file=None, 
                      temp=0.7, top_p=1.0, max_tokens=1000, stream_response=False):
    """Update conversation with new message"""
    # Get model ID from model_choice
    model_id = None
    for name, model_id_value, *_ in free_models:
        if name == model_choice or model_id_value == model_choice:
            model_id = model_id_value
            break
    
    if not model_id:
        # Fallback to a default model
        model_id = "google/gemini-2.0-pro-exp-02-05:free"
    
    # Build messages array from chat history
    messages = []
    for msg in chat_history:
        if isinstance(msg, dict):
            messages.append(msg)
        elif isinstance(msg, tuple) and len(msg) == 2:
            # Handle legacy tuple format
            user_msg, ai_msg = msg
            messages.append({"role": "user", "content": user_msg})
            messages.append({"role": "assistant", "content": ai_msg})
    
    # Prepare the new user message
    content = message
    
    # Handle file attachment
    if uploaded_file:
        file_content = encode_file(uploaded_file)
        content = f"{message}\n\nFile content:\n```\n{file_content}\n```"
    
    # Handle image
    if uploaded_image:
        base64_image = encode_image(uploaded_image)
        image_content = [
            {"type": "text", "text": content},
            {
                "type": "image_url",
                "image_url": {
                    "url": f"data:image/jpeg;base64,{base64_image}"
                }
            }
        ]
        messages.append({"role": "user", "content": image_content})
    else:
        messages.append({"role": "user", "content": content})
    
    # Add message to chat history
    user_message = {"role": "user", "content": content}
    assistant_message = {"role": "assistant", "content": ""}
    chat_history.append(user_message)
    chat_history.append(assistant_message)
    
    try:
        if stream_response:
            # Handle streaming response
            response = process_api_call(messages, model_id, temp, top_p, max_tokens, stream=True)
            
            full_response = ""
            buffer = ""
            
            for chunk in response.iter_content(chunk_size=1024, decode_unicode=False):
                if chunk:
                    buffer += chunk.decode('utf-8')
                    
                    while True:
                        line_end = buffer.find('\n')
                        if line_end == -1:
                            break
                            
                        line = buffer[:line_end].strip()
                        buffer = buffer[line_end + 1:]
                        
                        if line.startswith('data: '):
                            data = line[6:]
                            if data == '[DONE]':
                                break
                                
                            try:
                                data_obj = json.loads(data)
                                delta_content = data_obj["choices"][0]["delta"].get("content", "")
                                if delta_content:
                                    full_response += delta_content
                                    # Update the assistant message
                                    chat_history[-1]["content"] = full_response
                                    yield chat_history
                            except json.JSONDecodeError:
                                pass
        else:
            # Handle non-streaming response
            response = process_api_call(messages, model_id, temp, top_p, max_tokens, stream=False)
            response.raise_for_status()
            result = response.json()
            
            reply = result.get("choices", [{}])[0].get("message", {}).get("content", "No response")
            chat_history[-1]["content"] = reply
            yield chat_history
    
    except Exception as e:
        error_msg = f"Error: {str(e)}"
        chat_history[-1]["content"] = error_msg
        yield chat_history

# Create simpler UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("# 🔆 CrispChat - OpenRouter AI Models")
    
    with gr.Row():
        with gr.Column(scale=4):
            chatbot = gr.Chatbot(
                height=500,
                show_copy_button=True,
                show_share_button=False,
                layout="bubble",
                avatar_images=("👤", "🤖"),
                type="messages"
            )
            
            with gr.Row():
                user_message = gr.Textbox(
                    placeholder="Type your message here...",
                    show_label=False,
                    lines=3
                )
                
            with gr.Row():
                with gr.Column(scale=1):
                    image_upload = gr.Image(
                        type="pil", 
                        label="Upload Image",
                        show_label=True
                    )
                
                with gr.Column(scale=1):
                    file_upload = gr.File(
                        label="Upload Text File",
                        file_types=[".txt", ".md", ".py", ".js", ".html", ".css", ".json"]
                    )
                
                with gr.Column(scale=1):
                    submit_btn = gr.Button("Send", variant="primary")
        
        with gr.Column(scale=2):
            with gr.Accordion("Model Settings", open=True):
                model_selector = gr.Dropdown(
                    choices=[name for name, _ in free_models],
                    value=free_models[0][0],
                    label="Select Model"
                )
                
                temperature = gr.Slider(
                    minimum=0.1, 
                    maximum=2.0, 
                    value=0.7, 
                    step=0.1, 
                    label="Temperature"
                )
                
                top_p = gr.Slider(
                    minimum=0.1, 
                    maximum=1.0, 
                    value=1.0, 
                    step=0.1, 
                    label="Top P"
                )
                
                max_tokens = gr.Slider(
                    minimum=100, 
                    maximum=4000, 
                    value=1000, 
                    step=100, 
                    label="Max Tokens"
                )
                
                streaming = gr.Checkbox(
                    label="Enable Streaming", 
                    value=True
                )
                
                clear_btn = gr.Button("Clear Chat")
    
    # Set up event handlers
    msg_submit_event = user_message.submit(
        fn=update_conversation,
        inputs=[
            user_message, 
            chatbot, 
            model_selector, 
            image_upload, 
            file_upload,
            temperature,
            top_p,
            max_tokens,
            streaming
        ],
        outputs=chatbot
    )
    
    btn_submit_event = submit_btn.click(
        fn=update_conversation,
        inputs=[
            user_message, 
            chatbot, 
            model_selector, 
            image_upload, 
            file_upload,
            temperature,
            top_p,
            max_tokens,
            streaming
        ],
        outputs=chatbot
    )
    
    # Clear chat
    clear_btn.click(
        fn=lambda: [], 
        outputs=[chatbot]
    )
    
    # Clear input after submission
    msg_submit_event.then(
        fn=lambda: "", 
        outputs=[user_message]
    )
    
    btn_submit_event.then(
        fn=lambda: "", 
        outputs=[user_message]
    )

# Mount FastAPI for external access
from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()

class GenerateRequest(BaseModel):
    message: str
    model: str = None
    image_data: str = None

@app.post("/api/generate")
async def api_generate(request: GenerateRequest):
    """API endpoint for generating responses"""
    try:
        # Process request
        messages = [{"role": "user", "content": request.message}]
        
        # Handle image if provided
        if request.image_data:
            try:
                image_bytes = base64.b64decode(request.image_data)
                image = Image.open(BytesIO(image_bytes))
                base64_image = encode_image(image)
                
                messages = [{
                    "role": "user", 
                    "content": [
                        {"type": "text", "text": request.message},
                        {
                            "type": "image_url",
                            "image_url": {
                                "url": f"data:image/jpeg;base64,{base64_image}"
                            }
                        }
                    ]
                }]
            except Exception as e:
                return {"error": f"Image processing error: {str(e)}"}
        
        # Get model
        model_id = request.model or free_models[0][1]
        
        # Make API call
        response = process_api_call(messages, model_id, stream=False)
        response.raise_for_status()
        result = response.json()
        
        reply = result.get("choices", [{}])[0].get("message", {}).get("content", "No response")
        return {"response": reply}
        
    except Exception as e:
        return {"error": f"Error: {str(e)}"}

# Mount Gradio app
app = gr.mount_gradio_app(app, demo, path="/")

# Launch the app
if __name__ == "__main__":
    demo.launch()