Update app.py

app.py

@@ -10,6 +10,7 @@ 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),
@@ -71,33 +72,7 @@ free_models = [
     ("MythoMax 13B (free)", "gryphe/mythomax-l2-13b:free", 0, 0, 4096),
 ]
 
-# Filter for vision models
-vision_model_ids = [
-    "meta-llama/llama-3.2-11b-vision-instruct:free",
-    "qwen/qwen2.5-vl-72b-instruct:free", 
-    "qwen/qwen2.5-vl-3b-instruct:free",
-    "qwen/qwen2.5-vl-32b-instruct:free",
-    "qwen/qwen-2.5-vl-7b-instruct:free",
-    "google/gemini-2.0-pro-exp-02-05:free",
-    "google/gemini-2.5-pro-exp-03-25:free"
-]
-
-# Format model names to include context size
-def format_model_name(name, context_size):
-    if context_size >= 1000000:
-        context_str = f"{context_size/1000000:.1f}M tokens"
-    else:
-        context_str = f"{context_size/1000:.0f}K tokens"
-    return f"{name} ({context_str})"
-
-# Prefilter vision models
-vision_models = [(format_model_name(name, context_size), model_id, context_size) 
-                for name, model_id, _, _, context_size in free_models 
-                if model_id in vision_model_ids]
-
-text_models = [(format_model_name(name, context_size), model_id, context_size) 
-              for name, model_id, _, _, context_size in free_models]
-
+# Helper functions
 def encode_image(image):
     """Convert PIL Image to base64 string"""
     buffered = BytesIO()
@@ -112,79 +87,16 @@ def encode_file(file_path):
     except Exception as e:
         return f"Error reading file: {str(e)}"
 
-def process_message_stream(message, chat_history, model_name, uploaded_image=None, uploaded_file=None, 
-                           temperature=0.7, top_p=1.0, max_tokens=None, stream=True):
-    """Process message and stream the model response"""
-    # Extract model_id from the display name
-    model_id = model_name.split(' ')[1] if len(model_name.split(' ')) > 1 else model_name
-    
-    # Check if API key is set
-    if not OPENROUTER_API_KEY:
-        yield "Please set your OpenRouter API key in the environment variables.", chat_history
-        return
-    
-    # Setup headers and URL
+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",  # Replace with your actual space URL in production
+        "HTTP-Referer": "https://huggingface.co/spaces",
     }
     
     url = "https://openrouter.ai/api/v1/chat/completions"
     
-    # Build message content
-    messages = []
-    
-    # Add chat history
-    for item in chat_history:
-        if isinstance(item, tuple):
-            # Old format compatibility
-            human_msg, ai_msg = item
-            messages.append({"role": "user", "content": human_msg})
-            messages.append({"role": "assistant", "content": ai_msg})
-        else:
-            # New message format
-            messages.append(item)
-    
-    # Add current message with any attachments
-    if uploaded_image:
-        # Image processing for vision models
-        base64_image = encode_image(uploaded_image)
-        content = [
-            {"type": "text", "text": message}
-        ]
-        
-        # Add text from file if provided
-        if uploaded_file:
-            file_content = encode_file(uploaded_file)
-            content[0]["text"] = f"{message}\n\nFile content:\n```\n{file_content}\n```"
-            
-        # Add image
-        content.append({
-            "type": "image_url",
-            "image_url": {
-                "url": f"data:image/jpeg;base64,{base64_image}"
-            }
-        })
-        
-        messages.append({"role": "user", "content": content})
-    else:
-        if uploaded_file:
-            file_content = encode_file(uploaded_file)
-            content = f"{message}\n\nFile content:\n```\n{file_content}\n```"
-            messages.append({"role": "user", "content": content})
-        else:
-            messages.append({"role": "user", "content": message})
-    
-    # Get context length for the model
-    context_length = next((context for _, model_id, context in text_models if model_id == model_id), 4096)
-    
-    # Calculate default max tokens if not specified
-    if not max_tokens:
-        # Use 25% of context length as a reasonable default
-        max_tokens = min(4000, int(context_length * 0.25))
-    
-    # Build request data
     data = {
         "model": model_id,
         "messages": messages,
@@ -194,298 +106,244 @@ def process_message_stream(message, chat_history, model_name, uploaded_image=Non
         "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:
-        # Create a new message pair in the chat history
-        user_msg = {"role": "user", "content": message}
-        ai_msg = {"role": "assistant", "content": ""}
-        chat_history.append(user_msg)
-        chat_history.append(ai_msg)
-        
-        full_response = ""
-        
-        if stream:
-            # Make streaming API call
-            with requests.post(url, headers=headers, json=data, stream=True) as response:
-                response.raise_for_status()
-                buffer = ""
-                
-                for chunk in response.iter_content(chunk_size=1024, decode_unicode=False):
-                    if chunk:
-                        buffer += chunk.decode('utf-8')
+        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:]
                         
-                        while True:
-                            line_end = buffer.find('\n')
-                            if line_end == -1:
+                        if line.startswith('data: '):
+                            data = line[6:]
+                            if data == '[DONE]':
                                 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 last assistant message
-                                        chat_history[-1]["content"] = full_response
-                                        yield chat_history
-                                except json.JSONDecodeError:
-                                    pass
+                            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:
-            # Non-streaming API call
-            response = requests.post(url, headers=headers, json=data)
+            # 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()
-            full_response = result.get("choices", [{}])[0].get("message", {}).get("content", "No response")
-            chat_history[-1]["content"] = full_response
+            
+            reply = result.get("choices", [{}])[0].get("message", {}).get("content", "No response")
+            chat_history[-1]["content"] = reply
             yield chat_history
-        
-        return chat_history
     
     except Exception as e:
         error_msg = f"Error: {str(e)}"
         chat_history[-1]["content"] = error_msg
         yield chat_history
 
-# Create a nice CSS theme
-css = """
-.gradio-container {
-    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
-}
-.chat-message {
-    padding: 15px;
-    border-radius: 10px;
-    margin-bottom: 10px;
-}
-.user-message {
-    background-color: #f0f4f8;
-}
-.assistant-message {
-    background-color: #e9f5ff;
-}
-#chat-container {
-    height: 600px;
-    overflow-y: auto;
-}
-#chat-input {
-    min-height: 120px;
-    border-radius: 8px;
-    padding: 10px;
-}
-#model-select-container {
-    border-radius: 8px;
-    padding: 15px;
-    background-color: #f8fafc;
-}
-.app-header {
-    text-align: center;
-    margin-bottom: 20px;
-}
-.app-header h1 {
-    font-weight: 700;
-    color: #2C3E50;
-    margin-bottom: 5px;
-}
-.app-header p {
-    color: #7F8C8D;
-    margin-top: 0;
-}
-.parameter-container {
-    background-color: #f8fafc;
-    padding: 10px;
-    border-radius: 8px;
-    margin-top: 10px;
-}
-.file-upload-container {
-    margin-top: 10px;
-}
-"""
-
-with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
-    gr.HTML("""
-    <div class="app-header">
-        <h1>🔆 CrispChat</h1>
-        <p>Chat with free OpenRouter AI models - supports text, images, and files</p>
-    </div>
-    """)
+# 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=600,
+                height=500,
                 show_copy_button=True,
                 show_share_button=False,
-                elem_id="chatbot",
                 layout="bubble",
                 avatar_images=("👤", "🤖"),
-                bubble_full_width=False,
-                type="messages"  # Use new message format
+                type="messages"
             )
             
             with gr.Row():
-                with gr.Column(scale=10):
-                    user_message = gr.Textbox(
-                        placeholder="Type your message here...",
-                        show_label=False,
-                        elem_id="chat-input",
-                        lines=3
+                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.Row():
-                        image_upload = gr.Image(
-                            type="pil", 
-                            label="Image (optional)",
-                            show_label=True,
-                            scale=1
-                        )
-                        
-                        file_upload = gr.File(
-                            label="Text File (optional)",
-                            file_types=[".txt", ".md", ".py", ".js", ".html", ".css", ".json"],
-                            scale=1
-                        )
-                        
-                        submit_btn = gr.Button("Send", scale=1, variant="primary")
-            
-        with gr.Column(scale=2):
-            with gr.Accordion("Model Selection", open=True):
-                using_vision = gr.Checkbox(label="Using image", value=False)
                 
+                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 text_models],
-                    value=text_models[0][0],
-                    label="Select Model",
-                    elem_id="model-selector"
+                    choices=[name for name, _ in free_models],
+                    value=free_models[0][0],
+                    label="Select Model"
                 )
                 
-                context_info = gr.Markdown(value=f"Context: {text_models[0][2]:,} tokens")
+                temperature = gr.Slider(
+                    minimum=0.1, 
+                    maximum=2.0, 
+                    value=0.7, 
+                    step=0.1, 
+                    label="Temperature"
+                )
                 
-            with gr.Accordion("Parameters", open=False):
-                with gr.Group():
-                    temperature = gr.Slider(
-                        minimum=0.0, 
-                        maximum=2.0, 
-                        value=0.7, 
-                        step=0.1, 
-                        label="Temperature",
-                        info="Higher = more creative, Lower = more deterministic"
-                    )
-                    
-                    top_p = gr.Slider(
-                        minimum=0.1, 
-                        maximum=1.0, 
-                        value=1.0, 
-                        step=0.1, 
-                        label="Top P",
-                        info="Controls token diversity"
-                    )
-                    
-                    max_tokens = gr.Slider(
-                        minimum=100, 
-                        maximum=8000, 
-                        value=1000, 
-                        step=100, 
-                        label="Max Tokens",
-                        info="Maximum length of the response"
-                    )
-                    
-                    use_streaming = gr.Checkbox(
-                        label="Stream Response", 
-                        value=True,
-                        info="Show response as it's generated"
-                    )
-            
-            with gr.Accordion("Tips", open=False):
-                gr.Markdown("""
-                * Select a vision-capable model for images
-                * Upload text files to include their content
-                * Check model context window sizes
-                * Adjust temperature for creativity level
-                * Top P controls diversity of responses
-                """)
-    
-    # Define events
-    def update_model_selector(use_vision):
-        if use_vision:
-            return (
-                gr.Dropdown(choices=[name for name, _, _ in vision_models], value=vision_models[0][0]),
-                f"Context: {vision_models[0][2]:,} tokens"
-            )
-        else:
-            return (
-                gr.Dropdown(choices=[name for name, _, _ in text_models], value=text_models[0][0]),
-                f"Context: {text_models[0][2]:,} tokens"
-            )
-    
-    def update_context_info(model_name):
-        # Extract context size from model name
-        for name, _, context_size in text_models:
-            if name == model_name:
-                return f"Context: {context_size:,} tokens"
-        for name, _, context_size in vision_models:
-            if name == model_name:
-                return f"Context: {context_size:,} tokens"
-        return "Context size unknown"
+                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")
     
-    using_vision.change(
-        fn=update_model_selector,
-        inputs=using_vision,
-        outputs=[model_selector, context_info]
+    # 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
     )
     
-    model_selector.change(
-        fn=update_context_info,
-        inputs=model_selector,
-        outputs=context_info
+    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
     )
     
-    # Submit function
-    def on_submit(message, history, model, image, file, temp, top_p_val, max_tok, stream):
-        if not message and not image and not file:
-            return "", history
-        return "", process_message_stream(
-            message, 
-            history, 
-            model, 
-            image, 
-            file.name if file else None,
-            temperature=temp,
-            top_p=top_p_val,
-            max_tokens=max_tok,
-            stream=stream
-        )
+    # Clear chat
+    clear_btn.click(
+        fn=lambda: [], 
+        outputs=[chatbot]
+    )
     
-    # Set up submission events
-    submit_btn.click(
-        on_submit,
-        inputs=[
-            user_message, chatbot, model_selector, 
-            image_upload, file_upload,
-            temperature, top_p, max_tokens, use_streaming
-        ],
-        outputs=[user_message, chatbot]
+    # Clear input after submission
+    msg_submit_event.then(
+        fn=lambda: "", 
+        outputs=[user_message]
     )
     
-    user_message.submit(
-        on_submit,
-        inputs=[
-            user_message, chatbot, model_selector, 
-            image_upload, file_upload,
-            temperature, top_p, max_tokens, use_streaming
-        ],
-        outputs=[user_message, chatbot]
+    btn_submit_event.then(
+        fn=lambda: "", 
+        outputs=[user_message]
     )
 
-# Define FastAPI endpoint
-from fastapi import FastAPI, Request, HTTPException
-from fastapi.responses import JSONResponse
+# Mount FastAPI for external access
+from fastapi import FastAPI
 from pydantic import BaseModel
-from fastapi.middleware.cors import CORSMiddleware
 
 app = FastAPI()
 
@@ -498,108 +356,48 @@ class GenerateRequest(BaseModel):
 async def api_generate(request: GenerateRequest):
     """API endpoint for generating responses"""
     try:
-        message = request.message
-        model_name = request.model
-        image_data = request.image_data
+        # Process request
+        messages = [{"role": "user", "content": request.message}]
         
-        # Process image if provided
-        image = None
-        if image_data:
+        # Handle image if provided
+        if request.image_data:
             try:
-                # Decode base64 image
-                image_bytes = base64.b64decode(image_data)
+                image_bytes = base64.b64decode(request.image_data)
                 image = Image.open(BytesIO(image_bytes))
-            except Exception as e:
-                return JSONResponse(
-                    status_code=400,
-                    content={"error": f"Image processing error: {str(e)}"}
-                )
-        
-        # Generate response
-        try:
-            # Setup headers and URL
-            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"
-            
-            # Get model_id from model_name
-            model_id = None
-            if model_name:
-                for _, mid, _ in text_models + vision_models:
-                    if model_name in mid or model_name == mid:
-                        model_id = mid
-                        break
-            
-            if not model_id:
-                model_id = text_models[0][1]
-            
-            # Build messages
-            messages = []
-            
-            if image:
-                # Image processing for vision models
                 base64_image = encode_image(image)
-                content = [
-                    {"type": "text", "text": message},
-                    {
-                        "type": "image_url",
-                        "image_url": {
-                            "url": f"data:image/jpeg;base64,{base64_image}"
+                
+                messages = [{
+                    "role": "user", 
+                    "content": [
+                        {"type": "text", "text": request.message},
+                        {
+                            "type": "image_url",
+                            "image_url": {
+                                "url": f"data:image/jpeg;base64,{base64_image}"
+                            }
                         }
-                    }
-                ]
-                messages.append({"role": "user", "content": content})
-            else:
-                messages.append({"role": "user", "content": message})
-            
-            # Build request data
-            data = {
-                "model": model_id,
-                "messages": messages,
-                "temperature": 0.7
-            }
-            
-            # Make API call
-            response = requests.post(url, headers=headers, json=data)
-            response.raise_for_status()
-            
-            # Parse response
-            result = response.json()
-            reply = result.get("choices", [{}])[0].get("message", {}).get("content", "No response")
-            
-            return {"response": reply}
+                    ]
+                }]
+            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 JSONResponse(
-                status_code=500,
-                content={"error": f"Error generating response: {str(e)}"}
-            )
-            
     except Exception as e:
-        return JSONResponse(
-            status_code=500,
-            content={"error": f"Server error: {str(e)}"}
-        )
-
-# Add CORS middleware to allow cross-origin requests
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
+        return {"error": f"Error: {str(e)}"}
 
 # Mount Gradio app
-import gradio as gr
 app = gr.mount_gradio_app(app, demo, path="/")
 
-# Start the app
+# Launch the app
 if __name__ == "__main__":
-    # Use 'uvicorn' directly in HF Spaces
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    demo.launch()