Fix chat interface response issues: proper session handling, error handling, and UI display fixes

api/chat.py

@@ -1,44 +1,70 @@
 import json
 from fastapi import APIRouter, HTTPException
-from fastapi.responses import StreamingResponse
+from fastapi.responses import JSONResponse
 from core.llm import LLMClient
-from core.memory import save_user_state, load_user_state
 from core.session import session_manager
 import logging
 
 router = APIRouter()
 logger = logging.getLogger(__name__)
 
-llm_client = LLMClient(provider="ollama")  # Default to Ollama
+# Initialize LLM client with fallback support
+llm_client = LLMClient()
 
 @router.post("/chat")
 async def chat(user_id: str, message: str):
-    if not message:
+    """
+    Handle chat requests with proper session management and error handling.
+    """
+    if not message or not message.strip():
         raise HTTPException(status_code=400, detail="Message is required")
-    
+
     try:
-        # Use session manager for better session handling
+        # Get session using session manager
         session = session_manager.get_session(user_id)
         conversation_history = session.get("conversation", [])
         
         # Add user message to history
-        conversation_history.append({"role": "user", "content": message})
-        
-        # Generate AI response
-        full_response = ""
-        response_stream = llm_client.generate(prompt=message, stream=True)
-        
-        # Collect streamed response
-        for chunk in response_stream:
-            full_response += chunk
-        
-        # Save updated conversation using session manager
-        conversation_history.append({"role": "assistant", "content": full_response})
-        session_manager.update_session(user_id, {"conversation": conversation_history})
-        
+        conversation_history.append({
+            "role": "user", 
+            "content": message
+        })
+
+        # Generate AI response using factory pattern
+        try:
+            ai_response = llm_client.generate(
+                prompt=message,
+                conversation_history=conversation_history
+            )
+            
+            if not ai_response:
+                raise Exception("Empty response from LLM")
+                
+        except Exception as e:
+            logger.error(f"LLM generation failed: {e}")
+            # Try to provide a graceful fallback
+            ai_response = "I'm having trouble processing your request right now. Please try again."
+
+        # Add AI response to conversation history
+        conversation_history.append({
+            "role": "assistant",
+            "content": ai_response
+        })
+
+        # Update session with new conversation history
+        session_manager.update_session(user_id, {
+            "conversation": conversation_history
+        })
+
         logger.info(f"Successfully processed chat for user {user_id}")
-        return {"response": full_response}
-        
+        return JSONResponse(
+            content={"response": ai_response},
+            status_code=200
+        )
+
     except Exception as e:
-        logger.error(f"LLM generation failed for user {user_id}: {e}")
-        raise HTTPException(status_code=500, detail=f"LLM generation failed: {e}")
+        logger.error(f"Chat processing failed for user {user_id}: {e}")
+        raise HTTPException(
+            status_code=500, 
+            detail=f"Failed to process chat: {str(e)}"
+        )

app.py

@@ -252,7 +252,7 @@ def validate_user_input(text):
     
     return True, text.strip()
 
-# Chat input - FIXED VERSION (moved outside of tabs)
+# Chat input - FIXED VERSION
 user_input = st.chat_input("Type your message here...", disabled=st.session_state.is_processing)
 
 # Process message when received
@@ -278,8 +278,8 @@ if user_input and not st.session_state.is_processing:
         
         # Process AI response
         with st.chat_message("assistant"):
-            response_placeholder = st.empty()
             status_placeholder = st.empty()
+            response_placeholder = st.empty()
             
             try:
                 # Get conversation history
@@ -460,7 +460,7 @@ if user_input and not st.session_state.is_processing:
             finally:
                 st.session_state.is_processing = False
                 # Force UI update
-                st.experimental_rerun()
+                st.rerun()  # Changed from experimental_rerun() to rerun()
 
 # Add evaluation dashboard tab (separate from chat interface) - ONLY ABOUT TAB NOW
 st.divider()
@@ -506,4 +506,4 @@ if user_input and user_input.lower().strip() in ["tell me a story", "tell me a c
         "timestamp": datetime.now().strftime("%H:%M:%S")
     })
     st.session_state.is_processing = False
-    st.experimental_rerun()
+    st.rerun()

core/llm.py

@@ -5,47 +5,32 @@ from core.llm_factory import llm_factory, ProviderNotAvailableError
 logger = logging.getLogger(__name__)
 
 class LLMClient:
-    """High-level LLM client that uses the factory pattern"""
-    
-    def __init__(self, provider: Optional[str] = None):
-        self.provider_name = provider
+    """High-level LLM client that uses the factory pattern with improved error handling"""
+
+    def __init__(self):
         try:
-            self.provider = llm_factory.get_provider(provider)
+            self.provider = llm_factory.get_provider()
         except ProviderNotAvailableError:
             self.provider = None
             logger.error("No LLM providers available")
-    
-    def generate(self, prompt: str, conversation_history: List[Dict], stream: bool = False):
-        """Generate a response"""
+
+    def generate(self, prompt: str, conversation_history: List[Dict], stream: bool = False) -> Optional[str]:
+        """
+        Generate a response with robust error handling.
+        """
         if not self.provider:
             raise ProviderNotAvailableError("No LLM provider available")
-            
+
         try:
             if stream:
-                return self.provider.stream_generate(prompt, conversation_history)
+                result = self.provider.stream_generate(prompt, conversation_history)
+                # For streaming, combine chunks into single response
+                if isinstance(result, list):
+                    return "".join(result)
+                return result
             else:
                 return self.provider.generate(prompt, conversation_history)
+                
         except Exception as e:
             logger.error(f"LLM generation failed: {e}")
-            raise
-
-def send_to_ollama(prompt: str, conversation_history: List[Dict], ollama_url: str, model: str) -> Optional[str]:
-    """Legacy function for backward compatibility"""
-    try:
-        from core.providers.ollama import OllamaProvider
-        provider = OllamaProvider(model)
-        return provider.generate(prompt, conversation_history)
-    except Exception as e:
-        logger.error(f"Ollama call failed: {e}")
-        return None
-
-def send_to_hf(prompt: str, conversation_history: List[Dict]) -> Optional[str]:
-    """Legacy function for backward compatibility"""
-    try:
-        from utils.config import config
-        from core.providers.huggingface import HuggingFaceProvider
-        provider = HuggingFaceProvider("meta-llama/Llama-2-7b-chat-hf")
-        return provider.generate(prompt, conversation_history)
-    except Exception as e:
-        logger.error(f"Hugging Face call failed: {e}")
-        return None
+            raise  # Re-raise to let caller handle appropriately