Update app.py

app.py

@@ -203,6 +203,13 @@ def update_model_info(model_name):
             """
     return "<p>Model information not available</p>"
 
+def update_category_dropdown(category):
+    """Update the category model dropdown when a category is selected"""
+    models = get_models_for_category(category)
+    if not models:
+        return [], None
+    return models, models[0]
+
 def update_category_models_ui(category):
     """Completely regenerate the models dropdown based on selected category"""
     for cat in MODELS:
@@ -434,7 +441,7 @@ def extract_ai_response(result):
         return f"Error: {str(e)}"
 
 # streaming code:
-def streaming_handler(response, chatbot, message_idx):
+def streaming_handler(response, chatbot, message_idx, message):
     try:
         # First add the user message if needed
         if len(chatbot) == message_idx:
@@ -563,7 +570,7 @@ def ask_ai(message, history, model_choice, temperature, max_tokens, top_p,
             
             # Set up generator for streaming updates
             def streaming_generator():
-                for updated_history in streaming_handler(response, chat_history, len(chat_history) - 1):
+                for updated_history in streaming_handler(response, chat_history, len(chat_history) - 1, message):
                     yield updated_history
             
             return streaming_generator()
@@ -713,12 +720,13 @@ def create_app():
                         )
                     
                     with gr.Row(elem_classes="model-selection-row"):
+                        
+                        # Main model dropdown
                         model_choice = gr.Dropdown(
                             [model[0] for model in ALL_MODELS],
                             value=ALL_MODELS[0][0],
                             label="Model",
                             elem_id="model-choice",
-                            elem_classes="model-choice",
                             allow_custom_value=True
                         )
 
@@ -732,24 +740,18 @@ def create_app():
                     # Model category selection
                     with gr.Accordion("Browse by Category", open=False):
                         model_categories = gr.Dropdown(
-                            [category["category"] for category in MODELS],
+                            [model["category"] for model in MODELS],
                             label="Categories",
                             value=MODELS[0]["category"]
                         )
                         
-                        # Create a container for the category models dropdown
-                        with gr.Column(visible=True, elem_id="category-models-container") as category_models_container:
-                            # Create a hidden text component to store model choices as JSON
-                            category_model_choices = gr.Text(visible=False)
-
-                            # Create the dropdown with no initial choices
-                            category_models = gr.Dropdown(
-                                [],
-                                label="Models in Category",
-                                value=None,
-                                elem_classes="category-models",
-                                allow_custom_value=True
-                            )
+                        # Models in category dropdown
+                        category_models = gr.Dropdown(
+                            get_models_for_category(MODELS[0]["category"]),
+                            label="Models in Category",
+                            value=get_models_for_category(MODELS[0]["category"])[0] if get_models_for_category(MODELS[0]["category"]) else None,
+                            allow_custom_value=True
+                        )
                 
                 with gr.Accordion("Generation Parameters", open=False):
                     with gr.Group(elem_classes="parameter-grid"):
@@ -943,9 +945,9 @@ def create_app():
         
         # Update model list when category changes
         model_categories.change(
-            fn=lambda cat: json.dumps(get_models_for_category(cat)),
+            fn=update_category_dropdown,
             inputs=model_categories,
-            outputs=category_model_choices
+            outputs=[category_models, category_models]
         )
         
         # Update main model choice when category model is selected
@@ -955,54 +957,6 @@ def create_app():
             outputs=model_choice
         )
 
-        category_model_choices.change(
-            fn=None,
-            inputs=None,
-            outputs=None,
-            _js="""
-            function(choices_json) {
-                // Parse JSON string to array
-                const choices = JSON.parse(choices_json);
-                
-                // Find the dropdown element
-                const dropdown = document.querySelector('.category-models select');
-                
-                // Clear existing options
-                dropdown.innerHTML = '';
-                
-                // Add new options
-                choices.forEach(model => {
-                    const option = document.createElement('option');
-                    option.value = model;
-                    option.textContent = model;
-                    dropdown.appendChild(option);
-                });
-                
-                // Set the first option as selected if available
-                if (choices.length > 0) {
-                    dropdown.value = choices[0];
-                    
-                    // Update the main model dropdown
-                    const mainDropdown = document.querySelector('.model-choice select');
-                    mainDropdown.value = choices[0];
-                    
-                    // Trigger change events
-                    dropdown.dispatchEvent(new Event('change', { bubbles: true }));
-                    mainDropdown.dispatchEvent(new Event('change', { bubbles: true }));
-                }
-            }
-            """
-        )
-
-        # Function to initialize the category models dropdown
-        def init_category_models():
-            initial_category = MODELS[0]["category"]
-            initial_models = get_models_for_category(initial_category)
-            return json.dumps(initial_models)
-
-        # Set initial choices for category models dropdown
-        category_model_choices.value = init_category_models()
-
         # Process uploaded images
         image_upload_btn.upload(
             fn=lambda files: files,