Update app.py

app.py

@@ -82,9 +82,21 @@ vision_model_ids = [
     "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 = [(name, model_id) for name, model_id, _, _, _ in free_models if model_id in vision_model_ids]
-text_models = [(name, model_id) for name, model_id, _, _, _ in free_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]
 
 def encode_image(image):
     """Convert PIL Image to base64 string"""
@@ -92,9 +104,19 @@ def encode_image(image):
     image.save(buffered, format="JPEG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-def process_message_stream(message, chat_history, model_name, uploaded_image=None):
+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_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"""
-    model_id = next((model_id for name, model_id, _, _, _ in free_models if name == model_name), text_models[0][1])
+    # 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:
@@ -105,7 +127,7 @@ def process_message_stream(message, chat_history, model_name, uploaded_image=Non
     headers = {
         "Content-Type": "application/json",
         "Authorization": f"Bearer {OPENROUTER_API_KEY}",
-        "HTTP-Referer": "https://huggingface.co/spaces/cstr/CrispChat",  # Replace with your actual space URL in production
+        "HTTP-Referer": "https://huggingface.co/spaces",  # Replace with your actual space URL in production
     }
     
     url = "https://openrouter.ai/api/v1/chat/completions"
@@ -114,101 +136,173 @@ def process_message_stream(message, chat_history, model_name, uploaded_image=Non
     messages = []
     
     # Add chat history
-    for human_msg, ai_msg in chat_history:
-        messages.append({"role": "user", "content": human_msg})
-        messages.append({"role": "assistant", "content": ai_msg})
+    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
+    # 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},
-            {
-                "type": "image_url",
-                "image_url": {
-                    "url": f"data:image/jpeg;base64,{base64_image}"
-                }
-            }
+            {"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:
-        messages.append({"role": "user", "content": message})
+        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,
-        "stream": True,
-        "temperature": 0.7
+        "stream": stream,
+        "temperature": temperature,
+        "top_p": top_p,
+        "max_tokens": max_tokens
     }
     
     try:
         # Create a new message pair in the chat history
-        chat_history.append((message, ""))
+        user_msg = {"role": "user", "content": message}
+        ai_msg = {"role": "assistant", "content": ""}
+        chat_history.append(user_msg)
+        chat_history.append(ai_msg)
+        
         full_response = ""
         
-        # 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')
-                    
-                    while True:
-                        line_end = buffer.find('\n')
-                        if line_end == -1:
-                            break
-                            
-                        line = buffer[:line_end].strip()
-                        buffer = buffer[line_end + 1:]
+        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 line.startswith('data: '):
-                            data = line[6:]
-                            if data == '[DONE]':
+                        while True:
+                            line_end = buffer.find('\n')
+                            if line_end == -1:
                                 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] = (message, full_response)
-                                    yield full_response, chat_history
-                            except json.JSONDecodeError:
-                                pass
+                            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
+        else:
+            # Non-streaming API call
+            response = requests.post(url, headers=headers, json=data)
+            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
+            yield chat_history
         
-        return full_response, chat_history
+        return chat_history
     
     except Exception as e:
         error_msg = f"Error: {str(e)}"
-        chat_history[-1] = (message, error_msg)
-        yield error_msg, chat_history
+        chat_history[-1]["content"] = error_msg
+        yield chat_history
 
 # Create a nice CSS theme
 css = """
 .gradio-container {
-    font-family: 'Segoe UI', Arial, sans-serif;
+    font-family: 'Inter', -apple-system, BlinkMacSystemFont, sans-serif;
+}
+.chat-message {
+    padding: 15px;
+    border-radius: 10px;
+    margin-bottom: 10px;
 }
-#chat-message {
-    min-height: 100px;
+.user-message {
+    background-color: #f0f4f8;
 }
-#model-selector {
-    max-width: 100%;
+.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: 10px;
+    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;
 }
 """
 
@@ -216,182 +310,296 @@ with gr.Blocks(css=css, theme=gr.themes.Soft()) as demo:
     gr.HTML("""
     <div class="app-header">
         <h1>🔆 CrispChat</h1>
-        <p>Chat with AI models - supports text and images</p>
+        <p>Chat with free OpenRouter AI models - supports text, images, and files</p>
     </div>
     """)
     
     with gr.Row():
-        with gr.Column(scale=3):
+        with gr.Column(scale=4):
             chatbot = gr.Chatbot(
-                height=500,
+                height=600,
                 show_copy_button=True,
                 show_share_button=False,
                 elem_id="chatbot",
-                layout="panel",
+                layout="bubble",
+                avatar_images=("👤", "🤖"),
+                bubble_full_width=False,
                 type="messages"  # Use new message format
             )
             
             with gr.Row():
-                user_message = gr.Textbox(
-                    placeholder="Type your message here...",
-                    show_label=False,
-                    elem_id="chat-message",
-                    scale=10
-                )
-                image_upload = gr.Image(
-                    type="pil", 
-                    label="Image Upload (optional)",
-                    show_label=False,
-                    scale=2
-                )
-                submit_btn = gr.Button("Send", scale=1, variant="primary")
+                with gr.Column(scale=10):
+                    user_message = gr.Textbox(
+                        placeholder="Type your message here...",
+                        show_label=False,
+                        elem_id="chat-input",
+                        lines=3
+                    )
+                    
+                    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=1):
+        with gr.Column(scale=2):
             with gr.Accordion("Model Selection", open=True):
                 using_vision = gr.Checkbox(label="Using image", value=False)
                 
                 model_selector = gr.Dropdown(
-                    choices=[name for name, _ in text_models],
+                    choices=[name for name, _, _ in text_models],
                     value=text_models[0][0],
                     label="Select Model",
                     elem_id="model-selector"
                 )
+                
+                context_info = gr.Markdown(value=f"Context: {text_models[0][2]:,} tokens")
+                
+            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=True):
+            with gr.Accordion("Tips", open=False):
                 gr.Markdown("""
-                * For best results with images, select a vision-capable model
-                * Text models can handle up to 32k tokens
-                * Try different models for different tasks
-                * API output is in Markdown format for code highlighting
-                """)
-            
-            with gr.Accordion("API", open=False):
-                api_url = gr.Textbox(
-                    value="https://cstr-crispchat.hf.space/api/generate",
-                    label="API Endpoint",
-                    interactive=False
-                )
-                api_docs = gr.Markdown("""
-                ```json
-                POST /api/generate
-                {
-                    "message": "Your message here",
-                    "model": "model-id-here",
-                    "image_data": "optional-base64-encoded-image"
-                }
-                ```
+                * 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])
+            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])
+            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"
     
     using_vision.change(
         fn=update_model_selector,
         inputs=using_vision,
-        outputs=model_selector
+        outputs=[model_selector, context_info]
+    )
+    
+    model_selector.change(
+        fn=update_context_info,
+        inputs=model_selector,
+        outputs=context_info
     )
     
     # Submit function
-    def on_submit(message, history, model, image):
-        if not message and not image:
+    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)
+        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
+        )
     
     # Set up submission events
     submit_btn.click(
         on_submit,
-        inputs=[user_message, chatbot, model_selector, image_upload],
+        inputs=[
+            user_message, chatbot, model_selector, 
+            image_upload, file_upload,
+            temperature, top_p, max_tokens, use_streaming
+        ],
         outputs=[user_message, chatbot]
     )
     
     user_message.submit(
         on_submit,
-        inputs=[user_message, chatbot, model_selector, image_upload],
+        inputs=[
+            user_message, chatbot, model_selector, 
+            image_upload, file_upload,
+            temperature, top_p, max_tokens, use_streaming
+        ],
         outputs=[user_message, chatbot]
     )
 
-# API endpoint for external access
-@demo.queue()
-def api_generate(message, model=None, image_data=None):
+# Define FastAPI endpoint
+from fastapi import FastAPI, Request, HTTPException
+from fastapi.responses import JSONResponse
+from pydantic import BaseModel
+from fastapi.middleware.cors import CORSMiddleware
+
+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"""
-    model_name = model or text_models[0][0]
-    
-    # Process image if provided
-    image = None
-    if image_data:
-        try:
-            # Decode base64 image
-            image_bytes = base64.b64decode(image_data)
-            image = Image.open(BytesIO(image_bytes))
-        except Exception as e:
-            return {"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 = next((model_id for name, model_id, _, _, _ in free_models if name == model_name), None)
-        if not model_id and model:
-            # Check if model parameter is a direct model ID
-            model_id = model
+        message = request.message
+        model_name = request.model
+        image_data = request.image_data
         
-        if not model_id:
-            model_id = text_models[0][1]
-        
-        # Build messages
-        messages = []
+        # Process image if provided
+        image = None
+        if image_data:
+            try:
+                # Decode base64 image
+                image_bytes = base64.b64decode(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)}"}
+                )
         
-        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}"
+        # 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.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")
+                ]
+                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}
         
-        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 {"error": f"Error generating response: {str(e)}"}
+        return JSONResponse(
+            status_code=500,
+            content={"error": f"Server error: {str(e)}"}
+        )
 
-demo.queue()
-demo.launch(share=False)
+# Add CORS middleware to allow cross-origin requests
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Mount Gradio app
+import gradio as gr
+app = gr.mount_gradio_app(app, demo, path="/")
 
+# Start the app
 if __name__ == "__main__":
-    # Remove or comment out demo.launch() here if you added it above
-    pass
+    # Use 'uvicorn' directly in HF Spaces
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)