Update app.py

app.py

@@ -629,6 +629,10 @@ def update_context_display(provider, model_name):
 
 def is_vision_model(provider, model_name):
     """Check if a model supports vision/images"""
+    # Safety check for None model name
+    if model_name is None:
+        return False
+        
     if provider in VISION_MODELS:
         if model_name in VISION_MODELS[provider]:
             return True
@@ -1132,6 +1136,7 @@ def extract_ai_response(result, provider):
 # ==========================================================
 
 def openrouter_streaming_handler(response, history, message):
+    """Handle streaming responses from OpenRouter"""
     try:
         updated_history = history + [{"role": "user", "content": message}]
         assistant_response = ""
@@ -1163,66 +1168,54 @@ def openrouter_streaming_handler(response, history, message):
         # Add error message to the current response
         yield updated_history + [{"role": "assistant", "content": f"Error during streaming: {str(e)}"}]
 
-def openai_streaming_handler(response, chatbot, message_idx, message):
+def openai_streaming_handler(response, history, message):
+    """Handle streaming responses from OpenAI"""
     try:
-        # First add the user message if needed
-        if len(chatbot) == message_idx:
-            chatbot.append([message, ""])
-            
-        full_response = ""
+        updated_history = history + [{"role": "user", "content": message}]
+        assistant_response = ""
+        
         for chunk in response:
             if hasattr(chunk.choices[0].delta, "content") and chunk.choices[0].delta.content is not None:
                 content = chunk.choices[0].delta.content
-                full_response += content
-                chatbot[-1][1] = full_response
-                yield chatbot
-                
+                assistant_response += content
+                yield updated_history + [{"role": "assistant", "content": assistant_response}]
     except Exception as e:
         logger.error(f"Error in OpenAI streaming handler: {str(e)}")
         # Add error message to the current response
-        chatbot[-1][1] += f"\n\nError during streaming: {str(e)}"
-        yield chatbot
+        yield updated_history + [{"role": "assistant", "content": f"Error during streaming: {str(e)}"}]
 
-def groq_streaming_handler(response, chatbot, message_idx, message):
+def groq_streaming_handler(response, history, message):
+    """Handle streaming responses from Groq"""
     try:
-        # First add the user message if needed
-        if len(chatbot) == message_idx:
-            chatbot.append([message, ""])
-            
-        full_response = ""
+        updated_history = history + [{"role": "user", "content": message}]
+        assistant_response = ""
+        
         for chunk in response:
             if hasattr(chunk.choices[0].delta, "content") and chunk.choices[0].delta.content is not None:
                 content = chunk.choices[0].delta.content
-                full_response += content
-                chatbot[-1][1] = full_response
-                yield chatbot
-                
+                assistant_response += content
+                yield updated_history + [{"role": "assistant", "content": assistant_response}]
     except Exception as e:
         logger.error(f"Error in Groq streaming handler: {str(e)}")
         # Add error message to the current response
-        chatbot[-1][1] += f"\n\nError during streaming: {str(e)}"
-        yield chatbot
+        yield updated_history + [{"role": "assistant", "content": f"Error during streaming: {str(e)}"}]
 
-def together_streaming_handler(response, chatbot, message_idx, message):
+def together_streaming_handler(response, history, message):
+    """Handle streaming responses from Together"""
     try:
-        # First add the user message if needed
-        if len(chatbot) == message_idx:
-            chatbot.append([message, ""])
-            
-        full_response = ""
+        updated_history = history + [{"role": "user", "content": message}]
+        assistant_response = ""
+        
         for chunk in response:
             if hasattr(chunk.choices[0].delta, "content") and chunk.choices[0].delta.content is not None:
                 content = chunk.choices[0].delta.content
-                full_response += content
-                chatbot[-1][1] = full_response
-                yield chatbot
-                
+                assistant_response += content
+                yield updated_history + [{"role": "assistant", "content": assistant_response}]
     except Exception as e:
         logger.error(f"Error in Together streaming handler: {str(e)}")
         # Add error message to the current response
-        chatbot[-1][1] += f"\n\nError during streaming: {str(e)}"
-        yield chatbot
-
+        yield updated_history + [{"role": "assistant", "content": f"Error during streaming: {str(e)}"}]
+        
 # ==========================================================
 # MAIN FUNCTION TO ASK AI
 # ==========================================================
@@ -1236,11 +1229,8 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
     if not message.strip() and not images and not documents:
         return history
     
-    # Copy history to new list to avoid modifying the original
-    chat_history = list(history)
-    
-    # Create messages from chat history
-    messages = format_to_message_dict(chat_history)
+    # Create messages from chat history for API requests
+    messages = format_to_message_dict(history)
     
     # Add system message if provided
     if system_message and system_message.strip():
@@ -1252,7 +1242,7 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
     # Prepare message with images and documents if any
     content = prepare_message_with_media(message, images, documents)
     
-    # Add current message
+    # Add current message to API messages
     messages.append({"role": "user", "content": content})
     
     # Common parameters for all providers
@@ -1272,8 +1262,11 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in OpenRouter"
-                chat_history.append([message, error_message])
-                return chat_history
+                # Use proper message format
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
             # Build OpenRouter payload
             payload = {
@@ -1319,13 +1312,35 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
             
             # Handle streaming response
             if stream_output and response.status_code == 200:
-                # Add empty response slot to history
-                chat_history.append([message, ""])
+                # Add message to history
+                updated_history = history + [{"role": "user", "content": message}]
                 
                 # Set up generator for streaming updates
                 def streaming_generator():
-                    for updated_history in openrouter_streaming_handler(response, chat_history, len(chat_history) - 1, message):
-                        yield updated_history
+                    assistant_response = ""
+                    for line in response.iter_lines():
+                        if not line:
+                            continue
+                            
+                        line = line.decode('utf-8')
+                        if not line.startswith('data: '):
+                            continue
+                            
+                        data = line[6:]
+                        if data.strip() == '[DONE]':
+                            break
+                            
+                        try:
+                            chunk = json.loads(data)
+                            if "choices" in chunk and len(chunk["choices"]) > 0:
+                                delta = chunk["choices"][0].get("delta", {})
+                                if "content" in delta and delta["content"]:
+                                    # Update the current response
+                                    assistant_response += delta["content"]
+                                    # Yield updated history with current assistant response
+                                    yield updated_history + [{"role": "assistant", "content": assistant_response}]
+                        except json.JSONDecodeError:
+                            logger.error(f"Failed to parse JSON from chunk: {data}")
                 
                 return streaming_generator()
             
@@ -1337,9 +1352,11 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 # Extract AI response
                 ai_response = extract_ai_response(result, provider)
                 
-                # Add response to history
-                chat_history.append([message, ai_response])
-                return chat_history
+                # Add response to history with proper format
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": ai_response}
+                ]
             
             # Handle error response
             else:
@@ -1351,16 +1368,20 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                     error_message += f"\n\nResponse: {response.text}"
                 
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
         elif provider == "OpenAI":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in OpenAI"
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
             # Build OpenAI payload
             payload = {
@@ -1381,34 +1402,44 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 
                 # Handle streaming response
                 if stream_output:
-                    # Add empty response slot to history
-                    chat_history.append([message, ""])
+                    # Add message to history
+                    updated_history = history + [{"role": "user", "content": message}]
                     
                     # Set up generator for streaming updates
                     def streaming_generator():
-                        for updated_history in openai_streaming_handler(response, chat_history, len(chat_history) - 1, message):
-                            yield updated_history
+                        assistant_response = ""
+                        for chunk in response:
+                            if hasattr(chunk.choices[0].delta, "content") and chunk.choices[0].delta.content is not None:
+                                content = chunk.choices[0].delta.content
+                                assistant_response += content
+                                yield updated_history + [{"role": "assistant", "content": assistant_response}]
                     
                     return streaming_generator()
                 
                 # Handle normal response
                 else:
                     ai_response = extract_ai_response(response, provider)
-                    chat_history.append([message, ai_response])
-                    return chat_history
+                    return history + [
+                        {"role": "user", "content": message},
+                        {"role": "assistant", "content": ai_response}
+                    ]
             except Exception as e:
                 error_message = f"OpenAI API Error: {str(e)}"
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
         elif provider == "HuggingFace":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in HuggingFace"
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
             # Build HuggingFace payload
             payload = {
@@ -1426,21 +1457,27 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 
                 # Extract response
                 ai_response = extract_ai_response(response, provider)
-                chat_history.append([message, ai_response])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": ai_response}
+                ]
             except Exception as e:
                 error_message = f"HuggingFace API Error: {str(e)}"
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
         elif provider == "Groq":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in Groq"
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
             # Build Groq payload
             payload = {
@@ -1460,34 +1497,44 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 
                 # Handle streaming response
                 if stream_output:
-                    # Add empty response slot to history
-                    chat_history.append([message, ""])
+                    # Add message to history
+                    updated_history = history + [{"role": "user", "content": message}]
                     
                     # Set up generator for streaming updates
                     def streaming_generator():
-                        for updated_history in groq_streaming_handler(response, chat_history, len(chat_history) - 1, message):
-                            yield updated_history
+                        assistant_response = ""
+                        for chunk in response:
+                            if hasattr(chunk.choices[0].delta, "content") and chunk.choices[0].delta.content is not None:
+                                content = chunk.choices[0].delta.content
+                                assistant_response += content
+                                yield updated_history + [{"role": "assistant", "content": assistant_response}]
                     
                     return streaming_generator()
                 
                 # Handle normal response
                 else:
                     ai_response = extract_ai_response(response, provider)
-                    chat_history.append([message, ai_response])
-                    return chat_history
+                    return history + [
+                        {"role": "user", "content": message},
+                        {"role": "assistant", "content": ai_response}
+                    ]
             except Exception as e:
                 error_message = f"Groq API Error: {str(e)}"
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
         elif provider == "Cohere":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in Cohere"
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
             # Build Cohere payload (doesn't support streaming the same way)
             payload = {
@@ -1505,21 +1552,27 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 
                 # Extract response
                 ai_response = extract_ai_response(response, provider)
-                chat_history.append([message, ai_response])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": ai_response}
+                ]
             except Exception as e:
                 error_message = f"Cohere API Error: {str(e)}"
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
         elif provider == "Together":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in Together"
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
             
             # Build Together payload
             payload = {
@@ -1538,34 +1591,44 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 
                 # Handle streaming response
                 if stream_output:
-                    # Add empty response slot to history
-                    chat_history.append([message, ""])
+                    # Add message to history
+                    updated_history = history + [{"role": "user", "content": message}]
                     
                     # Set up generator for streaming updates
                     def streaming_generator():
-                        for updated_history in together_streaming_handler(response, chat_history, len(chat_history) - 1, message):
-                            yield updated_history
+                        assistant_response = ""
+                        for chunk in response:
+                            if hasattr(chunk.choices[0].delta, "content") and chunk.choices[0].delta.content is not None:
+                                content = chunk.choices[0].delta.content
+                                assistant_response += content
+                                yield updated_history + [{"role": "assistant", "content": assistant_response}]
                     
                     return streaming_generator()
                 
                 # Handle normal response
                 else:
                     ai_response = extract_ai_response(response, provider)
-                    chat_history.append([message, ai_response])
-                    return chat_history
+                    return history + [
+                        {"role": "user", "content": message},
+                        {"role": "assistant", "content": ai_response}
+                    ]
             except Exception as e:
                 error_message = f"Together API Error: {str(e)}"
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
         elif provider == "OVH":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in OVH"
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
             
             # Build OVH payload
             payload = {
@@ -1583,21 +1646,27 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 
                 # Extract response
                 ai_response = extract_ai_response(response, provider)
-                chat_history.append([message, ai_response])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": ai_response}
+                ]
             except Exception as e:
                 error_message = f"OVH API Error: {str(e)}"
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
         elif provider == "Cerebras":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in Cerebras"
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
             
             # Build Cerebras payload
             payload = {
@@ -1615,21 +1684,27 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 
                 # Extract response
                 ai_response = extract_ai_response(response, provider)
-                chat_history.append([message, ai_response])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": ai_response}
+                ]
             except Exception as e:
                 error_message = f"Cerebras API Error: {str(e)}"
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
                 
         elif provider == "GoogleAI":
             # Get model ID from registry
             model_id, _ = get_model_info(provider, model_choice)
             if not model_id:
                 error_message = f"Error: Model '{model_choice}' not found in GoogleAI"
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
             
             # Build GoogleAI payload
             payload = {
@@ -1648,24 +1723,32 @@ def ask_ai(message, history, provider, model_choice, temperature, max_tokens, to
                 
                 # Extract response
                 ai_response = extract_ai_response(response, provider)
-                chat_history.append([message, ai_response])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": ai_response}
+                ]
             except Exception as e:
                 error_message = f"GoogleAI API Error: {str(e)}"
                 logger.error(error_message)
-                chat_history.append([message, error_message])
-                return chat_history
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
             
         else:
             error_message = f"Error: Unsupported provider '{provider}'"
-            chat_history.append([message, error_message])
-            return chat_history
+            return history + [
+                {"role": "user", "content": message},
+                {"role": "assistant", "content": error_message}
+            ]
             
     except Exception as e:
         error_message = f"Error: {str(e)}"
         logger.error(f"Exception during API call: {error_message}")
-        chat_history.append([message, error_message])
-        return chat_history
+        return history + [
+            {"role": "user", "content": message},
+            {"role": "assistant", "content": error_message}
+        ]
 
 def clear_chat():
     """Reset all inputs"""
@@ -2160,14 +2243,20 @@ def create_app():
             
         def update_vision_indicator(provider, model_choice):
             """Update the vision capability indicator"""
+            # Safety check for None model
+            if model_choice is None:
+                return False
             return is_vision_model(provider, model_choice)
             
         def update_image_upload_visibility(provider, model_choice):
             """Show/hide image upload based on model vision capabilities"""
+            # Safety check for None model
+            if model_choice is None:
+                return gr.update(visible=False)
             is_vision = is_vision_model(provider, model_choice)
             return gr.update(visible=is_vision)
         
-        # Search model function - FIXED FUNCTION
+        # Search model function
         def search_openrouter_models(search_term):
             """Filter OpenRouter models based on search term"""
             all_models = [model[0] for model in OPENROUTER_ALL_MODELS]
@@ -2588,9 +2677,11 @@ def create_app():
             
             # Check if model is selected
             if not model_choice:
-                new_history = history.copy()
-                new_history.append([message, f"Error: No model selected for provider {provider}"])
-                return new_history
+                error_message = f"Error: No model selected for provider {provider}"
+                return history + [
+                    {"role": "user", "content": message},
+                    {"role": "assistant", "content": error_message}
+                ]
             
             # Select the appropriate API key based on the provider
             api_key_override = None