Spaces:

SAAHMATHWORKS
/

MultiCountryRAG

Sleeping

App Files Files Community

SAAHMATHWORKS commited on Oct 7

Commit

fbdfc24

1 Parent(s): 44d7656

Initial deployment: Legal Assistant application

Browse files

Files changed (43) hide show

.gitignore +10 -0
Dockerfile +30 -0
api/main.py +310 -0
api/requirements.txt +4 -0
app.py +7 -0
config/__init__.py +22 -0
config/constants.py +107 -0
config/settings.py +55 -0
core/assistance/__init__.py +5 -0
core/assistance/email_service.py +40 -0
core/assistance/workflow_nodes.py +186 -0
core/chat_manager.py +289 -0
core/conversation_repair.py +88 -0
core/email_tool.py +187 -0
core/graph_builder.py +266 -0
core/human_approval_node.py +222 -0
core/nodes/__init__.py +14 -0
core/nodes/base_node.py +79 -0
core/nodes/helper_nodes.py +147 -0
core/nodes/response_nodes.py +219 -0
core/nodes/retrieval_nodes.py +83 -0
core/nodes/routing_nodes.py +193 -0
core/prompts/__init__.py +3 -0
core/prompts/prompt_templates.py +94 -0
core/retriever.py +386 -0
core/router.py +238 -0
core/routing/__init__.py +3 -0
core/routing/routing_logic.py +158 -0
core/system_initializer.py +103 -0
database/__init__py +0 -0
database/mongodb_client.py +153 -0
database/postgres_checkpointer.py +97 -0
generate_graph.py +66 -0
interfaces/__init__.py +0 -0
interfaces/monitoring.py +109 -0
interfaces/web_interface.py +96 -0
main.py +629 -0
models/__init__py +0 -0
models/state_models.py +112 -0
requirements.txt +13 -0
utils/__init__.py +0 -0
utils/helpers.py +68 -0
utils/logger.py +74 -0

.gitignore ADDED Viewed

	@@ -0,0 +1,10 @@

+__pycache__/
+*.pyc
+*.log
+*.tmp
+.cache/
+.DS_Store
+*.pkl
+*.bin
+*.ipynb
+.ipynb_checkpoints/

Dockerfile ADDED Viewed

	@@ -0,0 +1,30 @@

+FROM python:3.11-slim
+WORKDIR /app
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    gcc \
+    curl \
+    && rm -rf /var/lib/apt/lists/*
+# Copy requirements and install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+# Copy application code
+COPY . .
+# Create non-root user
+RUN useradd -m -u 1000 user
+USER user
+# Expose port (Hugging Face uses 7860)
+EXPOSE 7860
+# Health check
+HEALTHCHECK --interval=30s --timeout=30s --start-period=5s --retries=3 \
+    CMD curl -f http://localhost:7860/health || exit 1
+# Start command
+CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "7860"]

api/main.py ADDED Viewed

	@@ -0,0 +1,310 @@

+# api/main.py
+from typing import Optional
+from contextlib import asynccontextmanager
+from fastapi import FastAPI, Query, HTTPException
+from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi.middleware.cors import CORSMiddleware
+from langchain_core.messages import AIMessageChunk
+import json
+from uuid import uuid4
+import logging
+import os
+import asyncio
+# Import your existing system
+from core.system_initializer import setup_system
+from models.state_models import MultiCountryLegalState
+# Setup logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+# Global variables
+chat_manager = None
+graph = None
+system_initialized = False
+async def initialize_system():
+    global chat_manager, graph, system_initialized
+    try:
+        # Check for required environment variables based on YOUR settings
+        required_vars = ['OPENAI_API_KEY', 'MONGO_URI', 'NEON_DB_URL', 'NEON_END_POINT']
+        missing_vars = [var for var in required_vars if not os.getenv(var)]
+        if missing_vars:
+            logger.warning(f"⚠️ Missing environment variables: {missing_vars}")
+            logger.warning("System will start but may not function properly")
+        system = await setup_system()
+        chat_manager = system["chat_manager"]
+        graph = system["graph"]
+        system_initialized = True
+        logger.info("✅ Legal assistant system initialized for Hugging Face")
+    except Exception as e:
+        logger.error(f"❌ Failed to initialize system: {e}")
+        system_initialized = False
+@asynccontextmanager
+async def lifespan(app: FastAPI):
+    """Modern lifespan event handler"""
+    # Startup logic
+    logger.info("🚀 Starting Legal Assistant API...")
+    # Initialize system in background
+    initialization_task = asyncio.create_task(initialize_system())
+    yield  # App runs here
+    # Shutdown logic
+    logger.info("🛑 Shutting down Legal Assistant API...")
+    initialization_task.cancel()
+    try:
+        await initialization_task
+    except asyncio.CancelledError:
+        pass
+app = FastAPI(
+    title="Legal Assistant API",
+    version="1.0.0",
+    description="Multi-country legal RAG system for Benin and Madagascar",
+    docs_url="/docs",
+    redoc_url="/redoc",
+    lifespan=lifespan
+)
+# Add CORS middleware
+app.add_middleware(
+    CORSMiddleware,
+    allow_origins=["*"],
+    allow_credentials=True,
+    allow_methods=["*"],
+    allow_headers=["*"],
+)
+@app.get("/", response_class=HTMLResponse)
+async def read_root():
+    """Simple homepage for better UX"""
+    return """
+    <html>
+        <head>
+            <title>Legal Assistant API</title>
+            <style>
+                body { font-family: Arial, sans-serif; margin: 40px; }
+                .container { max-width: 800px; margin: 0 auto; }
+                .card { border: 1px solid #ddd; padding: 20px; margin: 10px 0; border-radius: 8px; }
+                .status-ready { color: green; }
+                .status-starting { color: orange; }
+                .status-error { color: red; }
+            </style>
+        </head>
+        <body>
+            <div class="container">
+                <h1>🧑‍⚖️ Legal Assistant API</h1>
+                <p>Multi-country legal RAG system for Benin and Madagascar</p>
+                <div class="card">
+                    <h3>📚 Available Endpoints</h3>
+                    <ul>
+                        <li><a href="/docs">API Documentation</a></li>
+                        <li><a href="/health">Health Check</a></li>
+                        <li><strong>GET /chat</strong> - Streaming chat</li>
+                        <li><strong>GET /sessions/{id}/history</strong> - Conversation history</li>
+                    </ul>
+                </div>
+                <div class="card">
+                    <h3>🔧 System Status</h3>
+                    <div id="status">
+                        <p>Loading system status...</p>
+                    </div>
+                </div>
+                <script>
+                    async function updateStatus() {
+                        try {
+                            const response = await fetch('/health');
+                            const data = await response.json();
+                            const statusEl = document.getElementById('status');
+                            let statusClass = 'status-starting';
+                            let statusText = '🔄 Starting...';
+                            if (data.system_initialized) {
+                                statusClass = 'status-ready';
+                                statusText = '✅ System Ready';
+                            } else if (data.status === 'error') {
+                                statusClass = 'status-error';
+                                statusText = '❌ System Error';
+                            }
+                            statusEl.innerHTML = `
+                                <p class="${statusClass}"><strong>${statusText}</strong></p>
+                                <p><strong>MongoDB:</strong> ${data.mongodb_connected ? '✅ Connected' : '❌ Disconnected'}</p>
+                                <p><strong>Countries:</strong> ${data.available_countries?.join(', ') || 'Loading...'}</p>
+                                <p><strong>OpenAI:</strong> ${data.openai_configured ? '✅ Configured' : '❌ Not Configured'}</p>
+                            `;
+                        } catch (error) {
+                            document.getElementById('status').innerHTML =
+                                '<p class="status-error">❌ Failed to load system status</p>';
+                        }
+                    }
+                    updateStatus();
+                    setInterval(updateStatus, 5000);
+                </script>
+            </div>
+        </body>
+    </html>
+    """
+@app.get("/health")
+async def health_check():
+    """Enhanced health check with your specific environment variables"""
+    return {
+        "status": "healthy" if system_initialized else "starting",
+        "system_initialized": system_initialized,
+        "service": "Legal Assistant API",
+        "available_countries": ["benin", "madagascar"] if system_initialized else [],
+        "mongodb_connected": system_initialized and bool(os.getenv("MONGO_URI")),
+        "openai_configured": bool(os.getenv("OPENAI_API_KEY")),
+        "neon_postgres_configured": bool(os.getenv("NEON_END_POINT")),
+        "missing_variables": [var for var in ['OPENAI_API_KEY', 'MONGO_URI', 'NEON_DB_URL', 'NEON_END_POINT'] if not os.getenv(var)],
+    }
+def serialize_ai_message_chunk(chunk):
+    """Serialize AI message chunks for streaming"""
+    if isinstance(chunk, AIMessageChunk):
+        return chunk.content
+    else:
+        raise TypeError(
+            f"Object of type {type(chunk).__name__} is not correctly formatted for serialisation"
+        )
+async def generate_legal_chat_responses(message: str, session_id: Optional[str] = None):
+    """Generate streaming responses for legal chat"""
+    if not system_initialized:
+        yield f"data: {json.dumps({'type': 'error', 'message': 'System is still starting up. Please try again in a moment.'})}\n\n"
+        yield f"data: {json.dumps({'type': 'end'})}\n\n"
+        return
+    is_new_conversation = session_id is None
+    if is_new_conversation:
+        session_id = f"api_{uuid4()}"
+        logger.info(f"🆕 New conversation session: {session_id}")
+        yield f"data: {json.dumps({'type': 'session', 'session_id': session_id})}\n\n"
+    else:
+        logger.info(f"🔄 Continuing session: {session_id}")
+    try:
+        input_state = {
+            "messages": [{"role": "user", "content": message, "meta": {}}],
+            "legal_context": {
+                "jurisdiction": "Unknown",
+                "user_type": "general",
+                "document_type": "legal",
+                "detected_country": "unknown"
+            },
+            "session_id": session_id,
+            "router_decision": None,
+            "search_results": None,
+            "route_explanation": None,
+            "last_search_query": None,
+            "detected_articles": [],
+        }
+        config = {
+            "configurable": {
+                "thread_id": session_id
+            }
+        }
+        events = graph.astream_events(
+            MultiCountryLegalState(**input_state),
+            version="v2",
+            config=config
+        )
+        current_content = ""
+        current_node = ""
+        async for event in events:
+            event_type = event["event"]
+            node_name = event.get("name", "")
+            if node_name != current_node:
+                current_node = node_name
+                yield f"data: {json.dumps({'type': 'node_transition', 'node': node_name})}\n\n"
+            if event_type == "on_chat_model_stream":
+                chunk_content = serialize_ai_message_chunk(event["data"]["chunk"])
+                current_content += chunk_content
+                yield f"data: {json.dumps({'type': 'content', 'content': chunk_content})}\n\n"
+            elif event_type == "on_chat_model_end":
+                yield f"data: {json.dumps({'type': 'content_end'})}\n\n"
+            elif event_type == "on_chain_start" and "retrieval" in node_name:
+                country = node_name.replace("_retrieval", "")
+                yield f"data: {json.dumps({'type': 'search_start', 'country': country})}\n\n"
+            elif event_type == "on_chain_end" and "retrieval" in node_name:
+                country = node_name.replace("_retrieval", "")
+                yield f"data: {json.dumps({'type': 'search_end', 'country': country})}\n\n"
+            elif event_type == "on_tool_end":
+                tool_name = event["name"]
+                yield f"data: {json.dumps({'type': 'tool_complete', 'tool': tool_name})}\n\n"
+            elif event_type == "on_graph_end":
+                yield f"data: {json.dumps({'type': 'graph_end'})}\n\n"
+    except Exception as e:
+        logger.error(f"Error in streaming: {e}")
+        yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
+    yield f"data: {json.dumps({'type': 'end'})}\n\n"
+@app.get("/chat")
+async def chat_stream(
+    message: str = Query(..., description="User message"),
+    session_id: Optional[str] = Query(None, description="Existing session ID")
+):
+    """Streaming chat endpoint with initialization check"""
+    if not system_initialized:
+        raise HTTPException(
+            status_code=503,
+            detail="System is still starting up. Please try again in a moment."
+        )
+    return StreamingResponse(
+        generate_legal_chat_responses(message, session_id),
+        media_type="text/event-stream",
+        headers={
+            "Cache-Control": "no-cache",
+            "Connection": "keep-alive",
+        }
+    )
+@app.get("/sessions/{session_id}/history")
+async def get_conversation_history(session_id: str):
+    """Get conversation history for a session"""
+    if not chat_manager:
+        return {"error": "System not initialized"}
+    try:
+        history = await chat_manager.get_conversation_history(session_id)
+        return {
+            "session_id": session_id,
+            "history": [
+                {
+                    "role": msg.role if hasattr(msg, 'role') else msg.get('role', 'unknown'),
+                    "content": msg.content if hasattr(msg, 'content') else msg.get('content', ''),
+                    "timestamp": getattr(msg, 'timestamp', None)
+                }
+                for msg in history
+            ]
+        }
+    except Exception as e:
+        return {"error": str(e)}

api/requirements.txt ADDED Viewed

	@@ -0,0 +1,4 @@

+fastapi==0.104.1
+uvicorn==0.24.0
+python-multipart==0.0.6
+pydantic==2.5.0

app.py ADDED Viewed

	@@ -0,0 +1,7 @@

+# app.py - Main entry point for Hugging Face Spaces
+from api.main import app
+# Hugging Face Spaces will automatically use this 'app' variable
+if __name__ == "__main__":
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)  # HF uses port 7860

config/__init__.py ADDED Viewed

	@@ -0,0 +1,22 @@

+# config/__init__.py
+from .settings import settings
+from .constants import (
+    COUNTRY_PATTERNS,
+    ARTICLE_PATTERNS,
+    CATEGORY_KEYWORDS,
+    DOCUMENT_TYPE_KEYWORDS,
+    DOCUMENT_TYPE_DESCRIPTIONS,
+    LEGAL_CONTEXTS,
+    USER_TYPE_CONTEXTS
+)
+__all__ = [
+    'settings',
+    'COUNTRY_PATTERNS',
+    'ARTICLE_PATTERNS',
+    'CATEGORY_KEYWORDS',
+    'DOCUMENT_TYPE_KEYWORDS',
+    'DOCUMENT_TYPE_DESCRIPTIONS',
+    'LEGAL_CONTEXTS',
+    'USER_TYPE_CONTEXTS'
+]

config/constants.py ADDED Viewed

	@@ -0,0 +1,107 @@

+from typing import Dict, List
+# Country patterns for routing
+COUNTRY_PATTERNS = {
+    "benin": [
+        r"\bbénin\b", r"\bbeninois\b", r"\bbéninoise\b", r"\bbenin\b",
+        r"\bdahomey\b", r"\bporto-novo\b", r"\bcotonou\b",
+        r"\bdroit béninois\b", r"\bloi béninoise\b"
+    ],
+    "madagascar": [
+        r"\bmadagascar\b", r"\bmalgache\b", r"\bmalagasy\b",
+        r"\bantananarivo\b", r"\bmadagasikara\b",
+        r"\bdroit malgache\b", r"\bloi malgache\b"
+    ]
+}
+# Article detection patterns
+ARTICLE_PATTERNS = [
+    r"article[s]?\s+(\d+(?:\s+(?:et|à|\-)\s+\d+)*)",
+    r"art\.?\s*(\d+(?:\s+(?:et|à|\-)\s+\d+)*)",
+    r"articles?\s+(\d+)\s*à\s*(\d+)",
+    r"art\.?\s*(\d+)\s*au\s*(\d+)",
+]
+# Legal domain categories
+CATEGORY_KEYWORDS = {
+    "mariage": "Code des personnes et de la famille",
+    "divorce": "Code des personnes et de la famille",
+    "héritage": "Code des personnes et de la famille",
+    "succession": "Code des personnes et de la famille",
+    "adoption": "Code des personnes et de la famille",
+    "enfant": "Code des personnes et de la famille",
+    "pension": "Code des personnes et de la famille",
+    "infraction": "Droit pénal",
+    "délit": "Droit pénal",
+    "crime": "Droit pénal",
+    "peine": "Droit pénal",
+    "prison": "Droit pénal",
+    "entreprise": "Droit commercial",
+    "commerce": "Droit commercial",
+    "contrat": "Droit commercial",
+    "société": "Droit commercial",
+    "administration": "Droit administratif",
+    "fonctionnaire": "Droit administratif",
+    "service public": "Droit administratif"
+}
+# Document type detection keywords
+DOCUMENT_TYPE_KEYWORDS = {
+    "case_study": [
+        "jurisprudence", "arrêt", "décision", "tribunal", "cours", "jugement",
+        "affaire", "procès", "litige", "contentieux", "précédent", "cas",
+        "cour d'appel", "cour suprême", "conseil d'état"
+    ],
+    "articles": [
+        "article", "loi", "code", "décret", "texte", "disposition",
+        "règlement", "ordonnance", "prescription", "norme"
+    ]
+}
+# Document type descriptions
+DOCUMENT_TYPE_DESCRIPTIONS = {
+    "articles": "Textes législatifs et réglementaires (lois, codes, décrets)",
+    "case_study": "Jurisprudence et décisions de justice (arrêts, jugements)"
+}
+# Legal context templates
+LEGAL_CONTEXTS = {
+    "benin": {
+        "jurisdiction": "Bénin",
+        "user_type": "citizen",
+        "document_type": "Code des personnes et de la famille",
+        "language": "français",
+        "legal_system": "civil_law"
+    },
+    "madagascar": {
+        "jurisdiction": "Madagascar",
+        "user_type": "citizen",
+        "document_type": "legal",
+        "language": "français",
+        "legal_system": "mixed_civil_customary"
+    }
+}
+# User type contexts
+USER_TYPE_CONTEXTS = {
+    "citizen": {
+        "expertise_level": "basic",
+        "response_style": "accessible",
+        "include_procedures": True
+    },
+    "lawyer": {
+        "expertise_level": "advanced",
+        "response_style": "technical",
+        "include_precedents": True
+    },
+    "student": {
+        "expertise_level": "intermediate",
+        "response_style": "educational",
+        "include_examples": True
+    }
+}
+# LAW_KEYWORDS a été supprimé comme demandé - le filtre "titre" n'est plus utilisé

config/settings.py ADDED Viewed

	@@ -0,0 +1,55 @@

+import os
+from dotenv import load_dotenv
+# Change to:
+try:
+    load_dotenv("../.env", override=True)
+except:
+    pass  # Ignore if .env file doesn't exist (like on Hugging Face)
+class Settings:
+    # API Keys
+    OPENAI_API_KEY = os.getenv("OPENAI_API_KEY")
+    MONGO_URI = os.environ.get("MONGO_URI")
+    NEON_DB_URL = os.environ.get("NEON_DB_URL")
+    NEON_END_POINT = os.getenv("NEON_END_POINT")
+    # Database
+    DATABASE_URL = NEON_END_POINT
+    # Model Configurations
+    EMBEDDING_MODEL = "text-embedding-ada-002"
+    CHAT_MODEL = "gpt-4o-mini"
+    CHAT_MODEL_2 = "gpt-3.5-turbo"
+    CHAT_TEMPERATURE = 0.1
+    CHAT_MAX_TOKENS = 2000
+    # Vector Search
+    VECTOR_INDEX_NAME = "vector_index"
+    TEXT_KEY = "contenu"
+    EMBEDDING_KEY = "vecteur_embedding"
+    # Collections
+    BENIN_COLLECTION = "legal_documents"
+    MADAGASCAR_COLLECTION = "legal_documents_madagascar"
+    DATABASE_NAME = "legal_db"
+    # Search Parameters
+    MAX_SEARCH_RESULTS = 10
+    MAX_CONVERSATION_HISTORY = 8
+    def validate(self):
+        missing = []
+        if not self.OPENAI_API_KEY:
+            missing.append("OPENAI_API_KEY")
+        if not self.MONGO_URI:
+            missing.append("MONGO_URI")
+        if not self.NEON_DB_URL:
+            missing.append("NEON_DB_URL")
+        if not self.NEON_END_POINT:
+            missing.append("NEON_END_POINT")
+        if missing:
+            raise ValueError(f"Missing environment variables: {', '.join(missing)}")
+settings = Settings()

core/assistance/__init__.py ADDED Viewed

	@@ -0,0 +1,5 @@

+# [file name]: core/assistance/__init__.py
+from .workflow_nodes import AssistanceWorkflowNodes
+from .email_service import AssistanceEmailService
+__all__ = ["AssistanceWorkflowNodes", "AssistanceEmailService"]

core/assistance/email_service.py ADDED Viewed

	@@ -0,0 +1,40 @@

+# [file name]: core/assistance/email_service.py
+"""
+Wrapper for email functionality - provides a consistent interface
+"""
+import re
+import logging
+from typing import Optional, Dict
+from core.email_tool import LegalAssistanceEmailer
+logger = logging.getLogger(__name__)
+class AssistanceEmailService:
+    """Service wrapper for email operations"""
+    def __init__(self):
+        self.emailer = LegalAssistanceEmailer()
+    def extract_email_from_text(self, text: str) -> Optional[str]:
+        """Extract email from text"""
+        return self.emailer.extract_email_from_text(text)
+    def validate_email(self, email: str) -> bool:
+        """Validate email format"""
+        return self.emailer.validate_email(email)
+    def send_assistance_request(
+        self,
+        user_email: str,
+        user_query: str,
+        assistance_description: str,
+        country: str
+    ) -> Dict:
+        """Send assistance request emails"""
+        return self.emailer.send_assistance_request(
+            user_email=user_email,
+            user_query=user_query,
+            assistance_description=assistance_description,
+            country=country
+        )

core/assistance/workflow_nodes.py ADDED Viewed

	@@ -0,0 +1,186 @@

+# core/assistance/workflow_nodes.py
+import logging
+import re
+from typing import Dict, Any, List
+from langchain_core.runnables import RunnableConfig
+from models.state_models import MultiCountryLegalState
+logger = logging.getLogger(__name__)
+class AssistanceWorkflowNodes:
+    def __init__(self):
+        self.email_pattern = re.compile(r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$')
+    async def collect_assistance_info_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Collect assistance information (email, description)"""
+        s = state.model_dump()
+        assistance_step = s.get("assistance_step", "collecting_email")
+        user_input = s.get("messages", [{}])[-1].get("content", "") if s.get("messages") else ""
+        logger.info(f"📝 Collecting assistance info - step: {assistance_step}")
+        logger.debug(f"User input: {user_input}")
+        # 🔥 NEW: Check for cancellation commands
+        cancellation_keywords = ["annuler", "cancel", "stop", "arrêter", "je ne veux plus", "plus besoin", "abandonner"]
+        if any(keyword in user_input for keyword in cancellation_keywords):
+            logger.info("🚫 User requested cancellation of assistance workflow")
+            return {
+                "assistance_step": "cancelled",
+                "assistance_requested": False,
+                "user_email": None,
+                "assistance_description": None,
+                "messages": [{
+                    "role": "assistant",
+                    "content": "✅ Votre demande d'assistance a été annulée. Comment puis-je vous aider autrement ?",
+                    "meta": {"assistance_cancelled": True}
+                }]
+            }
+        if assistance_step == "collecting_email":
+            if not user_input:
+                logger.info(f"ℹ️  Waiting for email input")
+                return {
+                    "assistance_step": "collecting_email",
+                    "messages": []  # Response node will generate the message
+                }
+            if self.email_pattern.match(user_input):
+                logger.info(f"📧 Email collected: {user_input}")
+                return {
+                    "assistance_step": "collecting_description",
+                    "user_email": user_input,
+                    "assistance_requested": True,
+                    "messages": []  # Response node will generate the message
+                }
+            else:
+                logger.warning(f"Invalid email: {user_input}")
+                return {
+                    "assistance_step": "collecting_email",
+                    "messages": [{
+                        "role": "assistant",
+                        "content": """⚠️ L'adresse email fournie semble invalide. Veuillez fournir une adresse email valide.
+📧 **Veuillez me fournir votre adresse email :**""",
+                        "meta": {"assistance_step": "collecting_email"}
+                    }]
+                }
+        elif assistance_step == "collecting_description":
+            if not user_input or len(user_input.strip()) < 10:
+                logger.info(f"ℹ️  Waiting for description input")
+                return {
+                    "assistance_step": "collecting_description",
+                    "messages": []  # Response node will generate the message
+                }
+            # Detect country from the description
+            detected_country = MultiCountryLegalState.detect_country(user_input)
+            logger.info(f"📝 Description collected: {user_input[:50]}...")
+            logger.info(f"🌍 Detected country: {detected_country}")
+            # Return the update - move to confirmation step
+            return {
+                "assistance_description": user_input,
+                "assistance_step": "confirming_send",
+                "country": detected_country,
+                "legal_context": {
+                    **state.legal_context,
+                    "detected_country": detected_country
+                },
+                "messages": []  # Response node will generate the confirmation message
+            }
+        return {}
+    async def confirm_assistance_send_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Confirm assistance request before sending to legal team"""
+        s = state.model_dump()
+        user_input = s.get("messages", [{}])[-1].get("content", "").lower().strip() if s.get("messages") else ""
+        logger.info(f"✅ Confirmation node - user input: {user_input[:30]}...")
+        if user_input in ["oui", "yes", "ok", "confirmer"]:
+            logger.info("✅ User confirmed assistance request")
+            return {
+                "assistance_step": "confirmed",
+                "messages": []  # Let response node or approval node handle the message
+            }
+        elif user_input in ["non", "no", "cancel", "annuler"]:
+            logger.info("❌ User cancelled assistance request")
+            return {
+                "assistance_step": "cancelled",
+                "assistance_requested": False,
+                "messages": [{
+                    "role": "assistant",
+                    "content": """❌ Votre demande a été annulée.
+Si vous changez d'avis, vous pouvez relancer une demande en disant "Je veux parler à un avocat".""",
+                    "meta": {"assistance_step": "cancelled"}
+                }]
+            }
+        else:
+            logger.info("ℹ️  Awaiting valid confirmation")
+            return {
+                "assistance_step": "confirming_send",
+                "messages": [{
+                    "role": "assistant",
+                    "content": f"""⚠️ Veuillez confirmer avec "oui" ou "non".
+📋 **RÉCAPITULATIF DE VOTRE DEMANDE :**
+📧 **Email :** {s.get("user_email")}
+📝 **Description :** {s.get("assistance_description")}
+✅ **Confirmez-vous l'envoi de cette demande à notre équipe juridique ?**
+Répondez par :
+- **"oui"** pour confirmer et envoyer
+- **"non"** pour annuler et modifier""",
+                    "meta": {"assistance_step": "confirming_send"}
+                }]
+            }
+    def route_assistance(self, state: MultiCountryLegalState) -> str:
+        """Route assistance workflow based on current state"""
+        s = state.model_dump()
+        assistance_step = s.get("assistance_step", "collecting_email")
+        logger.info(f"📋 Assistance step: {assistance_step}")
+        logger.info(f"   - Has email: {s.get('user_email') is not None} ({s.get('user_email')})")
+        logger.info(f"   - Has description: {s.get('assistance_description') is not None} ({s.get('assistance_description')})")
+        if assistance_step == "collecting_email" and not s.get("user_email"):
+            logger.info("→ Routing to: need_email (waiting for email)")
+            return "need_email"
+        elif assistance_step == "collecting_description" and not s.get("assistance_description"):
+            logger.info("→ Routing to: need_description (waiting for description)")
+            return "need_description"
+        elif assistance_step == "confirming_send" and s.get("user_email") and s.get("assistance_description"):
+            logger.info("→ Routing to: ready_to_confirm (awaiting user confirmation)")
+            return "ready_to_confirm"
+        elif assistance_step == "cancelled":
+            logger.info("→ Routing to: cancelled")
+            return "cancelled"
+        logger.info("→ Routing to: need_email (default)")
+        return "need_email"
+    def route_after_confirmation(self, state: MultiCountryLegalState) -> str:
+        """Route after confirmation step"""
+        s = state.model_dump()
+        assistance_step = s.get("assistance_step")
+        last_message = s.get("messages", [{}])[-1] if s.get("messages") else {}
+        user_input = last_message.get("content", "").lower().strip() if last_message.get("role") == "user" else ""
+        logger.info(f"📋 Confirmation step: {assistance_step}")
+        logger.info(f"   - Last user message: '{user_input}'")
+        if assistance_step == "confirmed":
+            logger.info("→ Routing to: confirmed (human approval)")
+            return "confirmed"
+        elif assistance_step == "cancelled":
+            logger.info("→ Routing to: cancelled")
+            return "cancelled"
+        else:
+            logger.info("→ Routing to: needs_correction (need clarification)")
+            return "needs_correction"

core/chat_manager.py ADDED Viewed

	@@ -0,0 +1,289 @@

+# [file name]: core/chat_manager.py
+import asyncio
+import logging
+from datetime import datetime
+from typing import Dict, List, Optional
+from langchain_core.runnables import RunnableConfig
+from langchain_core.messages import BaseMessage
+from langgraph.types import Command
+from config.settings import settings
+from models.state_models import MultiCountryLegalState
+from utils.helpers import dict_to_message_obj
+logger = logging.getLogger(__name__)
+class LegalChatManager:
+    def __init__(self, graph, checkpointer):
+        self.graph = graph
+        self.checkpointer = checkpointer
+        self.active_sessions = {}
+        self.routing_stats = {
+            "benin": 0,
+            "madagascar": 0,
+            "unclear": 0,
+            "total_queries": 0
+        }
+        # 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"""
+        if not self.graph:
+            raise RuntimeError("System not initialized. Call setup_system() first.")
+        # Initialize or update session
+        self._initialize_session(session_id)
+        # Check if we have a pending interrupt for this session
+        if session_id in self.pending_interrupts:
+            return await self._handle_pending_interrupt(session_id, message)
+        # Prepare input state
+        input_state = self._prepare_input_state(message, session_id, legal_context)
+        config = RunnableConfig(
+            configurable={"thread_id": session_id},
+            recursion_limit=100
+        )
+        try:
+            # Track performance
+            start_time = datetime.now()
+            # Process through graph
+            result = await self.graph.ainvoke(MultiCountryLegalState(**input_state), config)
+            # 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)
+            # Track performance
+            processing_time = (datetime.now() - start_time).total_seconds()
+            self._update_session_stats(session_id, processing_time)
+            # Extract and return response
+            response = self._extract_response(result)
+            self._update_routing_stats(response)
+            return response
+        except Exception as e:
+            logger.exception(f"Chat error for session {session_id}")
+            self._log_error(session_id, str(e))
+            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=...)"""
+        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}")
+            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(
+                Command(resume=message),
+                config
+            )
+            # 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)
+            logger.info(f"✅ Graph resumed successfully for session {session_id}")
+            return response
+        except Exception as e:
+            logger.error(f"Error resuming from interrupt: {str(e)}")
+            # Clean up on error
+            if session_id in self.pending_interrupts:
+                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"""
+🔒 **APPROBATION HUMAINE REQUISE**
+📧 **Utilisateur**: {user_email}
+🌍 **Pays**: {country}
+📝 **Description**: {description}
+**Veuillez répondre avec:**
+- "approve [raison]" pour approuver la demande
+- "reject [raison]" pour rejeter la demande
+**Exemples:**
+- "approve Demande légitime de consultation"
+- "reject Email invalide ou description trop vague"
+**Votre décision:**
+"""
+    # === EXISTING METHODS (unchanged) ===
+    async def get_conversation_history(self, session_id: str) -> List[BaseMessage]:
+        """Get conversation history for a session"""
+        if not self.graph:
+            return []
+        config = RunnableConfig(configurable={"thread_id": session_id})
+        try:
+            state = await self.graph.aget_state(config)
+            if not state or not state.values:
+                return []
+            s = state.values
+            if isinstance(s, MultiCountryLegalState):
+                s = s.model_dump()
+            elif isinstance(s, dict):
+                pass
+            else:
+                s = {}
+            raw_messages = s.get("messages", [])
+            return [dict_to_message_obj(m) for m in raw_messages if isinstance(m, dict)]
+        except Exception as e:
+            logger.exception(f"Error getting conversation history for session {session_id}")
+            return []
+    def get_session_stats(self, session_id: str) -> Dict:
+        """Get statistics for a specific session"""
+        return self.active_sessions.get(session_id, {})
+    def get_global_stats(self) -> Dict:
+        """Get global system statistics"""
+        return {
+            "routing_stats": self.routing_stats,
+            "active_sessions": len(self.active_sessions),
+            "total_queries": self.routing_stats["total_queries"],
+            "pending_interrupts": len(self.pending_interrupts)
+        }
+    def _initialize_session(self, session_id: str):
+        """Initialize or update session tracking"""
+        if session_id not in self.active_sessions:
+            self.active_sessions[session_id] = {
+                "created": datetime.now(),
+                "query_count": 0,
+                "total_processing_time": 0,
+                "average_processing_time": 0,
+                "detected_countries": set(),
+                "last_activity": datetime.now()
+            }
+        session_info = self.active_sessions[session_id]
+        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"""
+        ctx = legal_context or {
+            "jurisdiction": "Unknown",
+            "user_type": "general",
+            "document_type": "legal",
+            "detected_country": "unknown"
+        }
+        if ctx.get("detected_country") is None:
+            ctx["detected_country"] = "unknown"
+        return {
+            "messages": [{"role": "user", "content": message, "meta": {}}],
+            "legal_context": ctx,
+            "session_id": session_id,
+            "router_decision": None,
+            "search_results": None,
+            "route_explanation": None,
+            "last_search_query": None,
+            "detected_articles": [],
+        }
+    def _extract_response(self, result) -> str:
+        """Extract response text from graph result"""
+        if isinstance(result, MultiCountryLegalState):
+            r = result.model_dump()
+        elif isinstance(result, dict):
+            r = result
+        else:
+            r = {}
+        msgs = r.get("messages", [])
+        for m in reversed(msgs):
+            if (m.get("role") or "").lower() in ("assistant", "ai"):
+                return m.get("content", "")
+        return "Désolé, je n'ai pas pu générer de réponse."
+    def _update_session_stats(self, session_id: str, processing_time: float):
+        """Update session statistics with processing time"""
+        if session_id in self.active_sessions:
+            session_info = self.active_sessions[session_id]
+            session_info["total_processing_time"] += processing_time
+            session_info["average_processing_time"] = (
+                session_info["total_processing_time"] / session_info["query_count"]
+            )
+    def _update_routing_stats(self, response: str):
+        """Update routing statistics based on response content"""
+        self.routing_stats["total_queries"] += 1
+        response_lower = response.lower()
+        if any(keyword in response_lower for keyword in ["bénin", "béninois", "béninoise"]):
+            self.routing_stats["benin"] += 1
+        elif any(keyword in response_lower for keyword in ["madagascar", "malgache", "malagasy"]):
+            self.routing_stats["madagascar"] += 1
+        else:
+            self.routing_stats["unclear"] += 1
+    def _log_error(self, session_id: str, error: str):
+        """Log error for monitoring"""
+        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"""
+        cutoff_time = datetime.now().timestamp() - (max_age_hours * 3600)
+        inactive_sessions = [
+            session_id for session_id, info in self.active_sessions.items()
+            if info["last_activity"].timestamp() < cutoff_time
+        ]
+        # Also clean up pending interrupts for inactive sessions
+        for session_id in inactive_sessions:
+            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}")

core/conversation_repair.py ADDED Viewed

	@@ -0,0 +1,88 @@

+# [file name]: core/conversation_repair.py
+import logging
+from typing import Dict, List, Optional, Any
+import re
+logger = logging.getLogger(__name__)
+class ConversationRepair:
+    def __init__(self):
+        self.meta_keywords = [
+            "pas compris", "mal compris", "reformuler", "autrement",
+            "différemment", "répéter", "redire", "expliquer autrement",
+            "plus simple", "plus clair", "clarifier", "précisez",
+            "explique mieux", "développe", "approfondis", "que veux-tu dire",
+            "c'est-à-dire", "concrètement", "en pratique", "recommence",
+            "ce n'est pas ça", "tu n'as pas compris", "erreur", "faux"
+        ]
+    def detect_repair_intent(self, query: str, conversation_history: List[Dict]) -> bool:
+        """Simple detection - just check if this is a repair request"""
+        query_lower = query.lower()
+        return any(keyword in query_lower for keyword in self.meta_keywords)
+    async def generate_repair_response(self, query: str, conversation_history: List[Dict], llm) -> str:
+        """Unified LLM-powered repair handling"""
+        try:
+            # Build conversation context
+            context = self._build_conversation_context(conversation_history)
+            repair_prompt = self._build_repair_prompt(query, context)
+            # Use LLM for intelligent repair response
+            from langchain_core.messages import HumanMessage
+            response = await llm.ainvoke([HumanMessage(content=repair_prompt)])
+            return response.content if hasattr(response, 'content') else str(response)
+        except Exception as e:
+            logger.error(f"LLM repair generation failed: {e}")
+            return self._generate_fallback_response()
+    def _build_conversation_context(self, conversation_history: List[Dict]) -> str:
+        """Build conversation context for LLM"""
+        if not conversation_history:
+            return "Aucun contexte de conversation"
+        # Get relevant conversation history
+        relevant_messages = conversation_history[-6:]  # Last 6 messages
+        context_lines = []
+        for msg in relevant_messages:
+            role = "Utilisateur" if msg.get("role") == "user" else "Assistant"
+            content = msg.get("content", "")
+            context_lines.append(f"{role}: {content}")
+        return "\n".join(context_lines)
+    def _build_repair_prompt(self, current_query: str, conversation_context: str) -> str:
+        """Build intelligent repair prompt"""
+        return f"""
+Vous êtes un assistant juridique expert. L'utilisateur exprime un problème de compréhension ou demande une clarification.
+**CONTEXTE DE LA CONVERSATION:**
+{conversation_context}
+**REQUÊTE ACTUELLE DE L'UTILISATEUR:**
+"{current_query}"
+**ANALYSE REQUISE:**
+1. Identifiez le type de problème : incompréhension, besoin de clarification, reformulation, correction d'erreur
+2. Analysez quel aspect précis pose problème dans la conversation
+3. Adaptez votre réponse au contexte juridique si pertinent
+**INSTRUCTIONS POUR LA RÉPONSE:**
+- Accusez réception du problème de compréhension
+- Fournissez une clarification adaptée et utile
+- Si c'est juridique, simplifiez sans perdre la précision légale
+- Utilisez des exemples concrets si pertinent
+- Proposez des pistes pour avancer
+- Gardez un ton professionnel et empathique
+- Maximum 5-7 phrases
+**RÉPONSE:**
+"""
+    def _generate_fallback_response(self) -> str:
+        """Fallback if LLM fails"""
+        return "Je m'excuse pour ce malentendu. Pouvez-vous reformuler votre demande ou préciser ce qui n'était pas clair ?"

core/email_tool.py ADDED Viewed

	@@ -0,0 +1,187 @@

+# acfai_project/core/email_tool.py
+import os
+import smtplib
+import logging
+from email.mime.text import MIMEText  # Correction: MIMEText au lieu de MimeText
+from email.mime.multipart import MIMEMultipart  # Correction: MIMEMultipart au lieu de MimeMultipart
+from typing import Dict, Optional
+import re
+import datetime  # Ajout pour la date
+from config.settings import settings
+logger = logging.getLogger(__name__)
+class LegalAssistanceEmailer:
+    def __init__(self):
+        self.email_address = os.getenv("EMAIL_ADDRESS")
+        self.email_password = os.getenv("EMAIL_APP_PASSWORD")
+        self.lawyer_email = os.getenv("LAWYER_EMAIL", "fitahiana@acfai.org")
+        self.smtp_server = "smtp.gmail.com"
+        self.smtp_port = 587
+    def is_assistance_request(self, query: str) -> bool:
+        """Détecte si l'utilisateur demande une assistance humaine"""
+        assistance_keywords = [
+            "parler à un avocat", "avocat humain", "assistance humaine",
+            "contactez-moi", "rappelez-moi", "assistance téléphonique",
+            "besoin d'un avocat", "consultation juridique", "avocat réel",
+            "aide humaine", "contact humain", "échange avec un avocat",
+            "assisté", "assisté par", "être assisté"
+        ]
+        query_lower = query.lower()
+        return any(keyword in query_lower for keyword in assistance_keywords)
+    def extract_email_from_text(self, text: str) -> Optional[str]:
+        """Extrait un email d'un texte"""
+        email_pattern = r'\b[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Z|a-z]{2,}\b'
+        matches = re.findall(email_pattern, text)
+        return matches[0] if matches else None
+    def validate_email(self, email: str) -> bool:
+        """Valide le format d'un email"""
+        pattern = r'^[a-zA-Z0-9._%+-]+@[a-zA-Z0-9.-]+\.[a-zA-Z]{2,}$'
+        return re.match(pattern, email) is not None
+    def send_assistance_request(self, user_email: str, user_query: str,
+                              assistance_description: str, country: str) -> Dict[str, any]:
+        """Envoie les emails de confirmation à l'utilisateur et à l'avocat"""
+        try:
+            # Validation des emails
+            if not self.validate_email(user_email):
+                return {"success": False, "error": "Format d'email utilisateur invalide"}
+            if not self.validate_email(self.lawyer_email):
+                return {"success": False, "error": "Format d'email avocat invalide"}
+            # Connexion SMTP
+            server = smtplib.SMTP(self.smtp_server, self.smtp_port)
+            server.starttls()
+            server.login(self.email_address, self.email_password)
+            # Email à l'utilisateur
+            user_email_sent = self._send_user_confirmation(server, user_email, user_query, country)
+            # Email à l'avocat
+            lawyer_email_sent = self._send_lawyer_notification(server, user_email, user_query,
+                                                             assistance_description, country)
+            server.quit()
+            if user_email_sent and lawyer_email_sent:
+                logger.info(f"✅ Emails envoyés avec succès pour {user_email}")
+                return {
+                    "success": True,
+                    "message": "Demande d'assistance envoyée avec succès",
+                    "user_email": user_email,
+                    "lawyer_email": self.lawyer_email
+                }
+            else:
+                return {"success": False, "error": "Échec de l'envoi des emails"}
+        except Exception as e:
+            logger.error(f"❌ Erreur d'envoi d'email: {e}")
+            return {"success": False, "error": f"Erreur SMTP: {str(e)}"}
+    def _send_user_confirmation(self, server, user_email: str, user_query: str, country: str) -> bool:
+        """Envoie l'email de confirmation à l'utilisateur"""
+        try:
+            message = MIMEMultipart()  # Correction: MIMEMultipart
+            message["From"] = self.email_address
+            message["To"] = user_email
+            message["Subject"] = "📧 Confirmation de votre demande d'assistance juridique"
+            body = f"""
+            <html>
+            <body>
+                <h2 style="color: #2E86AB;">Confirmation de votre demande d'assistance juridique</h2>
+                <p>Bonjour,</p>
+                <p>Nous accusons réception de votre demande d'assistance juridique concernant :</p>
+                <div style="background-color: #f8f9fa; padding: 15px; border-left: 4px solid #2E86AB;">
+                    <strong>Question initiale :</strong> {user_query}<br>
+                    <strong>Juridiction concernée :</strong> {country}<br>
+                    <strong>Votre email :</strong> {user_email}
+                </div>
+                <p>��� <strong>Notre équipe juridique a été notifiée</strong> et vous contactera dans les plus brefs délais.</p>
+                <h3>📞 Prochaines étapes :</h3>
+                <ul>
+                    <li>Un avocat spécialisé vous contactera à l'adresse {user_email}</li>
+                    <li>Préparez les documents relatifs à votre situation</li>
+                    <li>Durée de réponse estimée : 24-48 heures</li>
+                </ul>
+                <p>Pour toute urgence, vous pouvez répondre directement à cet email.</p>
+                <hr>
+                <p style="color: #6c757d;">
+                    <small>
+                        ACFAI - Assistance Juridique Intelligente<br>
+                        Email : {self.lawyer_email}<br>
+                        Ceci est un message automatique, merci de ne pas y répondre directement.
+                    </small>
+                </p>
+            </body>
+            </html>
+            """
+            message.attach(MIMEText(body, "html"))  # Correction: MIMEText
+            server.send_message(message)
+            return True
+        except Exception as e:
+            logger.error(f"Erreur envoi email utilisateur: {e}")
+            return False
+    def _send_lawyer_notification(self, server, user_email: str, user_query: str,
+                                assistance_description: str, country: str) -> bool:
+        """Envoie la notification à l'avocat"""
+        try:
+            message = MIMEMultipart()  # Correction: MIMEMultipart
+            message["From"] = self.email_address
+            message["To"] = self.lawyer_email
+            message["Subject"] = f"🔔 Nouvelle demande d'assistance juridique - {country}"
+            body = f"""
+            <html>
+            <body>
+                <h2 style="color: #A23B72;">Nouvelle demande d'assistance juridique</h2>
+                <div style="background-color: #fff3cd; padding: 15px; border-left: 4px solid #ffc107;">
+                    <h3>📋 Informations de la demande :</h3>
+                    <p><strong>Utilisateur :</strong> {user_email}</p>
+                    <p><strong>Pays/Juridiction :</strong> {country}</p>
+                    <p><strong>Question initiale :</strong> {user_query}</p>
+                    <p><strong>Description de l'assistance demandée :</strong><br>{assistance_description}</p>
+                </div>
+                <h3>🚀 Action requise :</h3>
+                <ul>
+                    <li>Contacter l'utilisateur à : {user_email}</li>
+                    <li>Spécialité requise : Droit {country}</li>
+                    <li>Priorité : Normale</li>
+                </ul>
+                <hr>
+                <p style="color: #6c757d;">
+                    <small>
+                        Système Automatique ACFAI - {settings.CHAT_MODEL}<br>
+                        Généré le : {datetime.datetime.now().strftime('%Y-%m-%d %H:%M')}
+                    </small>
+                </p>
+            </body>
+            </html>
+            """
+            message.attach(MIMEText(body, "html"))  # Correction: MIMEText
+            server.send_message(message)
+            return True
+        except Exception as e:
+            logger.error(f"Erreur envoi email avocat: {e}")
+            return False

core/graph_builder.py ADDED Viewed

	@@ -0,0 +1,266 @@

+# [file name]: core/graph_builder.py
+from langgraph.graph import StateGraph, START, END
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+import logging
+from typing import Dict, List, Any
+from langchain_core.runnables import RunnableConfig
+from models.state_models import MultiCountryLegalState
+from core.router import CountryRouter
+from core.retriever import LegalRetriever
+from core.conversation_repair import ConversationRepair
+from core.human_approval_node import HumanApprovalNode
+# Import modular components
+from core.nodes.routing_nodes import RoutingNodes
+from core.assistance.workflow_nodes import AssistanceWorkflowNodes
+from core.nodes.retrieval_nodes import RetrievalNodes
+from core.nodes.response_nodes import ResponseNodes
+from core.nodes.helper_nodes import HelperNodes
+from core.routing.routing_logic import RoutingLogic
+logger = logging.getLogger(__name__)
+class GraphBuilder:
+    def __init__(
+        self,
+        router: CountryRouter,
+        llm,
+        checkpointer: AsyncPostgresSaver,
+        # Country retrievers as a dictionary for easy extension
+        country_retrievers: Dict[str, LegalRetriever] = None
+    ):
+        self.router = router
+        self.llm = llm
+        self.checkpointer = checkpointer
+        # Initialize country retrievers - easily extensible!
+        self.country_retrievers = country_retrievers or {}
+        # Initialize modular components
+        self.conversation_repair = ConversationRepair()
+        self.human_approval = HumanApprovalNode()
+        self.routing_logic = RoutingLogic()
+        # Initialize node groups
+        self.routing_nodes = RoutingNodes(router, self.conversation_repair, llm)
+        self.assistance_nodes = AssistanceWorkflowNodes()
+        # Dynamic retrieval nodes based on available countries
+        self.retrieval_nodes = RetrievalNodes(self.country_retrievers)
+        self.response_nodes = ResponseNodes(llm)
+        self.helper_nodes = HelperNodes(llm)
+        logger.info(f"GraphBuilder initialized with countries: {list(self.country_retrievers.keys())}")
+    def add_country(self, country_code: str, retriever: LegalRetriever):
+        """Dynamically add a new country to the system"""
+        self.country_retrievers[country_code] = retriever
+        self.retrieval_nodes = RetrievalNodes(self.country_retrievers)  # Re-initialize
+        logger.info(f"Added country: {country_code}")
+    def build_graph(self) -> StateGraph:
+        """Build simplified flow with all routing categories"""
+        workflow = StateGraph(MultiCountryLegalState)
+        # Core nodes
+        workflow.add_node("router", self.routing_nodes.router_node)
+        workflow.add_node("response", self.response_nodes.response_generation_node)
+        # Country retrieval nodes - dynamically created
+        country_nodes = {}
+        for country_code in self.country_retrievers.keys():
+            node_name = f"{country_code}_retrieval"
+            workflow.add_node(node_name, self._create_country_retrieval_node(country_code))
+            country_nodes[country_code] = node_name
+        # Handler nodes
+        workflow.add_node("greeting_handler", self.routing_nodes.greeting_small_talk_node)
+        workflow.add_node("repair_handler", self.routing_nodes.conversation_repair_node)
+        workflow.add_node("summary_handler", self.helper_nodes.conversation_summarization_node)
+        workflow.add_node("unclear_handler", self.helper_nodes.unclear_route_node)
+        workflow.add_node("out_of_scope_handler", self.helper_nodes.out_of_scope_node)
+        # Assistance nodes - Using wrapper methods to ensure correct signatures
+        workflow.add_node("assistance_collect_info", self._create_assistance_collect_wrapper())
+        workflow.add_node("assistance_confirm", self._create_assistance_confirm_wrapper())
+        workflow.add_node("human_approval", self.human_approval.process_approval)
+        workflow.add_node("process_assistance", self._create_process_assistance_node)
+        # Main flow
+        workflow.add_edge(START, "router")
+        # Router directly routes to appropriate nodes
+        workflow.add_conditional_edges(
+            "router",
+            self._route_after_router,
+            {
+                **country_nodes,  # benin_retrieval, madagascar_retrieval, etc.
+                "greeting_small_talk": "greeting_handler",
+                "conversation_repair": "repair_handler",
+                "conversation_summarization": "summary_handler",
+                "unclear": "unclear_handler",
+                "out_of_scope": "out_of_scope_handler",
+                "assistance_request": "assistance_collect_info"
+            }
+        )
+        # All handlers go to response
+        workflow.add_edge("greeting_handler", "response")
+        workflow.add_edge("repair_handler", "response")
+        workflow.add_edge("summary_handler", "response")
+        workflow.add_edge("unclear_handler", "response")
+        workflow.add_edge("out_of_scope_handler", "response")
+        # Country nodes go to response
+        for country_code in self.country_retrievers.keys():
+            workflow.add_edge(f"{country_code}_retrieval", "response")
+        # Assistance sub-flow
+        workflow.add_conditional_edges(
+            "assistance_collect_info",
+            self.routing_logic.route_after_info_collection,
+            {
+                "need_email": "response",  # Ask for email
+                "need_description": "response",  # Ask for description
+                "ready_to_confirm": "assistance_confirm",
+                "cancelled": "response"
+            }
+        )
+        # CRITICAL FIX: After response, only continue assistance if we have new user input
+        workflow.add_conditional_edges(
+            "response",
+            self._route_after_response,
+            {
+                "continue_assistance": "assistance_collect_info",
+                "end": END
+            }
+        )
+        workflow.add_conditional_edges(
+            "assistance_confirm",
+            self.routing_logic.route_after_confirmation,
+            {
+                "confirmed": "human_approval",
+                "cancelled": "response",
+                "needs_correction": "response"
+            }
+        )
+        workflow.add_conditional_edges(
+            "human_approval",
+            self.routing_logic.route_after_human_approval,
+            {
+                "approved": "process_assistance",
+                "rejected": "response",
+                "interrupt": "response"
+            }
+        )
+        workflow.add_edge("process_assistance", "response")
+        logger.info(f"Scalable graph built for {len(self.country_retrievers)} countries: {list(self.country_retrievers.keys())}")
+        return workflow
+    def _create_assistance_collect_wrapper(self):
+        """Wrapper to ensure proper method signature for assistance collection"""
+        async def wrapper(state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+            result = await self.assistance_nodes.collect_assistance_info_node(state, config)
+            # Ensure supplemental_message is included if not present
+            if "supplemental_message" not in result:
+                result["supplemental_message"] = ""
+            return result
+        return wrapper
+    def _create_assistance_confirm_wrapper(self):
+        """Wrapper to ensure proper method signature for assistance confirmation"""
+        async def wrapper(state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+            result = await self.assistance_nodes.confirm_assistance_send_node(state, config)
+            # Ensure supplemental_message is included if not present
+            if "supplemental_message" not in result:
+                result["supplemental_message"] = ""
+            return result
+        return wrapper
+    def _route_after_router(self, state: MultiCountryLegalState) -> str:
+        """Route directly from router - single source of truth"""
+        router_decision = state.router_decision or "unclear"
+        logger.debug(f"Routing from router: {router_decision}")
+        return router_decision
+    def _route_after_response(self, state: MultiCountryLegalState) -> str:
+        """Route after response - check if we should continue assistance workflow"""
+        # Check if we're in the middle of an assistance workflow
+        assistance_step = state.assistance_step
+        if assistance_step and assistance_step not in [None, "cancelled", "completed"]:
+            # CRITICAL FIX: Only continue if we have new user input to process
+            # This prevents infinite loops when no new user input is available
+            has_new_user_input = self._has_new_user_input(state)
+            if has_new_user_input:
+                logger.info(f"🔄 Continuing assistance workflow from response: {assistance_step}")
+                return "continue_assistance"
+            else:
+                logger.info("⏸️  No new user input - waiting for user response")
+                return "end"
+        # Normal end of conversation
+        logger.debug("✅ Ending conversation - no assistance workflow active")
+        return "end"
+    def _has_new_user_input(self, state: MultiCountryLegalState) -> bool:
+        """Check if there's new user input to process in assistance workflow"""
+        if not state.messages:
+            return False
+        # Get the last message
+        last_message = state.messages[-1] if state.messages else None
+        # Check if the last message is from user and not already processed
+        if last_message and last_message.get("role") == "user":
+            # Check message metadata to see if it's been processed in current assistance step
+            message_meta = last_message.get("meta", {})
+            processed_in_step = message_meta.get("processed_in_assistance_step")
+            current_step = state.assistance_step
+            # If this message hasn't been processed in the current assistance step, it's new input
+            if processed_in_step != current_step:
+                logger.info(f"📥 New user input detected for assistance step: {current_step}")
+                return True
+        logger.info("📭 No new user input detected")
+        return False
+    def _create_country_retrieval_node(self, country_code: str):
+        """Create a dynamic country retrieval node (closure factory)"""
+        async def country_retrieval_node(state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+            logger.info(f"Country retrieval for: {country_code}")
+            return await self.retrieval_nodes.country_retrieval_node(state, config, country_code)
+        return country_retrieval_node
+    async def _create_process_assistance_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Process assistance after approval"""
+        logger.info("Processing assistance request")
+        # Mark assistance as completed with supplemental message
+        return {
+            "email_status": "sent",
+            "approval_status": "approved",
+            "assistance_step": "completed",
+            "messages": [],
+            # "supplemental_message": "Votre demande d'assistance a été traitée avec succès et envoyée à notre équipe juridique."
+        }
+    def debug_state(self, state: MultiCountryLegalState, step: str) -> None:
+        """Debug state information"""
+        if logger.isEnabledFor(logging.DEBUG):
+            logger.debug(f"=== STATE DEBUG at {step} ===")
+            logger.debug(f"Router decision: {getattr(state, 'router_decision', 'None')}")
+            logger.debug(f"Assistance step: {getattr(state, 'assistance_step', 'None')}")
+            logger.debug(f"User email: {getattr(state, 'user_email', 'None')}")
+            logger.debug(f"Assistance description: {getattr(state, 'assistance_description', 'None')}")
+            logger.debug(f"Supplemental message: {getattr(state, 'supplemental_message', 'None')}")
+            logger.debug(f"Available countries: {list(self.country_retrievers.keys())}")
+            logger.debug("=== END STATE DEBUG ===")

core/human_approval_node.py ADDED Viewed

	@@ -0,0 +1,222 @@

+# core/human_approval_node.py
+import logging
+from typing import Literal
+from langchain_core.runnables import RunnableConfig
+from langgraph.types import interrupt, Command
+from models.state_models import MultiCountryLegalState
+from core.assistance.email_service import AssistanceEmailService
+from datetime import datetime
+logger = logging.getLogger(__name__)
+class HumanApprovalNode:
+    def __init__(self):
+        self.email_service = AssistanceEmailService()
+    async def process_approval(
+        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)
+            return Command(
+                goto="response",
+                update={
+                    "approval_status": "error",
+                    "messages": [{
+                        "role": "assistant",
+                        "content": f"❌ Erreur lors de l'approbation: {str(e)}",
+                        "meta": {}
+                    }]
+                }
+            )
+    async def _handle_approval(
+        self,
+        state: MultiCountryLegalState,
+        decision: dict
+    ) -> Command[Literal["response"]]:
+        """Handle approved request (sends email and routes to response)"""
+        logger.info(f"✅ Request APPROVED for {state.user_email}")
+        # Send email
+        email_result = self.email_service.send_assistance_request(
+            user_email=state.user_email,
+            user_query=state.last_search_query or "Demande d'assistance",
+            assistance_description=state.assistance_description,
+            country=self._get_country_display(state)
+        )
+        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",
+            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"}
+                }]
+            }
+        )
+    async def _handle_rejection(
+        self,
+        state: MultiCountryLegalState,
+        decision: dict
+    ) -> Command[Literal["response"]]:  # Updated: Removed "process_assistance"
+        """Handle rejected request"""
+        logger.info(f"❌ Request REJECTED for {state.user_email}")
+        message_content = f"""❌ **DEMANDE REFUSÉE**
+Votre demande d'assistance n'a pas été approuvée.
+**Raison:** {decision['reason']}
+Si vous pensez qu'il s'agit d'une erreur, veuillez reformuler votre demande avec plus de détails.
+"""
+        return Command(
+            goto="response",
+            update={
+                "approval_status": "rejected",
+                "approval_reason": decision["reason"],
+                "approved_by": decision["moderator_id"],
+                "approval_timestamp": datetime.now().isoformat(),
+                "messages": [{
+                    "role": "assistant",
+                    "content": message_content,
+                    "meta": {"approval": "rejected"}
+                }]
+            }
+        )
+    def _format_approval_request(self, state: MultiCountryLegalState) -> str:
+        """Format the approval request message"""
+        return f"""
+🔒 **APPROBATION HUMAINE REQUISE**
+📧 **Email:** {state.user_email}
+🌍 **Pays:** {self._get_country_display(state)}
+📝 **Description:** {state.assistance_description}
+🔍 **Requête initiale:** {state.last_search_query or 'Non spécifiée'}
+**Instructions:**
+- Tapez "approve [raison]" pour approuver
+- Tapez "reject [raison]" pour rejeter
+**Exemples:**
+- "approve Demande légitime"
+- "reject Email invalide"
+"""
+    def _parse_decision(self, user_input: str) -> dict:
+        """Parse moderator decision from input"""
+        if not user_input or not isinstance(user_input, str):
+            return {
+                "approved": False,
+                "reason": "Input invalide",
+                "moderator_id": "system"
+            }
+        input_lower = user_input.lower().strip()
+        # Check for approval keywords
+        approve_keywords = ["approve", "approuver", "oui", "yes", "ok", "accept"]
+        is_approved = any(kw in input_lower for kw in approve_keywords)
+        # Extract reason (text after the decision keyword)
+        reason = user_input.strip()
+        for keyword in approve_keywords + ["reject", "rejeter", "non", "no"]:
+            if keyword in input_lower:
+                parts = user_input.split(keyword, 1)
+                if len(parts) > 1 and parts[1].strip():
+                    reason = parts[1].strip()
+                    break
+        if not reason or reason == user_input:
+            reason = "Approuvé par modérateur" if is_approved else "Refusé par modérateur"
+        return {
+            "approved": is_approved,
+            "reason": reason,
+            "moderator_id": "human_moderator"
+        }
+    def _get_country_display(self, state: MultiCountryLegalState) -> str:
+        """Get country display name"""
+        country = state.country or state.legal_context.get("detected_country", "unknown")
+        if country == "unknown" and state.assistance_description:
+            country = MultiCountryLegalState.detect_country(state.assistance_description)
+        country_map = {
+            "benin": "Bénin",
+            "madagascar": "Madagascar"
+        }
+        logger.debug(f"Country from state: {state.country}, legal_context: {state.legal_context.get('detected_country')}, description: {country}")
+        return country_map.get(country, "Non spécifié")

core/nodes/__init__.py ADDED Viewed

	@@ -0,0 +1,14 @@

+# [file name]: core/nodes/__init__.py
+from .routing_nodes import RoutingNodes
+from .retrieval_nodes import RetrievalNodes
+from .response_nodes import ResponseNodes
+from .helper_nodes import HelperNodes
+# Remove AssistanceNodes from exports since it's moved to core/assistance/
+__all__ = [
+    "RoutingNodes",
+    "RetrievalNodes",
+    "ResponseNodes",
+    "HelperNodes"
+]

core/nodes/base_node.py ADDED Viewed

	@@ -0,0 +1,79 @@

+# [file name]: core/nodes/base_node.py
+import logging
+from typing import Dict, List, Optional, Any
+from datetime import datetime
+from langchain_core.runnables import RunnableConfig
+from models.state_models import MultiCountryLegalState
+from utils.helpers import dict_to_message_obj
+logger = logging.getLogger(__name__)
+class BaseNode:
+    """Base class with common utilities for all nodes"""
+    def _get_last_human_message(self, messages: List[Dict]) -> Optional[Dict]:
+        """Get the last human message from conversation"""
+        if not messages:
+            return None
+        for msg in reversed(messages):
+            if msg.get("role", "").lower() in ("user", "human"):
+                return msg
+        return None
+    def _has_complete_response(self, messages: List[Dict]) -> bool:
+        """Check if there's already an assistant response in recent messages"""
+        if not messages:
+            return False
+        for msg in reversed(messages):
+            if msg.get("role") == "assistant" and msg.get("content"):
+                return True
+        return False
+    def _create_error_message(self, error: str) -> Dict[str, Any]:
+        """Create standardized error message"""
+        return {
+            "role": "assistant",
+            "content": f"Désolé, une erreur s'est produite lors du traitement de votre demande: {error}",
+            "meta": {
+                "is_error": True,
+                "timestamp": self._get_timestamp()
+            }
+        }
+    def _create_error_state(self, error: str) -> Dict[str, Any]:
+        """Create error state with message"""
+        return {
+            "messages": [self._create_error_message(error)],
+            "search_results": f"Error: {error}"
+        }
+    def _get_timestamp(self) -> str:
+        """Get current timestamp for message metadata"""
+        return datetime.now().isoformat()
+    def _update_legal_context(self, legal_context: Dict, country: str) -> Dict:
+        """Update legal context with country information"""
+        updated = legal_context.copy() if legal_context else {}
+        if country in ["benin", "madagascar"]:
+            updated["detected_country"] = country
+            if country == "benin":
+                updated["jurisdiction"] = "Bénin"
+            elif country == "madagascar":
+                updated["jurisdiction"] = "Madagascar"
+        else:
+            updated["jurisdiction"] = "Unknown"
+            updated["detected_country"] = "unknown"
+        return updated
+    def _create_router_response(self, country: str, explanation: str, legal_context: Dict) -> Dict[str, Any]:
+        """Create standardized router response"""
+        updated_context = self._update_legal_context(legal_context, country)
+        return {
+            "router_decision": country,
+            "route_explanation": explanation,
+            "legal_context": updated_context,
+            "primary_intent": country
+        }

core/nodes/helper_nodes.py ADDED Viewed

	@@ -0,0 +1,147 @@

+# [file name]: core/nodes/helper_nodes.py
+import logging
+from typing import Dict, Any, List
+from langchain_core.runnables import RunnableConfig
+from langchain_core.messages import HumanMessage
+from models.state_models import MultiCountryLegalState
+from .base_node import BaseNode
+from core.prompts.prompt_templates import PromptTemplates
+logger = logging.getLogger(__name__)
+class HelperNodes(BaseNode):
+    """Helper nodes for unclear routes and summarization"""
+    def __init__(self, llm):
+        self.llm = llm
+        self.prompts = PromptTemplates()
+    async def out_of_scope_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Handle out-of-scope questions - redirect to legal domain"""
+        try:
+            logger.info("🚫 Out of scope question detected")
+            redirect_message = {
+                "role": "assistant",
+                "content": (
+                    "Je suis un assistant juridique spécialisé dans le droit du Bénin et de Madagascar. "
+                    "Je ne peux répondre qu'aux questions relatives au droit et aux procédures juridiques.\n\n"
+                    "Comment puis-je vous aider avec vos questions juridiques ?"
+                ),
+                "meta": {
+                    "is_out_of_scope": True,
+                    "timestamp": self._get_timestamp()
+                }
+            }
+            return {
+                "messages": [redirect_message],
+                "current_country": "out_of_scope",
+                "search_results": "Out of scope query - no legal search performed"
+            }
+        except Exception as e:
+            logger.error(f"Error in out_of_scope handler: {str(e)}")
+            return self._create_error_state(f"Error in out_of_scope: {str(e)}")
+    async def unclear_route_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Handle unclear routing cases - for ambiguous legal queries"""
+        try:
+            s = state.model_dump()
+            route_explanation = s.get("route_explanation", "")
+            # This is now only for unclear LEGAL queries
+            clarification_msg = {
+                "role": "assistant",
+                "content": self.prompts.get_clarification_message(),
+                "meta": {
+                    "requires_clarification": True,
+                    "timestamp": self._get_timestamp()
+                }
+            }
+            return {
+                "messages": [clarification_msg],
+                "search_results": "Country clarification needed"
+            }
+        except Exception as e:
+            logger.error(f"Error in unclear route handling: {str(e)}")
+            return self._create_error_state(f"Error in unclear route: {str(e)}")
+    async def conversation_summarization_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Generate summary of conversation history"""
+        try:
+            s = state.model_dump()
+            messages = s.get("messages", [])
+            logger.info(f"📋 Generating conversation summary for {len(messages)} messages")
+            summary = await self._generate_conversation_summary(messages)
+            return {
+                "messages": [{
+                    "role": "assistant",
+                    "content": summary,
+                    "meta": {
+                        "is_summary": True,
+                        "conversation_length": len(messages),
+                        "timestamp": self._get_timestamp()
+                    }
+                }],
+                "search_results": "Conversation summary generated - no legal search performed"
+            }
+        except Exception as e:
+            logger.error(f"Error in conversation summarization: {str(e)}")
+            return self._create_error_state(f"Error in summarization: {str(e)}")
+    async def _generate_conversation_summary(self, messages: List[Dict]) -> str:
+        """Use LLM to generate conversation summary"""
+        conversation_messages = [
+            msg for msg in messages
+            if msg.get("role") in ["user", "assistant"]
+        ]
+        if len(conversation_messages) <= 2:
+            return "Notre conversation vient juste de commencer. Nous n'avons pas encore beaucoup échangé."
+        conversation_text = ""
+        for i, msg in enumerate(conversation_messages):
+            role = "Utilisateur" if msg.get("role") == "user" else "Assistant"
+            content = msg.get("content", "")
+            conversation_text += f"{role}: {content}\n\n"
+        summary_prompt = f"""
+Vous êtes un assistant juridique. Résumez la conversation suivante entre l'utilisateur et vous-même.
+**CONVERSATION:**
+{conversation_text}
+**INSTRUCTIONS:**
+- Faites un résumé concis et clair
+- Mettez en évidence les points juridiques principaux discutés
+- Mentionnez les pays concernés (Bénin/Madagascar) si pertinents
+- Gardez un ton professionnel mais accessible
+- Maximum 5-7 phrases
+**RÉSUMÉ:**
+"""
+        try:
+            response = await self.llm.ainvoke([HumanMessage(content=summary_prompt)])
+            return response.content if hasattr(response, 'content') else str(response)
+        except Exception as e:
+            logger.error(f"LLM summarization failed: {e}")
+            user_messages = [m for m in conversation_messages if m.get("role") == "user"]
+            assistant_messages = [m for m in conversation_messages if m.get("role") == "assistant"]
+            return f"""**Résumé de notre conversation:**
+- **Échanges totaux**: {len(conversation_messages)} messages
+- **Questions de l'utilisateur**: {len(user_messages)}
+- **Réponses fournies**: {len(assistant_messages)}
+- **Dernier échange**: {conversation_messages[-1].get('content', '')[:100]}...
+*Pour un résumé détaillé, veuillez reposer votre question.*"""

core/nodes/response_nodes.py ADDED Viewed

	@@ -0,0 +1,219 @@

+# [file name]: core/nodes/response_nodes.py
+import logging
+import time
+from datetime import datetime
+from typing import Dict, Any
+from langchain_core.runnables import RunnableConfig
+from models.state_models import MultiCountryLegalState
+from utils.helpers import dict_to_message_obj, message_obj_to_dict
+logger = logging.getLogger(__name__)
+class ResponseNodes:
+    def __init__(self, llm):
+        self.llm = llm
+    async def response_generation_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Generate appropriate responses based on current state"""
+        assistance_step = state.assistance_step
+        # Handle assistance workflow responses
+        if assistance_step == "collecting_email":
+            response_content = """
+Je vois que vous souhaitez parler à un avocat. Pour vous aider, j'ai besoin de votre adresse email pour que notre équipe puisse vous contacter.
+📧 **Veuillez me fournir votre adresse email :**
+"""
+            return {
+                "messages": [{
+                    "role": "assistant",
+                    "content": response_content,
+                    "meta": {"assistance_step": "collecting_email"}
+                }],
+                "supplemental_message": ""  # Clear any previous supplemental messages
+            }
+        elif assistance_step == "collecting_description":
+            response_content = f"""
+Merci ! Votre email ({state.user_email}) a été enregistré.
+📝 **Veuillez maintenant décrire brièvement votre situation :**
+- Quelle est votre question juridique ?
+- De quel pays s'agit-il ?
+- Quel type d'assistance recherchez-vous ?
+Cette description aidera notre équipe à mieux vous orienter.
+"""
+            return {
+                "messages": [{
+                    "role": "assistant",
+                    "content": response_content,
+                    "meta": {"assistance_step": "collecting_description"}
+                }],
+                "supplemental_message": ""  # Clear any previous supplemental messages
+            }
+        elif assistance_step == "confirming_send":
+            response_content = f"""
+📋 **RÉCAPITULATIF DE VOTRE DEMANDE :**
+📧 **Email :** {state.user_email}
+📝 **Description :** {state.assistance_description}
+✅ **Confirmez-vous l'envoi de cette demande à notre équipe juridique ?**
+Répondez par :
+- **"oui"** pour confirmer et envoyer
+- **"non"** pour annuler et modifier
+"""
+            return {
+                "messages": [{
+                    "role": "assistant",
+                    "content": response_content,
+                    "meta": {"assistance_step": "confirming_send"}
+                }],
+                "supplemental_message": ""  # Clear any previous supplemental messages
+            }
+        else:
+            # Default LLM response for non-assistance flows
+            return await self._generate_llm_response(state, config)
+    async def _generate_llm_response(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Generate LLM-based response for normal conversation flows"""
+        try:
+            # Include supplemental message in the response if present
+            supplemental_message = state.supplemental_message or ""
+            # Synthesize response using LLM
+            response_content = await self._synthesize_response(state, supplemental_message)
+            return {
+                "messages": [{
+                    "role": "assistant",
+                    "content": response_content,
+                    "meta": {
+                        "timestamp": datetime.now().isoformat(),
+                        "generated_by": "llm"
+                    }
+                }],
+                "supplemental_message": ""  # Clear after using
+            }
+        except Exception as e:
+            logger.error(f"Error generating LLM response: {str(e)}")
+            return {
+                "messages": [{
+                    "role": "assistant",
+                    "content": self._create_error_message(str(e)),
+                    "meta": {"is_error": True}
+                }],
+                "supplemental_message": f"Erreur: {str(e)}"
+            }
+    async def _synthesize_response(self, state: MultiCountryLegalState, supplemental_message: str = "") -> str:
+        """Synthesize final response based on graph execution"""
+        s = state.model_dump()
+        # Build context-aware system prompt
+        system_prompt = self._build_system_prompt(state, supplemental_message)
+        conversation_messages = self._build_conversation_messages(system_prompt, s.get("messages", []))
+        # Always use LLM to generate final response
+        logger.info("🧠 Generating final response with LLM")
+        ai_resp = await self.llm.ainvoke(conversation_messages)
+        return ai_resp.content if hasattr(ai_resp, 'content') else str(ai_resp)
+    def _build_system_prompt(self, state: MultiCountryLegalState, supplemental_message: str = "") -> str:
+        """Build context-aware system prompt"""
+        s = state.model_dump()
+        base_prompt = """Vous êtes un assistant juridique expert spécialisé dans le droit du Bénin et de Madagascar.
+TÂCHE: Fournir une réponse claire, précise et utile à l'utilisateur.
+"""
+        # Add supplemental message if available
+        if supplemental_message:
+            base_prompt += f"\nMESSAGE IMPORTANT: {supplemental_message}\n"
+        # Add legal context if available
+        country_name = s.get("legal_context", {}).get("jurisdiction", "Unknown")
+        if country_name != "Unknown":
+            base_prompt += f"\nCONTEXTE JURIDIQUE: Vous répondez dans le cadre du droit {country_name}.\n"
+        # Add search results if available
+        search_results = s.get("search_results", "")
+        if search_results and "RECHERCHE JURIDIQUE" in search_results:
+            base_prompt += f"\nINFORMATIONS JURIDIQUES DISPONIBLES:\n{search_results}\n"
+            base_prompt += """
+INSTRUCTIONS POUR LA RÉPONSE JURIDIQUE:
+- Basez-vous sur les informations juridiques disponibles
+- Citez les articles de loi pertinents si possible
+- Soyez précis mais accessible aux non-juristes
+- Indiquez si certaines informations manquent
+"""
+        else:
+            base_prompt += """
+INSTRUCTIONS GÉNÉRALES:
+- Répondez de manière naturelle et utile
+- Adaptez votre ton au contexte de la conversation
+- Soyez empathique et professionnel
+"""
+        # Add assistance context if relevant
+        if s.get("assistance_requested"):
+            base_prompt += "\nCONTEXTE ASSISTANCE: L'utilisateur a demandé à parler à un avocat.\n"
+        if s.get("approval_status") == "rejected":
+            base_prompt += "\nCONTEXTE: La demande d'assistance a été rejetée. Expliquez poliment et proposez des alternatives.\n"
+        elif s.get("approval_status") == "approved":
+            base_prompt += "\nCONTEXTE: La demande d'assistance a été approuvée. Confirmez et donnez les prochaines étapes.\n"
+        return base_prompt
+    def _build_conversation_messages(self, system_prompt: str, messages: list) -> list:
+        """Build conversation messages for LLM"""
+        from langchain_core.messages import SystemMessage
+        conversation_messages = [SystemMessage(content=system_prompt)]
+        # Include recent conversation history (last 6 messages)
+        recent_messages = messages[-6:] if len(messages) > 6 else messages
+        # Convert to message objects
+        conversation_messages.extend(dict_to_message_obj(m) for m in recent_messages)
+        return conversation_messages
+    async def human_approval_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Handle human approval interrupts"""
+        logger.info("👨‍⚖️ Human approval node - triggering interrupt")
+        return {
+            "approval_status": "pending",
+            "messages": [{
+                "role": "assistant",
+                "content": "⏳ Votre demande d'assistance nécessite une approbation manuelle. Un modérateur va examiner votre demande.",
+                "meta": {"requires_approval": True}
+            }],
+            "supplemental_message": ""
+        }
+    async def process_assistance_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Process assistance after approval - let LLM generate final message"""
+        logger.info("📧 Processing assistance request")
+        return {
+            "email_status": "sent",
+            "approval_status": "approved",
+            "assistance_step": "completed",
+            "messages": [],  # Empty messages so LLM generates the final response
+            "supplemental_message": "Votre demande d'assistance a été traitée avec succès."
+        }
+    def _create_error_message(self, error: str) -> str:
+        """Create error message"""
+        return f"❌ Désolé, une erreur s'est produite: {error}\n\nVeuillez réessayer ou reformuler votre demande."

core/nodes/retrieval_nodes.py ADDED Viewed

	@@ -0,0 +1,83 @@

+# [file name]: core/nodes/retrieval_nodes.py
+import logging
+from typing import Dict, Any
+from langchain_core.runnables import RunnableConfig
+from models.state_models import MultiCountryLegalState
+from core.retriever import LegalRetriever
+logger = logging.getLogger(__name__)
+class RetrievalNodes:
+    """Scalable legal retrieval nodes for any number of countries"""
+    def __init__(self, country_retrievers: Dict[str, LegalRetriever]):
+        self.country_retrievers = country_retrievers
+    async def country_retrieval_node(self, state: MultiCountryLegalState, config: RunnableConfig, country_code: str) -> Dict[str, Any]:
+        """Generic country retrieval for any country"""
+        try:
+            if country_code not in self.country_retrievers:
+                logger.error(f"❌ Country not configured: {country_code}")
+                return {
+                    "search_results": f"Country {country_code} not available",
+                    "detected_articles": [],
+                    "supplemental_message": f"Pays {country_code} non configuré dans le système."
+                }
+            retriever = self.country_retrievers[country_code]
+            s = state.model_dump()
+            last_human = self._get_last_human_message(s.get("messages", []))
+            if not last_human:
+                return {
+                    "search_results": f"No query for {country_code} retrieval",
+                    "detected_articles": [],
+                    "supplemental_message": "Aucune requête trouvée pour la recherche."
+                }
+            user_query = last_human.get("content", "").strip()
+            if not user_query:
+                return {
+                    "search_results": f"Empty query for {country_code} retrieval",
+                    "detected_articles": [],
+                    "supplemental_message": "Requête vide pour la recherche."
+                }
+            logger.info(f"🌍 Performing {country_code} retrieval for: '{user_query[:50]}...'")
+            enhanced_docs, detected_articles, applied_filters, supplemental_message = await retriever.smart_legal_query(user_query, country_code)
+            search_results = retriever.format_search_results(
+                user_query, enhanced_docs, detected_articles, applied_filters, country_code, supplemental_message
+            )
+            logger.info(f"📚 Retrieved {len(enhanced_docs)} documents for {country_code}")
+            return {
+                "search_results": search_results,
+                "detected_articles": detected_articles,
+                "last_search_query": user_query,
+                "supplemental_message": supplemental_message,  # Pass the supplemental message to state
+                # Store complex data in search_metadata instead of legal_context
+                "search_metadata": {
+                    "applied_filters": applied_filters,
+                    "documents_count": len(enhanced_docs),
+                    "supplemental_message": supplemental_message
+                }
+            }
+        except Exception as e:
+            logger.error(f"Error in {country_code} retrieval: {str(e)}")
+            return {
+                "search_results": f"Erreur lors de la recherche {country_code}: {str(e)}",
+                "detected_articles": [],
+                "supplemental_message": f"Erreur lors de la recherche: {str(e)}"
+            }
+    def _get_last_human_message(self, messages: list) -> Dict[str, Any]:
+        """Get the last human message"""
+        for msg in reversed(messages):
+            if msg.get("role") in ["user", "human"]:
+                return msg
+        return {}

core/nodes/routing_nodes.py ADDED Viewed

	@@ -0,0 +1,193 @@

+# [file name]: core/nodes/routing_nodes.py
+import logging
+from typing import Dict, Any
+from langchain_core.runnables import RunnableConfig
+from models.state_models import MultiCountryLegalState
+from core.router import CountryRouter
+from .base_node import BaseNode
+from core.prompts.prompt_templates import PromptTemplates
+logger = logging.getLogger(__name__)
+class RoutingNodes(BaseNode):
+    """Router, greeting, and conversation repair nodes"""
+    def __init__(self, router: CountryRouter, conversation_repair, llm):
+        self.router = router
+        self.conversation_repair = conversation_repair
+        self.llm = llm
+        self.prompts = PromptTemplates()
+    async def router_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Enhanced router that detects primary intent with state awareness"""
+        try:
+            s = state.model_dump()
+            # CRITICAL: Check if we're continuing an assistance workflow
+            # This prevents the router from misclassifying continuation messages
+            assistance_step = s.get("assistance_step")
+            if assistance_step and assistance_step not in [None, "cancelled", "completed"]:
+                logger.info(f"⏩ Bypassing router - continuing assistance at step: {assistance_step}")
+                return {
+                    "router_decision": "assistance_request",
+                    "route_explanation": f"Continuing assistance workflow: {assistance_step}",
+                    "assistance_step": assistance_step,  # Ensure step persists
+                    "assistance_requested": True
+                }
+            # Normal routing for new messages
+            return await self._perform_normal_routing(state, s)
+        except Exception as e:
+            logger.error(f"Router error: {str(e)}")
+            legal_context = state.legal_context if hasattr(state, 'legal_context') else {}
+            return self._create_router_response("unclear", f"Router error: {str(e)}", legal_context)
+    async def _perform_normal_routing(self, state: MultiCountryLegalState, state_dict: Dict) -> Dict[str, Any]:
+        """Perform normal routing for new user queries"""
+        if not state_dict.get("messages"):
+            logger.warning("No messages in state for router")
+            return self._create_router_response("unclear", "No messages in state", state_dict.get("legal_context", {}))
+        last_human = self._get_last_human_message(state_dict.get("messages", []))
+        if not last_human:
+            logger.warning("No user query found in router")
+            return self._create_router_response("unclear", "No user query found", state_dict.get("legal_context", {}))
+        user_query = last_human.get("content", "").strip()
+        if not user_query:
+            logger.warning("Empty user query in router")
+            return self._create_router_response("unclear", "Empty user query", state_dict.get("legal_context", {}))
+        logger.info(f"🔀 Routing query: '{user_query[:50]}...'")
+        routing_result = await self.router.route_query(user_query, state_dict["messages"])
+        primary_intent = routing_result.country
+        logger.info(f"🎯 Router decision: {primary_intent} ({routing_result.confidence}) - {routing_result.method}")
+        updated_context = self._update_legal_context(state_dict["legal_context"], primary_intent)
+        response = {
+            "router_decision": primary_intent,
+            "route_explanation": f"{routing_result.method}: {routing_result.explanation}",
+            "legal_context": updated_context,
+            "primary_intent": primary_intent
+        }
+        # If this is an assistance request, initialize the workflow
+        if primary_intent == "assistance_request":
+            response.update({
+                "assistance_step": "collecting_email",
+                "assistance_requested": True
+            })
+        return response
+    async def greeting_small_talk_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Handle greetings and small talk"""
+        try:
+            s = state.model_dump()
+            last_human = self._get_last_human_message(s.get("messages", []))
+            user_query = last_human.get("content", "").lower() if last_human else ""
+            logger.info(f"👋 Handling greeting/small_talk: '{user_query[:30]}...'")
+            greeting_response = self.prompts.generate_greeting_response(user_query)
+            return {
+                "messages": [{
+                    "role": "assistant",
+                    "content": greeting_response,
+                    "meta": {
+                        "is_greeting": True,
+                        "timestamp": self._get_timestamp()
+                    }
+                }],
+                "search_results": "Greeting handled - no legal search performed"
+            }
+        except Exception as e:
+            logger.error(f"Error in greeting node: {str(e)}")
+            return self._create_error_state(f"Error in greeting: {str(e)}")
+    async def conversation_repair_node(self, state: MultiCountryLegalState, config: RunnableConfig) -> Dict[str, Any]:
+        """Unified repair handling with LLM"""
+        try:
+            s = state.model_dump()
+            last_human = self._get_last_human_message(s.get("messages", []))
+            user_query = last_human.get("content", "") if last_human else ""
+            logger.info(f"🔧 Handling repair request: '{user_query[:30]}...'")
+            repair_response = await self.conversation_repair.generate_repair_response(
+                user_query, s.get("messages", []), self.llm
+            )
+            return {
+                "messages": [{
+                    "role": "assistant",
+                    "content": repair_response,
+                    "meta": {
+                        "is_repair_response": True,
+                        "timestamp": self._get_timestamp()
+                    }
+                }],
+                "search_results": "Repair handled - no legal search performed"
+            }
+        except Exception as e:
+            logger.error(f"Error in repair node: {str(e)}")
+            return self._create_error_state(f"Error in repair: {str(e)}")
+    def _create_router_response(self, decision: str, explanation: str, legal_context: Dict) -> Dict[str, Any]:
+        """Create a standardized router response"""
+        return {
+            "router_decision": decision,
+            "route_explanation": explanation,
+            "legal_context": legal_context,
+            "primary_intent": decision
+        }
+    def _get_last_human_message(self, messages: list) -> Dict[str, Any]:
+        """Get the last human message from conversation history"""
+        for msg in reversed(messages):
+            if msg.get("role") in ["user", "human"]:
+                return msg
+        return {}
+    def _update_legal_context(self, legal_context: Dict, primary_intent: str) -> Dict:
+        """Update legal context based on routing decision"""
+        updated_context = legal_context.copy()
+        # Map router decisions to detected_country
+        country_mapping = {
+            "benin": "benin",
+            "madagascar": "madagascar",
+            "assistance_request": updated_context.get("detected_country", "unknown"),
+            "greeting_small_talk": "unknown",
+            "conversation_repair": updated_context.get("detected_country", "unknown"),
+            "conversation_summarization": updated_context.get("detected_country", "unknown"),
+            "unclear": "unknown",
+            "out_of_scope": "unknown"
+        }
+        updated_context["detected_country"] = country_mapping.get(primary_intent, "unknown")
+        updated_context["primary_intent"] = primary_intent
+        return updated_context
+    def _get_timestamp(self) -> str:
+        """Get current timestamp"""
+        from datetime import datetime
+        return datetime.now().isoformat()
+    def _create_error_state(self, error_message: str) -> Dict[str, Any]:
+        """Create error state response"""
+        return {
+            "messages": [{
+                "role": "assistant",
+                "content": f"❌ Désolé, une erreur s'est produite. Veuillez réessayer.",
+                "meta": {"error": error_message}
+            }]
+        }

core/prompts/__init__.py ADDED Viewed

	@@ -0,0 +1,3 @@


1	+ from .prompt_templates import PromptTemplates
2	+
3	+ __all__ = ["PromptTemplates"]

core/prompts/prompt_templates.py ADDED Viewed

	@@ -0,0 +1,94 @@

+# [file name]: core/prompts/prompt_templates.py
+class PromptTemplates:
+    """All prompt templates used in the graph"""
+    @staticmethod
+    def get_email_request_message() -> str:
+        return """📧 **Demande d'assistance juridique**
+Pour vous mettre en relation avec un avocat, j'ai besoin de votre adresse email.
+**Votre email :**"""
+    @staticmethod
+    def get_description_prompt(email: str) -> str:
+        return f"""📝 **Description de votre besoin**
+Merci ! Email enregistré : {email}
+Maintenant, décrivez-moi **comment vous souhaitez être assisté(e)** :
+Exemples :
+• "Consultation téléphonique de 30 minutes sur le droit de la famille"
+• "Avis écrit sur un contrat de travail"
+• "Accompagnement pour une procédure de divorce"
+• "Explication sur mes droits successoraux"
+**Votre description :**"""
+    @staticmethod
+    def get_confirmation_prompt(data: dict) -> str:
+        email = data.get("email", "Non fourni")
+        description = data.get("description", "Non fournie")
+        return f"""✅ **Confirmation d'envoi**
+Veuillez confirmer l'envoi de votre demande d'assistance :
+📧 **Email** : {email}
+📋 **Description** : {description}
+**L'avocat vous contactera directement dans les 24-48 heures.**
+🔔 **Confirmez-vous l'envoi ?** (répondez par OUI/NON)"""
+    @staticmethod
+    def get_missing_info_prompt(current_step: str, has_email: bool) -> str:
+        if current_step == "collecting_email":
+            return "📧 **Email manquant** : Pourriez-vous me donner votre adresse email ?"
+        else:
+            return "📝 **Description manquante** : Pourriez-vous décrire comment vous souhaitez être assisté(e) ?"
+    @staticmethod
+    def get_non_legal_response() -> str:
+        return """🔍 **Hors de mon domaine d'expertise**
+Je suis un assistant juridique spécialisé pour le Bénin et Madagascar.
+**Je peux vous aider avec :**
+⚖️ **Questions juridiques** : lois, droits, procédures
+📚 **Textes de loi** : articles, codes, décrets
+🔧 **Assistance légale** : démarches, formalités
+👨‍⚖️ **Connexion avocat** : assistance humaine
+**Exemples de questions que je peux traiter :**
+• "Procédure de divorce au Bénin"
+• "Droits des enfants à Madagascar"
+• "Articles sur le droit du travail"
+• "Comment contacter un avocat ?"
+Posez-moi une question juridique !"""
+    @staticmethod
+    def get_clarification_message() -> str:
+        return """Je ne peux pas déterminer de quel pays vous parlez. Pourriez-vous préciser si votre question concerne le droit du **Bénin** ou de **Madagascar** ?"""
+    @staticmethod
+    def generate_greeting_response(query: str) -> str:
+        """Generate appropriate greeting responses"""
+        query_lower = query.lower()
+        if any(word in query_lower for word in ["bonjour", "hello", "hi"]):
+            return "👋 Bonjour ! Je suis votre assistant juridique spécialisé pour le Bénin et Madagascar. Comment puis-je vous aider aujourd'hui ?"
+        elif any(word in query_lower for word in ["salut", "coucou"]):
+            return "👋 Salut ! Je suis votre assistant juridique. Posez-moi vos questions sur le droit béninois ou malgache !"
+        elif any(word in query_lower for word in ["comment ça va", "ça va", "comment vas-tu"]):
+            return "😊 Je vais très bien, merci ! Je suis prêt à vous aider avec vos questions juridiques sur le Bénin ou Madagascar."
+        elif any(word in query_lower for word in ["merci", "thanks"]):
+            return "🤝 Je vous en prie ! N'hésitez pas si vous avez d'autres questions juridiques."
+        elif any(word in query_lower for word in ["au revoir", "bye", "à bientôt"]):
+            return "👋 Au revoir ! N'hésitez pas à revenir si vous avez besoin d'assistance juridique."
+        elif any(word in query_lower for word in ["qui es-tu", "ton nom", "te présenter"]):
+            return "⚖️ Je suis un assistant juridique IA spécialisé dans les droits du Bénin et de Madagascar. Je peux vous aider à trouver des informations sur les lois, articles, et procédures juridiques."
+        else:
+            return "👋 Bonjour ! Je suis votre assistant juridique. Posez-moi vos questions sur le droit béninois ou malgache !"

core/retriever.py ADDED Viewed

	@@ -0,0 +1,386 @@

+# core/retriever.py
+import re
+import logging
+import asyncio
+from typing import List, Dict, Any, Tuple
+from langchain_core.documents import Document
+from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
+from config.settings import settings
+from config.constants import ARTICLE_PATTERNS, CATEGORY_KEYWORDS, DOCUMENT_TYPE_KEYWORDS
+logger = logging.getLogger(__name__)
+class LegalRetriever:
+    def __init__(self, vectorstore: MongoDBAtlasVectorSearch, collection):
+        self.vectorstore = vectorstore
+        self.collection = collection
+    async def smart_legal_query(self, user_query: str, country: str) -> Tuple[List[Document], List[str], Dict[str, Any], str]:
+        """Perform smart legal search with automatic fallback and custom messages - ASYNC VERSION"""
+        try:
+            # Détection initiale du type de document
+            initial_doc_type = self._detect_document_type(user_query.lower())
+            pre_filter = self._build_pre_filters(user_query, country)
+            logger.info(f"📋 Filtre doc_type initial: {initial_doc_type}")
+            logger.info(f"🔍 Recherche {country} avec filtres: {pre_filter}")
+            # Première recherche
+            enhanced_docs, detected_articles, applied_filters = await self._perform_search_async(
+                user_query, country, pre_filter
+            )
+            message_supplementaire = ""
+            # Fallback automatique si aucun résultat pour case_study
+            if not enhanced_docs and initial_doc_type == "case_study":
+                logger.info("🔄 Fallback: Aucun case_study trouvé, recherche dans les articles")
+                # Create new filter for articles - DON'T rebuild, just modify
+                fallback_filter = pre_filter.copy()  # Copy the original filter
+                fallback_filter["doc_type"] = "articles"  # Force articles type
+                logger.info(f"🔄 Fallback filter: {fallback_filter}")  # Log the fallback filter
+                enhanced_docs, detected_articles, applied_filters = await self._perform_search_async(
+                    user_query, country, fallback_filter
+                )
+                # Mark that fallback was used
+                applied_filters["original_search"] = "case_study"
+                applied_filters["fallback_to"] = "articles"
+                applied_filters["fallback_used"] = True
+                # Message personnalisé pour le fallback
+                if enhanced_docs:
+                    message_supplementaire = (
+                        "⚠️ Nous nous excusons, mais aucune décision de justice n'a été trouvée pour votre requête. "
+                        "La base de données sera enrichie avec des décisions de justice prochainement. "
+                        "En attendant, voici des articles de loi pertinents qui peuvent vous aider."
+                    )
+                else:
+                    # Check if it's a MongoDB error
+                    if "mongodb_error" in applied_filters:
+                        message_supplementaire = (
+                            "⚠️ Nous nous excusons, mais une erreur technique s'est produite lors de la recherche. "
+                            "Nous travaillons à résoudre ce problème. Veuillez réessayer dans quelques instants."
+                        )
+                    else:
+                        message_supplementaire = (
+                            "⚠️ Nous nous excusons, mais aucune décision de justice n'a été trouvée pour votre requête. "
+                            "La base de données sera enrichie avec des décisions de justice prochainement. "
+                            "De plus, aucun article de loi correspondant n'a été trouvé. "
+                            "Essayez de reformuler votre question avec des termes plus généraux."
+                        )
+            logger.info(f"🔍 Search completed: {len(enhanced_docs)} documents found")
+            logger.info(f"📢 Supplemental message: {message_supplementaire[:100] if message_supplementaire else 'None'}")
+            return enhanced_docs, detected_articles, applied_filters, message_supplementaire
+        except Exception as e:
+            logger.error(f"Error in smart_legal_query: {str(e)}")
+            # Return empty results on error
+            return [], [], {"error": str(e)}, f"Erreur lors de la recherche: {str(e)}"
+    async def _perform_search_async(self, user_query: str, country: str, pre_filter: Dict) -> Tuple[List[Document], List[str], Dict[str, Any]]:
+        """Perform search with given filters - ASYNC VERSION"""
+        try:
+            detected_articles = self._detect_articles(user_query)
+            enhanced_query = self._enhance_query(user_query, detected_articles)
+            logger.info(f"🔢 Articles détectés: {detected_articles}")
+            logger.info(f"🔍 Requête enrichie: {enhanced_query[:100]}...")
+            # CRITICAL FIX: Run synchronous vectorstore operation in thread pool
+            docs = await asyncio.get_event_loop().run_in_executor(
+                None,  # Use default thread pool
+                lambda: self.vectorstore.similarity_search(
+                    enhanced_query,
+                    k=settings.MAX_SEARCH_RESULTS,
+                    pre_filter=pre_filter
+                )
+            )
+            logger.info(f"🎯 Vector search returned {len(docs)} raw documents")
+            if docs:
+                logger.info(f"📄 First result metadata: {docs[0].metadata}")
+            else:
+                logger.warning(f"⚠️ No documents found with filters: {pre_filter}")
+                await self._debug_search_issue(pre_filter)
+            enhanced_docs = self.enhance_with_article_context(docs)
+            return enhanced_docs, detected_articles, pre_filter
+        except Exception as e:
+            logger.error(f"Error in _perform_search_async: {str(e)}")
+            # Mark the filter with MongoDB error for better error handling
+            error_filter = pre_filter.copy()
+            error_filter["mongodb_error"] = str(e)
+            return [], [], error_filter
+    async def _debug_search_issue(self, pre_filter: Dict):
+        """Debug why search returned no results"""
+        try:
+            # Check total document count
+            total_count = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: self.collection.count_documents({})
+            )
+            logger.info(f"🔢 Total documents in collection: {total_count}")
+            # Check documents matching country filter
+            country_count = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: self.collection.count_documents({"pays": pre_filter.get("pays")})
+            )
+            logger.info(f"🌍 Documents for country {pre_filter.get('pays')}: {country_count}")
+            # Check documents with doc_type
+            doc_type_count = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: self.collection.count_documents({
+                    "pays": pre_filter.get("pays"),
+                    "doc_type": pre_filter.get("doc_type")
+                })
+            )
+            logger.info(f"📋 Documents with doc_type {pre_filter.get('doc_type')}: {doc_type_count}")
+            # Check documents with embeddings
+            embedding_count = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: self.collection.count_documents({
+                    "pays": pre_filter.get("pays"),
+                    "embedding": {"$exists": True, "$ne": None}
+                })
+            )
+            logger.info(f"🎯 Documents with embeddings: {embedding_count}")
+            # Sample document check
+            sample_doc = await asyncio.get_event_loop().run_in_executor(
+                None,
+                lambda: self.collection.find_one({"pays": pre_filter.get("pays")})
+            )
+            if sample_doc:
+                logger.info(f"📄 Sample document keys: {list(sample_doc.keys())}")
+                logger.info(f"📄 Sample doc_type: {sample_doc.get('doc_type', 'NOT_SET')}")
+            else:
+                logger.error("❌ No sample document found!")
+        except Exception as e:
+            logger.error(f"Error in debug: {str(e)}")
+    def _build_pre_filters(self, query: str, country: str) -> Dict[str, Any]:
+        """Build search filters based on query and country"""
+        # Filtre pays obligatoire - MAKE SURE EXACT MATCH
+        country_mapping = {
+            "benin": "Bénin",
+            "madagascar": "Madagascar"
+        }
+        pre_filter = {"pays": country_mapping.get(country.lower(), country)}
+        # Filtre doc_type pour différencier articles et études de cas
+        query_lower = query.lower()
+        detected_doc_type = self._detect_document_type(query_lower)
+        pre_filter["doc_type"] = detected_doc_type
+        logger.info(f"🏷️ Using country filter: {pre_filter['pays']}")
+        logger.info(f"📋 Using doc_type filter: {detected_doc_type}")
+        # Filtres par catégorie (optionnels)
+        logger.info("ℹ️  No category filter applied - using all available family law documents")
+        # for keyword, category in CATEGORY_KEYWORDS.items():
+        #     if keyword in query_lower:
+        #         pre_filter["categorie"] = category
+        #         logger.info(f"🏷️ Filtre catégorie: {category}")
+        #         break
+        return pre_filter
+    def _detect_document_type(self, query_lower: str) -> str:
+        """Détecte le type de document basé sur les mots-clés de la requête"""
+        # Mots-clés pour les études de cas
+        case_study_indicators = [
+            "jurisprudence", "arrêt", "décision", "tribunal", "cours", "jugement",
+            "affaire", "procès", "litige", "contentieux", "précédent", "cas",
+            "cour d'appel", "cour suprême", "conseil d'état", "juridiction"
+        ]
+        # Mots-clés pour les articles
+        articles_indicators = [
+            "article", "loi", "code", "décret", "texte", "disposition",
+            "règlement", "ordonnance", "prescription", "norme", "chapitre", "titre"
+        ]
+        case_study_score = sum(1 for keyword in case_study_indicators if keyword in query_lower)
+        articles_score = sum(1 for keyword in articles_indicators if keyword in query_lower)
+        if case_study_score > articles_score and case_study_score > 0:
+            return "case_study"
+        elif articles_score > 0:
+            return "articles"
+        else:
+            # Par défaut, on cherche les articles de loi
+            return "articles"
+    def _detect_articles(self, query: str) -> List[str]:
+        """Detect article references in query"""
+        detected_articles = []
+        for pattern in ARTICLE_PATTERNS:
+            matches = re.findall(pattern, query.lower())
+            for match in matches:
+                if isinstance(match, tuple):
+                    nums = [n for n in match if n.isdigit()]
+                    detected_articles.extend(nums)
+                else:
+                    nums = re.findall(r"\d+", match)
+                    detected_articles.extend(nums)
+        return sorted(list(set(detected_articles)))
+    def _enhance_query(self, query: str, detected_articles: List[str]) -> str:
+        """Enhance query with article context"""
+        if detected_articles:
+            enhanced = f"article {' '.join(detected_articles)} {query}"
+            logger.info(f"🔢 Requête enrichie avec articles: {detected_articles}")
+            return enhanced
+        return query
+    def enhance_with_article_context(self, results: List[Document]) -> List[Document]:
+        """Enhance search results with referenced article context"""
+        enhanced_results = []
+        for result in results:
+            enhanced_results.append(result)
+            # Pour les documents de type "articles", on peut ajouter les références
+            if result.metadata.get("doc_type") == "articles":
+                article_refs = result.metadata.get("article_references", [])
+                resolved_refs = result.metadata.get("resolved_references", {})
+                for article_num in article_refs[:3]:
+                    if article_num in resolved_refs:
+                        ref_doc = Document(
+                            page_content=f"Article {article_num} (Référencé): {resolved_refs[article_num][:500]}...",
+                            metadata={
+                                **result.metadata,
+                                "is_reference": True,
+                                "referenced_article": article_num,
+                                "doc_type": "article_reference"
+                            },
+                        )
+                        enhanced_results.append(ref_doc)
+        return enhanced_results
+    def format_search_results(self, query: str, enhanced_docs: List[Document],
+                            detected_articles: List[str], applied_filters: Dict[str, Any],
+                            country: str, supplemental_message: str = "") -> str:
+        """Format search results for system prompt"""
+        country_name = "Bénin" if country == "benin" else "Madagascar"
+        if not enhanced_docs:
+            doc_type = applied_filters.get("doc_type", "articles")
+            # Check if this was an error case
+            if "error" in applied_filters:
+                return f"""
+                    **🚨 ERREUR DE RECHERCHE - {country_name.upper()}**
+                    Une erreur s'est produite lors de la recherche: {applied_filters['error']}
+                    **Informations de débogage:**
+                    - **Requête**: "{query}"
+                    - **Pays**: {country_name}
+                    - **Type de document recherché**: {doc_type}
+                    - **Filtres**: {applied_filters}
+                    Veuillez réessayer ou contacter le support technique.
+                    """
+            if applied_filters.get("fallback_used"):
+                # Cas où le fallback a été utilisé mais n'a rien trouvé non plus
+                mongodb_error_note = ""
+                if "mongodb_error" in applied_filters:
+                    mongodb_error_note = f"\n\n**⚠️ Erreur technique**: {applied_filters['mongodb_error'][:200]}..."
+                return f"""
+                    **🔍 RECHERCHE JURIDIQUE - {country_name.upper()}**
+                    {supplemental_message}
+                    **💡 Informations :**
+                    - Votre recherche portait sur des **décisions de justice (jurisprudence)**
+                    - Aucune décision de justice n'a été trouvée dans la base de données
+                    - Aucun article de loi correspondant n'a été trouvé non plus
+                    {mongodb_error_note}
+                    **Suggestion**: Essayez de reformuler votre requête avec des termes plus généraux.
+                    **Recherche effectuée**:
+                    - Type initial: {applied_filters.get('original_search', 'N/A')}
+                    - Fallback vers: {applied_filters.get('fallback_to', 'N/A')}
+                    - Pays: {country_name}
+                    """
+            else:
+                # Cas normal sans fallback
+                return f"""
+                    **🔍 RECHERCHE JURIDIQUE - {country_name.upper()}**
+                    Aucun document trouvé avec les critères suivants:
+                    - **Type de document**: {doc_type}
+                    - **Catégorie**: {applied_filters.get('categorie', 'Toutes')}
+                    - **Requête**: "{query}"
+                    **Suggestion**: Essayez avec des termes plus généraux ou vérifiez l'orthographe.
+                    **Filtres appliqués**: {applied_filters}
+                    """
+        # Si des documents sont trouvés
+        doc_type = applied_filters.get("doc_type", "articles")
+        doc_type_fr = "articles de loi" if doc_type == "articles" else "études de cas/jurisprudence"
+        fallback_note = ""
+        if applied_filters.get("fallback_used"):
+            fallback_note = f"""
+                **💡 {supplemental_message}**
+                ---
+                """
+        search_results = f"""
+                **🔍 RECHERCHE JURIDIQUE - {country_name.upper()}**
+                **Type de documents**: {doc_type_fr}
+                **Requête**: "{query}"
+                **Juridiction**: {country_name}
+                **Articles détectés**: {', '.join(detected_articles) if detected_articles else 'Aucun'}
+                **Documents trouvés**: {len(enhanced_docs)}
+                {fallback_note}
+                """
+        # Formatage des documents trouvés
+        main_docs = [doc for doc in enhanced_docs if not doc.metadata.get("is_reference", False)]
+        for i, doc in enumerate(main_docs[:5]):
+            doc_type = doc.metadata.get("doc_type", "inconnu")
+            source = doc.metadata.get('source', 'Non spécifié')
+            content = doc.page_content[:600]
+            search_results += f"""
+                    **📄 DOCUMENT {i+1}** (Type: {doc_type})
+                    - **Source**: {source}
+                    - **Contenu**: {content}...
+                    """
+        return search_results
+    # BACKWARD COMPATIBILITY: Keep sync version for any remaining sync calls
+    def smart_legal_query_sync(self, user_query: str, country: str) -> Tuple[List[Document], List[str], Dict[str, Any], str]:
+        """Synchronous version for backward compatibility"""
+        logger.warning("Using sync version of smart_legal_query - consider migrating to async")
+        return asyncio.run(self.smart_legal_query(user_query, country))

core/router.py ADDED Viewed

	@@ -0,0 +1,238 @@

+# [file name]: core/router.py
+import re
+import logging
+import json
+from typing import Dict, List, Optional, Literal, Any
+from langchain_openai import ChatOpenAI
+from langchain_core.messages import SystemMessage, HumanMessage
+from config.settings import settings
+from models.state_models import RoutingResult
+logger = logging.getLogger(__name__)
+class CountryRouter:
+    def __init__(self):
+        self.llm = ChatOpenAI(
+            model=settings.CHAT_MODEL_2,
+            temperature=0.1,
+            max_tokens=200
+        )
+    async def route_query(self, query: str, conversation_history: List[Dict]) -> RoutingResult:
+        """Unified LLM-powered routing"""
+        try:
+            # Build conversation context
+            context = self._build_conversation_context(conversation_history)
+            # LLM routing prompt
+            routing_prompt = self._build_routing_prompt(query, context)
+            logger.info(f"🔀 Routing query: '{query[:50]}...'")
+            # Call LLM for routing decision
+            response = await self.llm.ainvoke([SystemMessage(content=routing_prompt)])
+            routing_result = self._parse_routing_response(response.content)
+            logger.info(f"🎯 Router decision: {routing_result.country} ({routing_result.confidence})")
+            return routing_result
+        except Exception as e:
+            logger.error(f"Router error: {e}")
+            # Fallback to unclear
+            return RoutingResult(
+                country="unclear",
+                confidence="low",
+                method="error_fallback",
+                explanation=f"Router error: {str(e)}"
+            )
+    def _build_routing_prompt(self, query: str, context: str) -> str:
+        """Build comprehensive routing prompt"""
+        return f"""
+Vous êtes un routeur intelligent pour un assistant juridique spécialisé dans le droit du Bénin et de Madagascar.
+**TÂCHE:** Analyser la requête utilisateur et déterminer la meilleure destination.
+**DESTINATIONS POSSIBLES:**
+- "benin": Questions juridiques concernant le Bénin (lois, procédures, droits)
+- "madagascar": Questions juridiques concernant Madagascar (lois, procédures, droits)
+- "assistance_request": Demande pour parler à un avocat humain
+- "greeting_small_talk": Salutations, présentations, remerciements (politesse uniquement)
+- "conversation_repair": Incompréhension, demande de clarification
+- "conversation_summarization": Demande de résumé de la conversation
+- "out_of_scope": Questions NON juridiques (café, météo, sports, recettes, etc.)
+- "unclear": Intention juridique incertaine
+**REQUÊTE:** "{query}"
+**CONTEXTE DE CONVERSATION:**
+{context}
+**RÈGLES DE CLASSIFICATION:**
+1. **greeting_small_talk** - UNIQUEMENT pour politesse basique:
+   - Salutations: "bonjour", "salut", "hello", "bonsoir", "au revoir"
+   - Présentations brèves: "je m'appelle X", "mon nom est X"
+   - Remerciements: "merci", "merci beaucoup"
+   - Politesses simples: "comment ça va", "ça va bien"
+   - Questions sur l'identité de l'assistant: "qui es-tu", "comment tu t'appelles"
+2. **benin** - Pour questions juridiques sur le Bénin:
+   - Mentions explicites: "bénin", "benin", "béninois"
+   - Villes: "cotonou", "porto-novo"
+   - Lois/procédures béninoises
+3. **madagascar** - Pour questions juridiques sur Madagascar:
+   - Mentions explicites: "madagascar", "malgache"
+   - Villes: "antananarivo", "toamasina"
+   - Lois/procédures malgaches
+4. **assistance_request** - Demande d'aide humaine:
+   - "parler à un avocat"
+   - "contacter un avocat"
+   - "assistance téléphonique"
+   - "besoin d'aide juridique personnalisée"
+5. **conversation_repair** - Problèmes de compréhension:
+   - "je n'ai pas compris"
+   - "répète s'il te plaît"
+   - "explique autrement"
+   - "qu'est-ce que tu veux dire"
+6. **conversation_summarization** - Demande de résumé:
+   - "résume notre conversation"
+   - "récapitulatif"
+   - "qu'avons-nous dit"
+7. **out_of_scope** - Questions clairement NON juridiques:
+   - Météo/Climat: "température à Douala", "il va pleuvoir?"
+   - Nourriture: "recette de ndolé", "fais-moi un café"
+   - Sport: "résultat du match", "qui a gagné?"
+   - Technologie: "comment réparer mon téléphone", "meilleur ordinateur"
+   - Divertissement: "raconte une blague", "parle-moi de musique"
+   - Santé non-juridique: "symptômes grippe", "remèdes traditionnels"
+   - **Règle clé**: AUCUN aspect juridique ou lien avec le droit
+8. **unclear** - Questions juridiques MAIS pays/détails manquants:
+   - "J'ai un problème de divorce" (quel pays?)
+   - "Comment créer une entreprise" (Bénin ou Madagascar?)
+   - "Besoin d'aide juridique" (trop vague)
+   - "Question sur l'héritage" (juridiction non précisée)
+   - **Règle clé**: Intention juridique évidente MAIS manque de précision sur le pays ou les détails
+**EXEMPLES COMPLETS:**
+- "Bonjour" → {{"destination": "greeting_small_talk", "confidence": "high", "reasoning": "Salutation simple"}}
+- "je m'appelle Thibaut" → {{"destination": "greeting_small_talk", "confidence": "high", "reasoning": "Présentation personnelle"}}
+- "comment est-ce que je m'appelle" → {{"destination": "greeting_small_talk", "confidence": "high", "reasoning": "Question personnelle de rappel"}}
+- "salut comment ça va" → {{"destination": "greeting_small_talk", "confidence": "high", "reasoning": "Salutation et politesse"}}
+- "merci beaucoup" → {{"destination": "greeting_small_talk", "confidence": "high", "reasoning": "Remerciement"}}
+- "qui es-tu" → {{"destination": "greeting_small_talk", "confidence": "high", "reasoning": "Question sur l'identité de l'assistant"}}
+- "procedure divorce Bénin" → {{"destination": "benin", "confidence": "high", "reasoning": "Question juridique explicite sur le Bénin"}}
+- "loi foncière Madagascar" → {{"destination": "madagascar", "confidence": "high", "reasoning": "Question juridique sur Madagascar"}}
+- "Je veux parler à un avocat" → {{"destination": "assistance_request", "confidence": "high", "reasoning": "Demande explicite d'assistance humaine"}}
+- "Je n'ai pas compris" → {{"destination": "conversation_repair", "confidence": "high", "reasoning": "Demande de clarification"}}
+- "résume notre conversation" → {{"destination": "conversation_summarization", "confidence": "high", "reasoning": "Demande de résumé"}}
+- "fais-moi un café" → {{"destination": "out_of_scope", "confidence": "high", "reasoning": "Demande sans rapport avec le droit"}}
+- "quelle est la météo" → {{"destination": "out_of_scope", "confidence": "high", "reasoning": "Question météorologique, non juridique"}}
+- "température à Douala" → {{"destination": "out_of_scope", "confidence": "high", "reasoning": "Question climatique, hors domaine juridique"}}
+- "raconte-moi une blague" → {{"destination": "out_of_scope", "confidence": "high", "reasoning": "Demande de divertissement, non juridique"}}
+- "J'ai un problème de divorce" → {{"destination": "unclear", "confidence": "medium", "reasoning": "Question juridique mais pays non précisé"}}
+- "Comment créer une entreprise" → {{"destination": "unclear", "confidence": "medium", "reasoning": "Question juridique mais juridiction manquante"}}
+**IMPORTANT:**
+- **out_of_scope**: Questions SANS aucun aspect juridique (météo, sport, nourriture, etc.)
+- **unclear**: Questions AVEC intention juridique MAIS manque de précision sur le pays
+- Les présentations, salutations et remerciements sont "greeting_small_talk"
+- Seules les questions JURIDIQUES avec pays identifié vont vers "benin" ou "madagascar"
+**FORMAT DE RÉPONSE:**
+Répondez UNIQUEMENT au format JSON valide:
+{{
+    "destination": "benin|madagascar|assistance_request|greeting_small_talk|conversation_repair|conversation_summarization|unclear",
+    "confidence": "high|medium|low",
+    "reasoning": "explication brève et claire"
+}}
+**RÉPONSE:**
+"""
+    def _parse_routing_response(self, response_text: str) -> RoutingResult:
+        """Parse LLM routing response"""
+        try:
+            # Extract JSON from response
+            json_match = re.search(r'\{.*\}', response_text, re.DOTALL)
+            if not json_match:
+                raise ValueError("No JSON found in response")
+            result = json.loads(json_match.group())
+            # Validate required fields
+            destination = result.get("destination", "unclear")
+            confidence = result.get("confidence", "low")
+            reasoning = result.get("reasoning", "No reasoning provided")
+            # Map destination to RoutingResult country field
+            valid_destinations = [
+                "benin", "madagascar", "unclear", "greeting_small_talk",
+                "conversation_repair", "assistance_request", "conversation_summarization",
+                "out_of_scope"
+            ]
+            if destination not in valid_destinations:
+                logger.warning(f"Invalid destination from LLM: {destination}, defaulting to unclear")
+                destination = "unclear"
+                confidence = "low"
+                reasoning = f"Destination invalide: {destination}"
+            return RoutingResult(
+                country=destination,
+                confidence=confidence,
+                method="llm_routing",
+                explanation=reasoning
+            )
+        except Exception as e:
+            logger.error(f"Error parsing routing response: {e}")
+            logger.error(f"Raw response: {response_text}")
+            return RoutingResult(
+                country="unclear",
+                confidence="low",
+                method="parse_error",
+                explanation=f"Parse error: {str(e)}"
+            )
+    def _build_conversation_context(self, conversation_history: List[Dict]) -> str:
+        """Build conversation context"""
+        if not conversation_history:
+            return "Aucun contexte de conversation"
+        # Get last 6 messages for context
+        recent_messages = conversation_history[-6:]
+        context_lines = []
+        for msg in recent_messages:
+            role = "Utilisateur" if msg.get("role") in ["user", "human"] else "Assistant"
+            content = msg.get("content", "")
+            context_lines.append(f"{role}: {content}")
+        return "\n".join(context_lines)
+    async def health_check(self) -> Dict[str, Any]:
+        """Router health check"""
+        try:
+            # Test with a simple query
+            test_result = await self.route_query("test", [])
+            return {
+                "status": "healthy",
+                "llm_responding": True,
+                "last_test_result": test_result.model_dump()
+            }
+        except Exception as e:
+            return {
+                "status": "unhealthy",
+                "llm_responding": False,
+                "error": str(e)
+            }

core/routing/__init__.py ADDED Viewed

	@@ -0,0 +1,3 @@


1	+ from .routing_logic import RoutingLogic
2	+
3	+ __all__ = ["RoutingLogic"]

core/routing/routing_logic.py ADDED Viewed

	@@ -0,0 +1,158 @@

+# [file name]: core/routing/routing_logic.py
+import logging
+from typing import Literal
+from models.state_models import MultiCountryLegalState
+logger = logging.getLogger(__name__)
+class RoutingLogic:
+    """Centralized routing logic for graph edges"""
+    def route_after_info_collection(
+        self,
+        state: MultiCountryLegalState
+    ) -> Literal["need_email", "need_description", "ready_to_confirm", "cancelled"]:
+        """Route based on current assistance step and collected data"""
+        step = state.assistance_step
+        has_email = bool(state.user_email)
+        has_description = bool(state.assistance_description)
+        logger.info(f"📋 Assistance step: {step}")
+        logger.info(f"   - Has email: {has_email} ({state.user_email})")
+        logger.info(f"   - Has description: {has_description} ({state.assistance_description})")
+        # 🔥 NEW: Handle cancellation first
+        if step == "cancelled":
+            logger.info("🔄 Assistance workflow cancelled by user")
+            return "cancelled"
+        # Route based on current step progression
+        if step == "collecting_email":
+            if not has_email:
+                logger.info("→ Routing to: need_email (waiting for email)")
+                return "need_email"
+            else:
+                # Email collected, move to description
+                logger.info("→ Routing to: need_description (email collected)")
+                return "need_description"
+        elif step == "collecting_description":
+            if not has_description:
+                logger.info("→ Routing to: need_description (waiting for description)")
+                return "need_description"
+            else:
+                # Description collected, ready for confirmation
+                logger.info("→ Routing to: ready_to_confirm (both collected)")
+                return "ready_to_confirm"
+        elif step == "confirming_send":
+            # We're already in confirmation step - stay here until user confirms
+            logger.info("→ Routing to: ready_to_confirm (awaiting user confirmation)")
+            return "ready_to_confirm"
+        else:
+            # Default fallback logic
+            if not has_email:
+                logger.info("→ Routing to: need_email (default)")
+                return "need_email"
+            elif not has_description:
+                logger.info("→ Routing to: need_description (default)")
+                return "need_description"
+            else:
+                logger.info("→ Routing to: ready_to_confirm (default)")
+                return "ready_to_confirm"
+    def route_after_confirmation(
+        self,
+        state: MultiCountryLegalState
+    ) -> Literal["confirmed", "cancelled", "needs_correction"]:
+        """Route based on user's confirmation response and current step"""
+        step = state.assistance_step
+        last_message = self._get_last_user_message(state)
+        logger.info(f"📋 Confirmation step: {step}")
+        logger.info(f"   - Last user message: '{last_message}'")
+        # 🔥 NEW: Handle cancellation from confirmation step
+        if step == "cancelled":
+            logger.info("→ Routing to: cancelled (workflow cancelled)")
+            return "cancelled"
+        elif step == "confirmed":
+            logger.info("→ Routing to: confirmed (human approval)")
+            return "confirmed"
+        elif step == "confirming_send":
+            # In confirmation step, check user response
+            user_response = last_message.lower().strip() if last_message else ""
+            if user_response in ["oui", "yes", "ok", "confirm", "confirmer", "c'est bon", "d'accord", "envoyer", "valider"]:
+                logger.info("→ Routing to: confirmed (user confirmed)")
+                return "confirmed"
+            elif user_response in ["non", "no", "cancel", "annuler", "pas maintenant", "arrêter", "stop", "je ne veux plus"]:
+                logger.info("→ Routing to: cancelled (user cancelled)")
+                return "cancelled"
+            else:
+                # User provided description or unclear response - go to response to ask again
+                logger.info("→ Routing to: needs_correction (need clarification)")
+                return "needs_correction"
+        else:
+            logger.info("→ Routing to: needs_correction (invalid state)")
+            return "needs_correction"
+    def route_after_human_approval(
+        self,
+        state: MultiCountryLegalState
+    ) -> Literal["approved", "rejected", "interrupt"]:
+        """Route based on human approval status"""
+        approval_status = state.approval_status
+        logger.info(f"📋 Approval status: {approval_status}")
+        if approval_status == "approved":
+            logger.info("→ Routing to: approved (process assistance)")
+            return "approved"
+        elif approval_status == "rejected":
+            logger.info("→ Routing to: rejected (response)")
+            return "rejected"
+        else:
+            # Still waiting for approval or error state
+            logger.info("→ Routing to: interrupt (waiting for decision)")
+            return "interrupt"
+    def _get_last_user_message(self, state: MultiCountryLegalState) -> str:
+        """Extract the last user message from state"""
+        if not state.messages:
+            return ""
+        for msg in reversed(state.messages):
+            if hasattr(msg, 'role'):
+                role = msg.role
+            else:
+                role = msg.get('role', '')
+            if role in ['user', 'human']:
+                if hasattr(msg, 'content'):
+                    return msg.content
+                else:
+                    return msg.get('content', '')
+        return ""
+    def _looks_like_description(self, text: str) -> bool:
+        """Check if text looks like a description rather than a confirmation"""
+        description_indicators = [
+            "j'ai besoin", "je veux", "je souhaite", "aide pour", "divorce",
+            "mariage", "héritage", "contrat", "travail", "familial", "bénin", "madagascar",
+            "problème", "situation", "question", "demande"
+        ]
+        text_lower = text.lower()
+        return any(indicator in text_lower for indicator in description_indicators)

core/system_initializer.py ADDED Viewed

	@@ -0,0 +1,103 @@

+# [file name]: core/system_initializer.py
+import logging
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+from core.graph_builder import GraphBuilder
+from core.chat_manager import LegalChatManager
+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
+logger = logging.getLogger(__name__)
+async def setup_system():
+    """Initialize the legal assistant system for API use"""
+    try:
+        # 1. Initialize MongoDB using your existing class
+        mongo_client = MongoDBClient()
+        if not mongo_client.connect():
+            raise Exception("MongoDB connection failed")
+        logger.info("✅ MongoDB connected successfully")
+        # 2. Use your existing vector stores directly from the client
+        vector_store_benin = mongo_client.benin_vectorstore
+        collection_benin = mongo_client.benin_collection
+        vector_store_madagascar = mongo_client.madagascar_vectorstore
+        collection_madagascar = mongo_client.madagascar_collection
+        # 3. Initialize retrievers
+        from core.retriever import LegalRetriever
+        benin_retriever = LegalRetriever(vector_store_benin, collection_benin)
+        madagascar_retriever = LegalRetriever(vector_store_madagascar, collection_madagascar)
+        country_retrievers = {
+            "benin": benin_retriever,
+            "madagascar": madagascar_retriever
+        }
+        # 4. Initialize LLM and router
+        llm = ChatOpenAI(
+            model="gpt-4o-mini",
+            temperature=0.1,
+            max_tokens=2000,
+            streaming=True
+        )
+        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")
+        # 6. Build graph
+        graph_builder = GraphBuilder(
+            router=router,
+            llm=llm,
+            checkpointer=checkpointer,
+            country_retrievers=country_retrievers
+        )
+        workflow = graph_builder.build_graph()
+        app = workflow.compile(checkpointer=checkpointer)
+        # 7. Initialize chat manager
+        chat_manager = LegalChatManager(app, checkpointer)
+        logger.info("✅ API System initialized successfully")
+        return {
+            "chat_manager": chat_manager,
+            "graph": app,
+            "checkpointer": checkpointer
+        }
+    except Exception as e:
+        logger.error(f"❌ Failed to initialize system: {e}")
+        raise

database/__init__py ADDED Viewed

File without changes

database/mongodb_client.py ADDED Viewed

	@@ -0,0 +1,153 @@

+from pymongo import MongoClient, ReadPreference
+from pymongo.errors import ServerSelectionTimeoutError, ConnectionFailure
+from langchain_mongodb.vectorstores import MongoDBAtlasVectorSearch
+from langchain_openai import OpenAIEmbeddings
+from typing import Dict
+import logging
+from config.settings import settings
+logger = logging.getLogger(__name__)
+class MongoDBClient:
+    def __init__(self):
+        self.client = None
+        self.db = None
+        self.benin_collection = None
+        self.madagascar_collection = None
+        self.benin_vectorstore = None
+        self.madagascar_vectorstore = None
+        self.embedding_model = None
+    def connect(self):
+        """Connect to MongoDB and initialize collections"""
+        try:
+            # CRITICAL FIX: Add read preference to allow reading from secondary nodes
+            self.client = MongoClient(
+                settings.MONGO_URI,
+                # Allow reading from secondary nodes when primary is unavailable
+                read_preference=ReadPreference.SECONDARY_PREFERRED,
+                # Reduce timeouts to fail faster (instead of 30s)
+                serverSelectionTimeoutMS=10000,  # 10 seconds
+                connectTimeoutMS=10000,
+                socketTimeoutMS=10000,
+                # Retry configuration
+                retryWrites=True,
+                retryReads=True,
+                # Connection pool settings
+                maxPoolSize=50,
+                minPoolSize=10,
+                # Write concern (for writes to still work)
+                w='majority',
+                journal=True
+            )
+            # Test the connection
+            self.client.admin.command('ping')
+            logger.info("✅ MongoDB connection test successful")
+            self.db = self.client[settings.DATABASE_NAME]
+            # Initialize collections
+            self.benin_collection = self.db[settings.BENIN_COLLECTION]
+            self.madagascar_collection = self.db[settings.MADAGASCAR_COLLECTION]
+            # Verify collections exist and have data
+            benin_count = self.benin_collection.count_documents({})
+            madagascar_count = self.madagascar_collection.count_documents({})
+            logger.info(f"📊 Bénin collection: {benin_count} documents")
+            logger.info(f"📊 Madagascar collection: {madagascar_count} documents")
+            # Initialize embedding model
+            self.embedding_model = OpenAIEmbeddings(
+                model=settings.EMBEDDING_MODEL,
+                openai_api_key=settings.OPENAI_API_KEY
+            )
+            # Initialize vector stores with read preference
+            self.benin_vectorstore = MongoDBAtlasVectorSearch(
+                collection=self.benin_collection,
+                embedding=self.embedding_model,
+                index_name=settings.VECTOR_INDEX_NAME,
+                text_key=settings.TEXT_KEY,
+                embedding_key=settings.EMBEDDING_KEY,
+            )
+            self.madagascar_vectorstore = MongoDBAtlasVectorSearch(
+                collection=self.madagascar_collection,
+                embedding=self.embedding_model,
+                index_name=settings.VECTOR_INDEX_NAME,
+                text_key=settings.TEXT_KEY,
+                embedding_key=settings.EMBEDDING_KEY,
+            )
+            print("✅ MongoDB connected successfully with SECONDARY_PREFERRED read preference")
+            return True
+        except (ServerSelectionTimeoutError, ConnectionFailure) as e:
+            logger.error(f"❌ MongoDB connection failed: {e}")
+            logger.error("🔍 Possible issues:")
+            logger.error("   1. MongoDB Atlas cluster is paused")
+            logger.error("   2. Network connectivity issues")
+            logger.error("   3. IP address not whitelisted in Atlas")
+            logger.error("   4. Cluster is undergoing maintenance")
+            print(f"❌ MongoDB connection failed: {e}")
+            return False
+        except Exception as e:
+            logger.error(f"❌ Unexpected error during MongoDB connection: {e}")
+            print(f"❌ MongoDB connection failed: {e}")
+            return False
+    def get_collection_stats(self) -> Dict:
+        """Get statistics for both collections"""
+        if not self.client:
+            return {}
+        try:
+            benin_count = self.benin_collection.count_documents({})
+            madagascar_count = self.madagascar_collection.count_documents({})
+            # Sample document to check schema
+            benin_sample = self.benin_collection.find_one()
+            madagascar_sample = self.madagascar_collection.find_one()
+            # Check for documents by doc_type
+            benin_case_study_count = self.benin_collection.count_documents({"doc_type": "case_study"})
+            benin_articles_count = self.benin_collection.count_documents({"doc_type": "articles"})
+            madagascar_case_study_count = self.madagascar_collection.count_documents({"doc_type": "case_study"})
+            madagascar_articles_count = self.madagascar_collection.count_documents({"doc_type": "articles"})
+            return {
+                "benin": {
+                    "total_documents": benin_count,
+                    "case_study_count": benin_case_study_count,
+                    "articles_count": benin_articles_count,
+                    "has_embeddings": bool(benin_sample and 'vecteur_embedding' in benin_sample),
+                    "sample_fields": list(benin_sample.keys()) if benin_sample else [],
+                    "sample_doc_type": benin_sample.get('doc_type', 'NOT_SET') if benin_sample else None
+                },
+                "madagascar": {
+                    "total_documents": madagascar_count,
+                    "case_study_count": madagascar_case_study_count,
+                    "articles_count": madagascar_articles_count,
+                    "has_embeddings": bool(madagascar_sample and 'vecteur_embedding' in madagascar_sample),
+                    "sample_fields": list(madagascar_sample.keys()) if madagascar_sample else [],
+                    "sample_doc_type": madagascar_sample.get('doc_type', 'NOT_SET') if madagascar_sample else None
+                }
+            }
+        except Exception as e:
+            logger.error(f"Error getting collection stats: {e}")
+            print(f"Error getting collection stats: {e}")
+            return {}
+    def close(self):
+        """Close MongoDB connection"""
+        if self.client:
+            self.client.close()
+            logger.info("✅ MongoDB connection closed")
+            print("✅ MongoDB connection closed")

database/postgres_checkpointer.py ADDED Viewed

	@@ -0,0 +1,97 @@

+# database/postgres_checkpointer.py - CORRECT VERSION
+from psycopg_pool import AsyncConnectionPool
+from psycopg.rows import dict_row
+from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver  # ✅ Correct import
+from langgraph.checkpoint.memory import MemorySaver
+import logging
+from typing import Optional
+logger = logging.getLogger(__name__)
+class PostgresCheckpointer:
+    def __init__(self, database_url: str, max_connections: int = 10, min_connections: int = 2):
+        self.database_url = database_url
+        self.max_connections = max_connections
+        self.min_connections = min_connections
+        self.pool: Optional[AsyncConnectionPool] = None
+        self.checkpointer: Optional[AsyncPostgresSaver] = None  # ✅ Correct type
+        self._is_initialized = False
+    async def initialize(self) -> bool:
+        """Initialize PostgreSQL connection pool and checkpointer"""
+        try:
+            # Create async connection pool
+            self.pool = AsyncConnectionPool(
+                conninfo=self.database_url,
+                max_size=self.max_connections,
+                min_size=self.min_connections,
+                kwargs={"row_factory": dict_row, "autocommit": True},
+                open=False,
+            )
+            await self.pool.open()
+            # ✅ CORRECT: Use AsyncPostgresSaver with AsyncConnectionPool
+            self.checkpointer = AsyncPostgresSaver(self.pool)
+            await self.checkpointer.setup()  # ✅ Async setup method
+            self._is_initialized = True
+            logger.info("✅ PostgreSQL checkpointer initialized successfully with AsyncPostgresSaver")
+            return True
+        except Exception as e:
+            logger.error(f"❌ PostgreSQL initialization failed: {e}")
+            # Fallback to in-memory
+            try:
+                from langgraph.checkpoint.memory_aio import AsyncMemorySaver  # ✅ Async memory saver
+                self.checkpointer = AsyncMemorySaver()
+                logger.warning("🔄 Falling back to async in-memory checkpointer")
+                self._is_initialized = True
+                return True
+            except ImportError:
+                # Fallback to sync MemorySaver if async not available
+                self.checkpointer = MemorySaver()
+                logger.warning("🔄 Falling back to sync in-memory checkpointer")
+                self._is_initialized = True
+                return True
+            except Exception as fallback_error:
+                logger.error(f"❌ Even fallback failed: {fallback_error}")
+                return False
+    async def close(self):
+        """Close connections with proper cleanup"""
+        if self.pool:
+            await self.pool.close()
+            logger.info("✅ PostgreSQL connection pool closed")
+        self._is_initialized = False
+    async def health_check(self) -> dict:
+        """Check the health of the PostgreSQL connection"""
+        if not self._is_initialized or not self.pool:
+            return {"status": "uninitialized", "healthy": False}
+        try:
+            async with self.pool.connection() as conn:
+                async with conn.cursor() as cur:
+                    await cur.execute("SELECT 1")
+                    result = await cur.fetchone()
+            return {
+                "status": "healthy",
+                "healthy": True,
+                "connection_count": self.pool.size if hasattr(self.pool, 'size') else "unknown"
+            }
+        except Exception as e:
+            return {"status": f"unhealthy: {str(e)}", "healthy": False}
+    def is_initialized(self) -> bool:
+        """Check if checkpointer is properly initialized"""
+        return self._is_initialized and self.checkpointer is not None
+    def get_checkpointer(self) -> AsyncPostgresSaver:
+        """Get the underlying checkpointer instance"""
+        if not self.is_initialized():
+            raise RuntimeError("Checkpointer not initialized. Call initialize() first.")
+        return self.checkpointer

generate_graph.py ADDED Viewed

	@@ -0,0 +1,66 @@

+# generate_graph.py (example)
+from graphviz import Digraph
+def generate_graph():
+    dot = Digraph(comment='Legal RAG System Workflow')
+    dot.attr(rankdir='TB', size='8,5')
+    dot.attr('node', shape='box', style='filled', fillcolor='#e6f3ff')
+    # Core nodes
+    dot.node("START", fillcolor="#90ee90")
+    dot.node("ROUTER", label="router")
+    dot.node("RESPONSE", label="response", fillcolor="#98fb98")
+    # Country nodes
+    dot.node("BENIN_RETRIEVAL", label="benin_retrieval")
+    dot.node("MADAGASCAR_RETRIEVAL", label="madagascar_retrieval")
+    # Handler nodes
+    dot.node("GREETING", label="greeting_handler")
+    dot.node("REPAIR", label="repair_handler")
+    dot.node("SUMMARY", label="summary_handler")
+    dot.node("UNCLEAR", label="unclear_handler")
+    dot.node("OUT_OF_SCOPE", label="out_of_scope_handler")
+    # Assistance nodes
+    dot.node("ASSIST_COLLECT", label="assistance_collect_info")
+    dot.node("ASSIST_CONFIRM", label="assistance_confirm")
+    dot.node("HUMAN_APPROVAL", label="human_approval", fillcolor="#ffa07a")
+    # End node
+    dot.node("END", fillcolor="#ff9999")
+    # Edges
+    dot.edge("START", "ROUTER")
+    dot.edge("ROUTER", "BENIN_RETRIEVAL", label="benin")
+    dot.edge("ROUTER", "MADAGASCAR_RETRIEVAL", label="madagascar")
+    dot.edge("ROUTER", "GREETING", label="greeting_small_talk")
+    dot.edge("ROUTER", "REPAIR", label="conversation_repair")
+    dot.edge("ROUTER", "SUMMARY", label="conversation_summarization")
+    dot.edge("ROUTER", "UNCLEAR", label="unclear")
+    dot.edge("ROUTER", "OUT_OF_SCOPE", label="out_of_scope")
+    dot.edge("ROUTER", "ASSIST_COLLECT", label="assistance_request")
+    dot.edge("GREETING", "RESPONSE")
+    dot.edge("REPAIR", "RESPONSE")
+    dot.edge("SUMMARY", "RESPONSE")
+    dot.edge("UNCLEAR", "RESPONSE")
+    dot.edge("OUT_OF_SCOPE", "RESPONSE")
+    dot.edge("ASSIST_COLLECT", "RESPONSE", label="need_email/need_description")
+    dot.edge("ASSIST_COLLECT", "ASSIST_CONFIRM", label="ready_to_confirm")
+    dot.edge("ASSIST_COLLECT", "RESPONSE", label="cancelled")
+    dot.edge("ASSIST_CONFIRM", "HUMAN_APPROVAL", label="confirmed")
+    dot.edge("ASSIST_CONFIRM", "RESPONSE", label="needs_correction/cancelled")
+    dot.edge("HUMAN_APPROVAL", "RESPONSE")
+    dot.edge("RESPONSE", "ASSIST_COLLECT", label="continue_assistance")
+    dot.edge("RESPONSE", "END", label="end")
+    dot.render('legal_rag_workflow', format='png', cleanup=True)
+    print("Graph visualization generated: legal_rag_workflow.png")
+if __name__ == "__main__":
+    generate_graph()

interfaces/__init__.py ADDED Viewed

File without changes

interfaces/monitoring.py ADDED Viewed

	@@ -0,0 +1,109 @@

+from datetime import datetime
+from typing import Dict, List
+import logging
+class LegalRAGMonitor:
+    """Monitoring and error tracking for the legal RAG system"""
+    def __init__(self):
+        self.error_log = []
+        self.performance_metrics = {
+            "query_times": [],
+            "routing_accuracy": [],
+            "retrieval_success_rate": 0
+        }
+        self.alerts = []
+    def log_error(self, error_type: str, message: str, context: Dict = None):
+        """Log errors for analysis"""
+        error_entry = {
+            "timestamp": datetime.now(),
+            "type": error_type,
+            "message": message,
+            "context": context or {}
+        }
+        self.error_log.append(error_entry)
+        logging.error(f"[{error_type}] {message}")
+        # Check for alert conditions
+        self._check_alerts(error_type, error_entry)
+    def track_query_performance(self, query_time: float, success: bool):
+        """Track query performance metrics"""
+        self.performance_metrics["query_times"].append(query_time)
+        # Update success rate
+        current_rate = self.performance_metrics["retrieval_success_rate"]
+        total_queries = len(self.performance_metrics["query_times"])
+        if success:
+            self.performance_metrics["retrieval_success_rate"] = (
+                (current_rate * (total_queries - 1) + 1) / total_queries
+            )
+    def get_health_report(self) -> Dict:
+        """Generate system health report"""
+        query_times = self.performance_metrics["query_times"]
+        return {
+            "error_count": len(self.error_log),
+            "recent_errors": self.error_log[-5:],
+            "avg_query_time": sum(query_times) / len(query_times) if query_times else 0,
+            "success_rate": self.performance_metrics["retrieval_success_rate"],
+            "total_queries": len(query_times),
+            "active_alerts": len(self.alerts)
+        }
+    def _check_alerts(self, error_type: str, error_entry: Dict):
+        """Check if error should trigger an alert"""
+        # Example alert conditions
+        if error_type == "database_connection":
+            self.alerts.append({
+                "type": "critical",
+                "message": "Database connection failure",
+                "timestamp": datetime.now(),
+                "error": error_entry
+            })
+        # Clean old alerts (keep only last 24 hours)
+        cutoff_time = datetime.now().timestamp() - (24 * 3600)
+        self.alerts = [
+            alert for alert in self.alerts
+            if alert["timestamp"].timestamp() > cutoff_time
+        ]
+class AlertManager:
+    """Manage system alerts and notifications"""
+    def __init__(self):
+        self.alerts = []
+        self.subscribers = []
+    def add_alert(self, alert_type: str, message: str, severity: str = "warning"):
+        """Add a new alert"""
+        alert = {
+            "type": alert_type,
+            "message": message,
+            "severity": severity,
+            "timestamp": datetime.now(),
+            "acknowledged": False
+        }
+        self.alerts.append(alert)
+        self._notify_subscribers(alert)
+    def acknowledge_alert(self, alert_index: int):
+        """Acknowledge an alert"""
+        if 0 <= alert_index < len(self.alerts):
+            self.alerts[alert_index]["acknowledged"] = True
+    def subscribe(self, callback):
+        """Subscribe to alert notifications"""
+        self.subscribers.append(callback)
+    def _notify_subscribers(self, alert):
+        """Notify all subscribers of a new alert"""
+        for subscriber in self.subscribers:
+            try:
+                subscriber(alert)
+            except Exception as e:
+                logging.error(f"Error notifying subscriber: {e}")

interfaces/web_interface.py ADDED Viewed

	@@ -0,0 +1,96 @@

+from fastapi import FastAPI, HTTPException, Depends
+from fastapi.middleware.cors import CORSMiddleware
+from pydantic import BaseModel
+from typing import Dict, Optional
+import uvicorn
+from datetime import datetime
+from core.chat_manager import LegalChatManager
+# Pydantic models for API
+class ChatRequest(BaseModel):
+    query: str
+    session_id: Optional[str] = None
+    context: Optional[Dict] = None
+class ChatResponse(BaseModel):
+    response: str
+    session_id: str
+    session_stats: Dict
+    error: Optional[str] = None
+class HealthResponse(BaseModel):
+    status: str
+    stats: Dict
+    timestamp: str
+class LegalRAGAPI:
+    def __init__(self, chat_manager: LegalChatManager):
+        self.app = FastAPI(title="Legal RAG API", version="1.0.0")
+        self.chat_manager = chat_manager
+        self._setup_middleware()
+        self._setup_routes()
+    def _setup_middleware(self):
+        """Setup CORS and other middleware"""
+        self.app.add_middleware(
+            CORSMiddleware,
+            allow_origins=["*"],
+            allow_credentials=True,
+            allow_methods=["*"],
+            allow_headers=["*"],
+        )
+    def _setup_routes(self):
+        """Setup API routes"""
+        @self.app.get("/")
+        async def root():
+            return {"message": "Legal RAG API is running"}
+        @self.app.post("/chat", response_model=ChatResponse)
+        async def chat_endpoint(request: ChatRequest):
+            try:
+                session_id = request.session_id or f"web_{datetime.now().timestamp()}"
+                response = await self.chat_manager.chat(
+                    request.query,
+                    session_id,
+                    request.context
+                )
+                session_stats = self.chat_manager.get_session_stats(session_id)
+                return ChatResponse(
+                    response=response,
+                    session_id=session_id,
+                    session_stats=session_stats,
+                    error=None
+                )
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+        @self.app.get("/health", response_model=HealthResponse)
+        async def health_check():
+            return HealthResponse(
+                status="healthy",
+                stats=self.chat_manager.get_global_stats(),
+                timestamp=datetime.now().isoformat()
+            )
+        @self.app.get("/sessions/{session_id}/history")
+        async def get_session_history(session_id: str):
+            try:
+                history = await self.chat_manager.get_conversation_history(session_id)
+                return {
+                    "session_id": session_id,
+                    "message_count": len(history),
+                    "messages": history
+                }
+            except Exception as e:
+                raise HTTPException(status_code=500, detail=str(e))
+    def run(self, host: str = "0.0.0.0", port: int = 8000):
+        """Run the API server"""
+        uvicorn.run(self.app, host=host, port=port)

main.py ADDED Viewed

	@@ -0,0 +1,629 @@

+#!/usr/bin/env python3
+"""
+Scalable Multi-Country Legal RAG System
+Supports dynamic addition of new countries with clean architecture
+"""
+import asyncio
+import logging
+import time
+from datetime import datetime
+from typing import List, Dict, Any, Optional
+from config.settings import settings
+from database.mongodb_client import MongoDBClient
+from database.postgres_checkpointer import PostgresCheckpointer
+from core.router import CountryRouter
+from core.retriever import LegalRetriever
+from core.graph_builder import GraphBuilder
+from core.chat_manager import LegalChatManager
+from utils.logger import setup_logging
+import uuid
+class MultiCountryLegalRAGSystem:
+    """Scalable system class supporting dynamic country addition"""
+    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.chat_manager = None
+        self.initialized = False
+    async def initialize(self) -> bool:
+        """Initialize the complete scalable 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")
+            # Initialize core components
+            self.router = CountryRouter()
+            # Initialize default countries - easily extensible!
+            self._initialize_default_countries()
+            # 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()
+            return True
+        except Exception as e:
+            logging.error(f"❌ System initialization failed: {e}")
+            import traceback
+            traceback.print_exc()
+            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
+            )
+        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
+        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)
+        # Test PostgreSQL connection
+        if hasattr(self.postgres_checkpointer, 'health_check'):
+            postgres_health = await self.postgres_checkpointer.health_check()
+            health_status["postgres_healthy"] = postgres_health
+        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")
+        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")
+            session_id = session_id or f"cli_{datetime.now().strftime('%Y%m%d_%H%M%S')}"
+            return await self.chat_manager.chat(message, session_id, ctx)
+        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}")
+        # session_id = f"test_{datetime.now().strftime('%H%M%S%f')}"
+        session_id = f"interactive_{uuid.uuid4().hex[:8]}"
+        current_response = ""
+        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]}...")
+            # 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]}...")
+            # 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]}...")
+            # 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]}...")
+            # Step 5: Check for interrupt
+            interrupt_detected = self._check_for_interrupt(current_response, session_id)
+            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
+                }
+            self.test_results.append(result)
+            return result
+        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
+    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']}")
+        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')}")
+        print("="*80)
+async def run_interrupt_tests():
+    """Run specialized tests for human approval interrupts"""
+    system = MultiCountryLegalRAGSystem()
+    tester = InterruptTester(system)
+    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()
+    try:
+        print("🚀 Initializing system...")
+        success = await system.initialize()
+        if not success:
+            print("❌ System initialization failed")
+            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()
+async def health_check_mode():
+    """Run system health check only"""
+    system = MultiCountryLegalRAGSystem()
+    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'}")
+            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'}")
+        else:
+            print("❌ System initialization failed")
+    finally:
+        await system.cleanup()
+async def quick_test_mode():
+    """Run a quick single test"""
+    system = MultiCountryLegalRAGSystem()
+    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())}")
+    except Exception as e:
+        print(f"❌ Quick test failed: {e}")
+    finally:
+        await system.cleanup()
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(
+        description="🚀 Scalable Multi-Country Legal RAG System"
+    )
+    parser.add_argument(
+        "--mode",
+        choices=["interactive", "health", "interrupt", "quick"],
+        default="interactive",
+        help="Run mode (default: interactive)"
+    )
+    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())

models/__init__py ADDED Viewed

File without changes

models/state_models.py ADDED Viewed

	@@ -0,0 +1,112 @@

+# [file name]: models/state_models.py
+from typing import List, Dict, Any, Optional, Annotated, Literal, Union
+from pydantic import BaseModel, Field
+import operator
+class MultiCountryLegalState(BaseModel):
+    messages: Annotated[List[Dict[str, Any]], operator.add] = Field(default_factory=list)
+    legal_context: Dict[str, Any] = Field(
+        default_factory=lambda: {
+            "jurisdiction": "Unknown",
+            "user_type": "general",
+            "document_type": "legal",
+            "detected_country": "unknown"
+        }
+    )
+    # FIX: Make supplemental_message handle concurrent updates
+    supplemental_message: Optional[str] = Field(
+        default="",
+        description="Supplemental message to display to user (e.g., fallback messages, apologies)"
+    )
+    session_id: Optional[str] = None
+    last_search_query: Optional[str] = None
+    detected_articles: Annotated[List[str], operator.add] = Field(default_factory=list)
+    router_decision: Optional[str] = None
+    search_results: Optional[str] = None
+    route_explanation: Optional[str] = None
+    country: Optional[str] = Field(default=None)
+    # Assistance email fields
+    assistance_requested: bool = Field(default=False)
+    user_email: Optional[str] = None
+    assistance_description: Optional[str] = None
+    email_status: Optional[str] = None  # "pending", "sent", "error"
+    assistance_step: Optional[str] = Field(default=None)  # "collecting_email", "collecting_description", "confirming_send"
+    pending_assistance_data: Dict[str, Any] = Field(default_factory=dict)
+    # Conversation repair tracking
+    repair_type: Optional[str] = None
+    original_query: Optional[str] = None
+    misunderstanding_count: int = Field(default=0)
+    # Enhanced routing support
+    primary_intent: Optional[str] = Field(default=None)
+    # NEW: Human approval fields
+    approval_status: Optional[str] = Field(default=None)  # "pending", "approved", "rejected"
+    approval_reason: Optional[str] = Field(default=None)
+    approved_by: Optional[str] = Field(default=None)
+    approval_timestamp: Optional[str] = Field(default=None)
+    # Conversation summary fields
+    summary_generated: bool = Field(default=False)
+    last_summary_timestamp: Optional[str] = Field(default=None)
+    # NEW: Search-related fields to prevent storing complex data in legal_context
+    search_metadata: Dict[str, Any] = Field(default_factory=dict)
+    @staticmethod
+    def detect_country(text: str) -> str:
+        """
+        Detect country from text based on keywords.
+        Args:
+            text: User input text to analyze
+        Returns:
+            Country code: "benin", "madagascar", or "unknown"
+        """
+        if not text:
+            return "unknown"
+        text_lower = text.lower()
+        # Benin keywords
+        benin_keywords = [
+            "bénin", "benin", "béninois", "béninoise",
+            "cotonou", "porto-novo", "porto novo",
+            "dahomey"  # Historical name
+        ]
+        # Madagascar keywords
+        madagascar_keywords = [
+            "madagascar", "malgache", "malagasy",
+            "antananarivo", "tananarive", "tana",
+            "toamasina", "tamatave"
+        ]
+        # Check for country mentions
+        benin_score = sum(1 for keyword in benin_keywords if keyword in text_lower)
+        madagascar_score = sum(1 for keyword in madagascar_keywords if keyword in text_lower)
+        if benin_score > madagascar_score and benin_score > 0:
+            return "benin"
+        elif madagascar_score > benin_score and madagascar_score > 0:
+            return "madagascar"
+        return "unknown"
+class RoutingResult(BaseModel):
+    country: Literal["benin", "madagascar", "unclear", "greeting_small_talk",
+                 "conversation_repair", "assistance_request", "conversation_summarization", "out_of_scope"]
+    confidence: Literal["high", "medium", "low"]
+    method: str
+    explanation: str
+class SearchResult(BaseModel):
+    documents: List[Any]
+    detected_articles: List[str]
+    applied_filters: Dict[str, Any]
+    query: str
+    country: str

requirements.txt ADDED Viewed

	@@ -0,0 +1,13 @@

+# Core dependencies
+fastapi==0.104.1
+uvicorn[standard]==0.24.0
+langchain-openai==0.0.8
+langchain-core==0.1.33
+langgraph==0.0.52
+langchain-mongodb==0.0.3
+pymongo==4.6.1
+openai==1.3.9
+pydantic==2.5.0
+python-dotenv==1.0.0
+psycopg[binary]==3.1.13
+langgraph-checkpoint-postgres==0.0.3

utils/__init__.py ADDED Viewed

File without changes

utils/helpers.py ADDED Viewed

	@@ -0,0 +1,68 @@

+from typing import Dict, Any
+from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
+def dict_to_message_obj(d: Dict[str, Any]) -> BaseMessage:
+    """Convert dictionary to LangChain message object"""
+    role = d.get("role", "").lower()
+    content = d.get("content", "")
+    meta = d.get("meta", {}) or {}
+    if role in ("user", "human", "humanmessage"):
+        return HumanMessage(content=content, metadata=meta)
+    if role in ("assistant", "ai", "aimessage"):
+        return AIMessage(content=content, metadata=meta)
+    return SystemMessage(content=content, metadata=meta)
+def message_obj_to_dict(msg: Any) -> Dict[str, Any]:
+    """Convert LangChain message object to dictionary"""
+    content = getattr(msg, "content", str(msg))
+    meta = getattr(msg, "metadata", {}) or {}
+    if isinstance(msg, HumanMessage):
+        role = "user"
+    elif isinstance(msg, AIMessage):
+        role = "assistant"
+    elif isinstance(msg, SystemMessage):
+        role = "system"
+    else:
+        role = meta.get("role", "assistant")
+    return {"role": role, "content": content, "meta": meta}
+def validate_country_code(country: str) -> str:
+    """Validate and normalize country code"""
+    country = country.lower().strip()
+    if country in ["benin", "bj", "bénin"]:
+        return "benin"
+    elif country in ["madagascar", "mg", "madagasikara"]:
+        return "madagascar"
+    else:
+        return "unclear"
+def format_legal_citation(article_number: str, law_title: str, country: str) -> str:
+    """Format legal citation in standard format"""
+    country_formats = {
+        "benin": f"Article {article_number} du {law_title} (Bénin)",
+        "madagascar": f"Article {article_number} du {law_title} (Madagascar)"
+    }
+    return country_formats.get(country, f"Article {article_number} du {law_title}")
+def safe_get(dictionary: Dict, key: str, default: Any = None) -> Any:
+    """Safely get value from dictionary with default"""
+    if isinstance(dictionary, dict):
+        return dictionary.get(key, default)
+    return default
+def truncate_text(text: str, max_length: int = 500) -> str:
+    """Truncate text to specified length"""
+    if len(text) <= max_length:
+        return text
+    return text[:max_length] + "..."
+def calculate_confidence_score(patterns_found: int, llm_confidence: str) -> float:
+    """Calculate a numerical confidence score"""
+    pattern_score = min(patterns_found * 0.3, 0.6)  # Max 0.6 from patterns
+    llm_scores = {"high": 0.8, "medium": 0.5, "low": 0.2}
+    llm_score = llm_scores.get(llm_confidence, 0.2)
+    return min(pattern_score + llm_score, 1.0)

utils/logger.py ADDED Viewed

	@@ -0,0 +1,74 @@

+import logging
+import sys
+from datetime import datetime
+from typing import Dict, Any
+def setup_logging(level=logging.INFO):
+    """Setup comprehensive logging configuration"""
+    # Create formatter
+    formatter = logging.Formatter(
+        '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+        datefmt='%Y-%m-%d %H:%M:%S'
+    )
+    # Console handler
+    console_handler = logging.StreamHandler(sys.stdout)
+    console_handler.setLevel(level)
+    console_handler.setFormatter(formatter)
+    # File handler
+    file_handler = logging.FileHandler(f'legal_rag_{datetime.now().strftime("%Y%m%d")}.log')
+    file_handler.setLevel(logging.DEBUG)
+    file_handler.setFormatter(formatter)
+    # Configure root logger
+    logging.basicConfig(
+        level=level,
+        handlers=[console_handler, file_handler],
+        force=True
+    )
+    # Specific logger configurations
+    legal_logger = logging.getLogger('legal_rag')
+    legal_logger.setLevel(logging.DEBUG)
+    mongodb_logger = logging.getLogger('pymongo')
+    mongodb_logger.setLevel(logging.WARNING)
+    print("✅ Logging setup completed")
+class PerformanceLogger:
+    """Logger for performance monitoring"""
+    def __init__(self):
+        self.metrics = {
+            "query_times": [],
+            "routing_times": [],
+            "retrieval_times": [],
+            "generation_times": []
+        }
+    def log_query_time(self, session_id: str, duration: float):
+        """Log query processing time"""
+        self.metrics["query_times"].append({
+            "session_id": session_id,
+            "duration": duration,
+            "timestamp": datetime.now()
+        })
+        logging.info(f"Query processed in {duration:.2f}s for session {session_id}")
+    def log_routing_decision(self, session_id: str, decision: str, confidence: str, method: str):
+        """Log routing decisions"""
+        logging.debug(f"Routing: session={session_id}, decision={decision}, confidence={confidence}, method={method}")
+    def get_performance_report(self) -> Dict[str, Any]:
+        """Generate performance report"""
+        query_times = [m["duration"] for m in self.metrics["query_times"]]
+        return {
+            "total_queries": len(query_times),
+            "average_query_time": sum(query_times) / len(query_times) if query_times else 0,
+            "max_query_time": max(query_times) if query_times else 0,
+            "min_query_time": min(query_times) if query_times else 0
+        }