Refactor iframe preview creation and enhance API key handling. Simplified iframe HTML structure and improved error messaging. Added model retrieval for non-Hugging Face providers in API key management.

app.py

@@ -141,126 +141,20 @@ def start_preview_app():
 
 
 def create_iframe_preview():
-    """Create an iframe HTML element for the preview."""
+    """Create a simple iframe for the preview."""
     print("🔄 Creating iframe preview...")
     # Start the preview app
     success, message = start_preview_app()
 
     if not success:
         print(f"❌ Failed to create preview iframe: {message}")
-        return f"""
-        <div style='padding: 20px; text-align: center; color: #d32f2f;'>
-            <h3>❌ Failed to start preview</h3>
-            <pre style='background: #f5f5f5; padding: 10px; \
-border-radius: 4px; text-align: left;'>{message}</pre>
-        </div>
-        """
-
-    # Add timestamp to force iframe refresh
-    timestamp = int(time.time() * 1000)
-    print("✅ Iframe preview created successfully")
-
-    return f"""
-    <div style='width: 100%; height: 70vh; border: 1px solid #ddd; border-radius: \
-    8px; overflow: hidden; position: relative; background: #f9f9f9;'>
-        <iframe
-            id="preview-iframe-{timestamp}"
-            src="{PREVIEW_URL}"
-            width="100%"
-            height="100%"
-            frameborder="0"
-            style="border: none; background: white;"
-            sandbox="allow-same-origin allow-scripts allow-forms allow-popups \
-            allow-top-navigation allow-modals"
-            loading="eager"
-            onload="document.getElementById('loading-{timestamp}').style.display='none'"
-            onerror="document.getElementById('iframe-{timestamp}').style.display='none';\
-             document.getElementById('fallback-{timestamp}').style.display='block'"
-        ></iframe>
-
-        <!-- Loading indicator -->
-        <div id="loading-{timestamp}" style="position: absolute; top: 0; left: 0; \
-            width: 100%; height: 100%; background: rgba(255,255,255,0.9); \
-            display: flex; align-items: center; justify-content: center; z-index: 10;">
-            <div style="text-align: center;">
-                <div style="border: 4px solid #f3f3f3; border-top: 4px solid #007bff; \
-                border-radius: 50%; width: 40px; height: 40px; animation: spin 1s \
-                linear infinite; margin: 0 auto 15px;"></div>
-                <p style="color: #666; margin: 0;">Loading preview...</p>
-            </div>
-        </div>
-
-        <!-- Fallback content -->
-        <div id="fallback-{timestamp}" style="display: none; padding: 20px; \
-            text-align: center; height: 100%; display: flex; flex-direction: column; \
-            justify-content: center;">
-            <div style='background: white; border-radius: 8px; padding: 30px; \
-                margin: 20px; box-shadow: 0 2px 10px rgba(0,0,0,0.1);'>
-                <h3 style='color: #28a745; margin-bottom: 20px;'>
-                    🚀 Preview App Running!
-                </h3>
-                <p style='color: #666; margin-bottom: 25px; font-size: 16px;'>
-                    Your Gradio app is running successfully.
-                </p>
-                <div style='margin: 20px 0;'>
-                    <a href="{PREVIEW_URL}" target="_blank"
-                       style='display: inline-block; padding: 12px 25px; background: \
-                        #007bff; color: white;
-                              text-decoration: none; border-radius: 6px; \
-                                font-weight: bold; font-size: 16px;
-                              box-shadow: 0 2px 4px rgba(0,0,0,0.2);'>
-                        🔗 Open Preview in New Tab
-                    </a>
-                </div>
-                <div style='margin-top: 20px; padding: 15px; background: #f8f9fa; \
-                border-radius: 4px; border-left: 4px solid #007bff;'>
-                    <p style='margin: 0; color: #495057; font-size: 14px;'>
-                        <strong>URL:</strong> <code style='background: #e9ecef; \
-                        padding: 2px 6px; border-radius: 3px;'>{PREVIEW_URL}</code>
-                    </p>
-                </div>
-            </div>
-        </div>
-
-        <!-- Open in new tab button -->
-        <div style='position: absolute; top: 10px; right: 10px; z-index: 20;'>
-            <a href="{PREVIEW_URL}" target="_blank"
-               style='display: inline-block; padding: 8px 12px; background: \
-                rgba(0,0,0,0.8); color: white;
-                      text-decoration: none; border-radius: 4px; font-size: 12px;'>
-                ↗ Open in new tab
-            </a>
-        </div>
-    </div>
-
-    <div style='padding: 10px; text-align: center; color: #666; font-size: 12px;'>
-        Preview running on <a href="{PREVIEW_URL}" target="_blank">{PREVIEW_URL}</a>
-        <span style='margin-left: 10px; color: #999;'>
-            Last updated: {time.strftime('%H:%M:%S')}
-        </span>
-    </div>
-
-    <style>
-    @keyframes spin {{
-        0% {{ transform: rotate(0deg); }}
-        100% {{ transform: rotate(360deg); }}
-    }}
-    </style>
-
-    <script>
-    // Auto-fallback after 10 seconds if iframe doesn't load
-    setTimeout(function() {{
-        const loading = document.getElementById('loading-{timestamp}');
-        const fallback = document.getElementById('fallback-{timestamp}');
-        if (loading && loading.style.display !== 'none') {{
-            loading.style.display = 'none';
-            if (fallback) {{
-                fallback.style.display = 'flex';
-            }}
-        }}
-    }}, 10000);
-    </script>
-    """
+        return f'<div style="padding: 20px; color: red;">\
+            ❌ Failed to start preview: {message}\
+            </div>'
+
+    print("✅ Simple iframe preview created")
+    return f'<iframe src="{PREVIEW_URL}" width="100%" height="600px" frameborder="0">\
+</iframe>'
 
 
 def is_preview_running():
@@ -283,8 +177,22 @@ def ensure_preview_running():
         start_preview_app()
 
 
+def get_default_model_for_provider(provider: str) -> str:
+    """Get the default model ID for a given provider."""
+    provider_model_map = {
+        "Anthropic": "anthropic/claude-sonnet-4-20250514",
+        "OpenAI": "openai/gpt-4o",
+        "Mistral": "mistral/codestral-latest",
+        "SambaNova": "sambanova/Meta-Llama-3.1-70B-Instruct",
+        "Hugging Face": "huggingface/together/Qwen/Qwen2.5-Coder-32B-Instruct",
+    }
+    return provider_model_map.get(
+        provider, "huggingface/together/Qwen/Qwen2.5-Coder-32B-Instruct"
+    )
+
+
 def save_api_key(provider, api_key):
-    """Save API key to environment variable."""
+    """Save API key to environment variable and update model accordingly."""
     if not api_key.strip():
         return f"⚠️ Please enter a valid API key for {provider}"
 
@@ -299,8 +207,20 @@ def save_api_key(provider, api_key):
 
     env_var_name = env_var_map.get(provider)
     if env_var_name:
+        # Always set the provider-specific API key
         os.environ[env_var_name] = api_key.strip()
-        return f"✅ {provider} API key saved successfully"
+
+        # For non-Hugging Face providers, also set the generic API_KEY and MODEL_ID
+        # This ensures the main agent uses the correct model and API key
+        if provider != "Hugging Face":
+            os.environ["API_KEY"] = api_key.strip()
+            os.environ["MODEL_ID"] = get_default_model_for_provider(provider)
+            return (
+                f"✅ {provider} API key saved successfully \n"
+                f"Model: {get_default_model_for_provider(provider)}"
+            )
+        else:
+            return f"✅ {provider} API key saved successfully"
     else:
         return f"❌ Unknown provider: {provider}"
 
@@ -331,9 +251,15 @@ def get_api_key_status(selected_llm_provider="Anthropic"):
     if llm_env_var and os.getenv(llm_env_var):
         key = os.getenv(llm_env_var)
         masked_key = f"{key[:8]}...{key[-4:]}" if len(key) > 12 else "***"
-        status.append(f"✅ {selected_llm_provider}: {masked_key}")
+        model = get_default_model_for_provider(selected_llm_provider)
+        status.append(f"✅ {selected_llm_provider}: {masked_key} (Model: {model})")
     else:
-        status.append(f"❌ {selected_llm_provider}: Not set")
+        model = get_default_model_for_provider(selected_llm_provider)
+        status.append(f"❌ {selected_llm_provider}: Not set (Would use: {model})")
+
+    # Show current active model
+    current_model = os.getenv("MODEL_ID", "Qwen/Qwen2.5-Coder-32B-Instruct")
+    status.append(f"🤖 Current Active Model: {current_model}")
 
     return "\n".join(status)
 
@@ -506,8 +432,8 @@ class GradioUI:
                             label="Current API Key Status",
                             value=get_api_key_status(),
                             interactive=False,
-                            lines=2,
-                            max_lines=2,
+                            lines=4,
+                            max_lines=4,
                         )
 
                         # with gr.Row():
@@ -557,22 +483,33 @@ class GradioUI:
 
             # Set up event handlers for API key saving
             def save_and_update_status(
-                provider, api_key, selected_llm_provider="Anthropic"
+                provider, api_key, selected_llm_provider="Anthropic", session_state=None
             ):
                 message = save_api_key(provider, api_key)
                 status = get_api_key_status(selected_llm_provider)
+
+                # If this is an LLM provider (not Hugging Face), recreate the agent
+                if provider != "Hugging Face" and session_state is not None:
+                    agent_message = self.recreate_agent_with_new_model(
+                        session_state, provider
+                    )
+                    if agent_message:
+                        message += f"\n{agent_message}"
+
                 return message, status, ""  # Clear the input field
 
             hf_save_btn.click(
                 lambda key, llm_prov: save_and_update_status(
-                    "Hugging Face", key, llm_prov
+                    "Hugging Face", key, llm_prov, session_state
                 ),
                 inputs=[hf_token, llm_provider],
                 outputs=[api_message, api_status, hf_token],
             ).then(lambda: gr.Textbox(visible=True), outputs=[api_message])
 
             llm_save_btn.click(
-                lambda provider, key: save_and_update_status(provider, key, provider),
+                lambda provider, key: save_and_update_status(
+                    provider, key, provider, session_state
+                ),
                 inputs=[llm_provider, llm_token],
                 outputs=[api_message, api_status, llm_token],
             ).then(lambda: gr.Textbox(visible=True), outputs=[api_message])
@@ -689,6 +626,35 @@ class GradioUI:
 
         return demo
 
+    def recreate_agent_with_new_model(self, session_state, provider=None):
+        """Recreate the agent with updated model configuration."""
+        from kiss_agent import KISSAgent
+
+        # Get the new model ID if provider is specified
+        if provider and provider != "Hugging Face":
+            model_id = get_default_model_for_provider(provider)
+
+            # Get API key from provider-specific environment variable
+            env_var_map = {
+                "Anthropic": "ANTHROPIC_API_KEY",
+                "OpenAI": "OPENAI_API_KEY",
+                "SambaNova": "SAMBANOVA_API_KEY",
+                "Mistral": "MISTRAL_API_KEY",
+            }
+
+            env_var_name = env_var_map.get(provider)
+            api_key = os.getenv(env_var_name) if env_var_name else None
+
+            if not api_key:
+                return f"❌ No API key found for {provider}"
+
+            # Create new agent with updated model
+            new_agent = KISSAgent(model_id=model_id, api_key=api_key)
+            session_state["agent"] = new_agent
+
+            return f"🔄 Agent updated to use {provider} model: {model_id}"
+        return ""
+
 
 if __name__ == "__main__":
     from kiss_agent import KISSAgent