from fastapi import WebSocket, WebSocketDisconnect from langchain_core.messages import HumanMessage, SystemMessage, AIMessage import logging import json import asyncio import re from typing import Dict, Any from hybrid_llm_service import HybridLLMService # Fixed import from voice_service import VoiceService from rag_service import search_documents from llm_service import create_graph, create_basic_graph from lancedb_service import lancedb_service from policy_chat_interface import PolicySimulatorChatInterface # Configure logging logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) # Initialize services hybrid_llm_service = HybridLLMService() # Create instance voice_service = VoiceService() policy_simulator = PolicySimulatorChatInterface() # Policy simulation detection patterns POLICY_PATTERNS = [ r"scenario.*analy", r"policy.*simulat", r"pension.*analy", r"simulate.*dr|dr.*simulat", r"simulate.*pension|pension.*simulat", r"impact.*analy", r"dearness.*relief", r"basic.*pension", r"medical.*allowance", r"chart.*pension|pension.*chart", r"visual.*analy|analy.*visual", r"show.*chart|chart.*show", r"explain.*chart|chart.*explain", r"using.*chart|chart.*using", r"dr.*\d+.*increase|increase.*dr.*\d+", r"analyze.*minimum.*pension", r"pension.*change" ] def is_policy_simulation_query(message: str) -> bool: """Check if the message is a policy simulation query""" message_lower = message.lower() return any(re.search(pattern, message_lower, re.IGNORECASE) for pattern in POLICY_PATTERNS) async def handle_websocket_connection(websocket: WebSocket): """Handle WebSocket connection for the voice bot""" await websocket.accept() logger.info("🔌 WebSocket client connected.") import uuid initial_data = await websocket.receive_json() messages = [] # Check if user authentication is provided flag = "user_id" in initial_data if flag: thread_id = initial_data.get("user_id") knowledge_base = initial_data.get("knowledge_base", "government_docs") # Create graph with RAG capabilities graph = await create_graph(kb_tool=True, mcp_config=None) config = { "configurable": { "thread_id": thread_id, "knowledge_base": knowledge_base, } } # Set system prompt for government document queries system_message = """You are a helpful assistant that can answer questions about government documents, policies, and procedures. Keep your responses clear and concise. When referencing specific documents or policies, mention the source. If you're uncertain about information, clearly state that and suggest where the user might find authoritative information.""" messages.append(SystemMessage(content=system_message)) else: # Basic graph for unauthenticated users graph = create_basic_graph() thread_id = str(uuid.uuid4()) config = {"configurable": {"thread_id": thread_id}} # Send initial greeting greeting_message = HumanMessage( content="Generate a brief greeting for the user, introduce yourself as a government document assistant, and explain how you can help them find information from government policies and documents." ) messages.append(greeting_message) try: response = await graph.ainvoke({"messages": messages}, config=config) greeting_response = response["messages"][-1].content messages.append(AIMessage(content=greeting_response)) await websocket.send_json({ "type": "connection_successful", "message": greeting_response }) except Exception as e: logger.error(f"❌ Error generating greeting: {e}") await websocket.send_json({ "type": "connection_successful", "message": "Hello! I'm your government document assistant. How can I help you today?" }) try: while True: data = await websocket.receive_json() if data["type"] == "text_message": # Handle text message user_message = data["message"] logger.info(f"💬 Received text message: {user_message}") messages.append(HumanMessage(content=user_message)) # Send acknowledgment await websocket.send_json({ "type": "message_received", "message": "Processing your message..." }) # Check if this is a policy simulation query if is_policy_simulation_query(user_message): logger.info("🎯 Detected policy simulation query") try: # Process with policy simulator policy_response = policy_simulator.process_policy_query(user_message) # Send policy simulation response await websocket.send_json({ "type": "policy_simulation", "data": policy_response }) messages.append(AIMessage(content=policy_response.get('message', 'Policy simulation completed'))) continue except Exception as policy_error: logger.error(f"❌ Policy simulation failed: {policy_error}") # Fall back to normal processing # First try to search for relevant documents search_results = None try: # Search for documents related to the user's query search_results = search_documents(user_message, limit=5) logger.info(f"🔍 Found {len(search_results) if search_results else 0} documents for query") except Exception as search_error: logger.warning(f"⚠️ Document search failed: {search_error}") # Get LLM response (with or without search context) try: if search_results and len(search_results) > 0: # Add search context to the message context_message = f"User query: {user_message}\n\nRelevant documents found:\n" for i, doc in enumerate(search_results[:3], 1): context_message += f"\n{i}. Source: {doc.get('filename', 'Unknown')}\nContent: {doc.get('content', '')[:400]}...\n" context_message += f"\nBased on the above documents, please provide a helpful response to the user's query: {user_message}" # Replace the user message with the enriched version messages[-1] = HumanMessage(content=context_message) result = await graph.ainvoke({"messages": messages}, config=config) llm_response = result["messages"][-1].content # Check if response contains scenario analysis images if "**SCENARIO_IMAGES_START**" in llm_response and "**SCENARIO_IMAGES_END**" in llm_response: # Extract images and text separately parts = llm_response.split("**SCENARIO_IMAGES_START**") text_response = parts[0].strip() image_part = parts[1].split("**SCENARIO_IMAGES_END**")[0].strip() try: import json images = json.loads(image_part) # Send text response first await websocket.send_json({ "type": "text_response", "message": text_response }) # Send images separately await websocket.send_json({ "type": "scenario_images", "images": images }) except json.JSONDecodeError: # If JSON parsing fails, send as regular text await websocket.send_json({ "type": "text_response", "message": llm_response }) else: # Send regular text response await websocket.send_json({ "type": "text_response", "message": llm_response }) # Add AI response to messages messages.append(AIMessage(content=llm_response)) logger.info(f"✅ Sent response to user: {thread_id}") except Exception as e: logger.error(f"❌ Error processing message: {e}") await websocket.send_json({ "type": "error", "message": "Sorry, I encountered an error processing your message." }) elif data["type"] == "ping": # Handle ping for connection keep-alive await websocket.send_json({"type": "pong"}) elif data["type"] == "get_knowledge_bases": # Send available knowledge bases try: kb_list = await lancedb_service.get_knowledge_bases() await websocket.send_json({ "type": "knowledge_bases", "knowledge_bases": kb_list }) except Exception as e: logger.error(f"❌ Error getting knowledge bases: {e}") await websocket.send_json({ "type": "error", "message": "Error retrieving knowledge bases" }) elif data["type"] == "end_session": logger.info("📞 Session ended by client") await websocket.close() break except WebSocketDisconnect: logger.info("🔌 WebSocket client disconnected.") except Exception as e: logger.error(f"❌ WebSocket error: {e}") try: await websocket.send_json({ "type": "error", "message": "Connection error occurred" }) except: pass finally: # Clean up when session ends logger.info(f"🔄 Session {thread_id} ended") async def send_welcome_message(websocket: WebSocket): """Send welcome message to the client""" try: welcome_text = """🇮🇳 Welcome to the Government Services AI Assistant! I'm here to help you with: • Government policies and procedures • Document information and guidance • Service-specific questions and redirects • Voice or text interaction (your choice!) How can I assist you today?""" await websocket.send_text(json.dumps({ "type": "bot_message", "content": welcome_text, "timestamp": asyncio.get_event_loop().time() })) except Exception as e: logger.error(f"❌ Error sending welcome message: {e}") async def handle_text_message(websocket: WebSocket, message_data: Dict[str, Any]): """Handle text-based messages""" try: user_message = message_data.get("content", "") logger.info(f"💬 Processing text message: {user_message}") # Search for relevant documents context = "" try: search_results = search_documents(user_message, limit=3) if search_results: context = "\n".join([doc.get("content", "") for doc in search_results]) logger.info(f"📚 Found {len(search_results)} relevant documents") except Exception as e: logger.warning(f"⚠️ Document search failed: {e}") # Get response from hybrid LLM response_text = "" try: # Check if this is a streaming request stream_response = message_data.get("stream", True) if stream_response: # Send streaming response await websocket.send_text(json.dumps({ "type": "bot_message_start", "timestamp": asyncio.get_event_loop().time() })) async for chunk in hybrid_llm_service.get_streaming_response(user_message, context): response_text += chunk await websocket.send_text(json.dumps({ "type": "bot_message_chunk", "content": chunk, "timestamp": asyncio.get_event_loop().time() })) await asyncio.sleep(0.01) # Small delay for better streaming await websocket.send_text(json.dumps({ "type": "bot_message_end", "timestamp": asyncio.get_event_loop().time() })) else: # Send complete response response_text = await hybrid_llm_service.get_response(user_message, context) await websocket.send_text(json.dumps({ "type": "bot_message", "content": response_text, "timestamp": asyncio.get_event_loop().time() })) except Exception as e: logger.error(f"❌ Error getting LLM response: {e}") await websocket.send_text(json.dumps({ "type": "bot_message", "content": f"I apologize, but I encountered an error processing your request: {str(e)}", "timestamp": asyncio.get_event_loop().time() })) # Add government service redirect suggestions try: redirect_suggestions = voice_service.generate_redirect_suggestions(user_message, "text") if redirect_suggestions: await websocket.send_text(json.dumps({ "type": "redirect_suggestions", "content": redirect_suggestions, "timestamp": asyncio.get_event_loop().time() })) except Exception as e: logger.warning(f"⚠️ Could not generate redirect suggestions: {e}") except Exception as e: logger.error(f"❌ Error handling text message: {e}") await websocket.send_text(json.dumps({ "type": "error", "content": f"Error processing your message: {str(e)}" })) async def handle_voice_message(websocket: WebSocket, message_data: Dict[str, Any]): """Handle voice-based messages""" try: # Check if voice features are enabled if not voice_service.voice_enabled: await websocket.send_text(json.dumps({ "type": "error", "content": "Voice features are currently disabled. Please use text input." })) return audio_data = message_data.get("audio_data", "") if not audio_data: await websocket.send_text(json.dumps({ "type": "error", "content": "No audio data received" })) return logger.info("🎤 Processing voice message") # Convert speech to text try: transcribed_text = await voice_service.speech_to_text(audio_data) logger.info(f"📝 Transcribed: {transcribed_text}") # Send transcription to client await websocket.send_text(json.dumps({ "type": "transcription", "content": transcribed_text, "timestamp": asyncio.get_event_loop().time() })) except Exception as e: logger.error(f"❌ Speech-to-text failed: {e}") await websocket.send_text(json.dumps({ "type": "error", "content": f"Speech recognition failed: {str(e)}" })) except Exception as e: logger.error(f"❌ Error handling voice message: {e}") await websocket.send_text(json.dumps({ "type": "error", "content": f"Error processing voice message: {str(e)}" }))