ready for hugging face space

api/config.py

@@ -0,0 +1,32 @@
+# api/config.py
+import os
+import logging
+from typing import Optional
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+)
+logger = logging.getLogger(__name__)
+
+
+class Settings:
+    """Application settings"""
+    APP_TITLE: str = "Legal Assistant API"
+    APP_VERSION: str = "2.0.0"
+    APP_DESCRIPTION: str = "Multi-country legal RAG with streaming & human-in-the-loop"
+    
+    # CORS
+    CORS_ORIGINS: list = ["*"]
+    
+    # API Settings
+    STREAM_DELAY: float = 0.02  # Delay between tokens in streaming
+    
+    # System
+    chat_manager: Optional[object] = None
+    graph: Optional[object] = None
+    system_initialized: bool = False
+
+
+settings = Settings()

api/main.py

@@ -1,704 +1,39 @@
-# api/main.py
-import sys
-from pathlib import Path
-sys.path.insert(0, str(Path(__file__).parent.parent))
-
-from typing import Optional, Any, Dict, AsyncGenerator
-from contextlib import asynccontextmanager
-from fastapi import FastAPI, Query, HTTPException, Body
-from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
+# api/main.py - MAIN FILE (Clean & Simple)
+from fastapi import FastAPI
 from fastapi.middleware.cors import CORSMiddleware
-from pydantic import BaseModel
-import json
-from uuid import uuid4
-import logging
-import os
-import asyncio
-from datetime import datetime
-
-# Import your existing system
-from core.system_initializer import setup_system
-
-# Setup logging
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-# Global variables
-chat_manager = None
-graph = None
-system_initialized = False
-
-
-# ============================================================================
-# Pydantic Models
-# ============================================================================
-class ChatRequest(BaseModel):
-    message: str
-    session_id: Optional[str] = None
-    stream: bool = False  # Option to enable/disable streaming
-
-class ApprovalRequest(BaseModel):
-    decision: str  # "approve" or "reject"
-    reason: Optional[str] = None
-
-class ChatResponse(BaseModel):
-    response: str
-    session_id: str
-    has_interrupt: bool = False
-    interrupt_type: Optional[str] = None
-    interrupt_data: Optional[Dict] = None
-
-
-# ============================================================================
-# System Initialization
-# ============================================================================
-async def initialize_system():
-    global chat_manager, graph, system_initialized
-    try:
-        system = await setup_system()
-        chat_manager = system["chat_manager"]
-        graph = system["graph"]
-        system_initialized = True
-        logger.info("✅ Legal assistant system initialized")
-    except Exception as e:
-        logger.error(f"❌ Failed to initialize system: {e}")
-        system_initialized = False
-
-@asynccontextmanager
-async def lifespan(app: FastAPI):
-    logger.info("🚀 Starting Legal Assistant API...")
-    initialization_task = asyncio.create_task(initialize_system())
-    yield
-    logger.info("🛑 Shutting down...")
-    initialization_task.cancel()
-    try:
-        await initialization_task
-    except asyncio.CancelledError:
-        pass
-
+from api.config import settings
+from api.utils import lifespan
+from api.routes import (
+    home_router,
+    chat_router,
+    sessions_router,
+    health_router
+)
 
-# ============================================================================
-# FastAPI App
-# ============================================================================
+# Create FastAPI app
 app = FastAPI(
-    title="Legal Assistant API",
-    version="2.0.0",
-    description="Multi-country legal RAG with streaming & human-in-the-loop",
+    title=settings.APP_TITLE,
+    version=settings.APP_VERSION,
+    description=settings.APP_DESCRIPTION,
     lifespan=lifespan
 )
 
+# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],
+    allow_origins=settings.CORS_ORIGINS,
     allow_credentials=True,
     allow_methods=["*"],
     allow_headers=["*"],
 )
 
+# Include routers
+app.include_router(home_router)
+app.include_router(chat_router)
+app.include_router(sessions_router)
+app.include_router(health_router)
 
-# ============================================================================
-# Streaming Helper
-# ============================================================================
-async def stream_chat_response(
-    message: str, 
-    session_id: str
-) -> AsyncGenerator[str, None]:
-    """
-    Stream chat responses token by token.
-    If interrupt occurs, yields special interrupt event.
-    """
-    try:
-        # Check for pending interrupt first
-        if chat_manager.has_pending_interrupt(session_id):
-            interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
-            interrupt_data_obj = interrupt_info.get("interrupt_data")
-            
-            if hasattr(interrupt_data_obj, 'value'):
-                interrupt_value = interrupt_data_obj.value
-            else:
-                interrupt_value = {}
-            
-            # Yield interrupt event
-            yield f"data: {json.dumps({'type': 'interrupt', 'data': interrupt_value})}\n\n"
-            return
-        
-        # Buffer to collect response
-        response_buffer = []
-        
-        # Process message
-        response = await chat_manager.chat(
-            message=message,
-            session_id=session_id,
-            interrupt_handler=None  # Async mode
-        )
-        
-        # Check if interrupt occurred during processing
-        if chat_manager.has_pending_interrupt(session_id):
-            interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
-            interrupt_data_obj = interrupt_info.get("interrupt_data")
-            
-            if hasattr(interrupt_data_obj, 'value'):
-                interrupt_value = interrupt_data_obj.value
-            else:
-                interrupt_value = {}
-            
-            # Yield interrupt event
-            yield f"data: {json.dumps({'type': 'interrupt', 'data': interrupt_value})}\n\n"
-            return
-        
-        # Stream response token by token (simulate streaming)
-        words = response.split()
-        for i, word in enumerate(words):
-            token = word + (" " if i < len(words) - 1 else "")
-            yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"
-            await asyncio.sleep(0.02)  # Small delay for visual effect
-        
-        # Send completion event
-        yield f"data: {json.dumps({'type': 'done', 'session_id': session_id})}\n\n"
-        
-    except Exception as e:
-        logger.exception(f"Error in streaming: {e}")
-        yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
-
-
-# ============================================================================
-# Routes
-# ============================================================================
-@app.get("/", response_class=HTMLResponse)
-async def read_root():
-    """Interactive chat interface with streaming support"""
-    return """
-    <html>
-        <head>
-            <title>Legal Assistant - Interactive Chat</title>
-            <meta charset="UTF-8">
-            <style>
-                * { margin: 0; padding: 0; box-sizing: border-box; }
-                body { 
-                    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
-                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-                    min-height: 100vh;
-                    padding: 20px;
-                }
-                .container { 
-                    max-width: 900px; 
-                    margin: 0 auto; 
-                    background: white;
-                    border-radius: 16px;
-                    box-shadow: 0 20px 60px rgba(0,0,0,0.3);
-                    overflow: hidden;
-                }
-                .header {
-                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
-                    color: white;
-                    padding: 30px;
-                    text-align: center;
-                }
-                .header h1 { margin-bottom: 10px; }
-                .header p { opacity: 0.9; }
-                .chat-container { padding: 20px; }
-                #messages {
-                    height: 500px;
-                    overflow-y: auto;
-                    padding: 20px;
-                    background: #f8f9fa;
-                    border-radius: 12px;
-                    margin-bottom: 20px;
-                }
-                .message {
-                    margin: 15px 0;
-                    padding: 12px 18px;
-                    border-radius: 18px;
-                    max-width: 75%;
-                    animation: slideIn 0.3s ease;
-                    word-wrap: break-word;
-                }
-                @keyframes slideIn {
-                    from { opacity: 0; transform: translateY(10px); }
-                    to { opacity: 1; transform: translateY(0); }
-                }
-                .user-message {
-                    background: #667eea;
-                    color: white;
-                    margin-left: auto;
-                    border-bottom-right-radius: 4px;
-                }
-                .assistant-message {
-                    background: white;
-                    color: #333;
-                    border: 1px solid #e0e0e0;
-                    border-bottom-left-radius: 4px;
-                }
-                .system-message {
-                    background: #fff3cd;
-                    color: #856404;
-                    border: 1px solid #ffc107;
-                    text-align: center;
-                    margin: 10px auto;
-                    max-width: 90%;
-                }
-                .interrupt-panel {
-                    background: #fff3cd;
-                    border: 2px solid #ffc107;
-                    border-radius: 12px;
-                    padding: 20px;
-                    margin: 20px 0;
-                    animation: pulse 2s infinite;
-                }
-                @keyframes pulse {
-                    0%, 100% { box-shadow: 0 0 0 0 rgba(255, 193, 7, 0.4); }
-                    50% { box-shadow: 0 0 0 10px rgba(255, 193, 7, 0); }
-                }
-                .interrupt-panel h3 {
-                    color: #856404;
-                    margin-bottom: 15px;
-                }
-                .interrupt-details {
-                    background: white;
-                    padding: 15px;
-                    border-radius: 8px;
-                    margin: 15px 0;
-                }
-                .interrupt-details p {
-                    margin: 8px 0;
-                    color: #333;
-                }
-                .approval-buttons {
-                    display: flex;
-                    gap: 10px;
-                    margin-top: 15px;
-                }
-                .btn {
-                    flex: 1;
-                    padding: 12px 24px;
-                    border: none;
-                    border-radius: 8px;
-                    font-weight: bold;
-                    cursor: pointer;
-                    transition: all 0.3s;
-                    font-size: 14px;
-                }
-                .btn-approve {
-                    background: #4caf50;
-                    color: white;
-                }
-                .btn-approve:hover {
-                    background: #45a049;
-                    transform: translateY(-2px);
-                    box-shadow: 0 4px 12px rgba(76, 175, 80, 0.4);
-                }
-                .btn-reject {
-                    background: #f44336;
-                    color: white;
-                }
-                .btn-reject:hover {
-                    background: #da190b;
-                    transform: translateY(-2px);
-                    box-shadow: 0 4px 12px rgba(244, 67, 54, 0.4);
-                }
-                .input-area {
-                    display: flex;
-                    gap: 10px;
-                    padding: 10px;
-                    background: #f8f9fa;
-                    border-radius: 12px;
-                }
-                #message-input {
-                    flex: 1;
-                    padding: 12px 16px;
-                    border: 2px solid #e0e0e0;
-                    border-radius: 8px;
-                    font-size: 14px;
-                    transition: border 0.3s;
-                }
-                #message-input:focus {
-                    outline: none;
-                    border-color: #667eea;
-                }
-                #send-btn {
-                    padding: 12px 32px;
-                    background: #667eea;
-                    color: white;
-                    border: none;
-                    border-radius: 8px;
-                    font-weight: bold;
-                    cursor: pointer;
-                    transition: all 0.3s;
-                }
-                #send-btn:hover:not(:disabled) {
-                    background: #5568d3;
-                    transform: translateY(-2px);
-                    box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4);
-                }
-                #send-btn:disabled {
-                    background: #ccc;
-                    cursor: not-allowed;
-                    transform: none;
-                }
-                .status-bar {
-                    padding: 10px;
-                    background: #f8f9fa;
-                    border-radius: 8px;
-                    margin-top: 10px;
-                    display: flex;
-                    justify-content: space-between;
-                    font-size: 12px;
-                    color: #666;
-                }
-                .typing-indicator {
-                    display: none;
-                    padding: 10px;
-                    color: #666;
-                    font-style: italic;
-                }
-                .typing-indicator.active {
-                    display: block;
-                }
-            </style>
-        </head>
-        <body>
-            <div class="container">
-                <div class="header">
-                    <h1>🧑‍⚖️ Legal Assistant</h1>
-                    <p>Multi-country legal support for Benin & Madagascar</p>
-                </div>
-                
-                <div class="chat-container">
-                    <div id="messages"></div>
-                    <div id="interrupt-zone"></div>
-                    <div class="typing-indicator" id="typing">Assistant is typing...</div>
-                    
-                    <div class="input-area">
-                        <input 
-                            type="text" 
-                            id="message-input" 
-                            placeholder="Type your legal question..."
-                            autocomplete="off"
-                        />
-                        <button id="send-btn" onclick="sendMessage()">
-                            Send
-                        </button>
-                    </div>
-                    
-                    <div class="status-bar">
-                        <span>Session: <strong id="session-id">Not started</strong></span>
-                        <span id="status">System ready</span>
-                    </div>
-                </div>
-            </div>
-
-            <script>
-                let sessionId = null;
-                let currentAssistantMessage = null;
-
-                function addMessage(content, type) {
-                    const messagesDiv = document.getElementById('messages');
-                    const messageDiv = document.createElement('div');
-                    messageDiv.className = `message ${type}-message`;
-                    messageDiv.textContent = content;
-                    messagesDiv.appendChild(messageDiv);
-                    messagesDiv.scrollTop = messagesDiv.scrollHeight;
-                    return messageDiv;
-                }
-
-                function showInterrupt(data) {
-                    const interruptZone = document.getElementById('interrupt-zone');
-                    interruptZone.innerHTML = `
-                        <div class="interrupt-panel">
-                            <h3>🔒 Human Approval Required</h3>
-                            <div class="interrupt-details">
-                                <p><strong>📧 Email:</strong> ${data.user_email || 'N/A'}</p>
-                                <p><strong>🌍 Country:</strong> ${data.country || 'N/A'}</p>
-                                <p><strong>📝 Description:</strong> ${data.description || 'N/A'}</p>
-                            </div>
-                            <div class="approval-buttons">
-                                <button class="btn btn-approve" onclick="handleApproval('approve')">
-                                    ✅ Approve Request
-                                </button>
-                                <button class="btn btn-reject" onclick="handleApproval('reject')">
-                                    ❌ Reject Request
-                                </button>
-                            </div>
-                        </div>
-                    `;
-                }
-
-                function hideInterrupt() {
-                    document.getElementById('interrupt-zone').innerHTML = '';
-                }
-
-                async function handleApproval(decision) {
-                    const reason = decision === 'approve' 
-                        ? 'Approved via web interface' 
-                        : 'Rejected via web interface';
-                    
-                    try {
-                        addMessage(`${decision === 'approve' ? '✅' : '❌'} Decision: ${decision}`, 'system');
-                        hideInterrupt();
-                        
-                        const response = await fetch(`/sessions/${sessionId}/approve`, {
-                            method: 'POST',
-                            headers: { 'Content-Type': 'application/json' },
-                            body: JSON.stringify({ decision, reason })
-                        });
-
-                        const data = await response.json();
-                        addMessage(data.response, 'assistant');
-                        
-                    } catch (error) {
-                        addMessage('Error: ' + error.message, 'system');
-                    }
-                }
-
-                async function sendMessage() {
-                    const input = document.getElementById('message-input');
-                    const message = input.value.trim();
-                    
-                    if (!message) return;
-                    
-                    // Initialize session
-                    if (!sessionId) {
-                        sessionId = 'web_' + Date.now();
-                        document.getElementById('session-id').textContent = sessionId;
-                    }
-
-                    // Add user message
-                    addMessage(message, 'user');
-                    input.value = '';
-                    
-                    // Disable input
-                    const sendBtn = document.getElementById('send-btn');
-                    sendBtn.disabled = true;
-                    input.disabled = true;
-                    
-                    // Show typing indicator
-                    document.getElementById('typing').classList.add('active');
-
-                    try {
-                        // Use streaming endpoint
-                        const response = await fetch('/chat/stream', {
-                            method: 'POST',
-                            headers: { 'Content-Type': 'application/json' },
-                            body: JSON.stringify({
-                                message: message,
-                                session_id: sessionId,
-                                stream: true
-                            })
-                        });
-
-                        const reader = response.body.getReader();
-                        const decoder = new TextDecoder();
-                        currentAssistantMessage = addMessage('', 'assistant');
-
-                        while (true) {
-                            const { value, done } = await reader.read();
-                            if (done) break;
-
-                            const chunk = decoder.decode(value);
-                            const lines = chunk.split('\\n\\n');
-
-                            for (const line of lines) {
-                                if (line.startsWith('data: ')) {
-                                    try {
-                                        const data = JSON.parse(line.slice(6));
-                                        
-                                        if (data.type === 'token') {
-                                            currentAssistantMessage.textContent += data.content;
-                                        } else if (data.type === 'interrupt') {
-                                            showInterrupt(data.data);
-                                        } else if (data.type === 'done') {
-                                            document.getElementById('status').textContent = 'Ready';
-                                        } else if (data.type === 'error') {
-                                            addMessage('Error: ' + data.message, 'system');
-                                        }
-                                    } catch (e) {
-                                        console.error('Parse error:', e);
-                                    }
-                                }
-                            }
-                        }
-                        
-                    } catch (error) {
-                        addMessage('Error: ' + error.message, 'system');
-                    } finally {
-                        sendBtn.disabled = false;
-                        input.disabled = false;
-                        input.focus();
-                        document.getElementById('typing').classList.remove('active');
-                    }
-                }
-
-                // Enter key to send
-                document.getElementById('message-input').addEventListener('keypress', (e) => {
-                    if (e.key === 'Enter' && !e.shiftKey) {
-                        e.preventDefault();
-                        sendMessage();
-                    }
-                });
-
-                // Focus input on load
-                document.getElementById('message-input').focus();
-            </script>
-        </body>
-    </html>
-    """
-
-@app.get("/health")
-async def health_check():
-    """Health check endpoint"""
-    return {
-        "status": "healthy" if system_initialized else "starting",
-        "system_initialized": system_initialized,
-        "timestamp": datetime.now().isoformat()
-    }
-
-@app.post("/chat")
-async def chat_simple(request: ChatRequest):
-    """
-    Simple non-streaming chat endpoint.
-    Best for handling interrupts - returns complete response.
-    """
-    if not system_initialized or not chat_manager:
-        raise HTTPException(status_code=503, detail="System initializing...")
-    
-    try:
-        session_id = request.session_id or f"api_{uuid4()}"
-        
-        # Check for pending interrupt
-        if chat_manager.has_pending_interrupt(session_id):
-            interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
-            interrupt_data_obj = interrupt_info.get("interrupt_data")
-            
-            if hasattr(interrupt_data_obj, 'value'):
-                interrupt_value = interrupt_data_obj.value
-            else:
-                interrupt_value = {}
-            
-            return ChatResponse(
-                response="⏸️ Pending approval request. Please approve/reject first.",
-                session_id=session_id,
-                has_interrupt=True,
-                interrupt_type="human_approval",
-                interrupt_data=interrupt_value
-            )
-        
-        # Process message
-        response_text = await chat_manager.chat(
-            message=request.message,
-            session_id=session_id
-        )
-        
-        # Check if interrupt occurred
-        has_interrupt = chat_manager.has_pending_interrupt(session_id)
-        interrupt_value = None
-        
-        if has_interrupt:
-            interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
-            interrupt_data_obj = interrupt_info.get("interrupt_data")
-            
-            if hasattr(interrupt_data_obj, 'value'):
-                interrupt_value = interrupt_data_obj.value
-            elif isinstance(interrupt_data_obj, dict):
-                interrupt_value = interrupt_data_obj.get("value", {})
-        
-        return ChatResponse(
-            response=response_text,
-            session_id=session_id,
-            has_interrupt=has_interrupt,
-            interrupt_type="human_approval" if has_interrupt else None,
-            interrupt_data=interrupt_value
-        )
-        
-    except Exception as e:
-        logger.exception(f"Error in chat: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.post("/chat/stream")
-async def chat_stream(request: ChatRequest):
-    """
-    Streaming chat endpoint using Server-Sent Events (SSE).
-    Streams tokens in real-time and handles interrupts.
-    """
-    if not system_initialized or not chat_manager:
-        raise HTTPException(status_code=503, detail="System initializing...")
-    
-    session_id = request.session_id or f"api_{uuid4()}"
-    
-    return StreamingResponse(
-        stream_chat_response(request.message, session_id),
-        media_type="text/event-stream",
-        headers={
-            "Cache-Control": "no-cache",
-            "Connection": "keep-alive",
-            "X-Accel-Buffering": "no"
-        }
-    )
-
-@app.post("/sessions/{session_id}/approve")
-async def approve_assistance(session_id: str, request: ApprovalRequest):
-    """Approve or reject assistance request"""
-    if not chat_manager:
-        raise HTTPException(status_code=503, detail="System not initialized")
-    
-    if not chat_manager.has_pending_interrupt(session_id):
-        raise HTTPException(status_code=404, detail="No pending interrupt")
-    
-    try:
-        decision_text = f"{request.decision} {request.reason or ''}"
-        response_text = await chat_manager.chat(
-            message=decision_text,
-            session_id=session_id
-        )
-        
-        return {
-            "status": "success",
-            "decision": request.decision,
-            "response": response_text,
-            "session_id": session_id
-        }
-    except Exception as e:
-        logger.exception(f"Error handling approval: {e}")
-        raise HTTPException(status_code=500, detail=str(e))
-
-@app.get("/sessions/{session_id}/status")
-async def get_session_status(session_id: str):
-    """Get session status including interrupt state"""
-    if not chat_manager:
-        raise HTTPException(status_code=503, detail="System not initialized")
-    
-    has_interrupt = chat_manager.has_pending_interrupt(session_id)
-    
-    interrupt_data = None
-    if has_interrupt:
-        interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
-        interrupt_data_obj = interrupt_info.get("interrupt_data")
-        
-        if hasattr(interrupt_data_obj, 'value'):
-            interrupt_data = interrupt_data_obj.value
-    
-    return {
-        "session_id": session_id,
-        "has_pending_interrupt": has_interrupt,
-        "interrupt_data": interrupt_data
-    }
 
-@app.get("/sessions/{session_id}/history")
-async def get_history(session_id: str):
-    """Get conversation history"""
-    if not chat_manager:
-        raise HTTPException(status_code=503, detail="System not initialized")
-    
-    try:
-        history = await chat_manager.get_conversation_history(session_id)
-        return {
-            "session_id": session_id,
-            "history": [
-                {
-                    "role": msg.type if hasattr(msg, 'type') else "unknown",
-                    "content": msg.content if hasattr(msg, 'content') else str(msg)
-                }
-                for msg in history
-            ]
-        }
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)

api/models/__init__.py

@@ -0,0 +1,4 @@
+# api/models/__init__.py
+from .schemas import ChatRequest, ApprovalRequest, ChatResponse
+
+__all__ = ["ChatRequest", "ApprovalRequest", "ChatResponse"]

api/models/schemas.py

@@ -0,0 +1,22 @@
+# api/models/schemas.py
+from pydantic import BaseModel
+from typing import Optional, Dict
+
+
+class ChatRequest(BaseModel):
+    message: str
+    session_id: Optional[str] = None
+    stream: bool = False
+
+
+class ApprovalRequest(BaseModel):
+    decision: str  # "approve" or "reject"
+    reason: Optional[str] = None
+
+
+class ChatResponse(BaseModel):
+    response: str
+    session_id: str
+    has_interrupt: bool = False
+    interrupt_type: Optional[str] = None
+    interrupt_data: Optional[Dict] = None

api/routes/__init__.py

@@ -0,0 +1,7 @@
+# api/routes/__init__.py
+from .home import router as home_router
+from .chat import router as chat_router
+from .sessions import router as sessions_router
+from .health import router as health_router
+
+__all__ = ["home_router", "chat_router", "sessions_router", "health_router"]

api/routes/chat.py

@@ -0,0 +1,48 @@
+from fastapi import APIRouter, HTTPException
+from fastapi.responses import StreamingResponse
+from api.models import ChatRequest, ChatResponse
+from api.services import ChatService, StreamService
+from api.config import settings
+from uuid import uuid4
+
+router = APIRouter(prefix="/chat", tags=["Chat"])
+
+
+@router.post("", response_model=ChatResponse)
+async def chat_simple(request: ChatRequest):
+    """
+    Simple non-streaming chat endpoint.
+    Best for handling interrupts - returns complete response.
+    """
+    if not settings.system_initialized or not settings.chat_manager:
+        raise HTTPException(status_code=503, detail="System initializing...")
+    
+    try:
+        return await ChatService.process_message(
+            message=request.message,
+            session_id=request.session_id
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.post("/stream")
+async def chat_stream(request: ChatRequest):
+    """
+    Streaming chat endpoint using Server-Sent Events (SSE).
+    Streams tokens in real-time and handles interrupts.
+    """
+    if not settings.system_initialized or not settings.chat_manager:
+        raise HTTPException(status_code=503, detail="System initializing...")
+    
+    session_id = request.session_id or f"api_{uuid4()}"
+    
+    return StreamingResponse(
+        StreamService.stream_chat_response(request.message, session_id),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+            "X-Accel-Buffering": "no"
+        }
+    )

api/routes/health.py

@@ -0,0 +1,16 @@
+# api/routes/health.py
+from fastapi import APIRouter
+from datetime import datetime
+from api.config import settings
+
+router = APIRouter(tags=["Health"])
+
+
+@router.get("/health")
+async def health_check():
+    """Health check endpoint"""
+    return {
+        "status": "healthy" if settings.system_initialized else "starting",
+        "system_initialized": settings.system_initialized,
+        "timestamp": datetime.now().isoformat()
+    }

api/routes/home.py

@@ -0,0 +1,521 @@
+from fastapi import APIRouter
+from fastapi.responses import HTMLResponse
+
+router = APIRouter(tags=["Home"])
+
+
+@router.get("/", response_class=HTMLResponse)
+async def root():
+    """Interactive chat interface - Main UI"""
+    return """
+    <html>
+        <head>
+            <title>Legal Assistant - Interactive Chat</title>
+            <meta charset="UTF-8">
+            <style>
+                * { margin: 0; padding: 0; box-sizing: border-box; }
+                body { 
+                    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                    min-height: 100vh;
+                    padding: 20px;
+                }
+                .container { 
+                    max-width: 900px; 
+                    margin: 0 auto; 
+                    background: white;
+                    border-radius: 16px;
+                    box-shadow: 0 20px 60px rgba(0,0,0,0.3);
+                    overflow: hidden;
+                }
+                .header {
+                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                    color: white;
+                    padding: 30px;
+                    text-align: center;
+                    position: relative;
+                }
+                .header h1 { margin-bottom: 10px; }
+                .header p { opacity: 0.9; }
+                .api-link {
+                    position: absolute;
+                    top: 20px;
+                    right: 20px;
+                    padding: 8px 16px;
+                    background: rgba(255,255,255,0.2);
+                    color: white;
+                    text-decoration: none;
+                    border-radius: 5px;
+                    font-size: 0.9em;
+                    transition: all 0.3s;
+                }
+                .api-link:hover {
+                    background: rgba(255,255,255,0.3);
+                    transform: translateY(-2px);
+                }
+                .chat-container { padding: 20px; }
+                #messages {
+                    height: 500px;
+                    overflow-y: auto;
+                    padding: 20px;
+                    background: #f8f9fa;
+                    border-radius: 12px;
+                    margin-bottom: 20px;
+                }
+                .message {
+                    margin: 15px 0;
+                    padding: 12px 18px;
+                    border-radius: 18px;
+                    max-width: 75%;
+                    animation: slideIn 0.3s ease;
+                    word-wrap: break-word;
+                }
+                @keyframes slideIn {
+                    from { opacity: 0; transform: translateY(10px); }
+                    to { opacity: 1; transform: translateY(0); }
+                }
+                .user-message {
+                    background: #667eea;
+                    color: white;
+                    margin-left: auto;
+                    border-bottom-right-radius: 4px;
+                }
+                .assistant-message {
+                    background: white;
+                    color: #333;
+                    border: 1px solid #e0e0e0;
+                    border-bottom-left-radius: 4px;
+                }
+                .system-message {
+                    background: #fff3cd;
+                    color: #856404;
+                    border: 1px solid #ffc107;
+                    text-align: center;
+                    margin: 10px auto;
+                    max-width: 90%;
+                }
+                .interrupt-panel {
+                    background: #fff3cd;
+                    border: 2px solid #ffc107;
+                    border-radius: 12px;
+                    padding: 20px;
+                    margin: 20px 0;
+                    animation: pulse 2s infinite;
+                }
+                @keyframes pulse {
+                    0%, 100% { box-shadow: 0 0 0 0 rgba(255, 193, 7, 0.4); }
+                    50% { box-shadow: 0 0 0 10px rgba(255, 193, 7, 0); }
+                }
+                .interrupt-panel h3 {
+                    color: #856404;
+                    margin-bottom: 15px;
+                }
+                .interrupt-details {
+                    background: white;
+                    padding: 15px;
+                    border-radius: 8px;
+                    margin: 15px 0;
+                }
+                .interrupt-details p {
+                    margin: 8px 0;
+                    color: #333;
+                }
+                .approval-buttons {
+                    display: flex;
+                    gap: 10px;
+                    margin-top: 15px;
+                }
+                .btn {
+                    flex: 1;
+                    padding: 12px 24px;
+                    border: none;
+                    border-radius: 8px;
+                    font-weight: bold;
+                    cursor: pointer;
+                    transition: all 0.3s;
+                    font-size: 14px;
+                }
+                .btn-approve {
+                    background: #4caf50;
+                    color: white;
+                }
+                .btn-approve:hover {
+                    background: #45a049;
+                    transform: translateY(-2px);
+                    box-shadow: 0 4px 12px rgba(76, 175, 80, 0.4);
+                }
+                .btn-reject {
+                    background: #f44336;
+                    color: white;
+                }
+                .btn-reject:hover {
+                    background: #da190b;
+                    transform: translateY(-2px);
+                    box-shadow: 0 4px 12px rgba(244, 67, 54, 0.4);
+                }
+                .input-area {
+                    display: flex;
+                    gap: 10px;
+                    padding: 10px;
+                    background: #f8f9fa;
+                    border-radius: 12px;
+                }
+                #message-input {
+                    flex: 1;
+                    padding: 12px 16px;
+                    border: 2px solid #e0e0e0;
+                    border-radius: 8px;
+                    font-size: 14px;
+                    transition: border 0.3s;
+                }
+                #message-input:focus {
+                    outline: none;
+                    border-color: #667eea;
+                }
+                #send-btn {
+                    padding: 12px 32px;
+                    background: #667eea;
+                    color: white;
+                    border: none;
+                    border-radius: 8px;
+                    font-weight: bold;
+                    cursor: pointer;
+                    transition: all 0.3s;
+                }
+                #send-btn:hover:not(:disabled) {
+                    background: #5568d3;
+                    transform: translateY(-2px);
+                    box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4);
+                }
+                #send-btn:disabled {
+                    background: #ccc;
+                    cursor: not-allowed;
+                    transform: none;
+                }
+                .status-bar {
+                    padding: 10px;
+                    background: #f8f9fa;
+                    border-radius: 8px;
+                    margin-top: 10px;
+                    display: flex;
+                    justify-content: space-between;
+                    font-size: 12px;
+                    color: #666;
+                }
+                .typing-indicator {
+                    display: none;
+                    padding: 10px;
+                    color: #666;
+                    font-style: italic;
+                }
+                .typing-indicator.active {
+                    display: block;
+                }
+            </style>
+        </head>
+        <body>
+            <div class="container">
+                <div class="header">
+                    <a href="/docs" class="api-link">📚 API Docs</a>
+                    <h1>🧑‍⚖️ Legal Assistant</h1>
+                    <p>Multi-country legal support for Benin & Madagascar</p>
+                </div>
+                
+                <div class="chat-container">
+                    <div id="messages"></div>
+                    <div id="interrupt-zone"></div>
+                    <div class="typing-indicator" id="typing">Assistant is typing...</div>
+                    
+                    <div class="input-area">
+                        <input 
+                            type="text" 
+                            id="message-input" 
+                            placeholder="Type your legal question..."
+                            autocomplete="off"
+                        />
+                        <button id="send-btn" onclick="sendMessage()">
+                            Send
+                        </button>
+                    </div>
+                    
+                    <div class="status-bar">
+                        <span>Session: <strong id="session-id">Not started</strong></span>
+                        <span id="status">System ready</span>
+                    </div>
+                </div>
+            </div>
+
+            <script>
+                let sessionId = null;
+                let currentAssistantMessage = null;
+
+                function addMessage(content, type) {
+                    const messagesDiv = document.getElementById('messages');
+                    const messageDiv = document.createElement('div');
+                    messageDiv.className = `message ${type}-message`;
+                    messageDiv.textContent = content;
+                    messagesDiv.appendChild(messageDiv);
+                    messagesDiv.scrollTop = messagesDiv.scrollHeight;
+                    return messageDiv;
+                }
+
+                function showInterrupt(data) {
+                    const interruptZone = document.getElementById('interrupt-zone');
+                    interruptZone.innerHTML = `
+                        <div class="interrupt-panel">
+                            <h3>🔒 Human Approval Required</h3>
+                            <div class="interrupt-details">
+                                <p><strong>📧 Email:</strong> ${data.user_email || 'N/A'}</p>
+                                <p><strong>🌍 Country:</strong> ${data.country || 'N/A'}</p>
+                                <p><strong>📝 Description:</strong> ${data.description || 'N/A'}</p>
+                            </div>
+                            <div class="approval-buttons">
+                                <button class="btn btn-approve" onclick="handleApproval('approve')">
+                                    ✅ Approve Request
+                                </button>
+                                <button class="btn btn-reject" onclick="handleApproval('reject')">
+                                    ❌ Reject Request
+                                </button>
+                            </div>
+                        </div>
+                    `;
+                }
+
+                function hideInterrupt() {
+                    document.getElementById('interrupt-zone').innerHTML = '';
+                }
+
+                async function handleApproval(decision) {
+                    const reason = decision === 'approve' 
+                        ? 'Approved via web interface' 
+                        : 'Rejected via web interface';
+                    
+                    try {
+                        addMessage(`${decision === 'approve' ? '✅' : '❌'} Decision: ${decision}`, 'system');
+                        hideInterrupt();
+                        
+                        const response = await fetch(`/sessions/${sessionId}/approve`, {
+                            method: 'POST',
+                            headers: { 'Content-Type': 'application/json' },
+                            body: JSON.stringify({ decision, reason })
+                        });
+
+                        const data = await response.json();
+                        addMessage(data.response, 'assistant');
+                        
+                    } catch (error) {
+                        addMessage('Error: ' + error.message, 'system');
+                    }
+                }
+
+                async function sendMessage() {
+                    const input = document.getElementById('message-input');
+                    const message = input.value.trim();
+                    
+                    if (!message) return;
+                    
+                    if (!sessionId) {
+                        sessionId = 'web_' + Date.now();
+                        document.getElementById('session-id').textContent = sessionId;
+                    }
+
+                    addMessage(message, 'user');
+                    input.value = '';
+                    
+                    const sendBtn = document.getElementById('send-btn');
+                    sendBtn.disabled = true;
+                    input.disabled = true;
+                    
+                    document.getElementById('typing').classList.add('active');
+
+                    try {
+                        const response = await fetch('/chat/stream', {
+                            method: 'POST',
+                            headers: { 'Content-Type': 'application/json' },
+                            body: JSON.stringify({
+                                message: message,
+                                session_id: sessionId,
+                                stream: true
+                            })
+                        });
+
+                        const reader = response.body.getReader();
+                        const decoder = new TextDecoder();
+                        currentAssistantMessage = addMessage('', 'assistant');
+
+                        while (true) {
+                            const { value, done } = await reader.read();
+                            if (done) break;
+
+                            const chunk = decoder.decode(value);
+                            const lines = chunk.split('\\n\\n');
+
+                            for (const line of lines) {
+                                if (line.startsWith('data: ')) {
+                                    try {
+                                        const data = JSON.parse(line.slice(6));
+                                        
+                                        if (data.type === 'token') {
+                                            currentAssistantMessage.textContent += data.content;
+                                        } else if (data.type === 'interrupt') {
+                                            showInterrupt(data.data);
+                                        } else if (data.type === 'done') {
+                                            document.getElementById('status').textContent = 'Ready';
+                                        } else if (data.type === 'error') {
+                                            addMessage('Error: ' + data.message, 'system');
+                                        }
+                                    } catch (e) {
+                                        console.error('Parse error:', e);
+                                    }
+                                }
+                            }
+                        }
+                        
+                    } catch (error) {
+                        addMessage('Error: ' + error.message, 'system');
+                    } finally {
+                        sendBtn.disabled = false;
+                        input.disabled = false;
+                        input.focus();
+                        document.getElementById('typing').classList.remove('active');
+                    }
+                }
+
+                document.getElementById('message-input').addEventListener('keypress', (e) => {
+                    if (e.key === 'Enter' && !e.shiftKey) {
+                        e.preventDefault();
+                        sendMessage();
+                    }
+                });
+
+                document.getElementById('message-input').focus();
+            </script>
+        </body>
+    </html>
+    """
+
+
+@router.get("/about", response_class=HTMLResponse)
+async def about():
+    """About page with API information"""
+    return """
+    <!DOCTYPE html>
+    <html lang="en">
+    <head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
+        <title>Legal Assistant API</title>
+        <style>
+            * { margin: 0; padding: 0; box-sizing: border-box; }
+            body {
+                font-family: 'Segoe UI', sans-serif;
+                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                min-height: 100vh;
+                display: flex;
+                align-items: center;
+                justify-content: center;
+                padding: 20px;
+            }
+            .container {
+                max-width: 900px;
+                width: 100%;
+                background: white;
+                border-radius: 20px;
+                box-shadow: 0 20px 60px rgba(0,0,0,0.3);
+                overflow: hidden;
+            }
+            .header {
+                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                color: white;
+                padding: 60px 40px;
+                text-align: center;
+            }
+            .header h1 { font-size: 2.5em; margin-bottom: 15px; }
+            .header p { font-size: 1.2em; opacity: 0.95; }
+            .version {
+                background: rgba(255,255,255,0.2);
+                padding: 5px 15px;
+                border-radius: 20px;
+                margin-top: 15px;
+                display: inline-block;
+            }
+            .content { padding: 40px; }
+            .features {
+                display: grid;
+                grid-template-columns: repeat(auto-fit, minmax(200px, 1fr));
+                gap: 20px;
+                margin-bottom: 40px;
+            }
+            .feature-card {
+                padding: 25px;
+                background: #f8f9fa;
+                border-radius: 12px;
+                border-left: 4px solid #667eea;
+            }
+            .feature-card h3 { color: #667eea; margin-bottom: 10px; }
+            .cta-section {
+                background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                padding: 40px;
+                border-radius: 12px;
+                text-align: center;
+            }
+            .cta-section h2 { color: white; margin-bottom: 25px; }
+            .buttons { display: flex; gap: 15px; justify-content: center; flex-wrap: wrap; }
+            .btn {
+                padding: 15px 35px;
+                border-radius: 8px;
+                text-decoration: none;
+                font-weight: 600;
+                transition: all 0.3s;
+            }
+            .btn-primary { background: white; color: #667eea; }
+            .btn-primary:hover { transform: translateY(-3px); }
+            .footer {
+                padding: 30px;
+                background: #f8f9fa;
+                text-align: center;
+                color: #666;
+            }
+        </style>
+    </head>
+    <body>
+        <div class="container">
+            <div class="header">
+                <h1>⚖️ Legal Assistant API</h1>
+                <p>Multi-country legal support for Benin & Madagascar</p>
+                <span class="version">v2.0.0</span>
+            </div>
+            <div class="content">
+                <div class="features">
+                    <div class="feature-card">
+                        <h3>🤖 AI-Powered</h3>
+                        <p>Advanced RAG technology</p>
+                    </div>
+                    <div class="feature-card">
+                        <h3>🌍 Multi-Country</h3>
+                        <p>Benin & Madagascar laws</p>
+                    </div>
+                    <div class="feature-card">
+                        <h3>🔒 Human Approval</h3>
+                        <p>Quality assurance</p>
+                    </div>
+                    <div class="feature-card">
+                        <h3>⚡ Real-time</h3>
+                        <p>Streaming responses</p>
+                    </div>
+                </div>
+                <div class="cta-section">
+                    <h2>Get Started</h2>
+                    <div class="buttons">
+                        <a href="/docs" class="btn btn-primary">📚 Swagger UI</a>
+                        <a href="/redoc" class="btn btn-primary">📖 ReDoc</a>
+                    </div>
+                </div>
+            </div>
+            <div class="footer">
+                <p>Built with FastAPI & LangGraph</p>
+            </div>
+        </div>
+    </body>
+    </html>
+    """

api/routes/sessions.py

@@ -0,0 +1,56 @@
+# api/routes/sessions.py
+from fastapi import APIRouter, HTTPException
+from api.models import ApprovalRequest
+from api.services import ChatService
+
+router = APIRouter(prefix="/sessions", tags=["Sessions"])
+
+
+@router.post("/{session_id}/approve")
+async def approve_assistance(session_id: str, request: ApprovalRequest):
+    """Approve or reject assistance request"""
+    try:
+        return await ChatService.approve_request(
+            session_id=session_id,
+            decision=request.decision,
+            reason=request.reason
+        )
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.get("/{session_id}/status")
+async def get_session_status(session_id: str):
+    """Get session status including interrupt state"""
+    try:
+        return await ChatService.get_session_status(session_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.get("/{session_id}/history")
+async def get_history(session_id: str):
+    """Get conversation history"""
+    try:
+        return await ChatService.get_history(session_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.get("")
+async def list_sessions():
+    """List all active sessions with their status"""
+    try:
+        return await ChatService.list_sessions()
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+
+
+@router.delete("/{session_id}")
+async def clear_session(session_id: str):
+    """Clear a session (useful for testing)"""
+    try:
+        return await ChatService.clear_session(session_id)
+    except Exception as e:
+        raise HTTPException(status_code=500, detail=str(e))
+

api/services/__init__.py

@@ -0,0 +1,6 @@
+# api/services/__init__.py
+# ============================================================================
+from .chat_service import ChatService
+from .stream_service import StreamService
+
+__all__ = ["ChatService", "StreamService"]

api/services/chat_service.py

@@ -0,0 +1,178 @@
+# api/services/chat_service.py
+from typing import Dict, Optional
+from api.config import settings, logger
+from api.models import ChatResponse
+from uuid import uuid4
+
+
+class ChatService:
+    """Service for handling chat operations"""
+    
+    @staticmethod
+    async def process_message(
+        message: str,
+        session_id: Optional[str] = None
+    ) -> ChatResponse:
+        """Process a chat message"""
+        if not settings.system_initialized or not settings.chat_manager:
+            raise Exception("System not initialized")
+        
+        session_id = session_id or f"api_{uuid4()}"
+        
+        # Check for pending interrupt
+        if settings.chat_manager.has_pending_interrupt(session_id):
+            interrupt_info = settings.chat_manager.pending_interrupts.get(session_id, {})
+            interrupt_data_obj = interrupt_info.get("interrupt_data")
+            
+            interrupt_value = (
+                interrupt_data_obj.value 
+                if hasattr(interrupt_data_obj, 'value') 
+                else {}
+            )
+            
+            return ChatResponse(
+                response="⏸️ Pending approval request. Please approve/reject first.",
+                session_id=session_id,
+                has_interrupt=True,
+                interrupt_type="human_approval",
+                interrupt_data=interrupt_value
+            )
+        
+        # Process message
+        response_text = await settings.chat_manager.chat(
+            message=message,
+            session_id=session_id
+        )
+        
+        # Check if interrupt occurred
+        has_interrupt = settings.chat_manager.has_pending_interrupt(session_id)
+        interrupt_value = None
+        
+        if has_interrupt:
+            interrupt_info = settings.chat_manager.pending_interrupts.get(session_id, {})
+            interrupt_data_obj = interrupt_info.get("interrupt_data")
+            
+            if hasattr(interrupt_data_obj, 'value'):
+                interrupt_value = interrupt_data_obj.value
+            elif isinstance(interrupt_data_obj, dict):
+                interrupt_value = interrupt_data_obj.get("value", {})
+        
+        return ChatResponse(
+            response=response_text,
+            session_id=session_id,
+            has_interrupt=has_interrupt,
+            interrupt_type="human_approval" if has_interrupt else None,
+            interrupt_data=interrupt_value
+        )
+    
+    @staticmethod
+    async def approve_request(
+        session_id: str,
+        decision: str,
+        reason: Optional[str] = None
+    ) -> Dict:
+        """Approve or reject an assistance request"""
+        if not settings.chat_manager:
+            raise Exception("System not initialized")
+        
+        if not settings.chat_manager.has_pending_interrupt(session_id):
+            raise Exception("No pending interrupt")
+        
+        decision_text = f"{decision} {reason or ''}"
+        response_text = await settings.chat_manager.chat(
+            message=decision_text,
+            session_id=session_id
+        )
+        
+        return {
+            "status": "success",
+            "decision": decision,
+            "response": response_text,
+            "session_id": session_id
+        }
+    
+    @staticmethod
+    async def get_session_status(session_id: str) -> Dict:
+        """Get session status"""
+        if not settings.chat_manager:
+            raise Exception("System not initialized")
+        
+        has_interrupt = settings.chat_manager.has_pending_interrupt(session_id)
+        interrupt_data = None
+        
+        if has_interrupt:
+            interrupt_info = settings.chat_manager.pending_interrupts.get(session_id, {})
+            interrupt_data_obj = interrupt_info.get("interrupt_data")
+            
+            if hasattr(interrupt_data_obj, 'value'):
+                interrupt_data = interrupt_data_obj.value
+        
+        return {
+            "session_id": session_id,
+            "has_pending_interrupt": has_interrupt,
+            "interrupt_data": interrupt_data
+        }
+    
+    @staticmethod
+    async def get_history(session_id: str) -> Dict:
+        """Get conversation history"""
+        if not settings.chat_manager:
+            raise Exception("System not initialized")
+        
+        history = await settings.chat_manager.get_conversation_history(session_id)
+        
+        return {
+            "session_id": session_id,
+            "history": [
+                {
+                    "role": msg.type if hasattr(msg, 'type') else "unknown",
+                    "content": msg.content if hasattr(msg, 'content') else str(msg)
+                }
+                for msg in history
+            ]
+        }
+    
+    @staticmethod
+    async def list_sessions() -> Dict:
+        """List all active sessions"""
+        if not settings.chat_manager:
+            raise Exception("System not initialized")
+        
+        sessions_info = []
+        for session_id in settings.chat_manager.pending_interrupts.keys():
+            interrupt_info = settings.chat_manager.pending_interrupts.get(session_id, {})
+            interrupt_data_obj = interrupt_info.get("interrupt_data")
+            
+            data = (
+                interrupt_data_obj.value 
+                if hasattr(interrupt_data_obj, 'value') 
+                else {}
+            )
+            
+            sessions_info.append({
+                "session_id": session_id,
+                "has_interrupt": True,
+                "user_email": data.get("user_email"),
+                "country": data.get("country"),
+                "description": data.get("description")
+            })
+        
+        return {
+            "total_pending": len(sessions_info),
+            "sessions": sessions_info
+        }
+    
+    @staticmethod
+    async def clear_session(session_id: str) -> Dict:
+        """Clear a session"""
+        if not settings.chat_manager:
+            raise Exception("System not initialized")
+        
+        if session_id in settings.chat_manager.pending_interrupts:
+            del settings.chat_manager.pending_interrupts[session_id]
+        
+        return {
+            "status": "success",
+            "message": f"Session {session_id} cleared"
+        }
+

api/services/stream_service.py

@@ -0,0 +1,66 @@
+# api/services/stream_service.py
+import json
+import asyncio
+from typing import AsyncGenerator
+from api.config import settings, logger
+
+
+class StreamService:
+    """Service for handling streaming responses"""
+    
+    @staticmethod
+    async def stream_chat_response(
+        message: str,
+        session_id: str
+    ) -> AsyncGenerator[str, None]:
+        """Stream chat responses token by token"""
+        try:
+            # Check for pending interrupt first
+            if settings.chat_manager.has_pending_interrupt(session_id):
+                interrupt_info = settings.chat_manager.pending_interrupts.get(session_id, {})
+                interrupt_data_obj = interrupt_info.get("interrupt_data")
+                
+                interrupt_value = (
+                    interrupt_data_obj.value 
+                    if hasattr(interrupt_data_obj, 'value') 
+                    else {}
+                )
+                
+                yield f"data: {json.dumps({'type': 'interrupt', 'data': interrupt_value})}\n\n"
+                return
+            
+            # Process message
+            response = await settings.chat_manager.chat(
+                message=message,
+                session_id=session_id,
+                interrupt_handler=None
+            )
+            
+            # Check if interrupt occurred during processing
+            if settings.chat_manager.has_pending_interrupt(session_id):
+                interrupt_info = settings.chat_manager.pending_interrupts.get(session_id, {})
+                interrupt_data_obj = interrupt_info.get("interrupt_data")
+                
+                interrupt_value = (
+                    interrupt_data_obj.value 
+                    if hasattr(interrupt_data_obj, 'value') 
+                    else {}
+                )
+                
+                yield f"data: {json.dumps({'type': 'interrupt', 'data': interrupt_value})}\n\n"
+                return
+            
+            # Stream response token by token
+            words = response.split()
+            for i, word in enumerate(words):
+                token = word + (" " if i < len(words) - 1 else "")
+                yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"
+                await asyncio.sleep(settings.STREAM_DELAY)
+            
+            # Send completion event
+            yield f"data: {json.dumps({'type': 'done', 'session_id': session_id})}\n\n"
+            
+        except Exception as e:
+            logger.exception(f"Error in streaming: {e}")
+            yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
+

api/utils/__init__.py

@@ -0,0 +1,4 @@
+# api/utils/__init__.py
+from .startup import initialize_system, lifespan
+
+__all__ = ["initialize_system", "lifespan"]

api/utils/startup.py

@@ -0,0 +1,37 @@
+# api/utils/startup.py
+import sys
+from pathlib import Path
+sys.path.insert(0, str(Path(__file__).parent.parent.parent))
+
+import asyncio
+from contextlib import asynccontextmanager
+from fastapi import FastAPI
+from core.system_initializer import setup_system
+from api.config import settings, logger
+
+
+async def initialize_system():
+    """Initialize the legal assistant system"""
+    try:
+        system = await setup_system()
+        settings.chat_manager = system["chat_manager"]
+        settings.graph = system["graph"]
+        settings.system_initialized = True
+        logger.info("✅ Legal assistant system initialized")
+    except Exception as e:
+        logger.error(f"❌ Failed to initialize system: {e}")
+        settings.system_initialized = False
+
+
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Application lifespan manager"""
+    logger.info("🚀 Starting Legal Assistant API...")
+    initialization_task = asyncio.create_task(initialize_system())
+    yield
+    logger.info("🛑 Shutting down...")
+    initialization_task.cancel()
+    try:
+        await initialization_task
+    except asyncio.CancelledError:
+        pass

config/settings.py

@@ -15,7 +15,7 @@ class Settings:
     NEON_END_POINT = os.getenv("NEON_END_POINT")
     
     # Database
-    DATABASE_URL = NEON_END_POINT
+    # DATABASE_URL = NEON_END_POINT
     
     # Model Configurations
     EMBEDDING_MODEL = "text-embedding-ada-002"

core/chat_manager.py

@@ -1,5 +1,4 @@
 # [file name]: core/chat_manager.py
-# Add this as the FIRST lines of code (after docstrings)
 import sys
 from pathlib import Path
 sys.path.insert(0, str(Path(__file__).parent.parent))
@@ -7,10 +6,11 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
 import asyncio
 import logging
 from datetime import datetime
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, Any, Callable
 from langchain_core.runnables import RunnableConfig
 from langchain_core.messages import BaseMessage
 from langgraph.types import Command
+from langgraph.errors import GraphInterrupt
 
 from config.settings import settings
 from models.state_models import MultiCountryLegalState
@@ -19,6 +19,11 @@ from utils.helpers import dict_to_message_obj
 logger = logging.getLogger(__name__)
 
 class LegalChatManager:
+    """
+    Chat manager with full human-in-the-loop interrupt support.
+    Handles both synchronous (callback-based) and asynchronous (stored interrupt) modes.
+    """
+    
     def __init__(self, graph, checkpointer):
         self.graph = graph
         self.checkpointer = checkpointer
@@ -32,9 +37,30 @@ class LegalChatManager:
         # Track pending interrupts by session
         self.pending_interrupts = {}
 
-    async def chat(self, message: str, session_id: str, 
-                  legal_context: Optional[Dict[str, str]] = None) -> str:
-        """Process a chat message with session management and interrupt handling"""
+    async def chat(
+        self, 
+        message: str, 
+        session_id: str,
+        legal_context: Optional[Dict[str, str]] = None,
+        interrupt_handler: Optional[Callable[[Dict], str]] = None
+    ) -> str:
+        """
+        Process a chat message with session management and interrupt handling.
+        
+        Args:
+            message: User message to process
+            session_id: Unique session identifier for conversation tracking
+            legal_context: Optional legal context (jurisdiction, user type, etc.)
+            interrupt_handler: Optional callback for handling interrupts synchronously.
+                             If provided, interrupts are handled immediately within this call.
+                             If None, interrupts are stored for later resolution via subsequent calls.
+                             
+        Returns:
+            Assistant's response text
+            
+        Raises:
+            RuntimeError: If system is not initialized
+        """
         if not self.graph:
             raise RuntimeError("System not initialized. Call setup_system() first.")
 
@@ -56,31 +82,89 @@ class LegalChatManager:
             # Track performance
             start_time = datetime.now()
             
-            # Process through graph
-            result = await self.graph.ainvoke(MultiCountryLegalState(**input_state), config)
+            # 🔥 CRITICAL: Use streaming to detect interrupts in real-time
+            final_response = None
             
-            # Check for interrupt
-            state_snapshot = await self.graph.aget_state(config)
-            if state_snapshot and state_snapshot.next:
-                # Graph is paused at an interrupt
-                logger.info(f"⏸️ Graph interrupted at: {state_snapshot.next}")
-                self.pending_interrupts[session_id] = {
-                    "type": "human_approval",
-                    "config": config,
-                    "created_at": datetime.now(),
-                    "paused_at": state_snapshot.next
-                }
-                return self._get_approval_prompt_message(result)
+            async for chunk in self.graph.astream(
+                MultiCountryLegalState(**input_state), 
+                config,
+                stream_mode="updates"
+            ):
+                # 🔥 Check if this chunk contains an interrupt
+                if "__interrupt__" in chunk:
+                    interrupt_data = chunk["__interrupt__"]
+                    logger.info(f"⏸️ Graph interrupted: {interrupt_data}")
+                    
+                    # Extract interrupt info - handle tuple/Interrupt object format
+                    # LangGraph returns interrupts as tuples containing Interrupt objects
+                    if isinstance(interrupt_data, (list, tuple)):
+                        interrupt_info = interrupt_data[0]
+                    else:
+                        interrupt_info = interrupt_data
+                    
+                    # Handle Interrupt object (has .value attribute)
+                    if hasattr(interrupt_info, 'value'):
+                        interrupt_value = interrupt_info.value
+                    elif isinstance(interrupt_info, dict):
+                        interrupt_value = interrupt_info.get("value", interrupt_info)
+                    else:
+                        interrupt_value = {}
+                    
+                    # Extract message from interrupt value
+                    interrupt_message = interrupt_value.get("message", "") if isinstance(interrupt_value, dict) else ""
+                    
+                    # 🔥 Two modes of operation:
+                    # 1. Synchronous: If interrupt_handler provided, handle immediately
+                    # 2. Asynchronous: Store interrupt and return, wait for next call
+                    
+                    if interrupt_handler:
+                        # SYNCHRONOUS MODE: Handle interrupt immediately
+                        logger.info("📞 Calling synchronous interrupt handler")
+                        moderator_response = interrupt_handler(interrupt_value)
+                        
+                        # Resume immediately with the moderator's response
+                        logger.info(f"🔄 Resuming graph with: {moderator_response}")
+                        async for resume_chunk in self.graph.astream(
+                            Command(resume=moderator_response),
+                            config,
+                            stream_mode="updates"
+                        ):
+                            # Continue processing resumed chunks
+                            for node_name, node_output in resume_chunk.items():
+                                if node_name != "__interrupt__":
+                                    logger.debug(f"📦 Resume chunk from {node_name}")
+                        
+                        # After resume completes, get final state
+                        state = await self.graph.aget_state(config)
+                        final_response = self._extract_response(state.values)
+                        break
+                    else:
+                        # ASYNCHRONOUS MODE: Store interrupt and return
+                        logger.info("💾 Storing interrupt for later resolution")
+                        self.pending_interrupts[session_id] = {
+                            "type": "human_approval",
+                            "config": config,
+                            "created_at": datetime.now(),
+                            "interrupt_data": interrupt_info
+                        }
+                        return interrupt_message or self._get_default_approval_prompt()
+                
+                # Process normal chunks (non-interrupt)
+                for node_name, node_output in chunk.items():
+                    if node_name != "__interrupt__":
+                        logger.debug(f"📦 Chunk from {node_name}")
+            
+            # If no interrupt occurred, get final state
+            if final_response is None:
+                state = await self.graph.aget_state(config)
+                final_response = self._extract_response(state.values)
             
             # Track performance
             processing_time = (datetime.now() - start_time).total_seconds()
             self._update_session_stats(session_id, processing_time)
+            self._update_routing_stats(final_response)
             
-            # Extract and return response
-            response = self._extract_response(result)
-            self._update_routing_stats(response)
-            
-            return response
+            return final_response
             
         except Exception as e:
             logger.exception(f"Chat error for session {session_id}")
@@ -88,32 +172,79 @@ class LegalChatManager:
             return f"Erreur lors du traitement: {str(e)}"
 
     async def _handle_pending_interrupt(self, session_id: str, message: str) -> str:
-        """Handle user response to a pending interrupt using Command(resume=...)"""
+        """
+        Handle user response to a pending interrupt using Command(resume=...).
+        This is called when there's a stored interrupt waiting for resolution.
+        
+        Args:
+            session_id: Session with pending interrupt
+            message: User's response (e.g., "approve" or "reject")
+            
+        Returns:
+            Final response after resuming from interrupt
+        """
         interrupt_data = self.pending_interrupts.get(session_id)
         if not interrupt_data:
             return "Erreur: Aucune interruption en attente."
         
         try:
-            logger.info(f"📥 Resuming graph with moderator decision: {message}")
+            logger.info(f"🔥 Resuming graph with moderator decision: {message}")
             
             config = interrupt_data["config"]
             
-            # CRITICAL FIX: Use Command(resume=...) to properly resume from interrupt
-            # This sends the user's message back to the interrupt() call
-            result = await self.graph.ainvoke(
+            # Use streaming to handle potential nested interrupts
+            final_response = None
+            async for chunk in self.graph.astream(
                 Command(resume=message),
-                config
-            )
+                config,
+                stream_mode="updates"
+            ):
+                if "__interrupt__" in chunk:
+                    # Another interrupt occurred during resume - store it
+                    new_interrupt = chunk["__interrupt__"]
+                    
+                    # Handle tuple/list format
+                    if isinstance(new_interrupt, (list, tuple)):
+                        interrupt_info = new_interrupt[0]
+                    else:
+                        interrupt_info = new_interrupt
+                    
+                    # Extract value from Interrupt object
+                    if hasattr(interrupt_info, 'value'):
+                        interrupt_value = interrupt_info.value
+                    elif isinstance(interrupt_info, dict):
+                        interrupt_value = interrupt_info.get("value", interrupt_info)
+                    else:
+                        interrupt_value = {}
+                    
+                    # Store the new interrupt
+                    self.pending_interrupts[session_id] = {
+                        "type": "human_approval",
+                        "config": config,
+                        "created_at": datetime.now(),
+                        "interrupt_data": interrupt_info
+                    }
+                    
+                    interrupt_message = interrupt_value.get("message", "") if isinstance(interrupt_value, dict) else ""
+                    return interrupt_message or self._get_default_approval_prompt()
+                
+                # Process normal chunks
+                for node_name, node_output in chunk.items():
+                    if node_name != "__interrupt__":
+                        logger.debug(f"📦 Resume chunk from {node_name}")
+            
+            # Get final state after successful resume
+            state = await self.graph.aget_state(config)
+            final_response = self._extract_response(state.values)
             
             # Clean up the pending interrupt
             del self.pending_interrupts[session_id]
             
-            # Extract and return final response
-            response = self._extract_response(result)
-            self._update_routing_stats(response)
+            # Update stats
+            self._update_routing_stats(final_response)
             
             logger.info(f"✅ Graph resumed successfully for session {session_id}")
-            return response
+            return final_response
             
         except Exception as e:
             logger.error(f"Error resuming from interrupt: {str(e)}")
@@ -122,26 +253,12 @@ class LegalChatManager:
                 del self.pending_interrupts[session_id]
             return f"Erreur lors du traitement de la décision: {str(e)}"
 
-    def _get_approval_prompt_message(self, state) -> str:
-        """Generate message asking for human approval"""
-        # Extract metadata from state
-        if isinstance(state, MultiCountryLegalState):
-            state_dict = state.model_dump()
-        elif isinstance(state, dict):
-            state_dict = state
-        else:
-            state_dict = {}
-        
-        user_email = state_dict.get("user_email", "Non spécifié")
-        country = state_dict.get("legal_context", {}).get("detected_country", "Non spécifié")
-        description = state_dict.get("assistance_description", "Non spécifié")
-        
-        return f"""
+    def _get_default_approval_prompt(self) -> str:
+        """Default approval prompt if interrupt message extraction fails"""
+        return """
 🔒 **APPROBATION HUMAINE REQUISE**
 
-📧 **Utilisateur**: {user_email}
-🌍 **Pays**: {country}
-📝 **Description**: {description}
+Une demande d'assistance juridique nécessite votre approbation.
 
 **Veuillez répondre avec:**
 - "approve [raison]" pour approuver la demande
@@ -154,10 +271,32 @@ class LegalChatManager:
 **Votre décision:**
 """
 
-    # === EXISTING METHODS (unchanged) ===
+    def get_checkpointer_info(self) -> Dict[str, Any]:
+        """Get information about the current checkpointer type"""
+        checkpointer_type = "unknown"
+        if hasattr(self.checkpointer, '__class__'):
+            class_name = self.checkpointer.__class__.__name__
+            if 'PostgresSaver' in class_name:
+                checkpointer_type = "postgres"
+            elif 'InMemorySaver' in class_name:
+                checkpointer_type = "memory"
+        
+        return {
+            "type": checkpointer_type,
+            "persistent": checkpointer_type == "postgres",
+            "description": "Persistent storage" if checkpointer_type == "postgres" else "In-memory (volatile)"
+        }
 
     async def get_conversation_history(self, session_id: str) -> List[BaseMessage]:
-        """Get conversation history for a session"""
+        """
+        Get conversation history for a session.
+        
+        Args:
+            session_id: Session identifier
+            
+        Returns:
+            List of message objects from conversation history
+        """
         if not self.graph:
             return []
         
@@ -183,18 +322,27 @@ class LegalChatManager:
             logger.exception(f"Error getting conversation history for session {session_id}")
             return []
 
-    def get_session_stats(self, session_id: str) -> Dict:
+    def get_session_stats(self, session_id: str) -> Dict[str, Any]:
         """Get statistics for a specific session"""
         return self.active_sessions.get(session_id, {})
 
-    def get_global_stats(self) -> Dict:
-        """Get global system statistics"""
-        return {
+    def get_global_stats(self) -> Dict[str, Any]:
+        """
+        Get global system statistics.
+        
+        Returns:
+            Dictionary with routing stats, active sessions, and storage info
+        """
+        stats = {
             "routing_stats": self.routing_stats,
             "active_sessions": len(self.active_sessions),
             "total_queries": self.routing_stats["total_queries"],
             "pending_interrupts": len(self.pending_interrupts)
         }
+        
+        # Add checkpointer info
+        stats.update(self.get_checkpointer_info())
+        return stats
 
     def _initialize_session(self, session_id: str):
         """Initialize or update session tracking"""
@@ -212,9 +360,23 @@ class LegalChatManager:
         session_info["query_count"] += 1
         session_info["last_activity"] = datetime.now()
 
-    def _prepare_input_state(self, message: str, session_id: str, 
-                           legal_context: Optional[Dict[str, str]]) -> Dict:
-        """Prepare input state for graph processing"""
+    def _prepare_input_state(
+        self, 
+        message: str, 
+        session_id: str, 
+        legal_context: Optional[Dict[str, str]]
+    ) -> Dict[str, Any]:
+        """
+        Prepare input state for graph processing.
+        
+        Args:
+            message: User message
+            session_id: Session identifier
+            legal_context: Optional legal context
+            
+        Returns:
+            Dictionary with complete input state for graph
+        """
         ctx = legal_context or {
             "jurisdiction": "Unknown", 
             "user_type": "general", 
@@ -237,7 +399,15 @@ class LegalChatManager:
         }
 
     def _extract_response(self, result) -> str:
-        """Extract response text from graph result"""
+        """
+        Extract response text from graph result.
+        
+        Args:
+            result: Graph execution result (state or dict)
+            
+        Returns:
+            Assistant's response text
+        """
         if isinstance(result, MultiCountryLegalState):
             r = result.model_dump()
         elif isinstance(result, dict):
@@ -246,6 +416,7 @@ class LegalChatManager:
             r = {}
 
         msgs = r.get("messages", [])
+        # Find the last assistant message
         for m in reversed(msgs):
             if (m.get("role") or "").lower() in ("assistant", "ai"):
                 return m.get("content", "")
@@ -278,7 +449,12 @@ class LegalChatManager:
         logger.error(f"Session {session_id}: {error}")
 
     def cleanup_inactive_sessions(self, max_age_hours: int = 24):
-        """Clean up sessions that have been inactive for too long"""
+        """
+        Clean up sessions that have been inactive for too long.
+        
+        Args:
+            max_age_hours: Maximum age in hours before cleanup
+        """
         cutoff_time = datetime.now().timestamp() - (max_age_hours * 3600)
         
         inactive_sessions = [
@@ -291,4 +467,35 @@ class LegalChatManager:
             if session_id in self.pending_interrupts:
                 del self.pending_interrupts[session_id]
             del self.active_sessions[session_id]
-            logger.info(f"Cleaned up inactive session: {session_id}")
+            logger.info(f"Cleaned up inactive session: {session_id}")
+
+    def has_pending_interrupt(self, session_id: str) -> bool:
+        """
+        Check if there's a pending interrupt for a session.
+        
+        Args:
+            session_id: Session identifier
+            
+        Returns:
+            True if session has pending interrupt, False otherwise
+        """
+        return session_id in self.pending_interrupts
+
+    def get_system_info(self) -> Dict[str, Any]:
+        """
+        Get comprehensive system information.
+        
+        Returns:
+            Dictionary with system status, storage info, and statistics
+        """
+        return {
+            "system": {
+                "initialized": self.graph is not None,
+                "active_sessions": len(self.active_sessions),
+                "pending_interrupts": len(self.pending_interrupts),
+                "total_queries": self.routing_stats["total_queries"]
+            },
+            "storage": self.get_checkpointer_info(),
+            "routing": self.routing_stats,
+            "timestamp": datetime.now().isoformat()
+        }

core/graph_builder.py

@@ -6,8 +6,9 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
 
 from langgraph.graph import StateGraph, START, END
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.checkpoint.memory import InMemorySaver
 import logging
-from typing import Dict, List, Any
+from typing import Dict, List, Any, Union
 from langchain_core.runnables import RunnableConfig
 
 from models.state_models import MultiCountryLegalState
@@ -31,7 +32,7 @@ class GraphBuilder:
         self, 
         router: CountryRouter,
         llm, 
-        checkpointer: AsyncPostgresSaver,
+        checkpointer: Union[AsyncPostgresSaver, InMemorySaver],
         # Country retrievers as a dictionary for easy extension
         country_retrievers: Dict[str, LegalRetriever] = None
     ):
@@ -57,7 +58,19 @@ class GraphBuilder:
         self.response_nodes = ResponseNodes(llm)
         self.helper_nodes = HelperNodes(llm)
         
-        logger.info(f"GraphBuilder initialized with countries: {list(self.country_retrievers.keys())}")
+        # Log checkpointer type
+        checkpointer_type = self._get_checkpointer_type()
+        logger.info(f"GraphBuilder initialized with {checkpointer_type} checkpointer and countries: {list(self.country_retrievers.keys())}")
+
+    def _get_checkpointer_type(self) -> str:
+        """Determine the type of checkpointer being used"""
+        if hasattr(self.checkpointer, '__class__'):
+            class_name = self.checkpointer.__class__.__name__
+            if 'PostgresSaver' in class_name:
+                return "PostgreSQL"
+            elif 'InMemorySaver' in class_name:
+                return "in-memory"
+        return "unknown"
 
     def add_country(self, country_code: str, retriever: LegalRetriever):
         """Dynamically add a new country to the system"""
@@ -166,7 +179,8 @@ class GraphBuilder:
         
         workflow.add_edge("process_assistance", "response")
         
-        logger.info(f"Scalable graph built for {len(self.country_retrievers)} countries: {list(self.country_retrievers.keys())}")
+        checkpointer_type = self._get_checkpointer_type()
+        logger.info(f"Scalable graph built with {checkpointer_type} checkpointer for {len(self.country_retrievers)} countries: {list(self.country_retrievers.keys())}")
         return workflow
     
     def _create_assistance_collect_wrapper(self):

core/human_approval_node.py

@@ -1,9 +1,4 @@
 # core/human_approval_node.py
-# Add this as the FIRST lines of code (after docstrings)
-import sys
-from pathlib import Path
-sys.path.insert(0, str(Path(__file__).parent.parent))
-
 import logging
 from typing import Literal
 from langchain_core.runnables import RunnableConfig
@@ -22,68 +17,59 @@ class HumanApprovalNode:
         self,
         state: MultiCountryLegalState,
         config: RunnableConfig
-    ) -> Command[Literal["response"]]:
-        """Process human approval with interrupt"""
-        try:
-            # Validate required fields
-            if not state.user_email or not state.assistance_description:
-                logger.warning("Missing required fields for approval")
-                return Command(
-                    goto="response",
-                    update={
-                        "messages": [{
-                            "role": "assistant",
-                            "content": "❌ Données incomplètes pour l'approbation.",
-                            "meta": {}
-                        }]
-                    }
-                )
-            
-            logger.info(f"🔒 Human approval node triggered for {state.user_email}")
-            
-            # Prepare interrupt message
-            interrupt_message = self._format_approval_request(state)
-            
-            # Trigger interrupt and wait for human input
-            moderator_input = interrupt({
-                "type": "human_approval",
-                "user_email": state.user_email,
-                "country": self._get_country_display(state),
-                "description": state.assistance_description,
-                "message": interrupt_message
-            })
-            
-            logger.info(f"📥 Received moderator input: {moderator_input}")
-            
-            # Parse moderator decision
-            decision = self._parse_decision(moderator_input)
-            
-            # Handle approval
-            if decision["approved"]:
-                return await self._handle_approval(state, decision)
-            else:
-                return await self._handle_rejection(state, decision)
-                
-        except Exception as e:
-            logger.error(f"Error in approval node: {str(e)}", exc_info=True)
+    ) -> Command[Literal["process_assistance", "response"]]:
+        """Process human approval with interrupt - uses Command(goto=...) pattern"""
+        
+        # Validate required fields BEFORE interrupt
+        if not state.user_email or not state.assistance_description:
+            logger.warning("Missing required fields for approval")
             return Command(
                 goto="response",
                 update={
                     "approval_status": "error",
+                    "approval_reason": "Données incomplètes",
                     "messages": [{
                         "role": "assistant",
-                        "content": f"❌ Erreur lors de l'approbation: {str(e)}",
+                        "content": "❌ Données incomplètes pour l'approbation.",
                         "meta": {}
                     }]
                 }
             )
-    
+        
+        logger.info(f"🔒 Human approval node triggered for {state.user_email}")
+        
+        # Prepare interrupt message
+        interrupt_message = self._format_approval_request(state)
+        
+        # 🔥 CRITICAL: DO NOT wrap interrupt() in try-except!
+        # Let GraphInterrupt propagate naturally to the graph executor
+        moderator_input = interrupt({
+            "type": "human_approval",
+            "user_email": state.user_email,
+            "country": self._get_country_display(state),
+            "description": state.assistance_description,
+            "message": interrupt_message
+        })
+        
+        # 🎯 Code below ONLY executes after graph resumes from interrupt
+        logger.info(f"🔥 Received moderator input: {moderator_input}")
+        
+        # Parse moderator decision
+        decision = self._parse_decision(moderator_input)
+        
+        # Handle approval or rejection
+        if decision["approved"]:
+            return await self._handle_approval(state, decision)
+        else:
+            return await self._handle_rejection(state, decision)
+
+            
     async def _handle_approval(
         self, 
         state: MultiCountryLegalState, 
         decision: dict
-    ) -> Command[Literal["response"]]:
-        """Handle approved request (sends email and routes to response)"""
+    ) -> Command[Literal["process_assistance"]]:
+        """Handle approved request - routes to process_assistance"""
         logger.info(f"✅ Request APPROVED for {state.user_email}")
         
         # Send email
@@ -95,38 +81,15 @@ class HumanApprovalNode:
         )
         logger.info(f"✅ Emails envoyés avec succès pour {state.user_email}")
         
-        # Build success message
-        if email_result.get("success"):
-            message_content = f"""✅ **DEMANDE APPROUVÉE ET ENVOYÉE**
-
-📧 Un email de confirmation a été envoyé à: {state.user_email}
-👨‍⚖️ Notre équipe juridique vous contactera sous 24-48 heures.
-
-**Raison de l'approbation:** {decision['reason']}
-**Approuvé par:** {decision['moderator_id']}
-"""
-        else:
-            message_content = f"""⚠️ **DEMANDE APPROUVÉE MAIS ERREUR D'ENVOI**
-
-La demande a été approuvée mais l'envoi d'email a échoué.
-**Erreur:** {email_result.get('error', 'Unknown')}
-
-Veuillez contacter directement: fitahiana@acfai.org
-"""
-        
         return Command(
-            goto="response",
+            goto="process_assistance",
             update={
                 "approval_status": "approved",
                 "approval_reason": decision["reason"],
                 "approved_by": decision["moderator_id"],
                 "approval_timestamp": datetime.now().isoformat(),
                 "email_status": "sent" if email_result.get("success") else "error",
-                "messages": [{
-                    "role": "assistant",
-                    "content": message_content,
-                    "meta": {"approval": "approved"}
-                }]
+                "email_result": email_result
             }
         )
     
@@ -134,8 +97,8 @@ Veuillez contacter directement: fitahiana@acfai.org
         self, 
         state: MultiCountryLegalState, 
         decision: dict
-    ) -> Command[Literal["response"]]:  # Updated: Removed "process_assistance"
-        """Handle rejected request"""
+    ) -> Command[Literal["response"]]:
+        """Handle rejected request - routes directly to response"""
         logger.info(f"❌ Request REJECTED for {state.user_email}")
         
         message_content = f"""❌ **DEMANDE REFUSÉE**
@@ -154,6 +117,7 @@ Si vous pensez qu'il s'agit d'une erreur, veuillez reformuler votre demande avec
                 "approval_reason": decision["reason"],
                 "approved_by": decision["moderator_id"],
                 "approval_timestamp": datetime.now().isoformat(),
+                "assistance_step": "completed",
                 "messages": [{
                     "role": "assistant",
                     "content": message_content,

core/system_initializer.py

@@ -6,6 +6,7 @@ sys.path.insert(0, str(Path(__file__).parent.parent))
 
 import logging
 from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from langgraph.checkpoint.memory import InMemorySaver
 
 from core.graph_builder import GraphBuilder
 from core.chat_manager import LegalChatManager
@@ -13,13 +14,12 @@ from core.router import CountryRouter
 from database.mongodb_client import MongoDBClient
 from database.postgres_checkpointer import PostgresCheckpointer
 from langchain_openai import ChatOpenAI
-from config import settings  # Make sure this import is correct
-
+from config import settings
 
 logger = logging.getLogger(__name__)
 
 async def setup_system():
-    """Initialize the legal assistant system for API use"""
+    """Initialize the legal assistant system with fallback to in-memory checkpointer"""
     
     try:
         # 1. Initialize MongoDB using your existing class
@@ -55,32 +55,8 @@ async def setup_system():
         
         router = CountryRouter()
         
-        # 5. Initialize PostgreSQL checkpointer - FIXED DATABASE URL
-        # Check what database URL setting you have
-        database_url = getattr(settings, 'DATABASE_URL', None)
-        
-        if not database_url:
-            # Try alternative setting names
-            database_url = getattr(settings, 'POSTGRES_URL', None) or \
-                          getattr(settings, 'POSTGRESQL_URL', None) or \
-                          getattr(settings, 'DB_URL', None)
-        
-        if not database_url:
-            raise Exception("No database URL found in settings")
-            
-        logger.info(f"🔗 Using database URL: {database_url.split('@')[-1] if '@' in database_url else 'local'}")  # Log safely
-        
-        postgres_checkpointer = PostgresCheckpointer(
-            database_url=database_url,  # Use actual database URL
-            max_connections=10,
-            min_connections=2
-        )
-        
-        if not await postgres_checkpointer.initialize():
-            raise Exception("PostgreSQL checkpointer initialization failed")
-            
-        checkpointer = postgres_checkpointer.get_checkpointer()
-        logger.info("✅ PostgreSQL checkpointer initialized for API")
+        # 5. Initialize checkpointer with fallback logic
+        checkpointer = await _initialize_checkpointer_with_fallback()
         
         # 6. Build graph
         graph_builder = GraphBuilder(
@@ -106,4 +82,60 @@ async def setup_system():
         
     except Exception as e:
         logger.error(f"❌ Failed to initialize system: {e}")
-        raise
+        raise
+
+async def _initialize_checkpointer_with_fallback():
+    """Initialize checkpointer with fallback to in-memory if PostgreSQL fails"""
+    
+    # First, try to initialize PostgreSQL checkpointer
+    postgres_checkpointer = None
+    database_url = getattr(settings, 'DATABASE_URL', None)
+    
+    if not database_url:
+        # Try alternative setting names
+        database_url = getattr(settings, 'POSTGRES_URL', None) or \
+                      getattr(settings, 'POSTGRESQL_URL', None) or \
+                      getattr(settings, 'DB_URL', None)
+    
+    if database_url:
+        try:
+            logger.info(f"🔗 Attempting PostgreSQL connection: {database_url.split('@')[-1] if '@' in database_url else 'local'}")
+            
+            postgres_checkpointer = PostgresCheckpointer(
+                database_url=database_url,
+                max_connections=10,
+                min_connections=2
+            )
+            
+            if await postgres_checkpointer.initialize():
+                checkpointer = postgres_checkpointer.get_checkpointer()
+                logger.info("✅ PostgreSQL checkpointer initialized successfully")
+                return checkpointer
+            else:
+                logger.warning("❌ PostgreSQL checkpointer initialization failed, will fall back to in-memory")
+                
+        except Exception as e:
+            logger.warning(f"❌ PostgreSQL connection failed: {e}, falling back to in-memory checkpointer")
+    
+    else:
+        logger.warning("❌ No database URL found in settings, using in-memory checkpointer")
+    
+    # Fall back to in-memory checkpointer
+    try:
+        checkpointer = InMemorySaver()
+        logger.info("✅ In-memory checkpointer initialized as fallback")
+        logger.warning("⚠️  Using in-memory checkpointer - conversation history will not persist across restarts")
+        return checkpointer
+        
+    except Exception as e:
+        logger.error(f"❌ Even in-memory checkpointer failed: {e}")
+        raise Exception("Failed to initialize any checkpointer")
+
+def get_checkpointer_type(checkpointer):
+    """Utility function to check what type of checkpointer is being used"""
+    if hasattr(checkpointer, '__class__'):
+        if 'PostgresSaver' in checkpointer.__class__.__name__:
+            return "postgres"
+        elif 'InMemorySaver' in checkpointer.__class__.__name__:
+            return "memory"
+    return "unknown"

main.py

@@ -1,633 +1,483 @@
 #!/usr/bin/env python3
 """
-Scalable Multi-Country Legal RAG System
-Supports dynamic addition of new countries with clean architecture
+Multi-Country Legal RAG System - Interactive Testing Mode with Human-in-the-Loop Support
 """
-# Add this as the FIRST lines of code (after docstrings)
 import sys
 from pathlib import Path
-sys.path.insert(0, str(Path(__file__).parent.parent))
+sys.path.insert(0, str(Path(__file__).parent))
 
 import asyncio
 import logging
 import time
 from datetime import datetime
-from typing import List, Dict, Any, Optional
-
-from MultiCountryRAG.config.settings import settings
-from MultiCountryRAG.database.mongodb_client import MongoDBClient
-from MultiCountryRAG.database.postgres_checkpointer import PostgresCheckpointer
-from MultiCountryRAG.core.router import CountryRouter
-from MultiCountryRAG.core.retriever import LegalRetriever
-from MultiCountryRAG.core.graph_builder import GraphBuilder
-from MultiCountryRAG.core.chat_manager import LegalChatManager
-from MultiCountryRAG.utils.logger import setup_logging
-
-import uuid
-
-
-class MultiCountryLegalRAGSystem:
-    """Scalable system class supporting dynamic country addition"""
+from typing import Optional
+
+from core.system_initializer import setup_system
+
+# Setup comprehensive logging
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+    handlers=[
+        logging.StreamHandler(sys.stdout),
+        logging.FileHandler('legal_rag_system.log', mode='a')
+    ]
+)
+logger = logging.getLogger(__name__)
+
+class LegalRAGTester:
+    """
+    Interactive tester for the Legal RAG system with support for:
+    - Human-in-the-loop interrupts
+    - Session management
+    - Statistics tracking
+    - Assistance workflow testing
+    """
     
     def __init__(self):
-        self.mongo_client = MongoDBClient()
-        self.postgres_checkpointer = PostgresCheckpointer(
-            database_url=settings.DATABASE_URL,
-            max_connections=10,
-            min_connections=2
-        )
-        self.router = None
-        # Dynamic country retrievers dictionary - easily extensible!
-        self.country_retrievers = {}
-        self.llm = None
-        self.graph = None
+        self.system = None
         self.chat_manager = None
-        self.initialized = False
-
-    async def initialize(self) -> bool:
-        """Initialize the complete scalable system"""
+        self.session_id = f"test_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        self.message_count = 0
+    
+    async def initialize(self):
+        """Initialize the Legal RAG system"""
+        print("🚀 Initializing Legal RAG System...")
         try:
-            setup_logging()
-            settings.validate()
-            
-            # Initialize databases
-            if not self.mongo_client.connect():
-                raise Exception("MongoDB connection failed")
-                
-            if not await self.postgres_checkpointer.initialize():
-                logging.warning("PostgreSQL initialization failed")
+            self.system = await setup_system()
+            self.chat_manager = self.system["chat_manager"]
             
-            # Initialize core components
-            self.router = CountryRouter()
+            print("✅ System initialized successfully!")
             
-            # Initialize default countries - easily extensible!
-            self._initialize_default_countries()
+            # Print system info
+            info = self.chat_manager.get_checkpointer_info()
+            stats = self.chat_manager.get_global_stats()
             
-            # Initialize LLM
-            from langchain_openai import ChatOpenAI
-            self.llm = ChatOpenAI(
-                model=settings.CHAT_MODEL,
-                temperature=settings.CHAT_TEMPERATURE,
-                max_tokens=settings.CHAT_MAX_TOKENS
-            )
-            
-            # Build scalable graph with country dictionary
-            graph_builder = GraphBuilder(
-                router=self.router,
-                llm=self.llm,
-                checkpointer=self.postgres_checkpointer.get_checkpointer(),
-                country_retrievers=self.country_retrievers  # Pass the dictionary
-            )
-            
-            workflow = graph_builder.build_graph()
-            
-            # Compile with interrupt support
-            self.graph = workflow.compile(
-                checkpointer=self.postgres_checkpointer.get_checkpointer(),
-                interrupt_before=["human_approval"]
-            )
-            
-            # Initialize chat manager
-            self.chat_manager = LegalChatManager(
-                self.graph, 
-                self.postgres_checkpointer.get_checkpointer()
-            )
-            
-            await self._perform_health_check()
-            
-            self.initialized = True
-            logging.info(f"✅ System initialized with {len(self.country_retrievers)} countries")
-            self._print_system_info()
+            print(f"📊 Checkpointer: {info['type']} ({info['description']})")
+            print(f"💾 Persistent: {info['persistent']}")
+            print(f"🎯 Session ID: {self.session_id}")
+            print(f"🌍 Available countries: Bénin, Madagascar")
             
             return True
             
         except Exception as e:
-            logging.error(f"❌ System initialization failed: {e}")
-            import traceback
-            traceback.print_exc()
+            logger.exception("Failed to initialize system")
+            print(f"❌ Initialization failed: {e}")
             return False
-
-    def _initialize_default_countries(self):
-        """Initialize default countries - easily extensible!"""
-        # Benin
-        if hasattr(self.mongo_client, 'benin_vectorstore'):
-            self.country_retrievers["benin"] = LegalRetriever(
-                self.mongo_client.benin_vectorstore,
-                self.mongo_client.benin_collection
-            )
-        
-        # Madagascar  
-        if hasattr(self.mongo_client, 'madagascar_vectorstore'):
-            self.country_retrievers["madagascar"] = LegalRetriever(
-                self.mongo_client.madagascar_vectorstore,
-                self.mongo_client.madagascar_collection
-            )
+    
+    def _handle_interrupt(self, interrupt_value: dict) -> str:
+        """
+        Handle human-in-the-loop interrupts synchronously.
+        This is called when the graph needs human approval.
         
-        logging.info(f"🌍 Initialized {len(self.country_retrievers)} default countries")
-
-    def add_country(self, country_code: str, vectorstore, collection) -> bool:
-        """Dynamically add a new country to the running system"""
-        try:
-            if country_code in self.country_retrievers:
-                logging.warning(f"Country {country_code} already exists")
-                return False
-            
-            new_retriever = LegalRetriever(vectorstore, collection)
-            self.country_retrievers[country_code] = new_retriever
-            
-            # Rebuild graph if system is already initialized
-            if self.initialized:
-                graph_builder = GraphBuilder(
-                    router=self.router,
-                    llm=self.llm,
-                    checkpointer=self.postgres_checkpointer.get_checkpointer(),
-                    country_retrievers=self.country_retrievers
-                )
-                workflow = graph_builder.build_graph()
-                self.graph = workflow.compile(
-                    checkpointer=self.postgres_checkpointer.get_checkpointer(),
-                    interrupt_before=["human_approval"]
-                )
-            
-            logging.info(f"🎉 Successfully added country: {country_code}")
-            return True
+        Args:
+            interrupt_value: Interrupt data containing message and context
             
-        except Exception as e:
-            logging.error(f"❌ Failed to add country {country_code}: {e}")
-            return False
-
-    async def _perform_health_check(self):
-        """Perform health check after initialization"""
-        try:
-            health_status = await self.health_check()
-            
-            unhealthy_components = [k for k, v in health_status.get('components', {}).items() if not v]
-            if unhealthy_components:
-                logging.warning(f"⚠️ Unhealthy components: {unhealthy_components}")
-                
-        except Exception as e:
-            logging.warning(f"⚠️ Health check failed: {e}")
-
-    async def health_check(self) -> Dict[str, Any]:
-        """Comprehensive system health check"""
-        health_status = {
-            "system_initialized": self.initialized,
-            "mongodb_connected": self.mongo_client.client is not None,
-            "postgres_healthy": {},
-            "interrupt_enabled": True,
-            "available_countries": list(self.country_retrievers.keys()),
-            "components": {
-                "router": self.router is not None,
-                "llm": self.llm is not None,
-                "graph": self.graph is not None,
-                "chat_manager": self.chat_manager is not None,
-                "country_retrievers": len(self.country_retrievers) > 0
-            },
-            "timestamp": datetime.now().isoformat(),
-            "settings": {
-                "chat_model": settings.CHAT_MODEL,
-                "embedding_model": settings.EMBEDDING_MODEL,
-                "max_search_results": settings.MAX_SEARCH_RESULTS
-            }
-        }
-        
-        # Test MongoDB connection
-        if health_status["mongodb_connected"]:
-            try:
-                self.mongo_client.client.admin.command('ping')
-                health_status["mongodb_ping"] = True
-            except Exception as e:
-                health_status["mongodb_ping"] = False
-                health_status["mongodb_error"] = str(e)
+        Returns:
+            Moderator's decision (approve/reject)
+        """
+        print("\n" + "="*70)
+        print("🔒 HUMAN APPROVAL REQUIRED")
+        print("="*70)
         
-        # Test PostgreSQL connection
-        if hasattr(self.postgres_checkpointer, 'health_check'):
-            postgres_health = await self.postgres_checkpointer.health_check()
-            health_status["postgres_healthy"] = postgres_health
+        # Extract interrupt information
+        message = interrupt_value.get("message", "")
+        email = interrupt_value.get("user_email", "N/A")
+        country = interrupt_value.get("country", "N/A")
+        description = interrupt_value.get("description", "N/A")
         
-        return health_status
-
-    async def chat(self, message: str, session_id: str = None, context: dict = None) -> str:
-        """Public chat interface"""
-        if not self.initialized:
-            raise RuntimeError("System not initialized. Call initialize() first.")
-        
-        if not message or not message.strip():
-            raise ValueError("Message cannot be empty")
+        # Display formatted approval request
+        print(message)
+        print()
         
-        try:
-            # Prepare context
-            ctx = context or {}
-            ctx.setdefault("jurisdiction", "Unknown")
-            ctx.setdefault("user_type", "general")
-            ctx.setdefault("document_type", "legal")
-            ctx.setdefault("detected_country", "unknown")
+        # Get moderator decision with validation
+        while True:
+            moderator_input = input("🔐 Moderator Decision: ").strip()
             
-            session_id = session_id or f"cli_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+            if not moderator_input:
+                print("⚠️  Please provide a decision (approve/reject)")
+                continue
             
-            return await self.chat_manager.chat(message, session_id, ctx)
+            # Validate input
+            input_lower = moderator_input.lower()
+            if any(keyword in input_lower for keyword in ["approve", "approuver", "accept"]):
+                print("✅ Request APPROVED")
+                break
+            elif any(keyword in input_lower for keyword in ["reject", "rejeter", "refuse"]):
+                print("❌ Request REJECTED")
+                break
+            else:
+                print("⚠️  Invalid decision. Use 'approve [reason]' or 'reject [reason]'")
+                continue
         
-        except Exception as e:
-            logging.error(f"❌ Chat error for session {session_id}: {e}")
-            return f"❌ Désolé, une erreur s'est produite lors du traitement de votre demande. Veuillez réessayer."
-
-    def get_session_info(self, session_id: str) -> Dict[str, Any]:
-        """Get information about a specific session"""
-        if not self.initialized:
-            raise RuntimeError("System not initialized")
-        return self.chat_manager.get_session_stats(session_id)
-
-    def get_global_stats(self) -> Dict[str, Any]:
-        """Get global system statistics"""
-        if not self.initialized:
-            raise RuntimeError("System not initialized")
-        return self.chat_manager.get_global_stats()
-
-    def get_available_countries(self) -> List[str]:
-        """Get list of available countries"""
-        return list(self.country_retrievers.keys())
-
-    async def cleanup(self):
-        """Cleanup resources"""
-        try:
-            if self.mongo_client:
-                self.mongo_client.close()
-            if self.postgres_checkpointer:
-                await self.postgres_checkpointer.close()
-            logging.info("✅ System cleanup completed")
-        except Exception as e:
-            logging.error(f"❌ Error during cleanup: {e}")
-
-    def _print_system_info(self):
-        """Print system configuration information"""
-        countries = list(self.country_retrievers.keys())
-        print("\n" + "="*60)
-        print("🚀 SCALABLE MULTI-COUNTRY LEGAL RAG SYSTEM")
-        print("="*60)
-        print(f"🌍 Available Countries: {', '.join(countries) if countries else 'None'}")
-        print(f"🤖 AI Model: {settings.CHAT_MODEL}")
-        print(f"💾 Database: MongoDB + PostgreSQL")
-        print(f"🔍 Vector Search: {settings.EMBEDDING_MODEL}")
-        print(f"⏸️  Interrupt Support: ENABLED")
-        print(f"🌡️  Temperature: {settings.CHAT_TEMPERATURE}")
-        print(f"📝 Max Tokens: {settings.CHAT_MAX_TOKENS}")
-        print("="*60)
-
-
-class InterruptTester:
-    """Specialized tester for human approval interrupts"""
-    
-    def __init__(self, system: MultiCountryLegalRAGSystem):
-        self.system = system
-        self.test_results = []
-    
-    async def test_assistance_workflow(self, test_name: str, 
-                                      user_query: str, 
-                                      user_email: str,
-                                      user_description: str,
-                                      moderator_response: str) -> Dict[str, Any]:
-        """Test the complete assistance workflow with interrupt"""
-        print(f"\n🧪 Interrupt Test: {test_name}")
-        print(f"📝 User Query: {user_query}")
+        print("="*70 + "\n")
+        return moderator_input
+    
+    async def chat(self, message: str) -> Optional[str]:
+        """
+        Send a chat message and get response with interrupt handling.
         
-        # session_id = f"test_{datetime.now().strftime('%H%M%S%f')}"
-        session_id = f"interactive_{uuid.uuid4().hex[:8]}"
-        current_response = ""
+        Args:
+            message: User message
+            
+        Returns:
+            Assistant response or None on error
+        """
+        if not self.chat_manager:
+            print("❌ System not initialized. Please restart.")
+            return None
         
         try:
-            # Step 1: Initial request
-            print("1️⃣  Step 1: Initial assistance request...")
-            current_response = await self.system.chat(user_query, session_id)
-            print(f"🤖 Response: {current_response[:150]}...")
+            start_time = time.time()
+            self.message_count += 1
             
-            # Step 2: Email collection
-            if user_email and any(keyword in current_response.lower() for keyword in ["email", "adresse", "@"]):
-                print(f"2️⃣  Step 2: Providing email: {user_email}")
-                current_response = await self.system.chat(user_email, session_id)
-                print(f"🤖 Response: {current_response[:150]}...")
+            # Send message with interrupt handler
+            response = await self.chat_manager.chat(
+                message=message,
+                session_id=self.session_id,
+                interrupt_handler=self._handle_interrupt  # Enable synchronous interrupt handling
+            )
             
-            # Step 3: Description collection  
-            if user_description and any(keyword in current_response.lower() for keyword in ["description", "décrire", "besoin"]):
-                print(f"3️⃣  Step 3: Providing description: {user_description[:50]}...")
-                current_response = await self.system.chat(user_description, session_id)
-                print(f"🤖 Response: {current_response[:150]}...")
+            response_time = time.time() - start_time
             
-            # Step 4: Confirmation
-            if any(keyword in current_response.lower() for keyword in ["confirmer", "confirmation", "oui/non"]):
-                print("4️⃣  Step 4: Confirming request...")
-                current_response = await self.system.chat("oui", session_id)
-                print(f"🤖 Response: {current_response[:150]}...")
+            # Display response
+            print(f"\n🤖 Assistant ({response_time:.2f}s):")
             
-            # Step 5: Check for interrupt
-            interrupt_detected = self._check_for_interrupt(current_response, session_id)
+            # Format multi-line responses
+            for line in response.split('\n'):
+                if line.strip():
+                    print(f"   {line}")
+                else:
+                    print()
             
-            if interrupt_detected:
-                print("⏸️  INTERRUPT DETECTED! Waiting for moderator...")
-                
-                # Step 6: Moderator decision
-                print(f"👨‍⚖️  Moderator: {moderator_response}")
-                final_response = await self.system.chat(moderator_response, session_id)
-                print(f"✅ Final Response: {final_response[:200]}...")
-                
-                result = {
-                    "test_name": test_name,
-                    "status": "PASS",
-                    "interrupt_detected": True,
-                    "moderator_decision": moderator_response,
-                    "final_response": final_response,
-                    "session_id": session_id
-                }
-            else:
-                print("⚠️  No interrupt detected in workflow")
-                result = {
-                    "test_name": test_name,
-                    "status": "FAIL",
-                    "interrupt_detected": False,
-                    "moderator_decision": None,
-                    "final_response": current_response,
-                    "error": "Interrupt not triggered",
-                    "session_id": session_id
-                }
+            # Check for pending interrupts (async mode fallback)
+            if self.chat_manager.has_pending_interrupt(self.session_id):
+                print("\n⏸️  💡 System paused - waiting for moderator decision")
+                print("   Your next message will be treated as approval/rejection")
             
-            self.test_results.append(result)
-            return result
+            return response
             
+        except KeyboardInterrupt:
+            raise  # Re-raise to allow graceful shutdown
         except Exception as e:
-            logging.error(f"❌ Test error: {e}")
-            error_result = {
-                "test_name": test_name,
-                "status": "ERROR",
-                "interrupt_detected": False,
-                "moderator_decision": None,
-                "final_response": current_response,
-                "error": str(e),
-                "session_id": session_id
-            }
-            self.test_results.append(error_result)
-            return error_result
-    
-    def _check_for_interrupt(self, response: str, session_id: str) -> bool:
-        """Enhanced interrupt detection"""
-        interrupt_indicators = [
-            "APPROBATION", "APPROVAL", "HUMAN", "MODERATOR", 
-            "DÉCISION", "DECISION", "APPROUVER", "REJETER"
-        ]
-        
-        if any(indicator in response.upper() for indicator in interrupt_indicators):
-            return True
-        
-        if (hasattr(self.system.chat_manager, 'pending_interrupts') and 
-            session_id in self.system.chat_manager.pending_interrupts):
-            return True
-            
-        return False
+            logger.exception(f"Error processing message: {message}")
+            print(f"❌ Error: {str(e)}")
+            return None
     
-    def print_summary(self):
-        """Print test summary"""
-        print("\n" + "="*80)
-        print("📊 INTERRUPT TEST SUMMARY")
-        print("="*80)
-        
-        total = len(self.test_results)
-        passed = len([r for r in self.test_results if r["status"] == "PASS"])
-        failed = len([r for r in self.test_results if r["status"] == "FAIL"])
-        errors = len([r for r in self.test_results if r["status"] == "ERROR"])
-        
-        print(f"📈 Total Tests: {total}")
-        print(f"✅ Passed: {passed}")
-        print(f"❌ Failed: {failed}")
-        print(f"🚨 Errors: {errors}")
-        
-        if passed > 0:
-            print(f"\n🎉 Successful Tests:")
-            for result in self.test_results:
-                if result["status"] == "PASS":
-                    print(f"  - {result['test_name']}")
+    def display_stats(self):
+        """Display current system statistics"""
+        if not self.chat_manager:
+            print("❌ System not initialized")
+            return
         
-        if failed > 0 or errors > 0:
-            print(f"\n💥 Failed/Error Tests:")
-            for result in self.test_results:
-                if result["status"] in ["FAIL", "ERROR"]:
-                    print(f"  - {result['test_name']}: {result.get('error', 'Unknown error')}")
+        stats = self.chat_manager.get_global_stats()
+        routing = stats.get('routing_stats', {})
         
-        print("="*80)
-
-
-async def run_interrupt_tests():
-    """Run specialized tests for human approval interrupts"""
-    system = MultiCountryLegalRAGSystem()
-    tester = InterruptTester(system)
+        print("\n" + "="*70)
+        print("📊 SYSTEM STATISTICS")
+        print("="*70)
+        print(f"💬 Session Messages: {self.message_count}")
+        print(f"🔢 Total System Queries: {stats.get('total_queries', 0)}")
+        print(f"👥 Active Sessions: {stats.get('active_sessions', 0)}")
+        print(f"⏸️  Pending Interrupts: {stats.get('pending_interrupts', 0)}")
+        print()
+        print("🔀 Routing Statistics:")
+        print(f"   📍 Bénin queries: {routing.get('benin', 0)}")
+        print(f"   📍 Madagascar queries: {routing.get('madagascar', 0)}")
+        print(f"   ❓ Unclear queries: {routing.get('unclear', 0)}")
+        print()
+        print(f"💾 Storage: {stats.get('type', 'unknown')} - {stats.get('description', '')}")
+        print("="*70 + "\n")
     
-    try:
-        print("🚀 Initializing system...")
-        success = await system.initialize()
-        if not success:
-            print("❌ System initialization failed")
-            return
-
-        print("\n🧪 STARTING INTERRUPT TESTS")
-        print("="*60)
-
-        test_scenarios = [
-            {
-                "name": "Complete Workflow - Approve",
-                "user_query": "Je veux parler a un avocat",
-                "user_email": "test@example.com",
-                "user_description": "Consultation pour divorce au Benin",
-                "moderator_response": "approve Demande legitime"
-            },
-            {
-                "name": "Complete Workflow - Reject", 
-                "user_query": "Contactez-moi",
-                "user_email": "test2@example.com",
-                "user_description": "J'ai besoin d'aide",
-                "moderator_response": "reject Description trop vague"
-            }
-        ]
-
-        for scenario in test_scenarios:
-            await tester.test_assistance_workflow(
-                scenario["name"],
-                scenario["user_query"],
-                scenario["user_email"], 
-                scenario["user_description"],
-                scenario["moderator_response"]
-            )
-            await asyncio.sleep(1)
-
-        tester.print_summary()
-
-    except Exception as e:
-        logging.error(f"❌ Error during testing: {e}")
-    finally:
-        await system.cleanup()
-
-
-async def interactive_mode():
-    """Run interactive chat mode"""
-    system = MultiCountryLegalRAGSystem()
+    def display_help(self):
+        """Display help information"""
+        print("\n" + "="*70)
+        print("📋 AVAILABLE COMMANDS")
+        print("="*70)
+        print("  quit, exit, q       - Exit the program")
+        print("  stats, s            - Show system statistics")
+        print("  clear, cls          - Clear the screen")
+        print("  help, h, ?          - Show this help message")
+        print("  history             - Show conversation history")
+        print("  reset               - Reset current session")
+        print()
+        print("💡 TESTING TIPS:")
+        print("  - Legal queries: Ask about laws in Bénin or Madagascar")
+        print("  - Assistance: Say 'je veux parler à un avocat'")
+        print("  - Follow the prompts for email and description")
+        print("  - You'll be prompted for approval when needed")
+        print("="*70 + "\n")
     
-    try:
-        print("🚀 Initializing system...")
-        success = await system.initialize()
-        if not success:
-            print("❌ System initialization failed")
+    async def display_history(self):
+        """Display conversation history"""
+        if not self.chat_manager:
+            print("❌ System not initialized")
             return
-
-        print("\n🎯 INTERACTIVE MODE - SCALABLE SYSTEM")
-        print("="*60)
-        print("Commands:")
-        print("  'quit' - Exit")
-        print("  'stats' - Show statistics") 
-        print("  'health' - Health check")
-        print("  'countries' - List available countries")
-        print("  'session' - Session info")
-        print("="*60)
-        
-        session_id = f"interactive_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
-        print(f"Session ID: {session_id}")
-        print(f"Available: {', '.join(system.get_available_countries())}\n")
         
-        while True:
-            try:
-                user_input = input("👤 You: ").strip()
-                
-                if user_input.lower() in ['quit', 'exit', 'q']:
-                    break
-                elif user_input.lower() == 'stats':
-                    stats = system.get_global_stats()
-                    print(f"\n📊 Statistics:")
-                    print(f"  Total Queries: {stats.get('total_queries', 0)}")
-                    print(f"  Active Sessions: {stats.get('active_sessions', 0)}")
-                    print(f"  Pending Interrupts: {stats.get('pending_interrupts', 0)}")
-                    continue
-                elif user_input.lower() == 'health':
-                    health = await system.health_check()
-                    print(f"\n❤️  System Health:")
-                    print(f"  Status: {'✅ HEALTHY' if health['system_initialized'] else '❌ UNHEALTHY'}")
-                    print(f"  Countries: {len(health['available_countries'])} available")
-                    print(f"  MongoDB: {'✅ Connected' if health['mongodb_connected'] else '❌ Disconnected'}")
-                    continue
-                elif user_input.lower() == 'countries':
-                    countries = system.get_available_countries()
-                    print(f"\n🌍 Available Countries: {', '.join(countries) if countries else 'None'}")
-                    continue
-                elif user_input.lower() == 'session':
-                    info = system.get_session_info(session_id)
-                    print(f"\n📋 Session Info:")
-                    print(f"  Queries: {info.get('query_count', 0)}")
-                    print(f"  Avg Time: {info.get('average_processing_time', 0):.2f}s")
-                    continue
-                elif not user_input:
-                    continue
-                
-                start_time = time.time()
-                response = await system.chat(user_input, session_id)
-                response_time = time.time() - start_time
-                
-                print(f"🤖 Assistant ({response_time:.2f}s): {response}\n")
-                
-                # Check for interrupt
-                if (hasattr(system.chat_manager, 'pending_interrupts') and 
-                    session_id in system.chat_manager.pending_interrupts):
-                    print("⏸️  💡 SYSTEM PAUSED - Next message treated as moderator decision\n")
-                
-            except KeyboardInterrupt:
-                print("\n👋 Goodbye!")
-                break
-            except Exception as e:
-                print(f"❌ Error: {str(e)}\n")
-                
-    finally:
-        await system.cleanup()
-
+        try:
+            history = await self.chat_manager.get_conversation_history(self.session_id)
+            
+            if not history:
+                print("\n💬 No conversation history yet\n")
+                return
+            
+            print("\n" + "="*70)
+            print("💬 CONVERSATION HISTORY")
+            print("="*70)
+            
+            for i, msg in enumerate(history, 1):
+                role = "👤 User" if msg.type == "human" else "🤖 Assistant"
+                content = msg.content[:100] + "..." if len(msg.content) > 100 else msg.content
+                print(f"{i}. {role}: {content}")
+            
+            print("="*70 + "\n")
+            
+        except Exception as e:
+            logger.exception("Error displaying history")
+            print(f"❌ Error: {e}\n")
+    
+    def reset_session(self):
+        """Reset the current session"""
+        old_session = self.session_id
+        self.session_id = f"test_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+        self.message_count = 0
+        print(f"\n🔄 Session reset")
+        print(f"   Old: {old_session}")
+        print(f"   New: {self.session_id}\n")
+    
+    def clear_screen(self):
+        """Clear the terminal screen"""
+        print("\033[H\033[J", end="")
 
-async def health_check_mode():
-    """Run system health check only"""
-    system = MultiCountryLegalRAGSystem()
+async def interactive_mode():
+    """Run the system in interactive mode with full interrupt support"""
+    tester = LegalRAGTester()
     
-    try:
-        print("🔍 Performing health check...")
-        success = await system.initialize()
-        
-        if success:
-            health = await system.health_check()
-            print("\n" + "="*50)
-            print("📋 SYSTEM HEALTH REPORT")
-            print("="*50)
-            print(f"✅ System Initialized: {health['system_initialized']}")
-            print(f"🌍 Available Countries: {len(health['available_countries'])}")
-            print(f"💾 MongoDB: {'✅ Connected' if health['mongodb_connected'] else '❌ Disconnected'}")
-            print(f"⏸️  Interrupt Support: {'✅ Enabled' if health['interrupt_enabled'] else '❌ Disabled'}")
+    # Initialize system
+    if not await tester.initialize():
+        print("❌ Failed to initialize. Exiting.")
+        return
+    
+    # Display welcome message
+    print("\n" + "="*70)
+    print("🎯 LEGAL RAG INTERACTIVE TESTING MODE")
+    print("="*70)
+    print("Type 'help' for available commands")
+    print("="*70)
+    
+    # Quick test questions
+    test_questions = [
+        "Bonjour, comment ça va?",
+        "Quelle est la procédure de divorce au Bénin?",
+        "Je veux parler à un avocat spécialisé",
+        "Quels sont les droits des enfants à Madagascar?",
+        "Quelles sont les conditions pour se marier au Bénin?",
+    ]
+    
+    print("\n💡 Quick test questions:")
+    for i, question in enumerate(test_questions, 1):
+        print(f"  {i}. {question}")
+    print()
+    
+    # Main interaction loop
+    while True:
+        try:
+            # Get user input
+            user_input = input("👤 You: ").strip()
             
-            print(f"\n🔧 Components:")
-            for component, status in health['components'].items():
-                print(f"  {component}: {'✅ OK' if status else '❌ Missing'}")
-                
-            all_healthy = (health['system_initialized'] and 
-                          health['mongodb_connected'] and 
-                          all(health['components'].values()))
-            print(f"\n🎯 Overall Status: {'✅ HEALTHY' if all_healthy else '❌ UNHEALTHY'}")
+            if not user_input:
+                continue
             
-        else:
-            print("❌ System initialization failed")
+            # Handle commands
+            cmd_lower = user_input.lower()
             
-    finally:
-        await system.cleanup()
-
+            # Exit commands
+            if cmd_lower in ['quit', 'exit', 'q']:
+                print("\n👋 Goodbye! Thank you for testing.")
+                break
+            
+            # Stats command
+            elif cmd_lower in ['stats', 's']:
+                tester.display_stats()
+                continue
+            
+            # Clear screen
+            elif cmd_lower in ['clear', 'cls']:
+                tester.clear_screen()
+                continue
+            
+            # Help command
+            elif cmd_lower in ['help', 'h', '?']:
+                tester.display_help()
+                continue
+            
+            # History command
+            elif cmd_lower == 'history':
+                await tester.display_history()
+                continue
+            
+            # Reset session
+            elif cmd_lower == 'reset':
+                tester.reset_session()
+                continue
+            
+            # Regular chat message
+            else:
+                await tester.chat(user_input)
+            
+        except KeyboardInterrupt:
+            print("\n\n⚠️  Interrupted by user")
+            print("Type 'quit' to exit or continue chatting\n")
+            continue
+        
+        except Exception as e:
+            logger.exception("Unexpected error in main loop")
+            print(f"\n❌ Unexpected error: {e}\n")
 
-async def quick_test_mode():
-    """Run a quick single test"""
-    system = MultiCountryLegalRAGSystem()
+async def demo_mode():
+    """Run automated demo of the system capabilities"""
+    tester = LegalRAGTester()
     
-    try:
-        print("🚀 Quick Test Mode")
-        print("Initializing system...")
-        success = await system.initialize()
-        if not success:
-            print("❌ System initialization failed")
-            return
-
-        test_query = "Bonjour, quelle est la procedure pour un divorce au Benin?"
-        session_id = "quick_test"
-        
-        print(f"\n🧪 Testing: {test_query}")
-        start_time = time.time()
-        response = await system.chat(test_query, session_id)
-        response_time = time.time() - start_time
-        
-        print(f"✅ Response ({response_time:.2f}s): {response}")
-        
-        print(f"\n📊 System Info:")
-        print(f"  Available Countries: {', '.join(system.get_available_countries())}")
+    print("🚀 Running Automated Demo...")
+    
+    if not await tester.initialize():
+        print("❌ Failed to initialize. Exiting demo.")
+        return
+    
+    demo_scenarios = [
+        {
+            "name": "Greeting",
+            "messages": ["Bonjour, je m'appelle Test"]
+        },
+        {
+            "name": "Legal Query - Bénin",
+            "messages": ["Quelle est la procédure de divorce au Bénin?"]
+        },
+        {
+            "name": "Legal Query - Madagascar",
+            "messages": ["Quels sont les droits de succession à Madagascar?"]
+        },
+        {
+            "name": "Conversation Repair",
+            "messages": ["Peux-tu répéter plus simplement?"]
+        },
+        {
+            "name": "Summary Request",
+            "messages": ["Résume notre conversation"]
+        }
+    ]
+    
+    print("\n" + "="*70)
+    print("🧪 DEMO SCENARIOS")
+    print("="*70 + "\n")
+    
+    for i, scenario in enumerate(demo_scenarios, 1):
+        print(f"\n{'='*70}")
+        print(f"📋 Scenario {i}/{len(demo_scenarios)}: {scenario['name']}")
+        print('='*70)
         
-    except Exception as e:
-        print(f"❌ Quick test failed: {e}")
-    finally:
-        await system.cleanup()
-
+        for message in scenario['messages']:
+            print(f"\n🧪 Testing: '{message}'")
+            await tester.chat(message)
+            await asyncio.sleep(1.5)  # Pause between messages
+    
+    print("\n" + "="*70)
+    print("✅ DEMO COMPLETED")
+    print("="*70)
+    tester.display_stats()
+
+async def test_assistance_workflow():
+    """Test the complete assistance workflow with interrupts"""
+    tester = LegalRAGTester()
+    
+    print("🧪 Testing Assistance Workflow with Human Approval...")
+    
+    if not await tester.initialize():
+        print("❌ Failed to initialize. Exiting test.")
+        return
+    
+    print("\n" + "="*70)
+    print("🔬 ASSISTANCE WORKFLOW TEST")
+    print("="*70 + "\n")
+    
+    # Simulated assistance workflow
+    workflow_steps = [
+        "Je veux parler à un avocat",
+        "test@example.com",
+        "J'ai besoin d'aide pour un problème de divorce au Bénin",
+        "oui"
+    ]
+    
+    step_names = [
+        "1. Initiate assistance request",
+        "2. Provide email",
+        "3. Describe situation",
+        "4. Confirm request"
+    ]
+    
+    for step_name, message in zip(step_names, workflow_steps):
+        print(f"\n📍 {step_name}")
+        print(f"   Message: '{message}'")
+        await tester.chat(message)
+        await asyncio.sleep(1)
+    
+    print("\n" + "="*70)
+    print("✅ WORKFLOW TEST COMPLETED")
+    print("="*70)
+    tester.display_stats()
 
-if __name__ == "__main__":
+def main():
+    """Main entry point with argument parsing"""
     import argparse
     
     parser = argparse.ArgumentParser(
-        description="🚀 Scalable Multi-Country Legal RAG System"
+        description="Legal RAG System - Interactive Testing with Human-in-the-Loop Support",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Examples:
+  python main.py                    # Interactive mode (default)
+  python main.py --mode demo        # Automated demo
+  python main.py --mode test        # Test assistance workflow
+  python main.py --debug            # Enable debug logging
+        """
     )
     
     parser.add_argument(
         "--mode",
-        choices=["interactive", "health", "interrupt", "quick"],
+        choices=["interactive", "demo", "test"],
         default="interactive",
-        help="Run mode (default: interactive)"
+        help="Run mode: interactive (default), demo, or test"
+    )
+    
+    parser.add_argument(
+        "--debug",
+        action="store_true",
+        help="Enable debug logging"
     )
     
     args = parser.parse_args()
     
-    if args.mode == "interactive":
-        asyncio.run(interactive_mode())
-    elif args.mode == "health":
-        asyncio.run(health_check_mode())
-    elif args.mode == "interrupt":
-        asyncio.run(run_interrupt_tests())
-    elif args.mode == "quick":
-        asyncio.run(quick_test_mode())
+    # Adjust logging level if debug
+    if args.debug:
+        logging.getLogger().setLevel(logging.DEBUG)
+        print("🐛 Debug logging enabled\n")
+    
+    # Run selected mode
+    try:
+        if args.mode == "interactive":
+            asyncio.run(interactive_mode())
+        elif args.mode == "demo":
+            asyncio.run(demo_mode())
+        elif args.mode == "test":
+            asyncio.run(test_assistance_workflow())
+    
+    except KeyboardInterrupt:
+        print("\n\n👋 Program interrupted by user. Goodbye!")
+    
+    except Exception as e:
+        logger.exception("Fatal error occurred")
+        print(f"\n❌ Fatal error: {e}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    main()

requirements.txt

@@ -6,6 +6,7 @@
 # Web Framework
 fastapi==0.118.2
 uvicorn[standard]==0.37.0
+slowapi==0.1.9
 
 # LangChain & LangGraph - Core only (dependencies will pull others)
 langgraph==0.6.8