Deploy clean Voice Bot backend to HF Spaces

🚀 Features:
- FastAPI application optimized for HF Spaces (port 7860)
- Voice processing with ASR and TTS
- LangChain-powered RAG system for document search
- WebSocket support for real-time communication
- JWT authentication
- Hybrid LLM service (Gemini + Groq)
- Docker configuration with health checks
- Clean project structure without deployment artifacts

✅ Ready for HF Spaces deployment

.gitignore

@@ -0,0 +1,62 @@
+__pycache__/
+*.py[cod]
+*$py.class
+*.so
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+*.manifest
+*.spec
+
+# Virtualenv
+venv/
+ENV/
+env/
+.venv/
+
+# IDE
+.vscode/
+.idea/
+*.swp
+*.swo
+*~
+
+# Environment files
+.env
+.env.local
+.env.production
+.env.development
+
+# Logs
+*.log
+logs/
+
+# Database
+*.db
+*.sqlite
+*.sqlite3
+
+# LanceDB data
+lancedb_data/
+
+# Temporary files
+*.tmp
+*.temp
+.DS_Store
+Thumbs.db

Dockerfile

@@ -0,0 +1,39 @@
+# Use Python 3.12 as specified
+FROM python:3.12-slim
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+
+# Create a non-root user
+RUN useradd -m -u 1000 user
+USER user
+ENV PATH="/home/user/.local/bin:$PATH"
+
+# Set working directory
+WORKDIR /app
+
+# Copy requirements first for better Docker layer caching
+COPY --chown=user ./requirements.txt requirements.txt
+
+# Install Python dependencies
+RUN pip install --no-cache-dir --upgrade pip && \
+    pip install --no-cache-dir --upgrade -r requirements.txt
+
+# Copy the application code
+COPY --chown=user . /app
+
+# Expose the port that HF Spaces requires
+EXPOSE 7860
+
+# Set environment variables
+ENV PYTHONPATH=/app
+ENV PYTHONUNBUFFERED=1
+
+# Health check
+HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:7860/health || exit 1
+
+# Run the application
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

README.md

@@ -1,10 +1,201 @@
----
-title: PensionBot
-emoji: 🔥
-colorFrom: gray
-colorTo: blue
+---------
+
+title: PensionBot - Voice Assistant
+
+emoji: 🎤title: Voice Bot Government Assistanttitle: Rajasthan Pension Assistant
+
+colorFrom: blue
+
+colorTo: green  emoji: 🎤emoji: 🏛️
+
 sdk: docker
-pinned: false
+
+pinned: falsecolorFrom: bluecolorFrom: blue
+
+license: mit
+
+app_port: 7860colorTo: green  colorTo: purple
+
 ---
 
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+sdk: dockersdk: gradio
+
+# PensionBot - Voice Assistant 🎤
+
+pinned: falsesdk_version: 4.44.0
+
+A sophisticated AI-powered voice assistant designed for government pension queries and document searches. Built with FastAPI, this backend provides comprehensive API endpoints for voice interaction, document processing, and intelligent responses.
+
+license: mitapp_file: gradio_app.py
+
+## 🚀 Features
+
+app_port: 8000pinned: false
+
+- **Voice Processing**: Advanced ASR and TTS capabilities
+
+- **Document Search**: RAG-based government document knowledge base---license: mit
+
+- **Hybrid AI**: Multiple LLM providers for optimal responses
+
+- **WebSocket Support**: Real-time communicationdisable_embedding: false
+
+- **Authentication**: JWT-based secure access
+
+- **Policy Analysis**: Visual charts and scenario analysis# Voice Bot Government Assistant 🎤---
+
+
+
+## 📡 API Endpoints
+
+
+
+- `GET /` - Service information and available endpointsA sophisticated AI-powered voice assistant designed for government policy queries and document searches. Built with FastAPI, this backend provides comprehensive API endpoints for voice interaction, document processing, and intelligent responses.# Rajasthan Government Assistant �️
+
+- `GET /health` - Health check with service status
+
+- `POST /chat` - Text-based conversation interface
+
+- `WebSocket /ws` - Real-time voice and text communication
+
+- `GET /docs` - Interactive API documentation## 🚀 FeaturesA sophisticated AI-powered assistant for Rajasthan government services that combines voice interaction, document search, and intelligent conversation capabilities. Built with FastAPI, LangChain, and advanced RAG (Retrieval-Augmented Generation) technology to help citizens access government information and services.
+
+
+
+## 🛠 Technology Stack
+
+
+
+- **FastAPI**: High-performance web framework- **Voice Processing**: Advanced ASR and TTS capabilities## Features
+
+- **LangChain**: AI orchestration and document processing
+
+- **LanceDB**: Vector database for document search- **Document Search**: RAG-based government document knowledge base
+
+- **Whisper**: Speech-to-text processing
+
+- **Edge-TTS**: Text-to-speech synthesis- **Hybrid AI**: Multiple LLM providers for optimal responses- 🎙️ **Voice Interaction**: Speech-to-text and text-to-speech capabilities
+
+- **WebSocket**: Real-time communication
+
+- **WebSocket Support**: Real-time communication- 📚 **Document Search**: Advanced RAG system with government document knowledge
+
+## 🔗 Usage
+
+- **Authentication**: JWT-based secure access- 🤖 **Hybrid LLM**: Combines multiple AI models for optimal responses
+
+The API is accessible at the base URL of this space. Use the `/docs` endpoint to explore the interactive API documentation.
+
+- **Policy Analysis**: Visual charts and scenario analysis- 🔍 **Scenario Analysis**: Policy impact simulation and analysis
+
+### Example Usage:
+
+- 📊 **Chart Generation**: Visual policy impact charts
+
+```bash
+
+# Health check## 📡 API Endpoints- 🌐 **WebSocket Support**: Real-time communication
+
+curl https://chabhishek28-pensionbot.hf.space/health
+
+- 🛡️ **Authentication**: JWT-based user authentication
+
+# Chat endpoint
+
+curl -X POST https://chabhishek28-pensionbot.hf.space/chat \- `GET /` - Service information and available endpoints
+
+  -H "Content-Type: application/json" \
+
+  -d '{"message": "Tell me about pension policies"}'- `GET /health` - Health check with service status## API Endpoints
+
+```
+
+- `POST /chat` - Text-based conversation interface
+
+## 📋 Environment Variables
+
+- `WebSocket /ws` - Real-time voice and text communication- `GET /`: Root endpoint with service information
+
+The following environment variables are required:
+
+- `GET /docs` - Interactive API documentation- `GET /health`: Health check for all services
+
+- `GOOGLE_API_KEY`: Google Gemini API key
+
+- `GROQ_API_KEY`: Groq API key for Whisper- `POST /chat`: Text-based chat interface
+
+- `TAVILY_API_KEY`: Tavily search API key
+
+- `JWT_SECRET_KEY`: JWT authentication secret## 🛠 Technology Stack- `POST /search`: Document search functionality
+
+
+
+## 🔒 Security- `WebSocket /ws`: Real-time voice and text communication
+
+
+
+This API includes JWT-based authentication for secure access to protected endpoints.- **FastAPI**: High-performance web framework
+
+
+
+## 📄 License- **LangChain**: AI orchestration and document processing## Technology Stack
+
+
+
+MIT License - see LICENSE for details.- **LanceDB**: Vector database for document search
+
+- **Whisper**: Speech-to-text processing- **Backend**: FastAPI, Python 3.11+
+
+- **Edge-TTS**: Text-to-speech synthesis- **AI/ML**: LangChain, Hugging Face Transformers, Sentence Transformers
+
+- **WebSocket**: Real-time communication- **Vector Database**: LanceDB
+
+- **Voice**: Whisper ASR, Edge TTS
+
+## 🔗 Usage- **Authentication**: JWT tokens
+
+
+
+The API is accessible at the base URL of this space. Use the `/docs` endpoint to explore the interactive API documentation.## Environment Variables
+
+
+
+### Example Usage:Set these in your Hugging Face Space secrets:
+
+
+
+```bash- `GOOGLE_API_KEY`: For Gemini AI model
+
+# Health check- `GROQ_API_KEY`: For Groq AI model
+
+curl https://your-space-name.hf.space/health- `TAVILY_API_KEY`: For search capabilities
+
+- `JWT_SECRET_KEY`: For authentication
+
+# Chat endpoint
+
+curl -X POST https://your-space-name.hf.space/chat \## Usage
+
+  -H "Content-Type: application/json" \
+
+  -d '{"message": "Tell me about government policies"}'Once deployed, the API will be available at your Hugging Face Space URL. Use the WebSocket endpoint for real-time voice interaction or the REST endpoints for text-based communication.
+
+```# Updated for HF Spaces deployment
+
+
+## 📋 Environment Variables
+
+The following environment variables are required:
+
+- `GOOGLE_API_KEY`: Google Gemini API key
+- `GROQ_API_KEY`: Groq API key for Whisper
+- `TAVILY_API_KEY`: Tavily search API key
+- `JWT_SECRET_KEY`: JWT authentication secret
+
+## 🔒 Security
+
+This API includes JWT-based authentication for secure access to protected endpoints.
+
+## 📄 License
+
+MIT License - see LICENSE for details.

app.py

@@ -0,0 +1,145 @@
+import os
+import logging
+from datetime import datetime
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, WebSocket, HTTPException
+from fastapi.middleware.cors import CORSMiddleware
+from fastapi.responses import JSONResponse
+from websocket_handler import handle_websocket_connection
+from enhanced_websocket_handler import handle_enhanced_websocket_connection
+from hybrid_llm_service import HybridLLMService
+from voice_service import VoiceService
+from rag_service import search_documents
+from lancedb_service import LanceDBService
+import config
+from dotenv import load_dotenv
+
+# MCP and Authentication imports
+from fastapi import Depends
+from pydantic import BaseModel
+from typing import Optional
+from auth import get_current_user
+
+# Load environment variables
+load_dotenv()
+
+# Configure logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s [%(levelname)s] %(message)s',
+    datefmt='%Y-%m-%d %H:%M:%S'
+)
+logger = logging.getLogger(__name__)
+
+# Get configuration
+config_dict = {
+    "ALLOWED_ORIGINS": config.ALLOWED_ORIGINS,
+    "ENABLE_VOICE_FEATURES": config.ENABLE_VOICE_FEATURES
+}
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan handler"""
+    # Startup
+    logger.info("🚀 Starting Voice Bot Application...")
+    logger.info("✅ Application started successfully")
+    yield
+    # Shutdown (if needed)
+    logger.info("🛑 Shutting down Voice Bot Application...")
+
+# Create FastAPI application
+app = FastAPI(
+    title="Voice Bot Government Assistant",
+    description="AI-powered voice assistant for government policies and services",
+    version="1.0.0",
+    lifespan=lifespan
+)
+
+# Configure CORS
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=config.ALLOWED_ORIGINS,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize services (lazy loading for HF Spaces)
+llm_service = None
+voice_service = None
+lancedb_service = None
+
+def get_llm_service():
+    global llm_service
+    if llm_service is None:
+        llm_service = HybridLLMService()
+    return llm_service
+
+def get_voice_service():
+    global voice_service
+    if voice_service is None:
+        voice_service = VoiceService()
+    return voice_service
+
+def get_lancedb_service():
+    global lancedb_service
+    if lancedb_service is None:
+        lancedb_service = LanceDBService()
+    return lancedb_service
+
+# Health check endpoint
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy",
+        "service": "voice-bot-api",
+        "timestamp": datetime.now().isoformat(),
+        "version": "1.0.0"
+    }
+
+# Root endpoint
+@app.get("/")
+async def root():
+    """Root endpoint with service information"""
+    return {
+        "message": "Voice Bot Government Assistant API",
+        "status": "running",
+        "version": "1.0.0",
+        "endpoints": {
+            "health": "/health",
+            "chat": "/chat",
+            "websocket": "/ws",
+            "docs": "/docs"
+        }
+    }
+
+# Chat endpoint
+@app.post("/chat")
+async def chat_endpoint(request: dict):
+    """Text-based chat endpoint"""
+    try:
+        message = request.get("message", "")
+        if not message:
+            raise HTTPException(status_code=400, detail="Message is required")
+        
+        llm = get_llm_service()
+        response = await llm.get_response(message)
+        
+        return {
+            "response": response,
+            "timestamp": datetime.now().isoformat()
+        }
+    except Exception as e:
+        logger.error(f"Chat error: {str(e)}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+# WebSocket endpoint
+@app.websocket("/ws")
+async def websocket_endpoint(websocket: WebSocket):
+    """WebSocket endpoint for real-time communication"""
+    await handle_enhanced_websocket_connection(websocket)
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

audio_services.py

@@ -0,0 +1,88 @@
+from groq import AsyncGroq
+from config import GROQ_API_KEY, ASR_MODEL, MURF_API_KEY
+import soundfile as sf
+import numpy as np
+from huggingface_hub import hf_hub_download
+from concurrent.futures import ThreadPoolExecutor
+from murf import AsyncMurf
+
+
+groq = AsyncGroq(api_key=GROQ_API_KEY)
+executor = ThreadPoolExecutor(max_workers=4)
+
+# kokoro_device = "cuda" if torch.cuda.is_available() else "cpu"
+# kokoro_model = KModel().to(kokoro_device).eval()
+# model_path = hf_hub_download(repo_id='hexgrad/Kokoro-82M', filename="kokoro-v1_0.pth")
+# kokoro_model.load_state_dict(torch.load(model_path, map_location=kokoro_device), strict=False)
+# kokoro_pipeline = KPipeline(lang_code='a', model=False)
+# voice_path = hf_hub_download("hexgrad/Kokoro-82M", "voices/af_heart.pt")
+# kokoro_voice = torch.load(voice_path, weights_only=True).to(kokoro_device)
+
+async def groq_asr_bytes(audio_bytes: bytes, model: str = ASR_MODEL, language: str = "en") -> str:
+    """Transcribes audio using Groq ASR."""
+    # Groq client is already async, so we can use it directly
+    resp = await groq.audio.transcriptions.create(
+        model=model,
+        file=("audio.wav", audio_bytes, "audio/wav"),
+        response_format="text",
+        language=language
+    )
+    return resp
+
+murf_client = AsyncMurf(api_key=MURF_API_KEY)
+
+async def murf_tts(text: str, voice_id: str = "en-IN-isha", format: str = "MP3") -> bytes:
+    resp = murf_client.text_to_speech.stream(
+        text=text,
+        voice_id=voice_id,
+        format=format,
+        sample_rate=44100.0
+    )
+    chunks = [chunk async for chunk in resp]
+    full_audio = b''.join(chunks)
+    return full_audio
+
+
+
+
+
+
+
+
+
+
+
+# def groq_tts(text: str, speed: float = 1.0) -> bytes:
+#     try:
+#         audio_segments = []
+#         for _, ps, _ in kokoro_pipeline(text, kokoro_voice, speed):
+#             ref_s = kokoro_voice[len(ps) - 1]
+#             audio = kokoro_model(ps, ref_s, speed)
+#             audio_np = audio.cpu().numpy().astype(np.float32)
+#             audio_segments.append(audio_np)
+
+#         full_audio = np.concatenate(audio_segments)
+
+#         # Write to WAV bytes
+#         buf = io.BytesIO()
+#         sf.write(buf, full_audio, samplerate=24000, format="WAV", subtype="PCM_16")
+#         buf.seek(0)
+#         return buf.read()
+
+#     except Exception as e:
+#         print("Kokoro TTS synthesis failed")
+#         raise RuntimeError(f"Kokoro TTS failed: {e}")
+
+
+'''def groq_tts(text: str, model: str = TTS_MODEL, voice: str = TTS_VOICE) -> bytes:
+    text = text[:1000]
+    resp = groq.audio.speech.create(
+        model=model,
+        voice=voice,
+        input=text,
+        response_format="wav"
+    )
+    print(resp.read()[:10])
+    return resp.read()
+    '''
+   

auth.py

@@ -0,0 +1,25 @@
+import jwt
+from fastapi import HTTPException, status, Header
+from jwt import PyJWTError
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+SUPABASE_JWT_SECRET = os.getenv("SUPABASE_JWT_SECRET")
+def verify_token(token: str):
+    try:
+        payload = jwt.decode(token, SUPABASE_JWT_SECRET, algorithms=["HS256"], audience="authenticated")
+        return payload
+    except PyJWTError:
+        raise HTTPException(
+            status_code=status.HTTP_401_UNAUTHORIZED,
+            detail="Invalid authentication credentials",
+        )
+    
+def get_current_user(authorization: str = Header(...)):
+    if not authorization.startswith("Bearer "):
+        raise HTTPException(status_code=401, detail="Invalid Authorization header")
+    token = authorization.split(" ")[1]
+    payload = verify_token(token)
+    return payload

config.py

@@ -0,0 +1,51 @@
+from dotenv import load_dotenv
+import os
+
+load_dotenv()
+
+# API Configuration
+GOOGLE_API_KEY = os.environ.get("GOOGLE_API_KEY")
+GEMINI_API_KEY = os.environ.get("GOOGLE_API_KEY")  # Backward compatibility
+
+# LangSmith Configuration (optional)
+LANGSMITH_API_KEY = os.environ.get("LANGSMITH_API_KEY")
+LANGCHAIN_TRACING_V2 = os.environ.get("LANGCHAIN_TRACING_V2", "false").lower() == "true"
+LANGCHAIN_PROJECT = os.environ.get("LANGCHAIN_PROJECT", "voice-bot-government-docs")
+
+# Hybrid LLM Configuration
+USE_HYBRID_LLM = os.environ.get("USE_HYBRID_LLM", "false").lower() == "true"
+FAST_LLM_PROVIDER = os.environ.get("FAST_LLM_PROVIDER", "groq")
+COMPLEX_LLM_PROVIDER = os.environ.get("COMPLEX_LLM_PROVIDER", "gemini")
+
+# Groq Configuration
+GROQ_API_KEY = os.environ.get("GROQ_API_KEY")
+GROQ_MODEL = os.environ.get("GROQ_MODEL", "llama-3.1-70b-versatile")
+
+# Gemini Model Configuration
+GEMINI_MODEL = os.environ.get("GEMINI_MODEL", "gemini-1.5-pro-latest")
+GEMINI_TEMPERATURE = float(os.environ.get("GEMINI_TEMPERATURE", "0.7"))
+
+# Voice Features Configuration
+ENABLE_VOICE_FEATURES = os.environ.get("ENABLE_VOICE_FEATURES", "false").lower() == "true"
+TTS_PROVIDER = os.environ.get("TTS_PROVIDER", "edge-tts")
+ASR_PROVIDER = os.environ.get("ASR_PROVIDER", "whisper")
+VOICE_LANGUAGE = os.environ.get("VOICE_LANGUAGE", "en-US")
+DEFAULT_VOICE_SPEED = float(os.environ.get("DEFAULT_VOICE_SPEED", "1.0"))
+
+# Embedding Model Configuration
+EMBEDDING_MODEL_NAME = os.environ.get("EMBEDDING_MODEL_NAME", "sentence-transformers/all-MiniLM-L6-v2")
+EMBEDDING_SIZE = 768
+
+# Text Processing Configuration
+CHUNK_SIZE = int(os.environ.get("CHUNK_SIZE", "1000"))
+CHUNK_OVERLAP = int(os.environ.get("CHUNK_OVERLAP", "200"))
+
+# CORS Configuration
+ALLOWED_ORIGINS = os.environ.get("ALLOWED_ORIGINS", "*").split(",") if os.environ.get("ALLOWED_ORIGINS") != "*" else ["*"]
+
+# LanceDB Configuration
+LANCEDB_PATH = os.environ.get("LANCEDB_PATH", "./lancedb_data")
+
+# JWT Configuration
+JWT_SECRET_KEY = os.environ.get("JWT_SECRET_KEY")
+JWT_ALGORITHM = os.environ.get("JWT_ALGORITHM", "HS256")

document_service.py

@@ -0,0 +1,171 @@
+from fastapi import UploadFile
+from langchain.text_splitter import RecursiveCharacterTextSplitter
+from langchain.docstore.document import Document
+import pdfplumber
+import os
+import asyncio
+from typing import List
+from lancedb_service import lancedb_service
+from config import CHUNK_SIZE, CHUNK_OVERLAP
+from datetime import datetime
+
+def read_pdf(file: UploadFile) -> str:
+    with pdfplumber.open(file.file) as pdf:
+        text = "\n".join(page.extract_text() or "" for page in pdf.pages)
+    return text
+
+async def process_document_upload(file: UploadFile, userid: str, knowledge_base: str):
+    try:
+        filename = file.filename
+        if not filename.lower().endswith(".pdf"):
+            return {"error": "Only PDF files are supported"}
+
+        # Read PDF
+        with pdfplumber.open(file.file) as pdf:
+            text = "\n".join(page.extract_text() or "" for page in pdf.pages)
+
+        # Chunk text
+        splitter = RecursiveCharacterTextSplitter(chunk_size=CHUNK_SIZE, chunk_overlap=CHUNK_OVERLAP)
+        chunks = splitter.split_text(text)
+
+        # Batch create Document objects with metadata including knowledge base
+        upload_date = datetime.now().isoformat()
+        docs = [
+            Document(
+                page_content=chunk,
+                metadata={
+                    "source": filename,
+                    "userid": userid,
+                    "knowledge_base": knowledge_base,
+                    "upload_date": upload_date
+                }
+            )
+            for chunk in chunks
+        ]
+
+        # ✅ Batch embed & insert using LanceDB
+        await lancedb_service.add_documents(docs, userid, knowledge_base, filename)
+
+        return {
+            "status": "uploaded",
+            "chunks": len(docs),
+            "file": filename,
+            "knowledge_base": knowledge_base
+        }
+
+    except Exception as e:
+        return {"error": str(e)}
+
+def read_pdf_from_path(pdf_path: str) -> str:
+    """Read PDF content from file path"""
+    try:
+        with pdfplumber.open(pdf_path) as pdf:
+            text = "\n".join(page.extract_text() or "" for page in pdf.pages)
+        return text
+    except Exception as e:
+        print(f"Error reading PDF {pdf_path}: {str(e)}")
+        return ""
+
+async def process_documents_from_folder(folder_path: str, userid: str = "system", knowledge_base: str = "government_docs"):
+    """Process all PDF documents from the specified folder"""
+    try:
+        if not os.path.exists(folder_path):
+            return {"error": f"Folder path {folder_path} does not exist"}
+        
+        pdf_files = [f for f in os.listdir(folder_path) if f.lower().endswith('.pdf')]
+        
+        if not pdf_files:
+            return {"error": "No PDF files found in the folder"}
+        
+        processed_files = []
+        total_chunks = 0
+        
+        for pdf_file in pdf_files:
+            pdf_path = os.path.join(folder_path, pdf_file)
+            
+            # Read PDF content
+            text = read_pdf_from_path(pdf_path)
+            
+            if not text.strip():
+                print(f"Skipping {pdf_file} - no text content extracted")
+                continue
+            
+            # Chunk text
+            splitter = RecursiveCharacterTextSplitter(
+                chunk_size=CHUNK_SIZE, 
+                chunk_overlap=CHUNK_OVERLAP
+            )
+            chunks = splitter.split_text(text)
+            
+            # Create Document objects with metadata
+            upload_date = datetime.now().isoformat()
+            docs = [
+                Document(
+                    page_content=chunk,
+                    metadata={
+                        "source": pdf_file,
+                        "userid": userid,
+                        "knowledge_base": knowledge_base,
+                        "upload_date": upload_date,
+                        "file_path": pdf_path
+                    }
+                )
+                for chunk in chunks
+            ]
+            
+            # Add documents to LanceDB
+            await lancedb_service.add_documents(docs, userid, knowledge_base, pdf_file)
+            
+            processed_files.append({
+                "file": pdf_file,
+                "chunks": len(chunks)
+            })
+            total_chunks += len(chunks)
+            
+            print(f"Processed {pdf_file}: {len(chunks)} chunks")
+        
+        return {
+            "status": "success",
+            "processed_files": len(processed_files),
+            "total_chunks": total_chunks,
+            "files": processed_files,
+            "knowledge_base": knowledge_base
+        }
+        
+    except Exception as e:
+        return {"error": str(e)}
+
+async def initialize_document_database():
+    """Initialize the document database with documents from the aa folder"""
+    # Path to the documents folder
+    documents_folder = "/Users/abhishekchoudhary/Abhi Project/aa/raw_documents/Documents"
+    
+    print("Starting document database initialization...")
+    result = await process_documents_from_folder(
+        folder_path=documents_folder,
+        userid="system",
+        knowledge_base="government_docs"
+    )
+    
+    if "error" in result:
+        print(f"Error initializing database: {result['error']}")
+    else:
+        print(f"Successfully initialized database with {result['total_chunks']} chunks from {result['processed_files']} files")
+    
+    return result
+
+async def get_available_knowledge_bases() -> List[str]:
+    """Get list of available knowledge bases"""
+    try:
+        return await lancedb_service.get_knowledge_bases()
+    except Exception as e:
+        print(f"Error getting knowledge bases: {str(e)}")
+        return []
+
+async def get_documents_by_knowledge_base(knowledge_base: str) -> List[dict]:
+    """Get list of documents in a specific knowledge base"""
+    try:
+        return await lancedb_service.get_documents_by_knowledge_base(knowledge_base)
+    except Exception as e:
+        print(f"Error getting documents for knowledge base {knowledge_base}: {str(e)}")
+        return []

enhanced_websocket_handler.py

@@ -0,0 +1,395 @@
+"""
+Enhanced WebSocket handler with hybrid LLM and optional voice features
+"""
+
+from fastapi import WebSocket, WebSocketDisconnect
+from langchain_core.messages import HumanMessage, SystemMessage, AIMessage
+import logging
+import json
+import asyncio
+import uuid
+import tempfile
+import base64
+from pathlib import Path
+
+from llm_service import create_graph, create_basic_graph
+from lancedb_service import lancedb_service
+from hybrid_llm_service import HybridLLMService
+from voice_service import voice_service
+from rag_service import search_government_docs
+
+# Initialize hybrid LLM service
+hybrid_llm_service = HybridLLMService()
+
+logger = logging.getLogger("voicebot")
+
+async def handle_enhanced_websocket_connection(websocket: WebSocket):
+    """Enhanced WebSocket handler with hybrid LLM and voice features"""
+    await websocket.accept()
+    logger.info("🔌 Enhanced WebSocket client connected.")
+
+    # Initialize session data
+    session_data = {
+        "messages": [],
+        "user_preferences": {
+            "voice_enabled": False,
+            "preferred_voice": "en-US-AriaNeural",
+            "response_mode": "text"  # text, voice, both
+        },
+        "context": ""
+    }
+
+    try:
+        # Get initial connection data
+        initial_data = await websocket.receive_json()
+        
+        # Extract user preferences
+        if "preferences" in initial_data:
+            session_data["user_preferences"].update(initial_data["preferences"])
+        
+        # Setup user session
+        flag = "user_id" in initial_data
+        graph = None  # Initialize graph variable
+        
+        if flag:
+            thread_id = initial_data.get("user_id")
+            knowledge_base = initial_data.get("knowledge_base", "government_docs")
+            
+            # Use hybrid LLM or traditional graph based on configuration
+            if hybrid_llm_service.use_hybrid:
+                logger.info("🤖 Using Hybrid LLM Service")
+                use_hybrid = True
+            else:
+                graph = await create_graph(kb_tool=True, mcp_config=None)
+                use_hybrid = False
+            
+            config = {
+                "configurable": {
+                    "thread_id": thread_id,
+                    "knowledge_base": knowledge_base,
+                }
+            }
+        else:
+            # Basic setup for unauthenticated users
+            thread_id = str(uuid.uuid4())
+            knowledge_base = "government_docs"
+            use_hybrid = hybrid_llm_service.use_hybrid
+            
+            if not use_hybrid:
+                graph = create_basic_graph()
+            
+            config = {"configurable": {"thread_id": thread_id}}
+
+        # Send initial greeting with voice/hybrid capabilities
+        await send_enhanced_greeting(websocket, session_data)
+
+        # Main message handling loop
+        while True:
+            try:
+                data = await websocket.receive_json()
+                
+                if data["type"] == "text_message":
+                    await handle_text_message(
+                        websocket, data, session_data, 
+                        use_hybrid, config, knowledge_base, graph
+                    )
+                
+                elif data["type"] == "voice_message":
+                    await handle_voice_message(
+                        websocket, data, session_data,
+                        use_hybrid, config, knowledge_base, graph
+                    )
+                
+                elif data["type"] == "preferences_update":
+                    await handle_preferences_update(websocket, data, session_data)
+                
+                elif data["type"] == "get_voice_status":
+                    await websocket.send_json({
+                        "type": "voice_status",
+                        "data": voice_service.get_voice_status()
+                    })
+                
+                elif data["type"] == "get_llm_status":
+                    await websocket.send_json({
+                        "type": "llm_status", 
+                        "data": hybrid_llm_service.get_provider_info()
+                    })
+
+            except WebSocketDisconnect:
+                logger.info("🔌 WebSocket client disconnected.")
+                break
+            except Exception as e:
+                logger.error(f"❌ Error handling message: {e}")
+                await websocket.send_json({
+                    "type": "error",
+                    "message": f"An error occurred: {str(e)}"
+                })
+
+    except WebSocketDisconnect:
+        logger.info("🔌 WebSocket client disconnected during setup.")
+    except Exception as e:
+        logger.error(f"❌ WebSocket error: {e}")
+        try:
+            await websocket.send_json({
+                "type": "error",
+                "message": f"Connection error: {str(e)}"
+            })
+        except:
+            pass
+
+async def send_enhanced_greeting(websocket: WebSocket, session_data: dict):
+    """Send enhanced greeting with system capabilities"""
+    
+    # Get system status
+    llm_info = hybrid_llm_service.get_provider_info()
+    voice_status = voice_service.get_voice_status()
+    
+    greeting_text = f"""🤖 Welcome to the Government Document Assistant!
+
+I'm powered by a hybrid AI system that can help you with:
+• Government policies and procedures
+• Document search and analysis 
+• Scenario analysis with visualizations
+• Quick answers and detailed explanations
+
+Current capabilities:
+• LLM: {'Hybrid (' + llm_info['fast_provider'] + '/' + llm_info['complex_provider'] + ')' if llm_info['hybrid_enabled'] else 'Single provider'}
+• Voice features: {'Enabled' if voice_status['voice_enabled'] else 'Disabled'}
+
+How can I assist you today? You can ask me about any government policies, procedures, or documents!"""
+
+    # Send text greeting
+    await websocket.send_json({
+        "type": "message_response",
+        "message": greeting_text,
+        "provider_used": "system",
+        "capabilities": {
+            "hybrid_llm": llm_info['hybrid_enabled'],
+            "voice_features": voice_status['voice_enabled'],
+            "scenario_analysis": True
+        }
+    })
+
+    # Send voice greeting if enabled
+    if session_data["user_preferences"]["voice_enabled"] and voice_status['voice_enabled']:
+        voice_greeting = "Welcome to the Government Document Assistant! I can help you with policies, procedures, and document analysis. How can I assist you today?"
+        audio_data = await voice_service.text_to_speech(voice_greeting)
+        
+        if audio_data:
+            await websocket.send_json({
+                "type": "audio_response",
+                "audio_data": base64.b64encode(audio_data).decode(),
+                "format": "mp3"
+            })
+
+async def handle_text_message(websocket: WebSocket, data: dict, session_data: dict, 
+                            use_hybrid: bool, config: dict, knowledge_base: str, graph=None):
+    """Handle text message with hybrid LLM"""
+    
+    user_message = data["message"]
+    logger.info(f"💬 Received text message: {user_message}")
+    
+    # Send acknowledgment
+    await websocket.send_json({
+        "type": "message_received",
+        "message": "Processing your message..."
+    })
+
+    try:
+        if use_hybrid:
+            # Use hybrid LLM service
+            response_text, provider_used = await get_hybrid_response(
+                user_message, session_data["context"], config, knowledge_base
+            )
+        else:
+            # Use traditional graph approach
+            session_data["messages"].append(HumanMessage(content=user_message))
+            result = await graph.ainvoke({"messages": session_data["messages"]}, config)
+            response_text = result["messages"][-1].content
+            provider_used = "traditional"
+
+        # Handle scenario analysis images
+        if "SCENARIO_ANALYSIS_IMAGE:" in response_text:
+            await handle_scenario_response(websocket, response_text, provider_used)
+        else:
+            await send_text_response(websocket, response_text, provider_used, session_data)
+
+    except Exception as e:
+        logger.error(f"❌ Error processing text message: {e}")
+        await websocket.send_json({
+            "type": "error", 
+            "message": f"Error processing your message: {str(e)}"
+        })
+
+async def handle_voice_message(websocket: WebSocket, data: dict, session_data: dict,
+                             use_hybrid: bool, config: dict, knowledge_base: str, graph=None):
+    """Handle voice message with ASR and TTS"""
+    
+    if not voice_service.is_voice_enabled():
+        await websocket.send_json({
+            "type": "error",
+            "message": "Voice features are not enabled"
+        })
+        return
+
+    try:
+        # Decode audio data
+        audio_data = base64.b64decode(data["audio_data"])
+        
+        # Save to temporary file
+        with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+            temp_file.write(audio_data)
+            temp_file_path = temp_file.name
+
+        # Convert speech to text
+        transcribed_text = await voice_service.speech_to_text(temp_file_path)
+        
+        # Clean up temp file
+        Path(temp_file_path).unlink()
+
+        if not transcribed_text:
+            await websocket.send_json({
+                "type": "error",
+                "message": "Could not transcribe audio"
+            })
+            return
+
+        logger.info(f"🎤 Transcribed: {transcribed_text}")
+
+        # Send transcription
+        await websocket.send_json({
+            "type": "transcription",
+            "text": transcribed_text
+        })
+
+        # Process as text message
+        if use_hybrid:
+            response_text, provider_used = await get_hybrid_response(
+                transcribed_text, session_data["context"], config, knowledge_base
+            )
+        else:
+            session_data["messages"].append(HumanMessage(content=transcribed_text))
+            result = await graph.ainvoke({"messages": session_data["messages"]}, config)
+            response_text = result["messages"][-1].content
+            provider_used = "traditional"
+
+        # Send text response
+        await send_text_response(websocket, response_text, provider_used, session_data)
+
+        # Send voice response if enabled
+        if session_data["user_preferences"]["response_mode"] in ["voice", "both"]:
+            voice_text = voice_service.create_voice_response_with_guidance(
+                response_text,
+                suggested_resources=["Government portal", "Local offices"],
+                redirect_info="contact your local government office for personalized assistance"
+            )
+            
+            audio_response = await voice_service.text_to_speech(
+                voice_text, 
+                session_data["user_preferences"]["preferred_voice"]
+            )
+            
+            if audio_response:
+                await websocket.send_json({
+                    "type": "audio_response",
+                    "audio_data": base64.b64encode(audio_response).decode(),
+                    "format": "mp3"
+                })
+
+    except Exception as e:
+        logger.error(f"❌ Error processing voice message: {e}")
+        await websocket.send_json({
+            "type": "error",
+            "message": f"Error processing voice message: {str(e)}"
+        })
+
+async def get_hybrid_response(user_message: str, context: str, config: dict, knowledge_base: str):
+    """Get response using hybrid LLM with document search"""
+    
+    # Search for relevant documents
+    try:
+        search_results = await search_government_docs.ainvoke(
+            {"query": user_message}, 
+            config=config
+        )
+        context = search_results if search_results else context
+    except:
+        logger.warning("Document search failed, using existing context")
+    
+    # Get hybrid LLM response
+    response_text = await hybrid_llm_service.get_response(
+        user_message, 
+        context=context,
+        system_prompt="""You are a helpful government document assistant. Provide accurate, helpful responses based on the context provided. When appropriate, suggest additional resources or redirect users to relevant departments for more assistance."""
+    )
+    
+    # Determine which provider was used
+    complexity = hybrid_llm_service.determine_task_complexity(user_message, context)
+    provider_used = hybrid_llm_service.choose_llm_provider(complexity)
+    
+    return response_text, provider_used
+
+async def send_text_response(websocket: WebSocket, response_text: str, provider_used: str, session_data: dict):
+    """Send text response to client"""
+    
+    await websocket.send_json({
+        "type": "message_response",
+        "message": response_text,
+        "provider_used": provider_used,
+        "timestamp": asyncio.get_event_loop().time()
+    })
+    
+    # Update session context
+    session_data["context"] = response_text[-1000:]  # Keep last 1000 chars as context
+
+async def handle_scenario_response(websocket: WebSocket, response_text: str, provider_used: str):
+    """Handle scenario analysis response with images"""
+    
+    parts = response_text.split("SCENARIO_ANALYSIS_IMAGE:")
+    text_part = parts[0].strip()
+    
+    # Send text part
+    if text_part:
+        await websocket.send_json({
+            "type": "message_response",
+            "message": text_part,
+            "provider_used": provider_used
+        })
+    
+    # Send image parts
+    for i, part in enumerate(parts[1:], 1):
+        try:
+            image_data = part.strip()
+            await websocket.send_json({
+                "type": "scenario_image",
+                "image_data": image_data,
+                "image_index": i,
+                "chart_type": "analysis"
+            })
+        except Exception as e:
+            logger.error(f"Error sending scenario image {i}: {e}")
+
+async def handle_preferences_update(websocket: WebSocket, data: dict, session_data: dict):
+    """Handle user preferences update"""
+    
+    try:
+        session_data["user_preferences"].update(data["preferences"])
+        
+        await websocket.send_json({
+            "type": "preferences_updated",
+            "preferences": session_data["user_preferences"]
+        })
+        
+        logger.info(f"🔧 Updated user preferences: {session_data['user_preferences']}")
+        
+    except Exception as e:
+        logger.error(f"❌ Error updating preferences: {e}")
+        await websocket.send_json({
+            "type": "error",
+            "message": f"Error updating preferences: {str(e)}"
+        })
+
+# Keep the original function for backward compatibility
+async def handle_websocket_connection(websocket: WebSocket):
+    """Original websocket handler for backward compatibility"""
+    await handle_enhanced_websocket_connection(websocket)

hybrid_llm_service.py

@@ -0,0 +1,261 @@
+"""
+Hybrid LLM Service that intelligently routes between Groq and Gemini APIs
+based on task complexity and user requirements.
+"""
+
+import os
+import asyncio
+from enum import Enum
+from typing import Dict, Any, Optional
+import logging
+from langchain_groq import ChatGroq
+from langchain_google_genai import ChatGoogleGenerativeAI
+from langchain_core.messages import HumanMessage, SystemMessage
+
+logger = logging.getLogger(__name__)
+
+class TaskComplexity(Enum):
+    SIMPLE = "simple"
+    COMPLEX = "complex"
+
+class LLMProvider(Enum):
+    GROQ = "groq"
+    GEMINI = "gemini"
+
+class HybridLLMService:
+    def __init__(self):
+        # Initialize Groq (Primary)
+        self.groq_api_key = os.getenv("GROQ_API_KEY")
+        self.groq_model = os.getenv("GROQ_MODEL", "llama-3.1-70b-versatile")
+        
+        if self.groq_api_key:
+            self.groq_llm = ChatGroq(
+                groq_api_key=self.groq_api_key,
+                model_name=self.groq_model,
+                temperature=0.7
+            )
+            logger.info(f"✅ Groq LLM initialized: {self.groq_model}")
+        else:
+            self.groq_llm = None
+            logger.warning("⚠️ Groq API key not found")
+        
+        # Initialize Gemini (Secondary/Fallback)
+        self.google_api_key = os.getenv("GOOGLE_API_KEY")
+        self.gemini_model = os.getenv("GEMINI_MODEL", "gemini-1.5-flash")  # Use flash model for free tier
+        
+        if self.google_api_key:
+            try:
+                self.gemini_llm = ChatGoogleGenerativeAI(
+                    model=self.gemini_model,
+                    google_api_key=self.google_api_key,
+                    temperature=0.7
+                )
+                logger.info(f"✅ Gemini LLM initialized: {self.gemini_model}")
+            except Exception as e:
+                self.gemini_llm = None
+                logger.warning(f"⚠️ Gemini initialization failed: {e}")
+        else:
+            self.gemini_llm = None
+            logger.warning("⚠️ Google API key not found")
+        
+        # Hybrid configuration
+        self.use_hybrid = os.getenv("USE_HYBRID_LLM", "true").lower() == "true"
+        self.primary_provider = LLMProvider.GROQ  # Always use Groq as primary
+        
+        logger.info(f"🤖 Hybrid LLM Service initialized (Primary: {self.primary_provider.value})")
+
+    def analyze_task_complexity(self, message: str) -> TaskComplexity:
+        """Analyze if a task requires complex reasoning or simple response"""
+        complex_keywords = [
+            'analyze', 'compare', 'evaluate', 'scenario', 'chart', 'graph', 
+            'visualization', 'complex', 'detailed analysis', 'multi-step',
+            'comprehensive', 'in-depth', 'elaborate', 'breakdown'
+        ]
+        
+        simple_keywords = [
+            'what is', 'who is', 'when', 'where', 'how to', 'define',
+            'explain', 'tell me', 'show me', 'list', 'summary'
+        ]
+        
+        message_lower = message.lower()
+        
+        # Count complex vs simple indicators
+        complex_score = sum(1 for keyword in complex_keywords if keyword in message_lower)
+        simple_score = sum(1 for keyword in simple_keywords if keyword in message_lower)
+        
+        # If message is very long (>200 chars) or has complex keywords, use complex
+        if len(message) > 200 or complex_score > simple_score:
+            return TaskComplexity.COMPLEX
+        
+        return TaskComplexity.SIMPLE
+
+    def choose_llm_provider(self, message: str) -> LLMProvider:
+        """Choose the best LLM provider based on task complexity and availability"""
+        
+        # If hybrid is disabled, always use primary (Groq)
+        if not self.use_hybrid:
+            return LLMProvider.GROQ if self.groq_llm else LLMProvider.GEMINI
+        
+        # Always prefer Groq for better speed and reliability
+        if self.groq_llm:
+            return LLMProvider.GROQ
+        
+        # Fallback to Gemini only if Groq is not available
+        if self.gemini_llm:
+            return LLMProvider.GEMINI
+        
+        # If neither is available, return Groq (will handle error gracefully)
+        return LLMProvider.GROQ
+
+    async def get_response(self, message: str, context: str = "") -> str:
+        """Get response from the chosen LLM provider"""
+        provider = self.choose_llm_provider(message)
+        complexity = self.analyze_task_complexity(message)
+        
+        logger.info(f"🎯 Using {provider.value} for {complexity.value} task")
+        
+        try:
+            if provider == LLMProvider.GROQ and self.groq_llm:
+                return await self._get_groq_response(message, context)
+            elif provider == LLMProvider.GEMINI and self.gemini_llm:
+                return await self._get_gemini_response(message, context)
+            else:
+                # Fallback logic
+                if self.groq_llm:
+                    logger.info("🔄 Falling back to Groq")
+                    return await self._get_groq_response(message, context)
+                elif self.gemini_llm:
+                    logger.info("🔄 Falling back to Gemini")
+                    return await self._get_gemini_response(message, context)
+                else:
+                    return "I apologize, but no AI providers are currently available. Please check your API keys."
+                    
+        except Exception as e:
+            logger.error(f"❌ Error with {provider.value}: {e}")
+            
+            # Try fallback provider
+            if provider == LLMProvider.GROQ and self.gemini_llm:
+                logger.info("🔄 Groq failed, trying Gemini")
+                try:
+                    return await self._get_gemini_response(message, context)
+                except Exception as gemini_error:
+                    logger.error(f"❌ Gemini fallback also failed: {gemini_error}")
+                    return f"I apologize, but I'm experiencing technical difficulties. Both AI providers are currently unavailable."
+            
+            elif provider == LLMProvider.GEMINI and self.groq_llm:
+                logger.info("🔄 Gemini failed, trying Groq")
+                try:
+                    return await self._get_groq_response(message, context)
+                except Exception as groq_error:
+                    logger.error(f"❌ Groq fallback also failed: {groq_error}")
+                    return f"I apologize, but I'm experiencing technical difficulties. Both AI providers are currently unavailable."
+            
+            return f"I apologize, but I encountered an error: {str(e)}"
+
+    async def _get_groq_response(self, message: str, context: str = "") -> str:
+        """Get response from Groq LLM"""
+        system_prompt = """You are a helpful AI assistant specializing in government policies and procedures. 
+        You have access to government documents and can provide accurate information based on them.
+        Provide clear, concise, and helpful responses."""
+        
+        if context:
+            system_prompt += f"\n\nRelevant context from documents:\n{context}"
+        
+        messages = [
+            SystemMessage(content=system_prompt),
+            HumanMessage(content=message)
+        ]
+        
+        response = await self.groq_llm.ainvoke(messages)
+        return response.content
+
+    async def _get_gemini_response(self, message: str, context: str = "") -> str:
+        """Get response from Gemini LLM"""
+        system_prompt = """You are a helpful AI assistant specializing in government policies and procedures. 
+        You have access to government documents and can provide accurate information based on them.
+        Provide detailed, analytical responses when needed."""
+        
+        if context:
+            system_prompt += f"\n\nRelevant context from documents:\n{context}"
+        
+        messages = [
+            SystemMessage(content=system_prompt),
+            HumanMessage(content=message)
+        ]
+        
+        response = await self.gemini_llm.ainvoke(messages)
+        return response.content
+
+    async def get_streaming_response(self, message: str, context: str = ""):
+        """Get streaming response from the chosen LLM provider"""
+        provider = self.choose_llm_provider(message)
+        
+        try:
+            if provider == LLMProvider.GROQ and self.groq_llm:
+                async for chunk in self._get_groq_streaming_response(message, context):
+                    yield chunk
+            elif provider == LLMProvider.GEMINI and self.gemini_llm:
+                async for chunk in self._get_gemini_streaming_response(message, context):
+                    yield chunk
+            else:
+                # Fallback to available provider
+                if self.groq_llm:
+                    async for chunk in self._get_groq_streaming_response(message, context):
+                        yield chunk
+                else:
+                    yield "No AI providers are currently available."
+                    
+        except Exception as e:
+            logger.error(f"❌ Streaming error with {provider.value}: {e}")
+            
+            # Try fallback
+            if provider == LLMProvider.GROQ and self.gemini_llm:
+                try:
+                    async for chunk in self._get_gemini_streaming_response(message, context):
+                        yield chunk
+                except:
+                    yield f"I apologize, but I'm experiencing technical difficulties."
+            elif provider == LLMProvider.GEMINI and self.groq_llm:
+                try:
+                    async for chunk in self._get_groq_streaming_response(message, context):
+                        yield chunk
+                except:
+                    yield f"I apologize, but I'm experiencing technical difficulties."
+            else:
+                yield f"Error: {str(e)}"
+
+    async def _get_groq_streaming_response(self, message: str, context: str = ""):
+        """Get streaming response from Groq"""
+        system_prompt = """You are a helpful AI assistant specializing in government policies and procedures."""
+        
+        if context:
+            system_prompt += f"\n\nRelevant context:\n{context}"
+        
+        messages = [
+            SystemMessage(content=system_prompt),
+            HumanMessage(content=message)
+        ]
+        
+        # Groq streaming
+        async for chunk in self.groq_llm.astream(messages):
+            if chunk.content:
+                yield chunk.content
+                await asyncio.sleep(0.01)
+
+    async def _get_gemini_streaming_response(self, message: str, context: str = ""):
+        """Get streaming response from Gemini"""
+        system_prompt = """You are a helpful AI assistant specializing in government policies and procedures."""
+        
+        if context:
+            system_prompt += f"\n\nRelevant context:\n{context}"
+        
+        messages = [
+            SystemMessage(content=system_prompt),
+            HumanMessage(content=message)
+        ]
+        
+        # Gemini streaming
+        async for chunk in self.gemini_llm.astream(messages):
+            if chunk.content:
+                yield chunk.content
+                await asyncio.sleep(0.01)

lancedb_service.py

@@ -0,0 +1,436 @@
+import lancedb
+import pandas as pd
+from langchain_huggingface import HuggingFaceEmbeddings
+from config import EMBEDDING_MODEL_NAME, LANCEDB_PATH
+from typing import List, Dict, Any, Optional
+import logging
+import os
+import uuid
+from datetime import datetime
+import json
+
+logger = logging.getLogger("voicebot")
+
+# Lazy load embedding model to reduce startup time and memory usage
+embedding_model = None
+
+def get_embedding_model():
+    """Lazy load the embedding model"""
+    global embedding_model
+    if embedding_model is None:
+        logger.info(f"Loading embedding model: {EMBEDDING_MODEL_NAME}")
+        embedding_model = HuggingFaceEmbeddings(
+            model_name=EMBEDDING_MODEL_NAME,
+            model_kwargs={
+                "device": "cpu",
+                "trust_remote_code": True
+            },
+            encode_kwargs={
+                "normalize_embeddings": True
+            }
+        )
+    return embedding_model
+
+class LanceDBService:
+    def __init__(self):
+        self.db_path = LANCEDB_PATH
+        self.db = None
+        self.embedding_model = get_embedding_model()
+        self._initialize_db()
+
+    def _initialize_db(self):
+        """Initialize LanceDB connection and create tables if they don't exist"""
+        try:
+            os.makedirs(self.db_path, exist_ok=True)
+            self.db = lancedb.connect(self.db_path)
+            
+            # Initialize tables
+            self._init_documents_table()
+            self._init_personas_table()
+            self._init_mcp_servers_table()
+            self._init_sessions_table()
+            
+            logger.info("✅ LanceDB initialized successfully")
+        except Exception as e:
+            logger.error(f"❌ Error initializing LanceDB: {e}")
+            raise
+
+    def _init_documents_table(self):
+        """Initialize documents table for vector storage"""
+        try:
+            if "documents" not in self.db.table_names():
+                # Create sample data to define schema
+                sample_data = pd.DataFrame({
+                    "id": [str(uuid.uuid4())],
+                    "content": ["sample"],
+                    "metadata": [json.dumps({})],
+                    "user_id": ["sample"],
+                    "knowledge_base": ["sample"],
+                    "filename": ["sample"],
+                    "upload_date": [datetime.now().isoformat()],
+                    "vector": [get_embedding_model().embed_query("sample")]
+                })
+                self.db.create_table("documents", sample_data)
+                # Delete sample data
+                tbl = self.db.open_table("documents")
+                tbl.delete("id = 'sample'")
+        except Exception as e:
+            logger.error(f"❌ Error initializing documents table: {e}")
+
+    def _init_personas_table(self):
+        """Initialize personas table"""
+        try:
+            if "personas" not in self.db.table_names():
+                sample_data = pd.DataFrame({
+                    "id": [str(uuid.uuid4())],
+                    "user_id": ["sample"],
+                    "name": ["sample"],
+                    "description": ["sample"],
+                    "icon": ["sample"],
+                    "custom_prompt": ["sample"],
+                    "knowledge_base": ["none"],
+                    "language": ["en"],
+                    "created_at": [datetime.now().isoformat()],
+                    "updated_at": [datetime.now().isoformat()]
+                })
+                self.db.create_table("personas", sample_data)
+                tbl = self.db.open_table("personas")
+                tbl.delete("id = 'sample'")
+        except Exception as e:
+            logger.error(f"❌ Error initializing personas table: {e}")
+
+    def _init_mcp_servers_table(self):
+        """Initialize MCP servers table"""
+        try:
+            if "mcp_servers" not in self.db.table_names():
+                sample_data = pd.DataFrame({
+                    "id": [str(uuid.uuid4())],
+                    "user_id": ["sample"],
+                    "name": ["sample"],
+                    "url": ["sample"],
+                    "bearer_token": ["sample"],
+                    "created_at": [datetime.now().isoformat()]
+                })
+                self.db.create_table("mcp_servers", sample_data)
+                tbl = self.db.open_table("mcp_servers")
+                tbl.delete("id = 'sample'")
+        except Exception as e:
+            logger.error(f"❌ Error initializing mcp_servers table: {e}")
+
+    def _init_sessions_table(self):
+        """Initialize sessions table"""
+        try:
+            if "sessions" not in self.db.table_names():
+                sample_data = pd.DataFrame({
+                    "id": [str(uuid.uuid4())],
+                    "user_id": ["sample"],
+                    "persona_id": ["sample"],
+                    "persona_source": ["sample"],
+                    "session_summary": ["sample"],
+                    "created_at": [datetime.now().isoformat()],
+                    "updated_at": [datetime.now().isoformat()]
+                })
+                self.db.create_table("sessions", sample_data)
+                tbl = self.db.open_table("sessions")
+                tbl.delete("id = 'sample'")
+        except Exception as e:
+            logger.error(f"❌ Error initializing sessions table: {e}")
+
+    async def add_documents(self, docs, user_id: str, knowledge_base: str, filename: str):
+        """Add documents to LanceDB vector store"""
+        try:
+            documents_to_insert = []
+            for doc in docs:
+                embedding = self.embedding_model.embed_query(doc.page_content)
+                
+                doc_data = {
+                    "id": str(uuid.uuid4()),
+                    "content": doc.page_content,
+                    "metadata": json.dumps(doc.metadata),
+                    "user_id": user_id,
+                    "knowledge_base": knowledge_base,
+                    "filename": filename,
+                    "upload_date": datetime.now().isoformat(),
+                    "vector": embedding
+                }
+                documents_to_insert.append(doc_data)
+
+            # Insert documents
+            tbl = self.db.open_table("documents")
+            df = pd.DataFrame(documents_to_insert)
+            tbl.add(df)
+            
+            logger.info(f"✅ Added {len(docs)} documents to LanceDB")
+            return len(docs)
+
+        except Exception as e:
+            logger.error(f"❌ Error adding documents to LanceDB: {e}")
+            raise e
+
+    async def similarity_search(self, query: str, user_id: str, knowledge_base: str, k: int = 4):
+        """Search for similar documents"""
+        try:
+            query_embedding = self.embedding_model.embed_query(query)
+            
+            tbl = self.db.open_table("documents")
+            
+            # Search with filters
+            results = tbl.search(query_embedding)\
+                        .where(f"user_id = '{user_id}' AND knowledge_base = '{knowledge_base}'")\
+                        .limit(k)\
+                        .to_list()
+
+            docs = []
+            for result in results:
+                docs.append(type('Document', (), {
+                    'page_content': result['content'],
+                    'metadata': json.loads(result['metadata']) if result['metadata'] else {}
+                })())
+
+            return docs
+
+        except Exception as e:
+            logger.error(f"❌ Error searching LanceDB: {e}")
+            return []
+
+    async def get_user_knowledge_bases(self, user_id: str) -> List[str]:
+        """Get all knowledge bases for a user"""
+        try:
+            tbl = self.db.open_table("documents")
+            df = tbl.search().where(f"user_id = '{user_id}'").to_pandas()
+            
+            if df.empty:
+                return []
+            
+            knowledge_bases = df['knowledge_base'].unique().tolist()
+            return [kb for kb in knowledge_bases if kb and kb != "none"]
+
+        except Exception as e:
+            logger.error(f"❌ Error fetching knowledge bases: {e}")
+            return []
+
+    async def get_kb_documents(self, user_id: str, kb_name: str):
+        """Get all documents in a knowledge base"""
+        try:
+            tbl = self.db.open_table("documents")
+            df = tbl.search().where(f"user_id = '{user_id}' AND knowledge_base = '{kb_name}'").to_pandas()
+            
+            documents = []
+            for _, row in df.iterrows():
+                documents.append({
+                    "id": row['id'],
+                    "filename": row['filename'],
+                    "knowledge_base": row['knowledge_base'],
+                    "upload_date": row['upload_date']
+                })
+            
+            return documents
+
+        except Exception as e:
+            logger.error(f"❌ Error fetching documents: {e}")
+            return []
+
+    async def delete_document_from_kb(self, user_id: str, kb_name: str, filename: str):
+        """Delete a document from knowledge base"""
+        try:
+            tbl = self.db.open_table("documents")
+            tbl.delete(f"user_id = '{user_id}' AND knowledge_base = '{kb_name}' AND filename = '{filename}'")
+            return True
+
+        except Exception as e:
+            logger.error(f"❌ Error deleting document: {e}")
+            return False
+
+    # Persona management methods
+    async def insert_persona(self, name: str, description: str, icon: str, custom_prompt: str, user_id: str):
+        """Insert a new persona"""
+        try:
+            persona_data = {
+                "id": str(uuid.uuid4()),
+                "user_id": user_id,
+                "name": name,
+                "description": description,
+                "icon": icon,
+                "custom_prompt": custom_prompt,
+                "knowledge_base": "none",
+                "language": "en",
+                "created_at": datetime.now().isoformat(),
+                "updated_at": datetime.now().isoformat()
+            }
+
+            tbl = self.db.open_table("personas")
+            df = pd.DataFrame([persona_data])
+            tbl.add(df)
+            
+            return persona_data
+
+        except Exception as e:
+            logger.error(f"❌ Error inserting persona: {e}")
+            raise e
+
+    async def get_user_personas(self, user_id: str):
+        """Get all personas for a user"""
+        try:
+            tbl = self.db.open_table("personas")
+            df = tbl.search().where(f"user_id = '{user_id}'").to_pandas()
+            return df.to_dict('records')
+
+        except Exception as e:
+            logger.error(f"❌ Error fetching personas: {e}")
+            return []
+
+    # MCP Server methods
+    async def create_mcp_server(self, user_id: str, name: str, url: str, bearer_token: str = None):
+        """Create MCP server entry"""
+        try:
+            server_data = {
+                "id": str(uuid.uuid4()),
+                "user_id": user_id,
+                "name": name,
+                "url": url,
+                "bearer_token": bearer_token,
+                "created_at": datetime.now().isoformat()
+            }
+
+            tbl = self.db.open_table("mcp_servers")
+            df = pd.DataFrame([server_data])
+            tbl.add(df)
+            
+            return server_data
+
+        except Exception as e:
+            logger.error(f"❌ Error creating MCP server: {e}")
+            raise e
+
+    async def get_mcp_servers_for_user(self, user_id: str):
+        """Get MCP servers for user"""
+        try:
+            tbl = self.db.open_table("mcp_servers")
+            df = tbl.search().where(f"user_id = '{user_id}'").to_pandas()
+            return df.to_dict('records')
+
+        except Exception as e:
+            logger.error(f"❌ Error fetching MCP servers: {e}")
+            return []
+
+    async def delete_mcp_server(self, user_id: str, server_id: str):
+        """Delete MCP server"""
+        try:
+            tbl = self.db.open_table("mcp_servers")
+            tbl.delete(f"user_id = '{user_id}' AND id = '{server_id}'")
+            return True
+
+        except Exception as e:
+            logger.error(f"❌ Error deleting MCP server: {e}")
+            return False
+
+    # Session management
+    async def upsert_session_summary(self, user_id: str, persona_id: str, persona_source: str, summary: str):
+        """Create or update session summary"""
+        try:
+            session_data = {
+                "id": str(uuid.uuid4()),
+                "user_id": user_id,
+                "persona_id": persona_id,
+                "persona_source": persona_source,
+                "session_summary": summary,
+                "created_at": datetime.now().isoformat(),
+                "updated_at": datetime.now().isoformat()
+            }
+
+            tbl = self.db.open_table("sessions")
+            df = pd.DataFrame([session_data])
+            tbl.add(df)
+            
+            return session_data
+
+        except Exception as e:
+            logger.error(f"❌ Error upserting session: {e}")
+            return None
+
+    async def get_knowledge_bases(self) -> List[str]:
+        """Get all unique knowledge bases"""
+        try:
+            tbl = self.db.open_table("documents")
+            df = tbl.search().to_pandas()
+            
+            if df.empty:
+                return []
+            
+            knowledge_bases = df['knowledge_base'].unique().tolist()
+            return [kb for kb in knowledge_bases if kb and kb != "none"]
+
+        except Exception as e:
+            logger.error(f"❌ Error getting knowledge bases: {e}")
+            return []
+
+    async def get_documents_by_knowledge_base(self, knowledge_base: str) -> List[dict]:
+        """Get list of documents in a specific knowledge base"""
+        try:
+            tbl = self.db.open_table("documents")
+            df = tbl.search().where(f"knowledge_base = '{knowledge_base}'").to_pandas()
+            
+            if df.empty:
+                return []
+            
+            # Group by filename and get document info
+            documents = []
+            for filename in df['filename'].unique():
+                file_docs = df[df['filename'] == filename]
+                documents.append({
+                    "filename": filename,
+                    "knowledge_base": knowledge_base,
+                    "chunks": len(file_docs),
+                    "upload_date": file_docs['upload_date'].iloc[0] if not file_docs.empty else None
+                })
+            
+            return documents
+
+        except Exception as e:
+            logger.error(f"❌ Error getting documents by knowledge base: {e}")
+            return []
+
+    async def delete_document(self, filename: str, knowledge_base: str, user_id: str = None):
+        """Delete a document from the knowledge base"""
+        try:
+            tbl = self.db.open_table("documents")
+            
+            where_clause = f"filename = '{filename}' AND knowledge_base = '{knowledge_base}'"
+            if user_id:
+                where_clause += f" AND user_id = '{user_id}'"
+            
+            # Delete the document chunks
+            tbl.delete(where_clause)
+            
+            logger.info(f"✅ Deleted document {filename} from knowledge base {knowledge_base}")
+            return True
+
+        except Exception as e:
+            logger.error(f"❌ Error deleting document: {e}")
+            return False
+
+    async def search_all_knowledge_bases(self, query: str, k: int = 4):
+        """Search across all knowledge bases"""
+        try:
+            query_embedding = self.embedding_model.embed_query(query)
+            
+            tbl = self.db.open_table("documents")
+            
+            # Search without user filters for system-wide search
+            results = tbl.search(query_embedding).limit(k).to_list()
+
+            docs = []
+            for result in results:
+                docs.append(type('Document', (), {
+                    'page_content': result['content'],
+                    'metadata': json.loads(result['metadata']) if result['metadata'] else {}
+                })())
+
+            return docs
+
+        except Exception as e:
+            logger.error(f"❌ Error searching all knowledge bases: {e}")
+            return []
+
+# Global instance
+lancedb_service = LanceDBService()

llm_service.py

@@ -0,0 +1,155 @@
+from langgraph.graph import StateGraph, END, START
+from langchain_core.messages import SystemMessage
+from typing import TypedDict
+from config import GOOGLE_API_KEY, GEMINI_MODEL, GEMINI_TEMPERATURE
+from rag_service import search_docs, search_government_docs, analyze_scenario
+from langchain_tavily import TavilySearch
+from langgraph.prebuilt import ToolNode, tools_condition
+from langgraph.checkpoint.memory import MemorySaver
+from typing import Annotated
+from typing_extensions import TypedDict
+from langgraph.graph.message import add_messages
+from langchain_google_genai import ChatGoogleGenerativeAI
+from pydantic import BaseModel, Field, field_validator
+from typing import List, Optional, Literal
+from langchain_core.tools import tool
+import asyncio
+
+# Optional MCP client - install with: pip install langchain-mcp-adapters
+try:
+    from langchain_mcp_adapters.client import MultiServerMCPClient
+    MCP_AVAILABLE = True
+except ImportError:
+    MCP_AVAILABLE = False
+
+# Optional Tavily search - requires TAVILY_API_KEY environment variable
+try:
+    tavily_search = TavilySearch(max_results=4)
+    TAVILY_AVAILABLE = True
+except Exception:
+    TAVILY_AVAILABLE = False
+
+
+@tool
+def search_tool(query: str):
+    """
+    Perform an advanced web search using the Tavily Search API with hardcoded options.
+
+    Parameters:
+    ----------
+    query : str
+        The search query string.
+
+    Returns:
+    -------
+    str
+        The search results as a string returned by the Tavily Search API.
+
+    Raises:
+    ------
+    Exception
+        Errors during the search are caught and returned as error strings.
+    """
+    if not TAVILY_AVAILABLE:
+        return "Web search is not available. Tavily API key is not configured."
+    
+    query_params = {"query": query, "auto_parameters": True}
+
+    try:
+        result = tavily_search.invoke(query_params)
+        return result
+    except Exception as e:
+        return f"Error during Tavily search: {str(e)}"
+
+
+# State definition
+class State(TypedDict):
+    # add_messages is known as a reducer, where it does not modify the list but adds messages to it
+    messages: Annotated[list, add_messages]
+    # messages: Annotated[list[BaseMessage], add_messages]
+    # both have same result no need to use BaseMessage
+
+
+async def create_graph(kb_tool: bool, mcp_config: dict):
+    if mcp_config and MCP_AVAILABLE:
+        server_config = {
+            "url": mcp_config["url"],
+            "transport": "streamable_http",
+        }
+
+        # Add headers if bearer token exists
+        if mcp_config.get("bearerToken"):
+            server_config["headers"] = {
+                "Authorization": f"Bearer {mcp_config['bearerToken']}"
+            }
+
+        client = MultiServerMCPClient({mcp_config["name"]: server_config})
+        mcp_tools = await client.get_tools()
+    else:
+        mcp_tools = []
+    llm = ChatGoogleGenerativeAI(
+        model=GEMINI_MODEL,
+        google_api_key=GOOGLE_API_KEY,
+        temperature=GEMINI_TEMPERATURE,
+    )
+    if kb_tool:
+        tools = [search_docs, search_government_docs, analyze_scenario, search_tool]
+    else:
+        tools = [search_tool, analyze_scenario]
+    tools = tools + mcp_tools
+    llm_with_tools = llm.bind_tools(tools)
+
+    async def llm_node(state: State):
+        messages = state["messages"]
+        response = await llm_with_tools.ainvoke(messages)
+        return {"messages": [response]}
+
+    builder = StateGraph(State)
+    builder.add_node("llm_with_tools", llm_node)
+    tool_node = ToolNode(tools=tools, handle_tool_errors=True)
+    builder.add_node("tools", tool_node)
+    builder.add_conditional_edges("llm_with_tools", tools_condition)
+    builder.add_edge("tools", "llm_with_tools")
+    builder.add_edge(START, "llm_with_tools")
+    builder.add_edge("llm_with_tools", END)
+    return builder.compile()
+
+
+# Build basic graph (no tools, no memory)
+def create_basic_graph():
+    llm = ChatGoogleGenerativeAI(
+        model=GEMINI_MODEL,
+        google_api_key=GOOGLE_API_KEY,
+        temperature=GEMINI_TEMPERATURE,
+    )
+
+    async def llm_basic_node(state: State):
+        messages = state["messages"]
+        system_prompt = SystemMessage(
+            content="""You are a helpful and friendly voice AI assistant. Your responses should be:
+
+    - Conversational and natural, as if speaking to a friend
+    - Concise but informative - aim for 1-3 sentences unless more detail is specifically requested
+    - Clear and easy to understand when spoken aloud
+    - Engaging and personable while remaining professional
+    - Avoid overly complex language or long lists that are hard to follow in audio format
+
+    When responding:
+    - Use a warm, approachable tone
+    - Speak in a natural rhythm suitable for text-to-speech
+    - If you need to provide multiple items or steps, break them into digestible chunks
+    - Ask clarifying questions when needed to better assist the user
+    - Acknowledge when you don't know something rather than guessing
+
+    Remember that users are interacting with you through voice, so structure your responses to be easily understood when heard rather than read.
+    Dont use abbreviations or numerical content in your responses."""
+        )
+        if not any(isinstance(m, SystemMessage) for m in messages):
+            messages.insert(0, system_prompt)
+        return {"messages": [llm.invoke(messages)]}
+
+    builder = StateGraph(State)
+    builder.add_node("llm_basic", llm_basic_node)
+    builder.set_entry_point("llm_basic")
+    builder.add_edge("llm_basic", END)
+    return builder.compile()  # No checkpointing

main.py

@@ -0,0 +1,21 @@
+"""
+Voice Bot Application - Entry Point
+
+This file has been refactored. The main application logic is now in app.py
+Please run: python app.py or use uvicorn app:app
+
+The modular structure:
+- config.py: Configuration and constants
+- audio_services.py: ASR and TTS functionality 
+- rag_service.py: Vector store and document search
+- llm_service.py: LangGraph and LLM handling
+- document_service.py: PDF processing and document upload
+- websocket_handler.py: WebSocket connection handling
+- app.py: FastAPI application and routes
+"""
+
+from app import app
+
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=8000)

rag_service.py

@@ -0,0 +1,322 @@
+from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_core.tools import tool
+from config import EMBEDDING_MODEL_NAME
+from langchain_core.runnables import RunnableConfig
+from typing import List, Dict, Any
+from lancedb_service import lancedb_service
+from scenario_analysis_service import scenario_service
+import logging
+import json
+import asyncio
+
+logger = logging.getLogger("voicebot")
+
+# Setup embedding model
+embedding_model = HuggingFaceEmbeddings(
+    model_name=EMBEDDING_MODEL_NAME,
+    model_kwargs={
+        "device": "cpu",
+        "trust_remote_code": True
+    },
+    encode_kwargs={
+        "normalize_embeddings": True
+    }
+)
+
+async def get_user_knowledge_bases(userid: str) -> List[str]:
+    """Get all knowledge bases for a user"""
+    try:
+        return await lancedb_service.get_user_knowledge_bases(userid)
+    except Exception as e:
+        logger.error(f"❌ Error fetching knowledge bases: {e}")
+        return []
+
+async def get_kb_documents(user_id: str, kb_name: str):
+    """Get all documents in a knowledge base"""
+    try:
+        return await lancedb_service.get_kb_documents(user_id, kb_name)
+    except Exception as e:
+        logger.error(f"❌ Error fetching documents: {e}")
+        return []
+
+async def delete_document_from_kb(user_id: str, kb_name: str, filename: str):
+    """Delete a document from knowledge base"""
+    try:
+        return await lancedb_service.delete_document_from_kb(user_id, kb_name, filename)
+    except Exception as e:
+        logger.error(f"❌ Error deleting document: {e}")
+        return False
+
+def search_documents(query: str, limit: int = 5) -> List[Dict[str, Any]]:
+    """
+    Synchronous wrapper for searching documents in government knowledge base.
+    Returns a list of documents with content for compatibility with existing code.
+    """
+    try:
+        # Run the async search function synchronously
+        loop = asyncio.new_event_loop()
+        asyncio.set_event_loop(loop)
+        
+        try:
+            # Determine which knowledge bases to search based on query content
+            knowledge_bases = ["government_docs"]  # Default
+            
+            # Add specific knowledge bases based on query keywords
+            query_lower = query.lower()
+            if any(keyword in query_lower for keyword in ["rajasthan", "pension", "circular", "pay", "rules"]):
+                # Use separate table for Rajasthan documents
+                return search_rajasthan_documents(query, limit)
+                
+            all_docs = []
+            
+            # Search across all relevant knowledge bases
+            for kb in knowledge_bases:
+                try:
+                    docs = loop.run_until_complete(
+                        lancedb_service.similarity_search(query, "system", kb, k=limit)
+                    )
+                    all_docs.extend(docs)
+                except Exception as e:
+                    logger.warning(f"Search failed for knowledge base {kb}: {e}")
+                    continue
+            
+            if not all_docs:
+                return []
+            
+            # Sort by relevance score if available and limit results
+            all_docs = sorted(all_docs, key=lambda x: getattr(x, 'score', 1.0), reverse=True)[:limit]
+            
+            # Convert to expected format
+            results = []
+            for doc in all_docs:
+                results.append({
+                    "content": doc.page_content,
+                    "source": doc.metadata.get('source', 'Unknown'),
+                    "score": getattr(doc, 'score', 1.0)
+                })
+            
+            logger.info(f"📚 Found {len(results)} documents for query: {query}")
+            return results
+            
+        finally:
+            loop.close()
+            
+    except Exception as e:
+        logger.error(f"❌ Error in search_documents: {e}")
+        return []
+
+def search_rajasthan_documents(query: str, limit: int = 5) -> List[Dict[str, Any]]:
+    """
+    Search specifically in the Rajasthan documents table using direct LanceDB query.
+    """
+    try:
+        import lancedb
+        
+        # Connect to LanceDB
+        db = lancedb.connect('./lancedb_data')
+        
+        # Check if rajasthan_documents table exists
+        if 'rajasthan_documents' not in db.table_names():
+            logger.warning("⚠️  Rajasthan documents table not found")
+            return []
+        
+        # Get the table
+        tbl = db.open_table('rajasthan_documents')
+        
+        # Create embedding for the query
+        query_embedding = embedding_model.embed_query(query)
+        
+        # Search using vector similarity
+        search_results = tbl.search(query_embedding).limit(limit).to_pandas()
+        
+        if search_results.empty:
+            logger.info(f"📚 No results found in Rajasthan documents for: {query}")
+            return []
+        
+        # Convert to expected format
+        results = []
+        for _, row in search_results.iterrows():
+            results.append({
+                "content": row['content'],
+                "source": row['filename'],
+                "score": float(row.get('_distance', 1.0))  # LanceDB returns _distance
+            })
+        
+        logger.info(f"📚 Found {len(results)} Rajasthan documents for query: {query}")
+        return results
+        
+    except Exception as e:
+        logger.error(f"❌ Error searching Rajasthan documents: {e}")
+        return []
+
+@tool
+async def search_docs(query: str, config: RunnableConfig) -> str:
+    """Search the knowledge base for relevant context within a specific knowledge base."""
+    userid = config["configurable"].get("thread_id")
+    knowledge_base = config["configurable"].get("knowledge_base", "government_docs")
+    
+    try:
+        # Search in the specified knowledge base
+        docs = await lancedb_service.similarity_search(query, userid, knowledge_base)
+        
+        if not docs:
+            return "No relevant documents found in the knowledge base."
+        
+        context = "\n\n".join([doc.page_content for doc in docs])
+        return f"📄 Found {len(docs)} relevant documents:\n\n{context}"
+        
+    except Exception as e:
+        logger.error(f"❌ Error searching documents: {e}")
+        return "Error occurred while searching documents."
+
+@tool
+async def search_government_docs(query: str, config: RunnableConfig) -> str:
+    """Search government documents for relevant information and policies."""
+    try:
+        # Search specifically in government_docs knowledge base
+        docs = await lancedb_service.similarity_search(query, "system", "government_docs")
+        
+        if not docs:
+            return "No relevant government documents found for your query."
+        
+        context = "\n\n".join([doc.page_content for doc in docs])
+        sources = list(set([doc.metadata.get('source', 'Unknown') for doc in docs]))
+        
+        result = f"📋 Found {len(docs)} relevant government documents:\n\n{context}"
+        if sources:
+            result += f"\n\n📁 Sources: {', '.join(sources)}"
+        
+        return result
+        
+    except Exception as e:
+        logger.error(f"❌ Error searching government documents: {e}")
+        return "Error occurred while searching government documents."
+
+@tool
+async def analyze_scenario(scenario_query: str, config: RunnableConfig) -> str:
+    """
+    Analyze government scenarios and create visualizations including charts, graphs, and diagrams.
+    Use this tool when users ask for scenario analysis, data visualization, charts, graphs, or diagrams
+    related to government processes, budgets, policies, organizational structures, or performance metrics.
+    
+    Args:
+        scenario_query: Description of the scenario to analyze (e.g., "budget analysis for health department",
+                       "policy implementation timeline", "organizational structure", "performance metrics")
+    """
+    try:
+        logger.info(f"🔍 Analyzing scenario: {scenario_query}")
+        
+        # Parse the scenario query to determine type and extract data
+        scenario_data = await _parse_scenario_query(scenario_query)
+        
+        # Perform scenario analysis
+        result = await scenario_service.analyze_government_scenario(scenario_data)
+        
+        if result.get("success", False):
+            # Format response with images
+            response = f"📊 **Scenario Analysis Complete!**\n\n"
+            response += result.get("analysis", "")
+            response += f"\n\n🖼️ **Generated {len(result.get('images', []))} visualization(s)**"
+            
+            # Add image information for frontend rendering
+            if result.get("images"):
+                response += "\n\n**SCENARIO_IMAGES_START**\n"
+                response += json.dumps(result["images"])
+                response += "\n**SCENARIO_IMAGES_END**"
+            
+            return response
+        else:
+            return f"❌ Error in scenario analysis: {result.get('error', 'Unknown error')}"
+            
+    except Exception as e:
+        logger.error(f"❌ Error in scenario analysis tool: {e}")
+        return f"Error occurred while analyzing scenario: {str(e)}"
+
+async def _parse_scenario_query(query: str) -> Dict[str, Any]:
+    """Parse scenario query to determine type and extract relevant data"""
+    query_lower = query.lower()
+    
+    # Determine scenario type based on keywords
+    if any(word in query_lower for word in ["budget", "financial", "expenditure", "allocation", "funding"]):
+        scenario_type = "budget"
+        # Extract budget data if mentioned in query
+        data = _extract_budget_data(query)
+    elif any(word in query_lower for word in ["policy", "implementation", "timeline", "plan", "strategy"]):
+        scenario_type = "policy"
+        data = _extract_policy_data(query)
+    elif any(word in query_lower for word in ["organization", "hierarchy", "structure", "reporting", "org"]):
+        scenario_type = "organization"
+        data = _extract_org_data(query)
+    elif any(word in query_lower for word in ["performance", "metrics", "kpi", "efficiency", "evaluation"]):
+        scenario_type = "performance"
+        data = _extract_performance_data(query)
+    elif any(word in query_lower for word in ["workflow", "process", "flow", "procedure", "steps"]):
+        scenario_type = "workflow"
+        data = _extract_workflow_data(query)
+    else:
+        scenario_type = "general"
+        data = {}
+    
+    return {
+        "type": scenario_type,
+        "title": f"Government {scenario_type.title()} Analysis",
+        "data": data
+    }
+
+def _extract_budget_data(query: str) -> Dict[str, Any]:
+    """Extract budget-related data from query"""
+    # This could be enhanced with NLP to extract actual numbers and departments
+    # For now, return sample data structure
+    return {}
+
+def _extract_policy_data(query: str) -> Dict[str, Any]:
+    """Extract policy-related data from query"""
+    return {}
+
+def _extract_org_data(query: str) -> Dict[str, Any]:
+    """Extract organizational data from query"""
+    return {}
+
+def _extract_performance_data(query: str) -> Dict[str, Any]:
+    """Extract performance data from query"""
+    return {}
+
+def _extract_workflow_data(query: str) -> Dict[str, Any]:
+    """Extract workflow data from query"""
+    return {}
+
+if __name__ == "__main__":
+    import asyncio
+
+    async def test_search():
+        print("🔍 Testing search_docs RAG tool with LanceDB vector store...\n")
+
+        test_user_id = "test_user_123"
+        test_knowledge_base = "test_kb"
+
+        while True:
+            user_input = input("Enter a query (or 'exit'): ").strip()
+            if user_input.lower() == "exit":
+                break
+
+            kb_input = input(f"Knowledge base (current: {test_knowledge_base}, press Enter to keep): ").strip()
+            if kb_input:
+                test_knowledge_base = kb_input
+
+            try:
+                result = await search_docs.ainvoke(
+                    {"query": user_input},
+                    config=RunnableConfig(
+                        configurable={
+                            "thread_id": test_user_id,
+                            "knowledge_base": test_knowledge_base
+                        }
+                    )
+                )
+                print(f"\n📄 Results from '{test_knowledge_base}' knowledge base:\n")
+                print(result)
+                print("\n" + "="*50 + "\n")
+            except Exception as e:
+                print(f"❌ Error: {e}")
+
+    asyncio.run(test_search())

requirements.txt

@@ -0,0 +1,32 @@
+dotenv>=0.9.9
+fastapi>=0.115.14
+gradio>=4.44.0
+requests>=2.31.0
+langchain>=0.3.26
+langchain-community>=0.3.27
+langchain-huggingface>=0.3.0
+langchain-google-genai>=2.0.1
+langchain-groq>=0.3.0
+langchain-tavily>=0.2.7
+langgraph>=0.5.1
+langsmith>=0.4.4
+lancedb>=0.13.0
+google-generativeai>=0.8.1
+pdfplumber>=0.11.7
+pip>=25.1.1
+pyjwt>=2.10.1
+python-multipart>=0.0.20
+sentence-transformers>=5.0.0
+uvicorn[standard]>=0.35.0
+pandas>=2.0.0
+pyarrow>=14.0.0
+einops>=0.8.0
+matplotlib>=3.7.0
+seaborn>=0.12.0
+plotly>=5.15.0
+networkx>=3.1
+pillow>=10.0.0
+edge-tts>=6.1.0
+whisper>=1.1.10
+pydub>=0.25.1
+websockets>=11.0.0

voice_service.py

@@ -0,0 +1,324 @@
+"""
+Voice Service for optional Text-to-Speech (TTS) and Automatic Speech Recognition (ASR)
+Provides voice interaction capabilities when enabled by user.
+"""
+
+import asyncio
+import logging
+import tempfile
+import os
+from typing import Optional, Dict, Any
+from pathlib import Path
+
+from config import (
+    ENABLE_VOICE_FEATURES, TTS_PROVIDER, ASR_PROVIDER, 
+    VOICE_LANGUAGE, DEFAULT_VOICE_SPEED
+)
+
+logger = logging.getLogger("voicebot")
+
+class VoiceService:
+    def __init__(self):
+        self.voice_enabled = ENABLE_VOICE_FEATURES
+        self.tts_provider = TTS_PROVIDER
+        self.asr_provider = ASR_PROVIDER
+        self.language = VOICE_LANGUAGE
+        self.voice_speed = DEFAULT_VOICE_SPEED
+        
+        # Initialize services if voice is enabled
+        if self.voice_enabled:
+            self._init_tts_service()
+            self._init_asr_service()
+            logger.info(f"🎤 Voice Service initialized - TTS: {self.tts_provider}, ASR: {self.asr_provider}")
+        else:
+            logger.info("🔇 Voice features disabled")
+
+    def _init_tts_service(self):
+        """Initialize Text-to-Speech service"""
+        try:
+            if self.tts_provider == "edge-tts":
+                import edge_tts
+                self.tts_available = True
+                logger.info("✅ Edge TTS initialized")
+            elif self.tts_provider == "openai-tts":
+                # OpenAI TTS would require OpenAI API key
+                self.tts_available = False
+                logger.info("⚠️ OpenAI TTS not configured")
+            else:
+                self.tts_available = False
+                logger.warning(f"⚠️ Unknown TTS provider: {self.tts_provider}")
+        except ImportError as e:
+            self.tts_available = False
+            logger.warning(f"⚠️ TTS dependencies not available: {e}")
+
+    def _init_asr_service(self):
+        """Initialize Automatic Speech Recognition service"""
+        try:
+            if self.asr_provider == "whisper":
+                import whisper
+                # Use base model for balance between speed and accuracy
+                self.whisper_model = whisper.load_model("base")
+                self.asr_available = True
+                logger.info("✅ Whisper ASR initialized (base model for accuracy)")
+            elif self.asr_provider == "browser-native":
+                # Browser-based ASR doesn't require server-side setup
+                self.asr_available = True
+                logger.info("✅ Browser ASR configured")
+            else:
+                self.asr_available = False
+                logger.warning(f"⚠️ Unknown ASR provider: {self.asr_provider}")
+        except ImportError as e:
+            self.asr_available = False
+            logger.warning(f"⚠️ ASR dependencies not available: {e}")
+
+    def _get_language_code(self, user_language: str = None) -> str:
+        """
+        Convert user language preference to Whisper language code
+        
+        Args:
+            user_language: User's language preference ('english', 'hindi', 'hi-IN', etc.)
+            
+        Returns:
+            Two-letter language code for Whisper (e.g., 'en', 'hi')
+        """
+        if not user_language:
+            # Fallback to default config language
+            return self.language.split('-')[0] if self.language else 'en'
+        
+        # Handle different language format inputs
+        user_lang_lower = user_language.lower()
+        
+        # Map common language names to codes
+        language_mapping = {
+            'english': 'en',
+            'hindi': 'hi',
+            'hinglish': 'hi',  # Treat Hinglish as Hindi for better results
+            'en': 'en',
+            'hi': 'hi',
+            'en-in': 'en',
+            'hi-in': 'hi',
+            'en-us': 'en'
+        }
+        
+        # Extract base language if it's a locale code (e.g., 'hi-IN' -> 'hi')
+        if '-' in user_lang_lower:
+            base_lang = user_lang_lower.split('-')[0]
+            return language_mapping.get(base_lang, 'en')
+        
+        return language_mapping.get(user_lang_lower, 'en')
+
+    def _get_default_voice(self) -> str:
+        """Get default voice based on language setting"""
+        language_voices = {
+            'hi-IN': 'hi-IN-SwaraNeural',  # Hindi (India) female voice
+            'en-IN': 'en-IN-NeerjaNeural',  # English (India) female voice
+            'en-US': 'en-US-AriaNeural',   # English (US) female voice
+            'es-ES': 'es-ES-ElviraNeural', # Spanish (Spain) female voice
+            'fr-FR': 'fr-FR-DeniseNeural', # French (France) female voice
+            'de-DE': 'de-DE-KatjaNeural',  # German (Germany) female voice
+            'ja-JP': 'ja-JP-NanamiNeural', # Japanese female voice
+            'ko-KR': 'ko-KR-SunHiNeural',  # Korean female voice
+            'zh-CN': 'zh-CN-XiaoxiaoNeural' # Chinese (Simplified) female voice
+        }
+        return language_voices.get(self.language, 'en-US-AriaNeural')
+
+    async def text_to_speech(self, text: str, voice: str = None) -> Optional[bytes]:
+        """
+        Convert text to speech audio
+        Returns audio bytes or None if TTS not available
+        """
+        if not self.voice_enabled or not self.tts_available:
+            return None
+
+        # Use default voice for the configured language if no voice specified
+        if voice is None:
+            voice = self._get_default_voice()
+
+        try:
+            if self.tts_provider == "edge-tts":
+                import edge_tts
+                
+                # Create TTS communication
+                communicate = edge_tts.Communicate(text, voice, rate=f"{int((self.voice_speed - 1) * 100):+d}%")
+                
+                # Generate audio
+                audio_data = b""
+                async for chunk in communicate.stream():
+                    if chunk["type"] == "audio":
+                        audio_data += chunk["data"]
+                
+                return audio_data
+                
+        except Exception as e:
+            logger.error(f"❌ TTS Error: {e}")
+            return None
+
+    async def speech_to_text(self, audio_file_path: str, user_language: str = None) -> Optional[str]:
+        """
+        Convert speech audio to text
+        Returns transcribed text or None if ASR not available
+        
+        Args:
+            audio_file_path: Path to the audio file
+            user_language: User's preferred language (e.g., 'english', 'hindi', 'hi-IN')
+        """
+        if not self.voice_enabled or not self.asr_available:
+            return None
+
+        try:
+            if self.asr_provider == "whisper":
+                # Determine language code based on user preference or default
+                language_code = self._get_language_code(user_language)
+                
+                logger.info(f"🎤 Using Whisper with language: {language_code} (user_pref: {user_language})")
+                
+                # Use enhanced transcription options for better accuracy
+                transcribe_options = {
+                    "fp16": False,  # Use FP32 for better accuracy on CPU
+                    "temperature": 0.0,  # Deterministic output
+                    "best_of": 1,  # Use best transcription
+                    "beam_size": 5,  # Better beam search
+                    "patience": 1.0,  # Wait for better results
+                }
+                
+                if language_code and language_code != 'en':
+                    transcribe_options["language"] = language_code
+                    result = self.whisper_model.transcribe(audio_file_path, **transcribe_options)
+                    logger.info(f"🎤 {language_code.upper()} transcription result: {result.get('text', '')}")
+                else:
+                    result = self.whisper_model.transcribe(audio_file_path, **transcribe_options)
+                    logger.info(f"🎤 English transcription result: {result.get('text', '')}")
+                
+                transcribed_text = result["text"].strip()
+                
+                # Log confidence/quality metrics if available
+                if "segments" in result and result["segments"]:
+                    avg_confidence = sum(seg.get("no_speech_prob", 0) for seg in result["segments"]) / len(result["segments"])
+                    logger.info(f"🎤 Average confidence: {1-avg_confidence:.2f}")
+                
+                return transcribed_text
+                
+        except Exception as e:
+            logger.error(f"❌ ASR Error: {e}")
+            return None
+
+    def get_available_voices(self) -> Dict[str, Any]:
+        """Get list of available TTS voices"""
+        if not self.voice_enabled or self.tts_provider != "edge-tts":
+            return {}
+            
+        # Common Edge TTS voices
+        voices = {
+            "english": {
+                "female": ["en-US-AriaNeural", "en-US-JennyNeural", "en-GB-SoniaNeural"],
+                "male": ["en-US-GuyNeural", "en-US-DavisNeural", "en-GB-RyanNeural"]
+            },
+            "multilingual": {
+                "spanish": ["es-ES-ElviraNeural", "es-MX-DaliaNeural"],
+                "french": ["fr-FR-DeniseNeural", "fr-CA-SylvieNeural"],
+                "german": ["de-DE-KatjaNeural", "de-AT-IngridNeural"],
+                "italian": ["it-IT-ElsaNeural", "it-IT-IsabellaNeural"],
+                "hindi": ["hi-IN-SwaraNeural", "hi-IN-MadhurNeural"]
+            }
+        }
+        return voices
+
+    def create_voice_response_with_guidance(self, 
+                                          answer: str, 
+                                          suggested_resources: list = None,
+                                          redirect_info: str = None) -> str:
+        """
+        Create a comprehensive voice response with guidance and redirection
+        """
+        response_parts = []
+        
+        # Main answer
+        response_parts.append(answer)
+        
+        # Add guidance for further information
+        if suggested_resources:
+            response_parts.append("\nFor more detailed information, I recommend checking:")
+            for resource in suggested_resources:
+                response_parts.append(f"• {resource}")
+        
+        # Add redirection information
+        if redirect_info:
+            response_parts.append(f"\nYou can also {redirect_info}")
+        
+        # Add helpful voice interaction tips
+        response_parts.append("\nIs there anything specific you'd like me to explain further? Just ask!")
+        
+        return " ".join(response_parts)
+
+    def generate_redirect_suggestions(self, topic: str, query_type: str) -> Dict[str, Any]:
+        """
+        Generate contextual redirect suggestions based on the topic and query type
+        """
+        suggestions = {
+            "documents": [],
+            "websites": [],
+            "departments": [],
+            "redirect_text": ""
+        }
+        
+        # Government policy topics
+        if "digital india" in topic.lower():
+            suggestions["documents"] = [
+                "Digital India Policy Framework 2023",
+                "E-Governance Implementation Guidelines"
+            ]
+            suggestions["websites"] = ["digitalindia.gov.in", "meity.gov.in"]
+            suggestions["departments"] = ["Ministry of Electronics & IT"]
+            suggestions["redirect_text"] = "visit the official Digital India portal or contact your local e-governance center"
+            
+        elif "education" in topic.lower():
+            suggestions["documents"] = [
+                "National Education Policy 2020",
+                "Sarva Shiksha Abhiyan Guidelines"
+            ]
+            suggestions["websites"] = ["education.gov.in", "mhrd.gov.in"]
+            suggestions["departments"] = ["Ministry of Education"]
+            suggestions["redirect_text"] = "contact your District Education Officer or visit the nearest education department office"
+            
+        elif "health" in topic.lower():
+            suggestions["documents"] = [
+                "National Health Policy 2017",
+                "Ayushman Bharat Implementation Guide"
+            ]
+            suggestions["websites"] = ["mohfw.gov.in", "pmjay.gov.in"]
+            suggestions["departments"] = ["Ministry of Health & Family Welfare"]
+            suggestions["redirect_text"] = "visit your nearest Primary Health Center or call the health helpline"
+            
+        elif "employment" in topic.lower() or "job" in topic.lower():
+            suggestions["documents"] = [
+                "Employment Generation Schemes",
+                "Skill Development Programs Guide"
+            ]
+            suggestions["websites"] = ["nrega.nic.in", "msde.gov.in"]
+            suggestions["departments"] = ["Ministry of Rural Development", "Ministry of Skill Development"]
+            suggestions["redirect_text"] = "visit your local employment exchange or skill development center"
+            
+        # Default for other topics
+        if not suggestions["redirect_text"]:
+            suggestions["redirect_text"] = "contact the relevant government department or visit your local district collector's office"
+        
+        return suggestions
+
+    def is_voice_enabled(self) -> bool:
+        """Check if voice features are enabled"""
+        return self.voice_enabled
+
+    def get_voice_status(self) -> Dict[str, Any]:
+        """Get current voice service status"""
+        return {
+            "voice_enabled": self.voice_enabled,
+            "tts_available": getattr(self, 'tts_available', False),
+            "asr_available": getattr(self, 'asr_available', False),
+            "tts_provider": self.tts_provider,
+            "asr_provider": self.asr_provider,
+            "language": self.language,
+            "voice_speed": self.voice_speed
+        }
+
+# Global instance
+voice_service = VoiceService()

voice_websocket_server.py

@@ -0,0 +1,492 @@
+#!/usr/bin/env python3
+"""
+Voice-enabled WebSocket server that combines the full voice backend with our document search
+"""
+from fastapi import FastAPI, WebSocket, WebSocketDisconnect
+from fastapi.middleware.cors import CORSMiddleware
+import uvicorn
+import json
+import logging
+import lancedb
+import pandas as pd
+import asyncio
+import os
+from dotenv import load_dotenv
+from dataclasses import asdict, is_dataclass
+
+# Try to import voice services, fallback if not available
+try:
+    from hybrid_llm_service import HybridLLMService
+    from voice_service import VoiceService
+    from settings_api import router as settings_router
+    from policy_simulator_api import router as policy_simulator_router
+    VOICE_AVAILABLE = True
+except ImportError:
+    VOICE_AVAILABLE = False
+    logging.warning("Voice services not available, text-only mode")
+
+load_dotenv()
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Simple response cache for common queries
+response_cache = {}
+MAX_CACHE_SIZE = 100
+
+app = FastAPI()
+
+# Include API routers
+if VOICE_AVAILABLE:
+    app.include_router(settings_router)
+    app.include_router(policy_simulator_router)
+
+# Enable CORS - Include both local development and production origins
+allowed_origins = [
+    "http://localhost:5176", "http://localhost:5177", 
+    "http://127.0.0.1:5176", "http://127.0.0.1:5177",
+    "http://localhost:3000", "http://localhost:5173",
+    "https://*.vercel.app", "https://*.hf.space"
+]
+
+# Add any custom origins from environment
+if os.getenv("ALLOWED_ORIGINS"):
+    try:
+        custom_origins = eval(os.getenv("ALLOWED_ORIGINS"))
+        if isinstance(custom_origins, list):
+            allowed_origins.extend(custom_origins)
+    except:
+        pass
+
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"] if "*" in str(allowed_origins) else allowed_origins,
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+
+# Initialize services if available
+if VOICE_AVAILABLE:
+    try:
+        hybrid_llm_service = HybridLLMService()
+        voice_service = VoiceService()
+        logger.info("✅ Voice services initialized")
+    except Exception as e:
+        logger.warning(f"⚠️ Voice services failed to initialize: {e}")
+        VOICE_AVAILABLE = False
+
+def serialize_for_json(obj):
+    """Custom JSON serializer for policy simulation objects"""
+    if is_dataclass(obj):
+        return asdict(obj)
+    elif hasattr(obj, '__dict__'):
+        return obj.__dict__
+    elif isinstance(obj, (list, tuple)):
+        return [serialize_for_json(item) for item in obj]
+    elif isinstance(obj, dict):
+        return {key: serialize_for_json(value) for key, value in obj.items()}
+    else:
+        return obj
+
+def search_documents_simple(query: str):
+    """Simple document search without embeddings"""
+    try:
+        db = lancedb.connect('./lancedb_data')
+        
+        # Check for Rajasthan documents first
+        if 'rajasthan_documents' in db.table_names():
+            tbl = db.open_table('rajasthan_documents')
+            df = tbl.to_pandas()
+            
+            # Enhanced search for Rajasthan/pension queries
+            query_lower = query.lower()
+            is_pension_query = any(keyword in query_lower for keyword in [
+                'pension', 'पेंशन', 'वृद्धावस्था', 'सामाजिक', 'भत्ता', 'allowance', 
+                'old age', 'social security', 'retirement', 'सेवानिवृत्ति'
+            ])
+            
+            if is_pension_query or 'rajasthan' in query_lower:
+                # Enhanced pension search with more keywords
+                pension_filter = df['content'].str.contains(
+                    'pension|Pension|पेंशन|वृद्धावस्था|सामाजिक|भत्ता|allowance|old.age|social.security|retirement|सेवानिवृत्ति|scheme|योजना', 
+                    case=False, na=False, regex=True
+                )
+                
+                relevant_docs = df[pension_filter]
+                
+                if not relevant_docs.empty:
+                    # Sort by relevance
+                    def score_relevance(content):
+                        keywords = ['pension', 'पेंशन', 'वृद्धावस्था', 'सामाजिक', 'भत्ता', 'allowance', 'old age']
+                        return sum(1 for keyword in keywords if keyword in content.lower())
+                    
+                    relevant_docs = relevant_docs.copy()
+                    relevant_docs['relevance_score'] = relevant_docs['content'].apply(score_relevance)
+                    relevant_docs = relevant_docs.sort_values('relevance_score', ascending=False)
+                    
+                    results = []
+                    for _, row in relevant_docs.head(5).iterrows():
+                        results.append({
+                            "content": row['content'][:800],
+                            "filename": row['filename']
+                        })
+                    return results, "rajasthan_pension_documents"
+        
+        return [], "none"
+        
+    except Exception as e:
+        logger.error(f"Search error: {e}")
+        return [], "error"
+
+async def get_llm_response(query: str, search_results: list):
+    """Get response using available LLM service with caching"""
+    # Create cache key based on query and search results
+    cache_key = f"{query}_{len(search_results) if search_results else 0}"
+    
+    # Check cache first
+    if cache_key in response_cache:
+        logger.info(f"📦 Cache hit for query: {query[:50]}...")
+        return response_cache[cache_key]
+    
+    try:
+        if VOICE_AVAILABLE and hybrid_llm_service:
+            # Use the hybrid LLM service
+            if search_results:
+                context = "\\n\\n".join([f"Document: {doc['filename']}\\nContent: {doc['content']}" for doc in search_results])
+                enhanced_query = f"Based on these Rajasthan government documents, please answer: {query}\\n\\nDocuments:\\n{context}"
+            else:
+                enhanced_query = query
+                
+            response = await hybrid_llm_service.get_response(enhanced_query)
+            
+            # Cache the response
+            if len(response_cache) >= MAX_CACHE_SIZE:
+                # Remove oldest entry
+                response_cache.pop(next(iter(response_cache)))
+            response_cache[cache_key] = response
+            
+            return response
+        else:
+            # Fallback to simple response
+            if search_results:
+                response = f"Based on the Rajasthan government documents, I found information about {query}. However, voice processing is currently limited. Please use text chat for detailed responses."
+            else:
+                response = f"I received your query about '{query}' but couldn't find specific documents. Please try using text chat for better results."
+            
+            # Cache fallback response too
+            response_cache[cache_key] = response
+            return response
+            
+    except Exception as e:
+        logger.error(f"LLM error: {e}")
+        return "I'm having trouble processing your request. Please try using the text chat."
+
+@app.websocket("/ws/stream")
+async def websocket_endpoint(websocket: WebSocket):
+    await websocket.accept()
+    logger.info("🔌 WebSocket client connected")
+    
+    # Store user session info
+    user_language = "english"  # Default language
+    
+    try:
+        # Send initial greeting
+        await websocket.send_json({
+            "type": "connection_successful",
+            "message": "Hello! I'm your Rajasthan government document assistant. I can help with text and voice queries about pension schemes and government policies."
+        })
+        
+        while True:
+            try:
+                # Receive message with better error handling
+                message = await websocket.receive()
+                
+                # Handle different message types
+                if message["type"] == "websocket.receive":
+                    if "text" in message:
+                        # Parse JSON text message
+                        try:
+                            data = json.loads(message["text"])
+                        except json.JSONDecodeError:
+                            logger.warning(f"⚠️ Invalid JSON received: {message['text']}")
+                            continue
+                            
+                        # Process text message
+                        if isinstance(data, dict) and data.get("type") == "text_message":
+                            user_message = data.get("message", "")
+                            if not user_message.strip():
+                                continue
+                                
+                            logger.info(f"💬 Text received: {user_message}")
+                            
+                            # Check for interactive scenario form triggers
+                            form_triggers = ["start scenario analysis", "scenario form", "interactive analysis", "step by step analysis", "guided analysis", "form analysis", "scenario chat form", "interactive scenario"]
+                            is_form_request = any(trigger in user_message.lower() for trigger in form_triggers)
+                            
+                            # Check if this is a policy simulation query (robust regex patterns)
+                            import re
+                            POLICY_PATTERNS = [
+                                r"policy.*simulation|simulation.*policy",
+                                r"policy.*scenario|scenario.*policy", 
+                                r"policy.*analysis|analysis.*policy",
+                                r"pension.*simulation|simulation.*pension",
+                                r"pension.*analysis|analysis.*pension",
+                                r"pension.*scenario|scenario.*pension",
+                                r"dearness.*relief|dr.*increase|dr.*adjustment",
+                                r"dearness.*allowance|da.*increase|da.*adjustment",
+                                r"minimum.*pension.*increase|increase.*minimum.*pension",
+                                r"calculate.*pension|pension.*calculation",
+                                r"impact.*dr|dr.*impact|impact.*da|da.*impact",
+                                r"show.*impact.*da|show.*impact.*dr",
+                                r"impact.*\d+.*da|impact.*\d+.*dr",
+                                r"\d+.*da.*increase|da.*\d+.*increase",
+                                r"\d+.*dr.*increase|dr.*\d+.*increase",
+                                r"inflation.*adjustment|adjustment.*inflation",
+                                r"scenario.*analysis|analysis.*scenario",
+                                r"what.*if.*dr|what.*if.*pension|what.*if.*da",
+                                r"compare.*scenario|scenario.*comparison",
+                                r"show.*chart|chart.*show",
+                                r"explain.*chart|chart.*explain", 
+                                r"using.*chart|chart.*using",
+                                r"dr.*\d+.*increase|increase.*dr.*\d+",
+                                r"da.*\d+.*increase|increase.*da.*\d+",
+                                r"analyze.*minimum.*pension",
+                                r"pension.*change",
+                                r"make.*chart|chart.*make",
+                                r"pension.*value|value.*pension",
+                                r"basic.*pension.*\d+|pension.*\d+",
+                                r"simulate.*dr|simulate.*pension|simulate.*da"
+                            ]
+                            
+                            def is_policy_simulation_query(message: str) -> bool:
+                                """Check if the message is a policy simulation query"""
+                                message_lower = message.lower()
+                                logger.info(f"🔍 Checking policy patterns for: '{message_lower}'")
+                                
+                                for i, pattern in enumerate(POLICY_PATTERNS):
+                                    if re.search(pattern, message_lower, re.IGNORECASE):
+                                        logger.info(f"✅ Pattern {i+1} matched: {pattern}")
+                                        return True
+                                
+                                logger.info("❌ No policy patterns matched")
+                                return False
+                            
+                            is_policy_query = is_policy_simulation_query(user_message)
+                            
+                            # Handle interactive scenario form request
+                            if is_form_request:
+                                logger.info("📋 Interactive scenario form requested")
+                                
+                                try:
+                                    from scenario_chat_form import start_scenario_analysis_form
+                                    form_response = start_scenario_analysis_form(data.get("user_id", "default"))
+                                    
+                                    # Format form response for chat
+                                    form_message = f"""🎯 **{form_response.get('title', 'Interactive Scenario Analysis')}**
+
+{form_response.get('message', '')}
+
+**{form_response.get('step_title', 'Step 1')}** ({form_response.get('current_step', 1)}/{form_response.get('total_steps', 4)})
+
+{form_response['form_data']['question']}
+
+**Available Options:**"""
+                                    
+                                    # Add form options
+                                    if form_response['form_data']['input_type'] == 'select':
+                                        for i, option in enumerate(form_response['form_data']['options'], 1):
+                                            form_message += f"\n{i}. {option['label']}"
+                                    
+                                    form_message += "\n\n**Quick Actions:**"
+                                    for action in form_response.get('quick_actions', []):
+                                        form_message += f"\n• {action['text']}"
+                                        
+                                    form_message += "\n\n💡 **Next:** Choose an option above or type your selection!"
+                                    
+                                    await websocket.send_json({
+                                        "type": "interactive_form",
+                                        "message": form_message,
+                                        "form_data": form_response
+                                    })
+                                    continue
+                                    
+                                except Exception as e:
+                                    logger.error(f"Form initialization failed: {str(e)}")
+                                    await websocket.send_json({
+                                        "type": "error_message",
+                                        "message": f"Sorry, I couldn't start the interactive scenario analysis. Error: {str(e)}"
+                                    })
+                                    continue
+                            
+                            # Handle policy queries
+                            elif is_policy_query:
+                                logger.info("🎯 Detected policy simulation query")
+                                
+                                try:
+                                    # Import policy chat interface
+                                    from policy_chat_interface import PolicySimulatorChatInterface
+                                    
+                                    # Send acknowledgment for policy simulation
+                                    await websocket.send_json({
+                                        "type": "message_received", 
+                                        "message": "🎯 Analyzing Rajasthan policy impact..."
+                                    })
+                                    
+                                    # Initialize and process policy simulation
+                                    policy_simulator = PolicySimulatorChatInterface()
+                                    policy_result = policy_simulator.process_policy_query(user_message)
+                                    
+                                    # Format policy response - use same format as working simple backend
+                                    if policy_result.get("type") == "policy_simulation":
+                                        # Serialize the response for JSON
+                                        serialized_response = serialize_for_json(policy_result)
+                                        
+                                        # Send policy simulation response
+                                        await websocket.send_json({
+                                            "type": "policy_simulation",
+                                            "data": serialized_response
+                                        })
+                                        logger.info("📤 Policy simulation response sent to client")
+                                    else:
+                                        # Handle other policy responses (errors, help, etc.)
+                                        await websocket.send_json({
+                                            "type": "text_response",
+                                            "message": policy_result.get('message', 'Policy analysis completed')
+                                        })
+                                    
+                                    continue
+                                    
+                                except Exception as e:
+                                    logger.error(f"Policy simulation failed: {str(e)}")
+                                    await websocket.send_json({
+                                        "type": "error_message",
+                                        "message": f"Sorry, policy analysis failed. Using document search instead."
+                                    })
+                                    # Fall through to regular document search
+                            
+                            # Regular document search (fallback)
+                            # Send acknowledgment
+                            await websocket.send_json({
+                                "type": "message_received",
+                                "message": "🔍 Searching Rajasthan government documents..."
+                            })
+                            
+                            # Search for relevant documents
+                            search_results, source = search_documents_simple(user_message)
+                            logger.info(f"🔍 Found {len(search_results)} documents from {source}")
+                            
+                            # Get LLM response
+                            llm_response = await get_llm_response(user_message, search_results)
+                            
+                            # Send response
+                            await websocket.send_json({
+                                "type": "text_response",
+                                "message": llm_response
+                            })
+                            
+                        elif isinstance(data, dict) and data.get("type") == "user_info":
+                            user_name = data.get("user_name", "Unknown")
+                            logger.info(f"👤 User connected: {user_name}")
+                            
+                        elif isinstance(data, dict) and data.get("lang"):
+                            new_language = data.get("lang", "english")
+                            if new_language != user_language:
+                                user_language = new_language
+                                logger.info(f"🌍 Language preference updated: {user_language}")
+                            # Avoid logging if language hasn't changed
+                            
+                    elif "bytes" in message:
+                        # Handle binary message (audio data)
+                        audio_data = message["bytes"]
+                        logger.info(f"🎤 Received audio data: {len(audio_data)} bytes")
+                        
+                        if VOICE_AVAILABLE and voice_service:
+                            try:
+                                # Save audio data to temporary file for processing
+                                import tempfile
+                                with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_file:
+                                    temp_file.write(audio_data)
+                                    temp_file_path = temp_file.name
+                                
+                                # Process audio with voice service using user's language preference
+                                text = await voice_service.speech_to_text(temp_file_path, user_language)
+                                
+                                # Clean up temp file
+                                os.unlink(temp_file_path)
+                                
+                                if text and text.strip():
+                                    logger.info(f"🎤 Transcribed: {text}")
+                                    
+                                    # Search documents
+                                    search_results, source = search_documents_simple(text)
+                                    logger.info(f"🔍 Found {len(search_results)} documents from {source}")
+                                    
+                                    # Get LLM response
+                                    llm_response = await get_llm_response(text, search_results)
+                                    
+                                    # Send text response
+                                    await websocket.send_json({
+                                        "type": "text_response",
+                                        "message": llm_response
+                                    })
+                                    
+                                    # Try to send voice response
+                                    try:
+                                        audio_response = await voice_service.text_to_speech(llm_response)
+                                        if audio_response:
+                                            await websocket.send_bytes(audio_response)
+                                    except Exception as tts_error:
+                                        logger.warning(f"TTS failed: {tts_error}")
+                                else:
+                                    await websocket.send_json({
+                                        "type": "text_response",
+                                        "message": "I couldn't understand what you said. Please try speaking more clearly or use text chat."
+                                    })
+                                        
+                            except Exception as voice_error:
+                                logger.error(f"Voice processing error: {voice_error}")
+                                await websocket.send_json({
+                                    "type": "text_response",
+                                    "message": "Sorry, I couldn't process your voice input. Please try speaking again or use text chat."
+                                })
+                        else:
+                            # Voice services not available
+                            await websocket.send_json({
+                                "type": "text_response", 
+                                "message": "Voice processing is currently unavailable. Please use the text chat to ask about Rajasthan pension schemes and government policies."
+                            })
+                            
+                elif message["type"] == "websocket.disconnect":
+                    break
+                    
+            except json.JSONDecodeError as e:
+                logger.warning(f"⚠️ JSON decode error: {e}")
+                continue
+            except KeyError as e:
+                logger.warning(f"⚠️ Missing key in message: {e}")
+                continue
+                
+    except WebSocketDisconnect:
+        logger.info("🔌 WebSocket client disconnected")
+    except Exception as e:
+        logger.error(f"❌ WebSocket error: {e}")
+
+@app.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    try:
+        db = lancedb.connect('./lancedb_data')
+        tables = db.table_names()
+        return {
+            "status": "healthy", 
+            "tables": tables,
+            "voice_available": VOICE_AVAILABLE
+        }
+    except Exception as e:
+        return {"status": "error", "error": str(e)}
+
+if __name__ == "__main__":
+    print("🚀 Starting voice-enabled WebSocket server on port 8000...")
+    uvicorn.run(app, host="0.0.0.0", port=8000)

websocket_handler.py

@@ -0,0 +1,403 @@
+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)}"
+        }))