Update app.py

app.py

@@ -273,9 +273,9 @@ def handle_test_example(example_text, history):
 # Initialize default API key if available
 default_initialized, default_status = initialize_default_api_if_available()
 
-# Enhanced Gradio interface with API key input
+# Gradio interface with API key input
 with gr.Blocks(
-    title="Enhanced MCP Agent Client",
+    title="MCP Agent Client",
     css="""
     .gradio-container {
         max-width: 100% !important;
@@ -326,6 +326,26 @@ with gr.Blocks(
         margin: 5px 0;
         background: rgba(0,0,0,0.1);
     }
+    .accordion-header {
+        cursor: pointer;
+        padding: 15px;
+        background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+        border-radius: 10px;
+        margin: 10px 0;
+        color: white;
+        display: flex;
+        justify-content: space-between;
+        align-items: center;
+    }
+    .accordion-header-optional {
+        background: linear-gradient(135deg, #28a745 0%, #20c997 100%);
+    }
+    .accordion-content {
+        background: rgba(255,255,255,0.05);
+        border-radius: 0 0 10px 10px;
+        padding: 20px;
+        margin-top: -10px;
+    }
     """
 ) as demo:
     gr.Markdown("# 🚀 Enhanced MCP Agent Client")
@@ -338,11 +358,9 @@ with gr.Blocks(
     
     # API Key Section - conditional visibility and messaging
     if default_initialized:
-        # Default key available - show optional upgrade section
-        with gr.Group(visible=True) as api_key_section:
+        # Default key available - show optional collapsible upgrade section
+        with gr.Accordion("🔓 Upgrade to Personal API Key (Optional)", open=False) as api_accordion:
             gr.Markdown("""
-            ## 🔓 Upgrade to Personal API Key (Optional)
-            
             **You're currently using a limited default API key. For unlimited access:**
             
             1. 🌐 Go to [Google AI Studio](https://aistudio.google.com/app/apikey)
@@ -351,7 +369,7 @@ with gr.Blocks(
             4. 📝 Paste it below and click "Upgrade to Personal Key"
             
             *Your personal API key will give you unlimited access and faster responses.*
-            """, elem_classes=["api-key-section-optional"])
+            """)
             
             with gr.Row():
                 api_key_input = gr.Textbox(
@@ -362,12 +380,12 @@ with gr.Blocks(
                     scale=4
                 )
                 api_key_submit = gr.Button("🚀 Upgrade to Personal Key", variant="primary", scale=1)
+            
+            api_key_status = gr.Markdown("", visible=True)
     else:
-        # No default key - show required section
-        with gr.Group(visible=True) as api_key_section:
+        # No default key - show required collapsible section (open by default)
+        with gr.Accordion("🔐 Setup Required: Gemini API Key", open=True) as api_accordion:
             gr.Markdown("""
-            ## 🔐 Setup Required: Gemini API Key
-            
             **To use this application, you need a free Gemini API key:**
             
             1. 🌐 Go to [Google AI Studio](https://aistudio.google.com/app/apikey)
@@ -376,7 +394,7 @@ with gr.Blocks(
             4. 📝 Paste it below and click "Validate & Start"
             
             Your API key is only stored locally in this session and is not saved anywhere.
-            """, elem_classes=["api-key-section"])
+            """)
             
             with gr.Row():
                 api_key_input = gr.Textbox(
@@ -387,8 +405,8 @@ with gr.Blocks(
                     scale=4
                 )
                 api_key_submit = gr.Button("🚀 Validate & Start", variant="primary", scale=1)
-    
-    api_key_status = gr.Markdown("", visible=True)
+            
+            api_key_status = gr.Markdown("", visible=True)
     
     # Main Chat Interface - visible if default key works, hidden otherwise
     with gr.Group(visible=default_initialized) as chat_interface:
@@ -526,18 +544,39 @@ with gr.Blocks(
             error_msg = [{"role": "assistant", "content": f"Test failed: {str(e)}"}]
             return history + error_msg, None, {"error": str(e)}, "", get_api_key_status_info()
 
+    def handle_api_key_submit(api_key):
+        """Handle API key submission."""
+        status_msg, is_valid = validate_and_set_api_key(api_key, is_user_provided=True)
+        
+        if is_valid:
+            return (
+                status_msg,
+                gr.update(open=False),  # Close accordion after successful setup
+                gr.update(visible=True),   # Show chat interface
+                get_api_key_status_info(),  # Update status
+                ""  # Clear API key input for security
+            )
+        else:
+            return (
+                status_msg,
+                gr.update(open=True),   # Keep accordion open on error
+                gr.update(visible=_api_key_set),  # Show chat if default key works
+                get_api_key_status_info(),  # Update status
+                api_key  # Keep the input value for correction
+            )
+        
     # API Key submission
     api_key_submit.click(
         fn=handle_api_key_submit,
         inputs=[api_key_input],
-        outputs=[api_key_status, api_key_section, chat_interface, api_status_display, api_key_input]
+        outputs=[api_key_status, api_accordion, chat_interface, api_status_display, api_key_input]
     )
     
     # Allow Enter key in API key input
     api_key_input.submit(
         fn=handle_api_key_submit,
         inputs=[api_key_input],
-        outputs=[api_key_status, api_key_section, chat_interface, api_status_display, api_key_input]
+        outputs=[api_key_status, api_accordion, chat_interface, api_status_display, api_key_input]
     )
 
     # Main chat handlers (only work when API key is set)