Update app.py

app.py

@@ -2,9 +2,10 @@ import os
 import base64
 import gradio as gr
 import requests
-from io import BytesIO
 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", "")
@@ -88,22 +89,23 @@ text_models = [(name, model_id) for name, model_id, _, _, _ in free_models]
 def encode_image(image):
     """Convert PIL Image to base64 string"""
     buffered = BytesIO()
-    image.save(buffered, format="PNG")
+    image.save(buffered, format="JPEG")
     return base64.b64encode(buffered.getvalue()).decode("utf-8")
 
-def process_message(message, chat_history, model_name, uploaded_image=None):
-    """Process message and return model response"""
+def process_message_stream(message, chat_history, model_name, uploaded_image=None):
+    """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])
     
     # Check if API key is set
     if not OPENROUTER_API_KEY:
-        return "Please set your OpenRouter API key in the environment variables.", chat_history
+        yield "Please set your OpenRouter API key in the environment variables.", chat_history
+        return
     
     # Setup headers and URL
     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/cstr/CrispChat",  # Replace with your actual space URL in production
     }
     
     url = "https://openrouter.ai/api/v1/chat/completions"
@@ -125,7 +127,7 @@ def process_message(message, chat_history, model_name, uploaded_image=None):
             {
                 "type": "image_url",
                 "image_url": {
-                    "url": f"data:image/png;base64,{base64_image}"
+                    "url": f"data:image/jpeg;base64,{base64_image}"
                 }
             }
         ]
@@ -136,24 +138,55 @@ def process_message(message, chat_history, model_name, uploaded_image=None):
     # Build request data
     data = {
         "model": model_id,
-        "messages": messages
+        "messages": messages,
+        "stream": True,
+        "temperature": 0.7
     }
     
     try:
-        # Make API call
-        response = requests.post(url, headers=headers, json=data)
-        response.raise_for_status()
+        # Create a new message pair in the chat history
+        chat_history.append((message, ""))
+        full_response = ""
         
-        # Parse response
-        result = response.json()
-        reply = result.get("choices", [{}])[0].get("message", {}).get("content", "No 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 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] = (message, full_response)
+                                    yield full_response, chat_history
+                            except json.JSONDecodeError:
+                                pass
         
-        # Update chat history
-        chat_history.append((message, reply))
-        return reply, chat_history
+        return full_response, chat_history
     
     except Exception as e:
-        return f"Error: {str(e)}", chat_history
+        error_msg = f"Error: {str(e)}"
+        chat_history[-1] = (message, error_msg)
+        yield error_msg, chat_history
 
 # Create a nice CSS theme
 css = """
@@ -179,11 +212,11 @@ css = """
 }
 """
 
-with gr.Blocks(css=css) as demo:
+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 and images</p>
+        <p>Chat with AI models - supports text and images</p>
     </div>
     """)
     
@@ -195,6 +228,7 @@ with gr.Blocks(css=css) as demo:
                 show_share_button=False,
                 elem_id="chatbot",
                 layout="panel",
+                type="messages"  # Use new message format
             )
             
             with gr.Row():
@@ -208,7 +242,6 @@ with gr.Blocks(css=css) as demo:
                     type="pil", 
                     label="Image Upload (optional)",
                     show_label=False,
-                    tool="select", 
                     scale=2
                 )
                 submit_btn = gr.Button("Send", scale=1, variant="primary")
@@ -234,7 +267,7 @@ with gr.Blocks(css=css) as demo:
             
             with gr.Accordion("API", open=False):
                 api_url = gr.Textbox(
-                    value="https://[your-space-name].hf.space/api/generate",
+                    value="https://cstr-crispchat.hf.space/api/generate",
                     label="API Endpoint",
                     interactive=False
                 )
@@ -266,7 +299,7 @@ with gr.Blocks(css=css) as demo:
     def on_submit(message, history, model, image):
         if not message and not image:
             return "", history
-        return "", process_message(message, history, model, image)
+        return "", process_message_stream(message, history, model, image)
     
     # Set up submission events
     submit_btn.click(
@@ -299,8 +332,60 @@ def api_generate(message, model=None, image_data=None):
     
     # Generate response
     try:
-        response, _ = process_message(message, [], model_name, image)
-        return {"response": response}
+        # 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
+        
+        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")
+        
+        return {"response": reply}
+    
     except Exception as e:
         return {"error": f"Error generating response: {str(e)}"}