Update ui_components.py

ui_components.py

@@ -1,5 +1,5 @@
 """
-UI Components for Universal MCP Client - Updated with ChatMessage support
+UI Components for Universal MCP Client - Enhanced with Inference Provider Support
 """
 import gradio as gr
 from gradio import ChatMessage
@@ -11,15 +11,20 @@ from server_manager import ServerManager
 from mcp_client import UniversalMCPClient
 
 class UIComponents:
-    """Manages Gradio UI components and event handlers with ChatMessage support"""
+    """Manages Gradio UI components and event handlers with Inference Provider support"""
     
     def __init__(self, mcp_client: UniversalMCPClient):
         self.mcp_client = mcp_client
         self.chat_handler = ChatHandler(mcp_client)
         self.server_manager = ServerManager(mcp_client)
+        
+        # State for current LLM backend
+        self.current_backend = "anthropic"  # "anthropic" or "hf_inference"
+        self.current_provider = None
+        self.current_model = None
     
     def create_interface(self) -> gr.Blocks:
-        """Create the main Gradio interface"""
+        """Create the main Gradio interface with provider selection"""
         with gr.Blocks(
             title="Universal MCP Client", 
             theme=getattr(gr.themes, AppConfig.GRADIO_THEME.title())(),
@@ -39,41 +44,129 @@ class UIComponents:
         return demo
     
     def _create_sidebar(self):
-        """Create the sidebar with server management"""
+        """Create the sidebar with LLM provider selection and server management"""
         with gr.Sidebar(elem_id="main-sidebar"):
             gr.Markdown("# Gradio.chat.app")
             
+            # LLM Backend Selection Section
+            self._create_llm_backend_selection()
+            
             # Collapsible information section
             with gr.Accordion("📚 Guide & Info", open=False):
                 gr.Markdown("""
                 ## 🎯 How To Use
-                - **Chat with LLM**: Select your fav LLM to chat to
-                - **Generation**: Select an MCP Server from hundreds of Gradio spaces
-                - **Subscribe**: PRO subscribers get higher usage on ZeroGPU spaces
+                - **Choose LLM Backend**: Select between Anthropic Claude or HuggingFace Inference Providers
+                - **Chat with LLM**: Interact with your selected model
+                - **MCP Integration**: Add MCP servers for enhanced capabilities
+                - **Subscribe**: PRO subscribers get higher usage on various services
                 
                 ## 💭 New UI Features
+                - **Provider Selection**: Choose from multiple inference providers
                 - **Tool Usage**: See tool calls in collapsible accordions with 🔧 icons
                 - **Results**: View tool results in nested thoughts with 📋 icons
                 - **Media Output**: Images, audio, and videos display separately from text
                 - **Real-time Status**: Watch tools execute with pending/done status indicators
                 """)
             
-            # Server status
-            gr.Markdown("## 🔧 Server Status")
-            self.server_count_display = gr.Markdown(f"**Connected Servers**: {len(self.mcp_client.servers)}")
+            # Current backend status
+            self.backend_status = gr.Markdown("**Current Backend**: Not configured")
             
-            if self.mcp_client.servers:
-                server_list = "\n".join([f"• **{name}**" for name in self.mcp_client.servers.keys()])
-                self.server_list_display = gr.Markdown(server_list)
+            # Server management
+            self._create_server_management_section()
+    
+    def _create_llm_backend_selection(self):
+        """Create LLM backend selection interface"""
+        gr.Markdown("## 🤖 LLM Backend Selection")
+        
+        # Radio buttons for backend selection
+        self.backend_radio = gr.Radio(
+            choices=[
+                ("Anthropic Claude Sonnet 4", "anthropic"),
+                ("HuggingFace Inference Providers", "hf_inference")
+            ],
+            value="anthropic",
+            label="Choose LLM Backend",
+            info="Select your preferred language model backend"
+        )
+        
+        # Anthropic configuration section
+        with gr.Column(visible=True, elem_classes="provider-selection anthropic-config") as self.anthropic_config:
+            gr.Markdown("### 🔹 Anthropic Claude Configuration")
+            
+            # Check if API key is available
+            if AppConfig.ANTHROPIC_API_KEY:
+                self.anthropic_status = gr.Markdown("✅ **Status**: API key configured")
             else:
-                self.server_list_display = gr.Markdown("*No servers connected*\n\nAdd servers below.")
+                self.anthropic_status = gr.Markdown("""
+                ❌ **Status**: API key not found
+                
+                **Setup Instructions**:
+                1. Go to Space Settings → Secrets
+                2. Add `ANTHROPIC_API_KEY` with your Anthropic API key
+                3. Restart the space
+                """)
+        
+        # HuggingFace Inference Provider configuration section
+        with gr.Column(visible=False, elem_classes="provider-selection hf-config") as self.hf_config:
+            gr.Markdown("### 🔸 HuggingFace Inference Provider Configuration")
             
-            # Server management in accordion
-            self._create_server_management_section()
+            # Check if HF token is available
+            if AppConfig.HF_TOKEN:
+                self.hf_status = gr.Markdown("✅ **Status**: HF token configured")
+            else:
+                self.hf_status = gr.Markdown("""
+                ❌ **Status**: HF token not found
+                
+                **Setup Instructions**:
+                1. Go to Space Settings → Secrets
+                2. Add `HF_TOKEN` with your HuggingFace token
+                3. Restart the space
+                """)
+            
+            # Provider dropdown
+            provider_choices = list(AppConfig.INFERENCE_PROVIDERS.keys())
+            self.provider_dropdown = gr.Dropdown(
+                choices=provider_choices,
+                label="🚀 Select Inference Provider",
+                value=provider_choices[0] if provider_choices else None,
+                info="Choose your preferred inference provider"
+            )
+            
+            # Model dropdown (initially empty, populated based on provider)
+            self.model_dropdown = gr.Dropdown(
+                choices=[],
+                label="🤖 Select Model",
+                value=None,
+                info="Choose the specific model to use"
+            )
+            
+            # Configure button
+            self.configure_btn = gr.Button(
+                "Configure Inference Provider", 
+                variant="primary",
+                size="sm"
+            )
+            
+            self.provider_config_status = gr.Textbox(
+                label="Configuration Status",
+                interactive=False,
+                visible=False
+            )
     
     def _create_server_management_section(self):
         """Create the server management section in sidebar"""
-        with gr.Accordion("⚙️ Manage Servers", open=False):
+        gr.Markdown("## 🔧 MCP Server Management")
+        
+        # Server status
+        self.server_count_display = gr.Markdown(f"**Connected Servers**: {len(self.mcp_client.servers)}")
+        
+        if self.mcp_client.servers:
+            server_list = "\n".join([f"• **{name}**" for name in self.mcp_client.servers.keys()])
+            self.server_list_display = gr.Markdown(server_list)
+        else:
+            self.server_list_display = gr.Markdown("*No servers connected*\n\nAdd servers below.")
+        
+        with gr.Accordion("⚙️ Manage MCP Servers", open=False):
             gr.Markdown("**Add MCP Servers**")
             
             # Get MCP spaces count for dropdown label
@@ -122,7 +215,7 @@ class UIComponents:
                 value=[
                     ChatMessage(
                         role="assistant",
-                        content="Welcome! I'm your MCP-powered AI assistant. I can help you with various tasks using connected MCP servers. Upload files, ask questions, or request content generation!"
+                        content="Welcome! I'm your MCP-powered AI assistant. Please configure your LLM backend in the sidebar to get started. I can help you with various tasks using connected MCP servers!"
                     )
                 ]
             )
@@ -141,11 +234,84 @@ class UIComponents:
         return chatbot
     
     def _setup_event_handlers(self, chatbot: gr.Chatbot):
-        """Set up all event handlers for the interface with ChatMessage support"""
+        """Set up all event handlers for the interface"""
+        
+        # Backend selection event handlers
+        def handle_backend_change(backend_choice):
+            """Handle LLM backend selection change"""
+            self.current_backend = backend_choice
+            
+            # Update visibility of configuration sections
+            anthropic_visible = (backend_choice == "anthropic")
+            hf_visible = (backend_choice == "hf_inference")
+            
+            # Update backend status
+            if backend_choice == "anthropic":
+                if AppConfig.ANTHROPIC_API_KEY:
+                    status = "**Current Backend**: Anthropic Claude Sonnet 4 ✅"
+                else:
+                    status = "**Current Backend**: Anthropic Claude (❌ API key needed)"
+            else:
+                if AppConfig.HF_TOKEN:
+                    status = "**Current Backend**: HF Inference Providers (⚙️ Configure provider)"
+                else:
+                    status = "**Current Backend**: HF Inference Providers (❌ HF token needed)"
+            
+            return (
+                gr.update(visible=anthropic_visible),  # anthropic_config
+                gr.update(visible=hf_visible),         # hf_config
+                status                                 # backend_status
+            )
+        
+        def handle_provider_change(provider):
+            """Handle inference provider change"""
+            if not provider:
+                return gr.update(choices=[], value=None)
+            
+            models = AppConfig.get_provider_models(provider)
+            return gr.update(choices=models, value=models[0] if models else None)
+        
+        def handle_provider_configuration(provider, model):
+            """Handle inference provider configuration"""
+            if not provider or not model:
+                return "❌ Please select both provider and model", gr.update(visible=True)
+            
+            if not AppConfig.HF_TOKEN:
+                return "❌ HF_TOKEN not configured. Please add it to space secrets.", gr.update(visible=True)
+            
+            # Configure the MCP client with the selected provider
+            success = self.mcp_client.configure_inference_provider(provider, model)
+            
+            if success:
+                self.current_provider = provider
+                self.current_model = model
+                
+                # Update chat handler
+                self.chat_handler.mcp_client = self.mcp_client
+                
+                status_msg = f"✅ Configured {provider} with {model}"
+                
+                # Update backend status
+                backend_status = f"**Current Backend**: {provider}/{model} ✅"
+                
+                return status_msg, gr.update(visible=True)
+            else:
+                return "❌ Failed to configure inference provider", gr.update(visible=True)
         
         # Chat event handlers
         def submit_message(message, history):
             """Handle message submission with ChatMessage support"""
+            # Check if backend is properly configured
+            if self.current_backend == "anthropic" and not AppConfig.ANTHROPIC_API_KEY:
+                error_msg = "❌ Please configure Anthropic API key in space settings first."
+                history.append(ChatMessage(role="assistant", content=error_msg))
+                return history, gr.MultimodalTextbox(value=None, interactive=False)
+            
+            if self.current_backend == "hf_inference" and (not self.current_provider or not self.current_model):
+                error_msg = "❌ Please configure HuggingFace inference provider first."
+                history.append(ChatMessage(role="assistant", content=error_msg))
+                return history, gr.MultimodalTextbox(value=None, interactive=False)
+            
             if message and (message.get("text", "").strip() or message.get("files", [])):
                 # Convert existing history to ChatMessage objects if needed
                 converted_history = []
@@ -169,6 +335,30 @@ class UIComponents:
             """Re-enable input after processing"""
             return gr.MultimodalTextbox(interactive=True)
         
+        # Connect backend selection events
+        self.backend_radio.change(
+            handle_backend_change,
+            inputs=[self.backend_radio],
+            outputs=[self.anthropic_config, self.hf_config, self.backend_status]
+        )
+        
+        # Connect provider selection events
+        self.provider_dropdown.change(
+            handle_provider_change,
+            inputs=[self.provider_dropdown],
+            outputs=[self.model_dropdown]
+        )
+        
+        # Connect provider configuration
+        self.configure_btn.click(
+            handle_provider_configuration,
+            inputs=[self.provider_dropdown, self.model_dropdown],
+            outputs=[self.provider_config_status, self.provider_config_status]
+        ).then(
+            lambda: self._update_backend_status(),
+            outputs=[self.backend_status]
+        )
+        
         # Set up the chat flow - using submit event
         chat_msg_enter = self.chat_input.submit(
             submit_message,
@@ -231,7 +421,7 @@ class UIComponents:
         # Connect server management events
         self.add_server_btn.click(
             handle_add_server,
-            inputs=[self.server_name, self.mcp_dropdown],  # Use dropdown directly
+            inputs=[self.server_name, self.mcp_dropdown],
             outputs=[
                 self.add_server_output, 
                 self.add_server_details, 
@@ -251,56 +441,16 @@ class UIComponents:
             handle_refresh_spaces,
             outputs=[self.mcp_dropdown]
         )
-        
-        # Optional: Add a clear chat button for better UX
-        def clear_chat():
-            """Clear chat history and return to welcome message"""
-            return [
-                ChatMessage(
-                    role="assistant",
-                    content="Chat cleared! I'm ready to help you with new tasks using MCP servers."
-                )
-            ]
-        
-        # You can add a clear button in the interface if needed
-        # clear_btn = gr.Button("🗑️ Clear Chat", variant="secondary", size="sm")
-        # clear_btn.click(clear_chat, outputs=[chatbot])
     
-    def add_sample_conversation(self, chatbot: gr.Chatbot):
-        """Add a sample conversation to demonstrate ChatMessage features (for testing)"""
-        sample_messages = [
-            ChatMessage(
-                role="user",
-                content="Generate an image of a sunset over mountains"
-            ),
-            ChatMessage(
-                role="assistant",
-                content="",
-                metadata={
-                    "title": "🔧 Using dalle_3_xl_lora_v2_generate",
-                    "id": "tool_1",
-                    "status": "done",
-                    "duration": 3.2,
-                    "log": "Server: text_to_image"
-                }
-            ),
-            ChatMessage(
-                role="assistant",
-                content="[{'image': {'path': '/tmp/generated_image.png', 'url': 'https://example.com/image.png'}}]",
-                metadata={
-                    "title": "📋 Tool Result",
-                    "parent_id": "tool_1",
-                    "status": "done"
-                }
-            ),
-            ChatMessage(
-                role="assistant",
-                content="I've generated a beautiful sunset image for you! The image shows a stunning sunset over mountains with warm orange and pink colors filling the sky."
-            ),
-            ChatMessage(
-                role="assistant",
-                content={"path": "https://example.com/image.png"}
-            )
-        ]
-        
-        return sample_messages
+    def _update_backend_status(self):
+        """Update backend status based on current configuration"""
+        if self.current_backend == "anthropic":
+            if AppConfig.ANTHROPIC_API_KEY:
+                return "**Current Backend**: Anthropic Claude Sonnet 4 ✅"
+            else:
+                return "**Current Backend**: Anthropic Claude (❌ API key needed)"
+        else:
+            if self.current_provider and self.current_model:
+                return f"**Current Backend**: {self.current_provider}/{self.current_model} ✅"
+            else:
+                return "**Current Backend**: HF Inference Providers (⚙️ Configure provider)"