feat: Migrate Literary Chat to WebSockets (Flask-SocketIO)

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

@@ -8,6 +8,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
@@ -46,7 +47,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
@@ -285,6 +286,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)
 
@@ -1249,17 +1253,29 @@ def handle_generate_note(validated_data):
 
 @app.route('/api/v1/chat', methods=['POST'])
 @limiter.limit("10 per minute")
-@validate_schema(ChatRequest)
-def handle_chat(validated_data):
-    """Handle chat messages and generate bookseller responses."""
+def handle_chat():
+    """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:
-
+        # Validate using existing ChatRequest rules (manual dict pass)
+        is_valid, validated_data = validate_request(ChatRequest, data)
+        if not is_valid:
+            emit('chat_error', {"error": "Validation failed", "details": validated_data})
+            return
+
         user_message = validated_data.message
         conversation_history = validated_data.history or []
 
@@ -1270,39 +1286,28 @@ def handle_chat(validated_data):
             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."})
 
 
 
@@ -1328,7 +1333,7 @@ def add_to_library(validated_data):
     from error_responses import handle_exception
 
     try:
-
+        current_user_id = get_jwt_identity()
         if str(validated_data.user_id) != str(current_user_id):
             return unauthorized_access_error("Cannot access another user's library")
 
@@ -1501,7 +1506,7 @@ def _get_yearly_stats(user_id, year):
 def update_library_item(item_id, validated_data):
     """Update a library item (e.g. move to different shelf)."""
     try:
-
+        current_user_id = get_jwt_identity()
         item = ShelfItem.query.with_for_update().get(item_id)
         if not item:
             return not_found_error("Library item")
@@ -2103,7 +2108,7 @@ def verify_auth_session():
 @validate_schema(SetGoalRequest)
 def set_reading_goal(validated_data):
     """Set or update annual reading goal."""
-
+    current_user_id = get_jwt_identity()
     if str(validated_data.user_id) != str(current_user_id):
         return forbidden_error("Unauthorized")
 
@@ -2221,7 +2226,7 @@ def get_leaderboard():
 @validate_schema(CollectionRequest)
 def create_collection(validated_data):
     """Create a new collection."""
-
+    current_user_id = get_jwt_identity()
     if str(validated_data.user_id) != str(current_user_id):
         return forbidden_error("Unauthorized")
 
@@ -2291,6 +2296,7 @@ def update_collection(collection_id, validated_data):
     """Update a collection."""
 
     try:
+        current_user_id = get_jwt_identity()
         collection = Collection.query.get(collection_id)
         if not collection:
             return jsonify({"error": "Collection not found"}), 404
@@ -2344,10 +2350,12 @@ def delete_collection(collection_id):
 
 @app.route('/api/v1/collections/<int:collection_id>/books', methods=['POST'])
 @jwt_required()
-def add_book_to_collection(collection_id):
+@validate_schema(AddToCollectionRequest)
+def add_book_to_collection(collection_id, validated_data):
     """Add a book to a collection."""
 
     try:
+        current_user_id = get_jwt_identity()
         collection = Collection.query.get(collection_id)
         if not collection:
             return jsonify({"error": "Collection not found"}), 404
@@ -2741,6 +2749,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)