Update app.py

app.py

@@ -317,97 +317,102 @@ def process_citations(complete_answer: str, ranked_docs: List[dict]) -> Tuple[st
 # ------------------------------------------------------------------------------
 @app.websocket("/chat")
 async def websocket_endpoint(websocket: WebSocket):
+    logger.info("Client connected to WebSocket")
     await websocket.accept()
+    
     try:
-        # Receive and validate the request
-        try:
-            data = await asyncio.wait_for(websocket.receive_json(), timeout=30)
-            chat_request = ChatRequest(**data)
-        except ValidationError as e:
-            logger.error(f"Validation error: {e}")
-            await safe_send(websocket, {"response": "Something went wrong with your request!", "sources": []})
-            return
-        except Exception as e:
-            logger.error(f"Error receiving data: {e}")
-            await safe_send(websocket, {"response": "Something went wrong with your request!", "sources": []})
-            return
-
-        question = chat_request.question
-        language = chat_request.language
-
-        # Retrieve documents using the retriever
-        try:
-            retrieved_docs = await asyncio.to_thread(retriever.invoke, question)
-        except Exception as e:
-            logger.error(f"Document retrieval error: {e}")
-            await safe_send(websocket, {"response": "Document retrieval failed", "sources": []})
-            return
-
-        docs = [{
-            "summary": ele.metadata.get("summary", ""),
-            "chunk": ele.page_content,
-            "page_source": ele.metadata.get("source", "")
-        } for ele in retrieved_docs]
-
-        if not docs:
-            await safe_send(websocket, {"response": "Cannot provide answer to this question", "sources": []})
-            return
-
-        # Rerank the documents (fallback to original docs if reranking fails)
-        try:
-            ranked_docs = await asyncio.to_thread(rerank_docs, question, docs, pc)
-        except Exception as e:
-            logger.error(f"Reranking error: {e}")
-            ranked_docs = docs
-
-        # Prepare the conversation messages
-        messages = [{"role": "system", "content": system_prompt}]
-        messages.extend(chat_request.previous_chats)
-        messages.append({"role": "user", "content": format_query(question, language, ranked_docs)})
-
-        complete_answer = ""
-        chunk_buffer = ""
-
-        # Generate and stream the chat response
-        try:
-            completion = await openai_client.chat.completions.create(
-                model="gpt-4o",
-                messages=messages,
-                temperature=0.2,
-                max_completion_tokens=1024,
-                stream=True
-            )
-            async for chunk in completion:
-                delta_content = chunk.choices[0].delta.content
-                if delta_content:
-                    complete_answer += delta_content
-                    # Remove inline citation markers from the streamed chunk before sending
-                    cleaned_content = re.sub(r'\[\d+\]', '', delta_content)
-                    chunk_buffer += cleaned_content
-                    if len(chunk_buffer) >= 1:
-                        await safe_send(websocket, {"response": chunk_buffer})
-                        chunk_buffer = ""
-            if chunk_buffer:
-                await safe_send(websocket, {"response": chunk_buffer})
-        except Exception as e:
-            logger.error(f"Streaming error: {e}")
-            await safe_send(websocket, {"response": "Response generation failed", "sources": []})
-            return
-
-        # Process and map citations in the final answer
-        complete_answer, citations = process_citations(complete_answer, ranked_docs)
-
-        await safe_send(websocket, {
-            "response": complete_answer,
-            "sources": citations
-        })
-
-    except WebSocketDisconnect:
-        logger.info("Client disconnected")
+        while True:
+            try:
+                # Wait indefinitely for the next query from the client.
+                data = await websocket.receive_json()
+            except WebSocketDisconnect:
+                logger.info("Client disconnected")
+                break
+            except Exception as e:
+                logger.error(f"Error receiving data: {e}")
+                await safe_send(websocket, {"response": "Error receiving data", "sources": []})
+                continue
+
+            # Validate the received query
+            try:
+                chat_request = ChatRequest(**data)
+            except Exception as e:
+                logger.error(f"Validation error: {e}")
+                await safe_send(websocket, {"response": "Invalid query data", "sources": []})
+                continue
+
+            question = chat_request.question
+            language = chat_request.language
+            previous_chats = chat_request.previous_chats
+
+            # Retrieve documents using the retriever
+            try:
+                retrieved_docs = await asyncio.to_thread(retriever.invoke, question)
+            except Exception as e:
+                logger.error(f"Document retrieval error: {e}")
+                await safe_send(websocket, {"response": "Document retrieval failed", "sources": []})
+                continue
+
+            docs = [{
+                "summary": ele.metadata.get("summary", ""),
+                "chunk": ele.page_content,
+                "page_source": ele.metadata.get("source", "")
+            } for ele in retrieved_docs]
+
+            if not docs:
+                await safe_send(websocket, {"response": "Cannot provide an answer to this question", "sources": []})
+                continue
+
+            # Rerank the documents; if the reranking fails, use the original docs
+            try:
+                ranked_docs = await asyncio.to_thread(rerank_docs, question, docs, pc)
+            except Exception as e:
+                logger.error(f"Reranking error: {e}")
+                ranked_docs = docs
+
+            # Prepare the conversation messages for the chat model.
+            messages = [{"role": "system", "content": system_prompt}]
+            messages.extend(previous_chats)
+            messages.append({"role": "user", "content": format_query(question, language, ranked_docs)})
+
+            complete_answer = ""
+            chunk_buffer = ""
+
+            # Generate the answer in streaming mode.
+            try:
+                completion = await openai_client.chat.completions.create(
+                    model="gpt-4o",
+                    messages=messages,
+                    temperature=0.2,
+                    max_completion_tokens=1024,
+                    stream=True
+                )
+                async for chunk in completion:
+                    delta_content = chunk.choices[0].delta.content
+                    if delta_content:
+                        complete_answer += delta_content
+                        # Remove inline citation markers from the streamed chunk before sending.
+                        cleaned_content = re.sub(r'\[\d+\]', '', delta_content)
+                        chunk_buffer += cleaned_content
+                        if len(chunk_buffer) >= 1:
+                            await safe_send(websocket, {"response": chunk_buffer})
+                            chunk_buffer = ""
+                if chunk_buffer:
+                    await safe_send(websocket, {"response": chunk_buffer})
+            except Exception as e:
+                logger.error(f"Streaming error: {e}")
+                await safe_send(websocket, {"response": "Response generation failed", "sources": []})
+                continue
+
+            # Process citations in the complete answer and send the final response.
+            complete_answer, citations = process_citations(complete_answer, ranked_docs)
+            await safe_send(websocket, {"response": complete_answer, "sources": citations})
+            
     except Exception as e:
         logger.error(f"Unexpected error: {e}")
         await safe_send(websocket, {"response": "Something went wrong! Please try again.", "sources": []})
 
+
 # ------------------------------------------------------------------------------
 # Simple health check endpoint
 # ------------------------------------------------------------------------------