Fix UI response display issues and enhance chat handling

app.py

@@ -35,6 +35,8 @@ if "show_welcome" not in st.session_state:
     st.session_state.show_welcome = True
 if "session_id" not in st.session_state:
     st.session_state.session_id = f"sess_{int(time.time())}_{os.urandom(4).hex()}"
+if "last_processed_message" not in st.session_state:
+    st.session_state.last_processed_message = ""
 
 # Start session tracking
 session_analytics.start_session_tracking("default_user", st.session_state.session_id)
@@ -130,6 +132,7 @@ with st.sidebar:
     if st.button("🗑️ Clear History"):
         session_analytics.track_interaction("default_user", st.session_state.session_id, "clear_history_click")
         st.session_state.messages = []
+        st.session_state.last_processed_message = ""
         # Also clear backend session
         session_manager.clear_session("default_user")
         st.success("History cleared!")
@@ -259,11 +262,13 @@ for message in st.session_state.messages:
 # Chat input with enhanced processing
 user_input = st.chat_input("Type your message here...", key="chat_input")
 
-if user_input:
-    session_analytics.track_interaction("default_user", st.session_state.session_id, "chat_message_sent", {
-        "message_length": len(user_input)
-    })
-    chat_handler.process_user_message(user_input, selected_model_name)
+# Process message when received
+if user_input and user_input.strip():
+    # Prevent processing if already processing
+    if not getattr(chat_handler, 'is_processing', False):
+        chat_handler.process_user_message(user_input, selected_model_name)
+    else:
+        st.warning("Still processing your previous request...")
 
 # About tab
 st.divider()

src/llm/ollama_provider.py

@@ -30,21 +30,39 @@ class OllamaProvider(EnhancedLLMProvider):
 
     def generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[str]:
         """Generate a response synchronously with context enrichment"""
+        return self._retry_with_backoff(self._generate_impl, prompt, conversation_history)
+
+    def stream_generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[Union[str, List[str]]]:
+        """Generate a response with streaming support"""
+        return self._retry_with_backoff(self._stream_generate_impl, prompt, conversation_history)
+
+    def _generate_impl(self, prompt: str, conversation_history: List[Dict]) -> str:
+        """Implementation of synchronous generation with enhanced debugging"""
         try:
             # Enrich context
             enriched_history = self._enrich_context(conversation_history)
             
             url = f"{self.host}/api/chat"
+            
+            # Prepare messages - ensure proper format
+            messages = []
+            for msg in enriched_history:
+                if isinstance(msg, dict) and "role" in msg and "content" in msg:
+                    messages.append({
+                        "role": msg["role"],
+                        "content": str(msg["content"])
+                    })
+
             payload = {
                 "model": self.model_name,
-                "messages": enriched_history,
+                "messages": messages,
                 "stream": False
             }
-            
+
             logger.info(f"Ollama request URL: {url}")
             logger.info(f"Ollama request payload: {payload}")
             logger.info(f"Ollama headers: {self.headers}")
-            
+
             response = requests.post(
                 url,
                 json=payload,
@@ -54,31 +72,32 @@ class OllamaProvider(EnhancedLLMProvider):
             
             logger.info(f"Ollama response status: {response.status_code}")
             logger.info(f"Ollama response headers: {dict(response.headers)}")
-            
+
             response.raise_for_status()
             result = response.json()
             logger.info(f"Ollama response body: {result}")
-            
-            # Extract content properly
+
+            # Extract content properly with multiple fallbacks
+            content = None
             if "message" in result and "content" in result["message"]:
                 content = result["message"]["content"]
-                logger.info(f"Extracted content: {content[:100] if content else 'None'}")
-                return content
             elif "response" in result:
                 content = result["response"]
-                logger.info(f"Extracted response: {content[:100] if content else 'None'}")
-                return content
             else:
                 content = str(result)
-                logger.info(f"Raw result as string: {content[:100]}")
-                return content
                 
+            logger.info(f"Extracted content length: {len(content) if content else 0}")
+            return content if content else ""
+
+        except requests.exceptions.RequestException as e:
+            logger.error(f"Ollama API request error: {str(e)}")
+            raise Exception(f"Ollama API error: {str(e)}")
         except Exception as e:
-            logger.error(f"Ollama generation failed: {e}")
-            raise
+            logger.error(f"Failed to parse Ollama response: {str(e)}")
+            raise Exception(f"Failed to parse Ollama response: {str(e)}")
 
-    def stream_generate(self, prompt: str, conversation_history: List[Dict]) -> Optional[Union[str, List[str]]]:
-        """Generate a response with streaming support"""
+    def _stream_generate_impl(self, prompt: str, conversation_history: List[Dict]) -> List[str]:
+        """Implementation of streaming generation"""
         try:
             # Enrich context
             enriched_history = self._enrich_context(conversation_history)

src/ui/chat_handler.py

@@ -3,7 +3,6 @@ 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__)
@@ -20,10 +19,12 @@ class ChatHandler:
             st.warning("Please enter a message")
             return
             
-        if self.is_processing:
-            st.warning("Still processing previous request...")
-            return
-            
+        # Prevent duplicate processing
+        if hasattr(st.session_state, 'last_processed_message'):
+            if st.session_state.last_processed_message == user_input:
+                return
+        
+        st.session_state.last_processed_message = user_input
         self.is_processing = True
         
         try:
@@ -33,7 +34,10 @@ class ChatHandler:
                 st.markdown(user_input)
                 st.caption(f"🕒 {timestamp}")
             
-            # Add to session history
+            # Add to session state history
+            if "messages" not in st.session_state:
+                st.session_state.messages = []
+                
             st.session_state.messages.append({
                 "role": "user",
                 "content": user_input,
@@ -50,10 +54,15 @@ class ChatHandler:
                     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)
+                    # Get response with timeout handling
+                    response = None
+                    try:
+                        response = self._get_ai_response(user_input, provider_name)
+                    except Exception as e:
+                        logger.error(f"AI response error: {e}")
+                        raise
                     
-                    if response:
+                    if response and response.strip():
                         status_placeholder.success("✅ Response received!")
                         response_placeholder.markdown(response)
                         
@@ -65,8 +74,8 @@ class ChatHandler:
                             "provider": provider_name
                         })
                     else:
-                        status_placeholder.error("❌ Empty response received")
-                        response_placeholder.markdown("I received your message but couldn't generate a proper response.")
+                        status_placeholder.warning("⚠️ Empty response received")
+                        response_placeholder.markdown("*No response generated. Please try again.*")
                         
                 except ProviderNotAvailableError as e:
                     status_placeholder.error("❌ No AI providers available")
@@ -76,21 +85,24 @@ class ChatHandler:
                 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}")
+                    logger.error(f"Chat processing error: {e}", exc_info=True)
                     
+        except Exception as e:
+            logger.error(f"Unexpected error in process_user_message: {e}", exc_info=True)
+            st.error("An unexpected error occurred. Please try again.")
         finally:
             self.is_processing = False
-            time.sleep(0.1)  # Small delay to ensure UI updates
-            # st.experimental_rerun()  # Removed to prevent automatic rerun
+            # Force UI refresh
+            st.experimental_rerun()
     
     def _get_provider_for_model(self, selected_model: str) -> str:
         """Determine which provider to use based on model selection"""
-        model_map = {
+        model_to_provider = {
             "Mistral 7B (Local)": "ollama",
-            "Llama 2 7B (Local)": "ollama", 
+            "Llama 2 7B (Local)": "ollama",
             "OpenChat 3.5 (Local)": "ollama"
         }
-        return model_map.get(selected_model, "ollama")
+        return model_to_provider.get(selected_model, "ollama")
     
     def _get_provider_display_name(self, provider_name: str) -> str:
         """Get display name for provider"""
@@ -113,8 +125,10 @@ class ChatHandler:
             # Get provider
             provider = llm_factory.get_provider(provider_name)
             
-            # Generate response
+            # Generate response with timeout
+            logger.info(f"Generating response with {provider_name} provider")
             response = provider.generate(user_input, conversation_history)
+            logger.info(f"Received response from {provider_name}: {response[:100] if response else 'None'}")
             
             # Update session with conversation
             if response:
@@ -126,7 +140,7 @@ class ChatHandler:
             return response
             
         except Exception as e:
-            logger.error(f"AI response generation failed: {e}")
+            logger.error(f"AI response generation failed: {e}", exc_info=True)
             raise
 
 # Global instance