feat(chat): migrate bookseller chat to WebSockets for real-time streaming

Author: diksha78dev

backend/ai_service.py

@@ -508,6 +508,113 @@ def _build_flat_prompt(system: str, msgs: list) -> str:
             logger.error(f"generate_chat failed: {type(e).__name__}: {e}", exc_info=True)
 
         return None
+
+    def generate_chat_stream(self, system_prompt: str, messages: list, max_tokens: Optional[int] = None):
+        """Yield chunks of text from the LLM for streaming."""
+        if not self.is_available():
+            yield "The candles are flickering and my connection to the ether seems troubled today. "
+            return
+
+        if max_tokens is None:
+            max_tokens = self.config.get('gemini_max_tokens', 600)
+
+        # Per-model context window limits (conservative estimates)
+        _MODEL_CONTEXT_LIMITS = {
+            'llama-3.1-8b-instant': 8192,
+            'llama-3.3-70b-versatile': 32768,
+            'gpt-3.5-turbo': 4096,
+            'gpt-4': 8192,
+            'gpt-4o': 16384,
+            'models/gemini-2.0-flash-lite': 32768,
+            'models/gemini-1.5-flash': 32768,
+        }
+
+        def _get_context_limit(model_name: str) -> int:
+            return _MODEL_CONTEXT_LIMITS.get(model_name, 4096)
+
+        try:
+            # --- Gemini ---
+            if self.gemini_client and (self.preferred_llm == 'gemini' or not self.groq_client):
+                try:
+                    from google.genai import types
+                    context_limit = _get_context_limit(self.config['gemini_model'])
+                    trimmed = self.trim_history_to_token_budget(
+                        system_prompt, messages, max_tokens, context_limit
+                    )
+                    
+                    flat_prompt = system_prompt + "\n\n"
+                    for m in trimmed:
+                        role = "Elara" if m.get("role") == "assistant" else "Customer"
+                        flat_prompt += f"{role}: {m.get('content', '')}\n"
+                    flat_prompt += "Elara:"
+                    
+                    response = self.gemini_client.models.generate_content_stream(
+                        model=self.config['gemini_model'],
+                        contents=flat_prompt,
+                    )
+                    for chunk in response:
+                        if chunk.text:
+                            yield chunk.text
+                    return
+                except Exception as e:
+                    logger.warning(f"Gemini streaming chat failed, falling back: {e}")
+
+            # --- Groq ---
+            if self.groq_client:
+                try:
+                    context_limit = _get_context_limit(self.config['groq_model'])
+                    trimmed = self.trim_history_to_token_budget(
+                        system_prompt, messages, max_tokens, context_limit
+                    )
+                    
+                    groq_messages = [{"role": "system", "content": system_prompt}] + [
+                        {"role": m.get("role", "user"), "content": m.get("content", "")}
+                        for m in trimmed
+                    ]
+                    response = self.groq_client.chat.completions.create(
+                        model=self.config['groq_model'],
+                        messages=groq_messages,
+                        max_tokens=min(max_tokens, self.config['groq_max_tokens']),
+                        temperature=self.config['groq_temperature'],
+                        stream=True
+                    )
+                    for chunk in response:
+                        content = chunk.choices[0].delta.content
+                        if content:
+                            yield content
+                    return
+                except Exception as e:
+                    logger.warning(f"Groq streaming chat failed, falling back: {e}")
+
+            # --- OpenAI ---
+            if self.openai_client:
+                from openai import OpenAI
+                client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
+                context_limit = _get_context_limit(self.config['openai_model'])
+                trimmed = self.trim_history_to_token_budget(
+                    system_prompt, messages, max_tokens, context_limit
+                )
+                
+                oai_messages = [{"role": "system", "content": system_prompt}] + [
+                    {"role": m.get("role", "user"), "content": m.get("content", "")}
+                    for m in trimmed
+                ]
+                response = client.chat.completions.create(
+                    model=self.config['openai_model'],
+                    messages=oai_messages,
+                    max_tokens=min(max_tokens, self.config['openai_max_tokens']),
+                    temperature=self.config['openai_temperature'],
+                    stream=True
+                )
+                for chunk in response:
+                    content = chunk.choices[0].delta.content
+                    if content:
+                        yield content
+                return
+
+        except Exception as e:
+            logger.error(f"generate_chat_stream failed: {type(e).__name__}: {e}", exc_info=True)
+            yield "An error occurred while peering into the books. Please ask again."
 
     def generate_text(self, prompt: str, max_tokens: Optional[int] = None, retry_count: int = 0) -> Optional[str]:
         """Generate text using available LLM service with retry logic."""
@@ -653,7 +760,7 @@ def _generate_with_gemini(self, prompt: str, max_tokens: int) -> Optional[str]:
 llm_service = LLMService()
 
 __all__ = ['generate_book_note', 'get_ai_recommendations', 'get_category_books',
-           'get_book_mood_tags_safe', 'generate_chat_response', 'llm_service', 
+           'get_book_mood_tags_safe', 'generate_chat_response', 'generate_chat_response_stream', 'llm_service', 
            'LLMService', 'PromptTemplates']
 
 
@@ -944,6 +1051,40 @@ def generate_chat_response(user_message: str, conversation_history: list = []) -
     # SMART FALLBACK: A variety of poetic responses that adapt to keywords
     msg_lower = user_message.lower()
 
+    if any(k in msg_lower for k in ['rain', 'melancholy', 'sad', 'quiet']):
+        return reply or "The spirits are quiet... Let's talk again when they wake."
+
+def generate_chat_response_stream(user_message: str, conversation_history: list = []):
+    """
+    Yield an emotionally rich, persona-driven chat response from Elara.
+    """
+    if not llm_service.is_available():
+        logger.warning("generate_chat_response_stream: No LLM available")
+        yield "The candles are flickering and my connection to the ether seems troubled today. "
+        yield "Come back in a moment — the books are waiting, and so am I."
+        return
+
+    messages = []
+    for msg in (conversation_history or []):
+        role = msg.get("type", msg.get("role", "user"))
+        if role in ("bookseller", "assistant"):
+            role = "assistant"
+        else:
+            role = "user"
+        content = msg.get("content", "")
+        if content:
+            messages.append({"role": role, "content": content})
+
+    messages.append({"role": "user", "content": user_message})
+
+    yield from llm_service.generate_chat_stream(
+        system_prompt=_WISE_BOOKSELLER_SYSTEM_PROMPT,
+        messages=messages,
+    )
+
+    # SMART FALLBACK: A variety of poetic responses that adapt to keywords
+    msg_lower = user_message.lower()
+    
     if any(k in msg_lower for k in ['rain', 'melancholy', 'sad', 'quiet']):
         return "The rain has a way of turning the heart into a library of its own. I've gathered these quiet, thoughtful volumes for your pensive mood."
     elif any(k in msg_lower for k in ['adventure', 'journey', 'travel', 'exciting']):

backend/app.py

@@ -3,6 +3,7 @@
 
 from flask import Flask, request, jsonify, redirect, url_for
 from flask_cors import CORS
+from flask_socketio import SocketIO, emit
 from flask_jwt_extended import (
     JWTManager, create_access_token, jwt_required, 
     get_jwt_identity, set_access_cookies, unset_jwt_cookies
@@ -42,7 +43,7 @@
 
 # Environment variables are now loaded centrally in backend/config.py
 from config import app_config, setup_logging, validate_required_env_vars
-from ai_service import generate_book_note, get_ai_recommendations, get_category_books, get_book_mood_tags_safe, generate_chat_response, llm_service, get_vibe_recommendations
+from ai_service import generate_book_note, get_ai_recommendations, get_category_books, get_book_mood_tags_safe, generate_chat_response, generate_chat_response_stream, llm_service, get_vibe_recommendations
 from models import db, User, Book, ShelfItem, BookNote, ReadingGoal, ReadingStats, Collection, CollectionItem, PriceHistory, PriceAlert, Review, register_user, login_user
 from price_tracker import get_price_tracker
 from cache_service import cache_service
@@ -224,6 +225,9 @@ def _apply_security_headers(response):
     resources={r"/api/*": {"origins": _cors_origins}},
 )
 
+# Initialize SocketIO
+socketio = SocketIO(app, cors_allowed_origins=_cors_origins)
+
 # Initialize cache service
 cache_service.init_app(app)
 
@@ -1100,20 +1104,28 @@ def handle_generate_note():
 @app.route('/api/v1/chat', methods=['POST'])
 @limiter.limit("10 per minute")
 def handle_chat():
-    """Handle chat messages and generate bookseller responses."""
+    """Legacy chat endpoint (deprecated in favor of WebSockets)."""
+    return jsonify({
+        "success": False,
+        "error": "This endpoint is deprecated. Please use the WebSockets interface for Literary Chat."
+    }), 426
+
+@socketio.on('chat_message')
+def handle_socket_chat(data):
+    """Handle chat messages via WebSockets and stream bookseller responses."""
     from exceptions import (
         LLMCircuitBreakerOpenError, AIServiceException,
         ValidationException, InvalidInputError
     )
     from error_responses import handle_exception
 
     try:
-        data = request.get_json()
-        
+        # Validate using existing ChatRequest rules (manual dict pass)
         is_valid, validated_data = validate_request(ChatRequest, data)
         if not is_valid:
-            return jsonify(validated_data), 400
-        
+            emit('chat_error', {"error": "Validation failed", "details": validated_data})
+            return
+            
         user_message = validated_data.message
         conversation_history = validated_data.history or []
 
@@ -1124,39 +1136,28 @@ def handle_chat():
             else:
                 validated_history.append(msg)
 
-        # Generate contextual response based on conversation history
-        response = generate_chat_response(user_message, validated_history)
-        
-        # Try to get book recommendations based on the message
+        # We need a streaming version of generate_chat_response from ai_service
+        # For now, if the streaming generator exists we use it, otherwise fallback
         recommendations = get_ai_recommendations(user_message)
 
-        # =========================================================================
-        # TIMESTAMP STANDARDIZATION
-        # =========================================================================
-        # Ensure that the timestamp returned to the client is explicitly set to
-        # UTC using timezone-aware objects. This prevents subtle bugs where server 
-        # locale or deployment environments might skew the time by relying on 
-        # naive datetime.now() calls. This is a critical fix for ensuring
-        # consistent client-side formatting regardless of geographical region.
-        # =========================================================================
-        
-        return success_response(
-            data={
-                "response": response,
-                "recommendations": recommendations,
-                "timestamp": datetime.now(timezone.utc).isoformat()
-            }
-        )
+        # Start streaming back the LLM output
+        full_response = ""
+        for chunk in generate_chat_response_stream(user_message, validated_history):
+            if chunk:
+                full_response += chunk
+                emit('chat_stream', {'chunk': chunk})
+        
+        # When generation is complete, emit the final payload with metadata
+        emit('chat_complete', {
+            "success": True,
+            "response": full_response,
+            "recommendations": recommendations,
+            "timestamp": datetime.now(timezone.utc).isoformat()
+        })
 
-    except (LLMCircuitBreakerOpenError, AIServiceException) as e:
-        logger.error(f"AI service error in handle_chat: {e}", exc_info=True)
-        return handle_exception(e, "handle_chat")
-    except (ValidationException, InvalidInputError) as e:
-        logger.warning(f"Validation error in handle_chat: {e}")
-        return handle_exception(e, "handle_chat")
     except Exception as e:
-        logger.error(f"Unexpected error in handle_chat: {type(e).__name__}: {e}", exc_info=True)
-        return handle_exception(e, "handle_chat")
+        logger.error(f"Error in WebSocket handle_chat: {type(e).__name__}: {e}", exc_info=True)
+        emit('chat_error', {"error": "An error occurred while generating a response. Please try again."})
 
 @app.route('/api/v1/health', methods=['GET'])
 def health_check():
@@ -2611,6 +2612,6 @@ def delete_price_alert(alert_id):
         logger.info("  POST /api/v1/chat - Chat with bookseller")
         logger.info("  GET  /api/v1/health - Health check")
 
-    app.run(debug=server_config.debug, port=server_config.port, host=server_config.host)
+    socketio.run(app, debug=server_config.debug, port=server_config.port, host=server_config.host)
 
 

backend/requirements.txt

@@ -10,11 +10,32 @@ class ChatInterface {
 
         this.conversationHistory = [];
         this.isProcessing = false;
+        
+        this.socket = null;
+        this.currentMessageContent = "";
+        this.currentMessageBubble = null;
 
         this.init();
     }
 
     init() {
+        // Init socket.io
+        if (typeof io !== 'undefined') {
+            const moodApiBase = window.MOOD_API_BASE || 'http://127.0.0.1:5000';
+            const baseUrl = moodApiBase.replace(/\/api\/v1\/?$/, '');
+            this.socket = io(baseUrl);
+            
+            this.socket.on('chat_stream', (data) => {
+                if (this.currentMessageBubble) {
+                    this.hideTypingIndicator(); // Hide indicator as soon as text flows
+                    this.currentMessageContent += data.chunk;
+                    this.currentMessageBubble.innerHTML = '';
+                    this.renderTextContent(this.currentMessageBubble, this.currentMessageContent, true);
+                    this.scrollToBottom();
+                }
+            });
+        }
+
         // Load conversation history from localStorage
         this.loadConversationHistory();
 
@@ -164,6 +185,63 @@ Tell me: what is stirring in you today?`,
         // Show typing indicator
         this.showTypingIndicator();
 
+        // --- WebSockets Path ---
+        if (this.socket && this.socket.connected) {
+            this.currentMessageContent = "";
+            const booksellerMessage = {
+                type: 'bookseller',
+                content: "",
+                books: null,
+                timestamp: new Date().toISOString()
+            };
+            this.conversationHistory.push(booksellerMessage);
+            const messageDiv = this.renderMessage(booksellerMessage);
+            this.currentMessageBubble = messageDiv.querySelector('.message-bubble');
+
+            const history = this._buildTokenBudgetHistory(message);
+            this.socket.emit('chat_message', { message: message, history: history });
+
+            const onComplete = async (data) => {
+                this.socket.off('chat_complete', onComplete);
+                this.socket.off('chat_error', onError);
+                
+                booksellerMessage.content = data.response;
+                // Try to get actual books from Google Books API
+                const books = await this.searchGoogleBooks(message);
+                if (books && books.length > 0) {
+                    booksellerMessage.books = books;
+                    const bookRec = this.createBookRecommendations(books);
+                    this.currentMessageBubble.appendChild(bookRec);
+                }
+                this.saveConversationHistory();
+                this.scrollToBottom();
+            };
+
+            const onError = (data) => {
+                this.socket.off('chat_complete', onComplete);
+                this.socket.off('chat_error', onError);
+                this.hideTypingIndicator();
+                
+                // Remove the empty message we created
+                this.conversationHistory.pop();
+                if (messageDiv) messageDiv.remove();
+                
+                const errorMessage = {
+                    type: 'bookseller',
+                    content: "The candles are flickering and something seems amiss with my connection to the literary spirits. Give me a moment — the books are waiting, and so is the perfect story for you.",
+                    timestamp: new Date().toISOString()
+                };
+                this.conversationHistory.push(errorMessage);
+                this.renderMessage(errorMessage);
+                this.saveConversationHistory();
+            };
+
+            this.socket.on('chat_complete', onComplete);
+            this.socket.on('chat_error', onError);
+            return;
+        }
+
+        // --- HTTP Fallback Path ---
         try {
             // Get AI response
             const response = await this.getBooksellerResponse(message);
@@ -501,6 +579,7 @@ Tell me: what is stirring in you today?`,
 
         this.chatMessages.appendChild(messageDiv);
         this.scrollToBottom();
+        return messageDiv;
     }
 
     createBookRecommendations(books) {