Implement enhanced UI with proper response display and feedback

app.py

@@ -7,12 +7,13 @@ from datetime import datetime
 from pathlib import Path
 sys.path.append(str(Path(__file__).parent))
 
+# Import our new handler
+from src.ui.chat_handler import chat_handler
 from utils.config import config
 from core.session import session_manager
 from core.memory import check_redis_health
 from core.errors import translate_error
 from core.personality import personality
-from core.providers.ollama import OllamaProvider
 import logging
 
 # Set up logging
@@ -41,7 +42,7 @@ with st.sidebar:
     # Model selection
     model_options = {
         "Mistral 7B (Local)": "mistral:latest",
-        "Llama 2 7B (Local)": "llama2:latest",
+        "Llama 2 7B (Local)": "llama2:latest", 
         "OpenChat 3.5 (Local)": "openchat:latest"
     }
     selected_model_name = st.selectbox(
@@ -70,6 +71,7 @@ with st.sidebar:
     
     if st.button("📡 Test Connection"):
         try:
+            from core.providers.ollama import OllamaProvider
             ollama_provider = OllamaProvider(st.session_state.selected_model)
             is_valid = ollama_provider.validate_model()
             if is_valid:
@@ -81,12 +83,14 @@ with st.sidebar:
     
     if st.button("🗑️ Clear History"):
         st.session_state.messages = []
+        # Also clear backend session
+        session_manager.clear_session("default_user")
         st.success("History cleared!")
     
     st.divider()
     
-    # System Status
-    with st.expander("🔍 System Status", expanded=False):
+    # System Status with enhanced HF monitoring
+    with st.expander("🔍 System Status", expanded=True):
         st.subheader("📊 Status")
         
         # Ollama Status
@@ -100,11 +104,45 @@ with st.sidebar:
         except:
             st.info("🦙 Ollama: Unknown")
         
+        # HF Endpoint Status (Enhanced)
+        try:
+            from src.services.hf_monitor import hf_monitor
+            status_message = hf_monitor.get_human_readable_status()
+            
+            # Display appropriate status icon
+            if "🟢" in status_message:
+                st.success(status_message)
+            elif "🟡" in status_message:
+                st.warning(status_message)
+            elif "🔴" in status_message or "❌" in status_message:
+                st.error(status_message)
+            elif "⏳" in status_message:
+                st.info(status_message)
+            else:
+                st.info(status_message)
+                
+            # Add wake-up button if scaled to zero
+            if "scaled to zero" in status_message.lower():
+                if st.button("⚡ Wake Up HF Endpoint", key="wake_up_hf"):
+                    with st.spinner("Waking up HF endpoint... This may take 2-4 minutes..."):
+                        if hf_monitor.attempt_wake_up():
+                            st.success("✅ HF endpoint is waking up! Try your request again in a moment.")
+                            time.sleep(2)
+                            st.experimental_rerun()
+                        else:
+                            st.error("❌ Failed to wake up HF endpoint. Please try again.")
+                            
+        except Exception as e:
+            st.info(f"🤗 HF Endpoint: Error checking status - {str(e)}")
+        
         # Redis Status
-        if check_redis_health():
-            st.success("💾 Redis: Connected")
-        else:
-            st.error("💾 Redis: Disconnected")
+        try:
+            if check_redis_health():
+                st.success("💾 Redis: Connected")
+            else:
+                st.error("💾 Redis: Disconnected")
+        except:
+            st.info("💾 Redis: Unknown")
     
     st.divider()
     
@@ -124,88 +162,19 @@ if st.session_state.show_welcome:
         st.markdown(greeting)
     st.session_state.show_welcome = False
 
-# Display messages
+# Display conversation history
 for message in st.session_state.messages:
     with st.chat_message(message["role"]):
         st.markdown(message["content"])
         if "timestamp" in message:
-            st.caption(f"🕒 {message['timestamp']}")
+            provider_info = f" (via {message.get('provider', 'unknown')})" if message["role"] == "assistant" else ""
+            st.caption(f"🕒 {message['timestamp']}{provider_info}")
 
-# Chat input
-user_input = st.chat_input("Type your message here...", disabled=st.session_state.is_processing)
+# Chat input with enhanced processing
+user_input = st.chat_input("Type your message here...", key="chat_input")
 
-# Process message
-if user_input and not st.session_state.is_processing:
-    st.session_state.is_processing = True
-    
-    # Display user message
-    with st.chat_message("user"):
-        st.markdown(user_input)
-    
-    # Add to history
-    timestamp = datetime.now().strftime("%H:%M:%S")
-    st.session_state.messages.append({
-        "role": "user",
-        "content": user_input,
-        "timestamp": timestamp
-    })
-    
-    # Process response
-    try:
-        # Get conversation history
-        user_session = session_manager.get_session("default_user")
-        conversation_history = user_session.get("conversation", []).copy()
-        conversation_history.append({"role": "user", "content": user_input})
-        
-        # Generate response
-        try:
-            ollama_provider = OllamaProvider(st.session_state.selected_model)
-            ai_response = ollama_provider.generate(user_input, conversation_history)
-            
-            if ai_response and ai_response.strip():
-                with st.chat_message("assistant"):
-                    st.markdown(ai_response)
-                status = "✅ Response received!"
-            else:
-                with st.chat_message("assistant"):
-                    st.warning("⚠️ Received empty response")
-                ai_response = "I received your message but couldn't generate a proper response."
-                
-        except Exception as e:
-            error_message = str(e)
-            with st.chat_message("assistant"):
-                st.error(f"❌ Error: {error_message[:100]}...")
-            ai_response = f"Error: {error_message[:100]}..."
-        
-        # Save to session
-        if ai_response:
-            try:
-                conversation = user_session.get("conversation", []).copy()
-                conversation.append({"role": "user", "content": user_input})
-                conversation.append({"role": "assistant", "content": str(ai_response)})
-                session_manager.update_session("default_user", {"conversation": conversation})
-                
-                st.session_state.messages.append({
-                    "role": "assistant",
-                    "content": str(ai_response),
-                    "timestamp": timestamp
-                })
-            except Exception as session_error:
-                logger.error(f"Session update error: {session_error}")
-                
-    except Exception as e:
-        error_msg = f"System error: {str(e)}"
-        logger.error(f"Processing error: {error_msg}")
-        with st.chat_message("assistant"):
-            st.error(error_msg)
-        st.session_state.messages.append({
-            "role": "assistant",
-            "content": error_msg,
-            "timestamp": timestamp
-        })
-    finally:
-        st.session_state.is_processing = False
-        st.experimental_rerun()
+if user_input:
+    chat_handler.process_user_message(user_input, selected_model_name)
 
 # About tab
 st.divider()
@@ -220,11 +189,13 @@ with tab1:
     - **Local AI processing** with Ollama models
     - **Persistent memory** using Redis
     - **Space-themed personality** for fun interactions
+    - **HF Endpoint integration** for advanced capabilities
     
     ### 🚀 Cosmic Mode
     When enabled, the AI responds with space-themed language and metaphors.
     
     ### 🛠️ Technical Architecture
     - **Primary model**: Ollama (local processing)
+    - **Secondary model**: HF Endpoint (advanced processing)
     - **Memory system**: Redis-based session management
     """)

src/ui/chat_handler.py

@@ -0,0 +1,133 @@
+import streamlit as st
+import time
+import logging
+from typing import Optional
+from src.llm.factory import llm_factory, ProviderNotAvailableError
+from src.services.hf_monitor import hf_monitor
+from core.session import session_manager
+
+logger = logging.getLogger(__name__)
+
+class ChatHandler:
+    """Handles chat interactions with better UI feedback"""
+    
+    def __init__(self):
+        self.is_processing = False
+        
+    def process_user_message(self, user_input: str, selected_model: str):
+        """Process user message with enhanced UI feedback"""
+        if not user_input or not user_input.strip():
+            st.warning("Please enter a message")
+            return
+            
+        if self.is_processing:
+            st.warning("Still processing previous request...")
+            return
+            
+        self.is_processing = True
+        
+        try:
+            # Show user message immediately
+            timestamp = time.strftime("%H:%M:%S")
+            with st.chat_message("user"):
+                st.markdown(user_input)
+                st.caption(f"🕒 {timestamp}")
+            
+            # Add to session history
+            st.session_state.messages.append({
+                "role": "user",
+                "content": user_input,
+                "timestamp": timestamp
+            })
+            
+            # Show processing status
+            with st.chat_message("assistant"):
+                status_placeholder = st.empty()
+                response_placeholder = st.empty()
+                
+                try:
+                    # Determine provider based on model selection
+                    provider_name = self._get_provider_for_model(selected_model)
+                    status_placeholder.info(f"🚀 Contacting {self._get_provider_display_name(provider_name)}...")
+                    
+                    # Get response
+                    response = self._get_ai_response(user_input, provider_name)
+                    
+                    if response:
+                        status_placeholder.success("✅ Response received!")
+                        response_placeholder.markdown(response)
+                        
+                        # Add to session history
+                        st.session_state.messages.append({
+                            "role": "assistant",
+                            "content": response,
+                            "timestamp": time.strftime("%H:%M:%S"),
+                            "provider": provider_name
+                        })
+                    else:
+                        status_placeholder.error("❌ Empty response received")
+                        response_placeholder.markdown("I received your message but couldn't generate a proper response.")
+                        
+                except ProviderNotAvailableError as e:
+                    status_placeholder.error("❌ No AI providers available")
+                    response_placeholder.markdown("No AI providers are configured. Please check your settings.")
+                    logger.error(f"Provider not available: {e}")
+                    
+                except Exception as e:
+                    status_placeholder.error(f"❌ Error: {str(e)[:100]}...")
+                    response_placeholder.markdown(f"Sorry, I encountered an error: {str(e)[:100]}...")
+                    logger.error(f"Chat processing error: {e}")
+                    
+        finally:
+            self.is_processing = False
+            time.sleep(0.1)  # Small delay to ensure UI updates
+            # st.experimental_rerun()  # Removed to prevent automatic rerun
+    
+    def _get_provider_for_model(self, selected_model: str) -> str:
+        """Determine which provider to use based on model selection"""
+        model_map = {
+            "Mistral 7B (Local)": "ollama",
+            "Llama 2 7B (Local)": "ollama", 
+            "OpenChat 3.5 (Local)": "ollama"
+        }
+        return model_map.get(selected_model, "ollama")
+    
+    def _get_provider_display_name(self, provider_name: str) -> str:
+        """Get display name for provider"""
+        display_names = {
+            "ollama": "🦙 Ollama",
+            "huggingface": "🤗 HF Endpoint"
+        }
+        return display_names.get(provider_name, provider_name)
+    
+    def _get_ai_response(self, user_input: str, provider_name: str) -> Optional[str]:
+        """Get AI response from specified provider"""
+        try:
+            # Get session and conversation history
+            user_session = session_manager.get_session("default_user")
+            conversation_history = user_session.get("conversation", []).copy()
+            
+            # Add current user message
+            conversation_history.append({"role": "user", "content": user_input})
+            
+            # Get provider
+            provider = llm_factory.get_provider(provider_name)
+            
+            # Generate response
+            response = provider.generate(user_input, conversation_history)
+            
+            # Update session with conversation
+            if response:
+                conversation = user_session.get("conversation", []).copy()
+                conversation.append({"role": "user", "content": user_input})
+                conversation.append({"role": "assistant", "content": response})
+                session_manager.update_session("default_user", {"conversation": conversation})
+            
+            return response
+            
+        except Exception as e:
+            logger.error(f"AI response generation failed: {e}")
+            raise
+
+# Global instance
+chat_handler = ChatHandler()