feat: implement WebSocket-based streaming chat

main.py

@@ -16,7 +16,7 @@ from fastapi.responses import RedirectResponse
 
 from src.materials.routes import router as materials_router
 from src.summary_generator.routes import router as summary_router
-from src.rag.routes import router as tutor_router
+from src.rag.routes import router as tutor_router, ws_router
 from src.quiz_generator.routes import router as quiz_router
 from src.auth.routes import router as auth_router
 from src.asr.routes import router as asr_router
@@ -132,6 +132,7 @@ app.add_middleware(
 app.include_router(materials_router)
 app.include_router(summary_router)
 app.include_router(tutor_router)
+app.include_router(ws_router)
 app.include_router(quiz_router)
 app.include_router(auth_router)
 app.include_router(asr_router)

rag/rag.py

@@ -450,6 +450,141 @@ def rag_answer(
             logger.error(f"Fallback LLM call also failed: {fallback_err}")
             raise fallback_err
 
+
+async def rag_answer_stream(
+    query: str,
+    material_id: Optional[str] = None,
+    chunks: Optional[list[str]] = None,
+    summaries: str = "",
+    memory = None,
+):
+    """
+    Async generator that streams LLM response tokens using Mistral AI (with Gemini fallback).
+    Yields individual token strings as they arrive.
+    Updates conversation memory upon completion if not a refusal.
+    """
+    if memory is None:
+        memory = ConversationBufferWindowMemory(
+            input_key="input", memory_key="chat_history", return_messages=True, k=MEMORY_WINDOW_SIZE
+        )
+
+    # Fetch material info if material_id is provided
+    mat = None
+    if material_id:
+        mat = await asyncio.to_thread(get_material, material_id)
+
+    is_topic = not (material_id and mat and mat.get("source_type") != "topic")
+
+    context_parts = []
+    has_chunks = False
+
+    # Inject Subject/Topic
+    if mat and mat.get("title"):
+        context_parts.append(f"Subject / Topic: {mat.get('title')}")
+
+    if not is_topic:
+        # --- Material-based query (PDF/URL): vector similarity search ---
+        results = await asyncio.to_thread(similarity_search, query, material_id, k=TOP_K_CHUNKS)
+        if results:
+            has_chunks = True
+            chunks = [r["content"] for r in results]
+            context_parts.append("Relevant Excerpts:\n" + "\n---\n".join(chunks))
+
+        # Fallback: summary
+        if not has_chunks and summaries:
+            context_parts.append(f"Material Summary (No specific excerpts found for your query):\n{summaries}")
+
+        # Fallback: sample head + tail chunks
+        if not has_chunks and not summaries:
+            all_chunks = await asyncio.to_thread(get_chunks, material_id)
+            if all_chunks:
+                head = all_chunks[:3]
+                tail = all_chunks[-2:] if len(all_chunks) > 3 else []
+                sampled = head + [c for c in tail if c not in head]
+                sampled_text = "\n---\n".join(c["content"] for c in sampled)
+                context_parts.append(f"Material Sample (No summary found; showing start and end of material):\n{sampled_text}")
+
+    # --- Wikipedia search: topics only ---
+    wiki_snippets = ""
+    if is_topic:
+        wiki_snippets = await asyncio.to_thread(direct_wiki_search, query)
+        if wiki_snippets:
+            context_parts.append(f"Wikipedia Results:\n{wiki_snippets}")
+
+    # --- DuckDuckGo search: ALL material types (topics, PDFs, URLs) ---
+    subject_title = mat.get("title") if mat and mat.get("title") else ""
+    ddg_query = f"{query} {subject_title}".strip() if subject_title else query
+    ddg_snippets = await asyncio.to_thread(direct_ddg_search, ddg_query)
+    if ddg_snippets:
+        context_parts.append(f"Web Search Results (DuckDuckGo):\n{ddg_snippets}")
+
+    context_str = "\n\n".join(context_parts) if context_parts else "No specific context provided."
+
+    has_knowledge = not is_topic
+    subject_title = mat.get("title") if mat and mat.get("title") else ""
+    prompt = _rag_prompt(
+        has_ddg=bool(ddg_snippets),
+        has_wiki=bool(wiki_snippets),
+        has_knowledge_retriever=has_knowledge,
+        subject=subject_title,
+    )
+
+    _REFUSAL_PREFIXES = (
+        "I can't respond on a gibberish",
+        "I can't respond on a NSFW",
+        "I can't respond on a political",
+        "I can't respond on a religious",
+    )
+
+    def _is_refusal(text: str) -> bool:
+        t = text.strip()
+        return any(t.startswith(p) for p in _REFUSAL_PREFIXES)
+
+    memory_vars = memory.load_memory_variables({"input": query})
+    chat_history = memory_vars.get("chat_history", [])
+
+    full_answer_parts = []
+
+    try:
+        primary_llm = get_llm()
+        chain = prompt | primary_llm
+
+        async for chunk in chain.astream({
+            "input": query,
+            "context": context_str,
+            "chat_history": chat_history,
+            "agent_scratchpad": "",
+        }):
+            token = _clean_llm_response(chunk.content)
+            if token:
+                full_answer_parts.append(token)
+                yield token
+
+    except Exception as e:
+        logger.warning(f"Primary LLM streaming failed or rate-limited: {e}. Falling back to secondary LLM.")
+        try:
+            fallback_llm = get_fallback_llm()
+            chain = prompt | fallback_llm
+
+            async for chunk in chain.astream({
+                "input": query,
+                "context": context_str,
+                "chat_history": chat_history,
+                "agent_scratchpad": "",
+            }):
+                token = _clean_llm_response(chunk.content)
+                if token:
+                    full_answer_parts.append(token)
+                    yield token
+        except Exception as fallback_err:
+            logger.error(f"Fallback LLM streaming failed: {fallback_err}")
+            raise fallback_err
+
+    full_answer = "".join(full_answer_parts)
+    if not _is_refusal(full_answer):
+        memory.save_context({"input": query}, {"output": full_answer})
+
+
 def extract_chat_title(query: str, material_title: Optional[str] = None) -> str:
     topic_context = ""
     if material_title:

rag/routes.py

@@ -1,12 +1,16 @@
 import time
+import json
 import asyncio
+import jwt as pyjwt
 from loguru import logger
-from fastapi import APIRouter, HTTPException, Depends
+from fastapi import APIRouter, HTTPException, Depends, WebSocket, WebSocketDisconnect, status
 from typing import Optional, Any
 
-from src.rag.rag import rag_answer, extract_chat_title
+from src.config import settings
+from src.database import get_supabase, get_auth_supabase
+from src.rag.rag import rag_answer, extract_chat_title, rag_answer_stream
 from src.rag.constants import REFUSAL_PREFIXES
-from src.dependencies import get_current_user, get_current_user_id
+from src.dependencies import get_current_user, get_current_user_id, DEV_USER_ID, _verify_token_cached
 from src.store import (
     get_material, get_chunks, get_summary, get_or_create_memory, append_memory_message,
     # Session-based chat
@@ -21,6 +25,194 @@ from src.summary_generator.summary import clean_summary
 from .schemas import TutorQuery, TutorResponse, SessionRequest, RenameSessionRequest, ExtractTitleRequest, SaveChatRequest
 
 router = APIRouter(prefix="/api/tutor", tags=["Tutor"])
+ws_router = APIRouter(tags=["WebSocket Chat"])
+
+
+def _authenticate_ws_token(token: Optional[str]) -> Optional[str]:
+    """
+    Authenticate a JWT token passed during WebSocket auth handshake.
+    Returns user_id string if valid, None otherwise.
+    """
+    client = get_auth_supabase() or get_supabase()
+
+    # Dev mode — no Supabase configured
+    if client is None:
+        return DEV_USER_ID
+
+    if not token:
+        return None
+
+    # 1. Stateless custom JWT
+    try:
+        from src.auth.jwt_utils import decode_access_token
+        payload = decode_access_token(token)
+        user_id = payload.get("sub")
+        if user_id:
+            return str(user_id)
+    except Exception:
+        pass
+
+    # 2. Supabase JWT secret verification
+    if settings.supabase_jwt_secret:
+        try:
+            payload = pyjwt.decode(
+                token,
+                settings.supabase_jwt_secret,
+                algorithms=["HS256", "HS384", "HS512"],
+                options={"verify_aud": False},
+            )
+            user_id = payload.get("sub")
+            if user_id:
+                return str(user_id)
+        except Exception:
+            pass
+
+    # 3. Cached Supabase user lookup
+    try:
+        user = _verify_token_cached(client, token)
+        if user and getattr(user, "id", None):
+            return str(user.id)
+    except Exception as e:
+        logger.debug(f"WS token validation error: {e}")
+
+    return None
+
+
+@ws_router.websocket("/ws/chat")
+async def websocket_chat(websocket: WebSocket):
+    """
+    WebSocket endpoint for real-time streaming LLM chat responses.
+    
+    Handshake Protocol:
+    1. Client connects to ws(s)://<host>/ws/chat (No tokens in URL).
+    2. Client MUST send an initial JSON message: {"type": "auth", "token": "<JWT>"}
+    3. Server verifies token. Sends {"type": "auth_ok"} on success or closes with 1008 on failure.
+    4. Client sends chat JSON messages: {"query": "...", "material_id": "...", "session_id": "...", "source_type": "..."}
+    5. Server streams token strings one by one, followed by "[DONE]".
+    """
+    await websocket.accept()
+
+    # Step 1: Two-step Auth Handshake
+    try:
+        raw_msg = await asyncio.wait_for(websocket.receive_text(), timeout=10.0)
+        auth_data = json.loads(raw_msg)
+    except Exception as e:
+        logger.warning(f"WS auth handshake timeout or invalid message: {e}")
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Auth timeout or invalid payload")
+        return
+
+    if not isinstance(auth_data, dict) or auth_data.get("type") != "auth":
+        logger.warning("WS auth handshake failed: first message was not of type 'auth'")
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="First message must be type 'auth'")
+        return
+
+    token = auth_data.get("token")
+    user_id = _authenticate_ws_token(token)
+
+    if not user_id:
+        logger.warning("WS auth handshake failed: invalid or expired token")
+        await websocket.close(code=status.WS_1008_POLICY_VIOLATION, reason="Invalid authentication token")
+        return
+
+    # Send auth confirmation
+    await websocket.send_json({"type": "auth_ok"})
+    logger.info(f"WebSocket client authenticated successfully for user_id={user_id}")
+
+    # Step 2: Message Loop
+    while True:
+        try:
+            msg_text = await websocket.receive_text()
+        except WebSocketDisconnect:
+            logger.info(f"WebSocket client disconnected gracefully (user: {user_id}).")
+            break
+        except Exception as e:
+            logger.error(f"WebSocket connection error: {e}")
+            break
+
+        try:
+            data = json.loads(msg_text)
+        except json.JSONDecodeError:
+            await websocket.send_json({"type": "error", "message": "Invalid JSON payload"})
+            continue
+
+        query = data.get("query", "").strip()
+        if not query:
+            await websocket.send_json({"type": "error", "message": "Query cannot be empty"})
+            continue
+
+        material_id = data.get("material_id")
+        session_id = data.get("session_id")
+        source_type = data.get("source_type", "topic")
+
+        # Material ownership check for pdf/url sources
+        if source_type in ("pdf", "url") and material_id:
+            mat = get_material(material_id)
+            if not mat:
+                await websocket.send_json({"type": "error", "message": f"No {source_type} material found."})
+                continue
+            if mat.get("user_id") != user_id:
+                await websocket.send_json({"type": "error", "message": "Access denied"})
+                continue
+
+        # Load / seed memory
+        mem_key = session_id or material_id
+        seed_msgs = None
+        if session_id:
+            try:
+                seed_msgs = get_session_messages(session_id)
+            except Exception:
+                seed_msgs = None
+
+        memory, memory_id = get_or_create_memory(mem_key, seed_messages=seed_msgs)
+
+        # Persist user message to session DB immediately
+        if session_id:
+            try:
+                append_session_message(session_id, "user", query)
+            except Exception as e:
+                logger.warning(f"Failed to append user session message: {e}")
+
+        # Summary fallback
+        summary_text = ""
+        if material_id:
+            mat_summary = get_summary(material_id)
+            summary_text = mat_summary.get("summary", "") if mat_summary else ""
+
+        # Stream response
+        full_answer_parts = []
+        try:
+            async for token_chunk in rag_answer_stream(
+                query=query,
+                material_id=material_id,
+                summaries=summary_text,
+                memory=memory,
+            ):
+                full_answer_parts.append(token_chunk)
+                await websocket.send_text(token_chunk)
+
+            # Send special completion signal
+            await websocket.send_text("[DONE]")
+
+            full_answer = "".join(full_answer_parts)
+            cleaned_answer = clean_summary(full_answer)
+
+            is_refusal = any(cleaned_answer.strip().startswith(p) for p in REFUSAL_PREFIXES)
+
+            # Persist assistant response if not a refusal
+            if session_id and not is_refusal:
+                try:
+                    append_session_message(session_id, "assistant", cleaned_answer)
+                except Exception as e:
+                    logger.warning(f"Failed to append assistant session message: {e}")
+
+            if session_id and not is_refusal:
+                append_memory_message(memory_id, "user", query)
+                append_memory_message(memory_id, "assistant", cleaned_answer)
+
+        except Exception as e:
+            logger.error(f"Error during WS chat streaming for user {user_id}: {e}", exc_info=True)
+            await websocket.send_json({"type": "error", "message": f"Error generating answer: {str(e)}"})
+
 
 
 @router.post("/ask", response_model=TutorResponse)