fix interrupt and streaming

api/main.py

@@ -1,28 +1,23 @@
 # api/main.py
-# Add project root to Python path
 import sys
 from pathlib import Path
 sys.path.insert(0, str(Path(__file__).parent.parent))
 
-from typing import Optional, Any
+from typing import Optional, Any, Dict, AsyncGenerator
 from contextlib import asynccontextmanager
-from fastapi import FastAPI, Query, HTTPException
-from fastapi.responses import StreamingResponse, HTMLResponse
+from fastapi import FastAPI, Query, HTTPException, Body
+from fastapi.responses import StreamingResponse, HTMLResponse, JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
-from langchain_core.messages import AIMessageChunk
-from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, SystemMessage
+from pydantic import BaseModel
 import json
 from uuid import uuid4
 import logging
 import os
 import asyncio
+from datetime import datetime
 
 # Import your existing system
 from core.system_initializer import setup_system
-from models.state_models import MultiCountryLegalState
-
-# utils functions
-from utils.helpers import message_obj_to_dict, dict_to_message_obj
 
 # Setup logging
 logging.basicConfig(level=logging.INFO)
@@ -35,375 +30,663 @@ system_initialized = False
 
 
 # ============================================================================
-# CRITICAL: Safe JSON Serialization Utilities
+# Pydantic Models
 # ============================================================================
-class SafeJSONEncoder(json.JSONEncoder):
-    """
-    Custom JSON encoder that safely handles Pydantic models and other non-serializable objects.
-    This is the ultimate fallback for any serialization issues.
-    """
-    def default(self, obj):
-        # Handle Pydantic models
-        if hasattr(obj, 'model_dump'):
-            return obj.model_dump()
-        if hasattr(obj, 'dict'):
-            return obj.dict()
-        
-        # Handle LangChain messages
-        if isinstance(obj, BaseMessage):
-            return {
-                "role": "assistant" if isinstance(obj, AIMessage) else "user",
-                "content": obj.content if hasattr(obj, 'content') else str(obj),
-                "meta": getattr(obj, "additional_kwargs", {}),
-            }
-        
-        # Handle sets
-        if isinstance(obj, set):
-            return list(obj)
-        
-        # Handle bytes
-        if isinstance(obj, bytes):
-            return obj.decode('utf-8', errors='ignore')
-        
-        # Fallback: convert to string
-        try:
-            return str(obj)
-        except Exception:
-            return f"<Unserializable: {type(obj).__name__}>"
+class ChatRequest(BaseModel):
+    message: str
+    session_id: Optional[str] = None
+    stream: bool = False  # Option to enable/disable streaming
 
+class ApprovalRequest(BaseModel):
+    decision: str  # "approve" or "reject"
+    reason: Optional[str] = None
 
-def safe_json_dumps(obj: Any) -> str:
-    """
-    Safely convert any object to JSON string with multiple fallback strategies.
-    """
-    try:
-        # Try standard JSON encoding first
-        return json.dumps(obj)
-    except (TypeError, ValueError):
-        try:
-            # Try with custom encoder
-            return json.dumps(obj, cls=SafeJSONEncoder)
-        except Exception:
-            try:
-                # Try with default=str fallback
-                return json.dumps(obj, default=str)
-            except Exception as e:
-                # Ultimate fallback: return error message
-                logger.error(f"Complete JSON serialization failure: {e}")
-                return json.dumps({"error": "serialization_failed", "message": str(e)})
-# ============================================================================
+class ChatResponse(BaseModel):
+    response: str
+    session_id: str
+    has_interrupt: bool = False
+    interrupt_type: Optional[str] = None
+    interrupt_data: Optional[Dict] = None
 
 
+# ============================================================================
+# System Initialization
+# ============================================================================
 async def initialize_system():
     global chat_manager, graph, system_initialized
     try:
-        # 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")
+        logger.info("✅ Legal assistant system initialized")
     except Exception as e:
         logger.error(f"❌ Failed to initialize system: {e}")
         system_initialized = False
 
 @asynccontextmanager
 async def lifespan(app: FastAPI):
-    """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...")
+    yield
+    logger.info("🛑 Shutting down...")
     initialization_task.cancel()
     try:
         await initialization_task
     except asyncio.CancelledError:
         pass
 
+
+# ============================================================================
+# FastAPI App
+# ============================================================================
 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",
+    version="2.0.0",
+    description="Multi-country legal RAG with streaming & human-in-the-loop",
     lifespan=lifespan
 )
 
-# Add CORS middleware
 app.add_middleware(
     CORSMiddleware,
-    allow_origins=["*"],  
+    allow_origins=["*"],
     allow_credentials=True,
-    allow_methods=["*"],  
-    allow_headers=["*"], 
+    allow_methods=["*"],
+    allow_headers=["*"],
 )
 
+
+# ============================================================================
+# Streaming Helper
+# ============================================================================
+async def stream_chat_response(
+    message: str, 
+    session_id: str
+) -> AsyncGenerator[str, None]:
+    """
+    Stream chat responses token by token.
+    If interrupt occurs, yields special interrupt event.
+    """
+    try:
+        # Check for pending interrupt first
+        if chat_manager.has_pending_interrupt(session_id):
+            interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
+            interrupt_data_obj = interrupt_info.get("interrupt_data")
+            
+            if hasattr(interrupt_data_obj, 'value'):
+                interrupt_value = interrupt_data_obj.value
+            else:
+                interrupt_value = {}
+            
+            # Yield interrupt event
+            yield f"data: {json.dumps({'type': 'interrupt', 'data': interrupt_value})}\n\n"
+            return
+        
+        # Buffer to collect response
+        response_buffer = []
+        
+        # Process message
+        response = await chat_manager.chat(
+            message=message,
+            session_id=session_id,
+            interrupt_handler=None  # Async mode
+        )
+        
+        # Check if interrupt occurred during processing
+        if chat_manager.has_pending_interrupt(session_id):
+            interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
+            interrupt_data_obj = interrupt_info.get("interrupt_data")
+            
+            if hasattr(interrupt_data_obj, 'value'):
+                interrupt_value = interrupt_data_obj.value
+            else:
+                interrupt_value = {}
+            
+            # Yield interrupt event
+            yield f"data: {json.dumps({'type': 'interrupt', 'data': interrupt_value})}\n\n"
+            return
+        
+        # Stream response token by token (simulate streaming)
+        words = response.split()
+        for i, word in enumerate(words):
+            token = word + (" " if i < len(words) - 1 else "")
+            yield f"data: {json.dumps({'type': 'token', 'content': token})}\n\n"
+            await asyncio.sleep(0.02)  # Small delay for visual effect
+        
+        # Send completion event
+        yield f"data: {json.dumps({'type': 'done', 'session_id': session_id})}\n\n"
+        
+    except Exception as e:
+        logger.exception(f"Error in streaming: {e}")
+        yield f"data: {json.dumps({'type': 'error', 'message': str(e)})}\n\n"
+
+
+# ============================================================================
+# Routes
+# ============================================================================
 @app.get("/", response_class=HTMLResponse)
 async def read_root():
-    """Simple homepage for better UX"""
+    """Interactive chat interface with streaming support"""
     return """
     <html>
         <head>
-            <title>Legal Assistant API</title>
+            <title>Legal Assistant - Interactive Chat</title>
+            <meta charset="UTF-8">
             <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; }
+                * { margin: 0; padding: 0; box-sizing: border-box; }
+                body { 
+                    font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif;
+                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                    min-height: 100vh;
+                    padding: 20px;
+                }
+                .container { 
+                    max-width: 900px; 
+                    margin: 0 auto; 
+                    background: white;
+                    border-radius: 16px;
+                    box-shadow: 0 20px 60px rgba(0,0,0,0.3);
+                    overflow: hidden;
+                }
+                .header {
+                    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
+                    color: white;
+                    padding: 30px;
+                    text-align: center;
+                }
+                .header h1 { margin-bottom: 10px; }
+                .header p { opacity: 0.9; }
+                .chat-container { padding: 20px; }
+                #messages {
+                    height: 500px;
+                    overflow-y: auto;
+                    padding: 20px;
+                    background: #f8f9fa;
+                    border-radius: 12px;
+                    margin-bottom: 20px;
+                }
+                .message {
+                    margin: 15px 0;
+                    padding: 12px 18px;
+                    border-radius: 18px;
+                    max-width: 75%;
+                    animation: slideIn 0.3s ease;
+                    word-wrap: break-word;
+                }
+                @keyframes slideIn {
+                    from { opacity: 0; transform: translateY(10px); }
+                    to { opacity: 1; transform: translateY(0); }
+                }
+                .user-message {
+                    background: #667eea;
+                    color: white;
+                    margin-left: auto;
+                    border-bottom-right-radius: 4px;
+                }
+                .assistant-message {
+                    background: white;
+                    color: #333;
+                    border: 1px solid #e0e0e0;
+                    border-bottom-left-radius: 4px;
+                }
+                .system-message {
+                    background: #fff3cd;
+                    color: #856404;
+                    border: 1px solid #ffc107;
+                    text-align: center;
+                    margin: 10px auto;
+                    max-width: 90%;
+                }
+                .interrupt-panel {
+                    background: #fff3cd;
+                    border: 2px solid #ffc107;
+                    border-radius: 12px;
+                    padding: 20px;
+                    margin: 20px 0;
+                    animation: pulse 2s infinite;
+                }
+                @keyframes pulse {
+                    0%, 100% { box-shadow: 0 0 0 0 rgba(255, 193, 7, 0.4); }
+                    50% { box-shadow: 0 0 0 10px rgba(255, 193, 7, 0); }
+                }
+                .interrupt-panel h3 {
+                    color: #856404;
+                    margin-bottom: 15px;
+                }
+                .interrupt-details {
+                    background: white;
+                    padding: 15px;
+                    border-radius: 8px;
+                    margin: 15px 0;
+                }
+                .interrupt-details p {
+                    margin: 8px 0;
+                    color: #333;
+                }
+                .approval-buttons {
+                    display: flex;
+                    gap: 10px;
+                    margin-top: 15px;
+                }
+                .btn {
+                    flex: 1;
+                    padding: 12px 24px;
+                    border: none;
+                    border-radius: 8px;
+                    font-weight: bold;
+                    cursor: pointer;
+                    transition: all 0.3s;
+                    font-size: 14px;
+                }
+                .btn-approve {
+                    background: #4caf50;
+                    color: white;
+                }
+                .btn-approve:hover {
+                    background: #45a049;
+                    transform: translateY(-2px);
+                    box-shadow: 0 4px 12px rgba(76, 175, 80, 0.4);
+                }
+                .btn-reject {
+                    background: #f44336;
+                    color: white;
+                }
+                .btn-reject:hover {
+                    background: #da190b;
+                    transform: translateY(-2px);
+                    box-shadow: 0 4px 12px rgba(244, 67, 54, 0.4);
+                }
+                .input-area {
+                    display: flex;
+                    gap: 10px;
+                    padding: 10px;
+                    background: #f8f9fa;
+                    border-radius: 12px;
+                }
+                #message-input {
+                    flex: 1;
+                    padding: 12px 16px;
+                    border: 2px solid #e0e0e0;
+                    border-radius: 8px;
+                    font-size: 14px;
+                    transition: border 0.3s;
+                }
+                #message-input:focus {
+                    outline: none;
+                    border-color: #667eea;
+                }
+                #send-btn {
+                    padding: 12px 32px;
+                    background: #667eea;
+                    color: white;
+                    border: none;
+                    border-radius: 8px;
+                    font-weight: bold;
+                    cursor: pointer;
+                    transition: all 0.3s;
+                }
+                #send-btn:hover:not(:disabled) {
+                    background: #5568d3;
+                    transform: translateY(-2px);
+                    box-shadow: 0 4px 12px rgba(102, 126, 234, 0.4);
+                }
+                #send-btn:disabled {
+                    background: #ccc;
+                    cursor: not-allowed;
+                    transform: none;
+                }
+                .status-bar {
+                    padding: 10px;
+                    background: #f8f9fa;
+                    border-radius: 8px;
+                    margin-top: 10px;
+                    display: flex;
+                    justify-content: space-between;
+                    font-size: 12px;
+                    color: #666;
+                }
+                .typing-indicator {
+                    display: none;
+                    padding: 10px;
+                    color: #666;
+                    font-style: italic;
+                }
+                .typing-indicator.active {
+                    display: block;
+                }
             </style>
         </head>
         <body>
             <div class="container">
-                <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 class="header">
+                    <h1>🧑‍⚖️ Legal Assistant</h1>
+                    <p>Multi-country legal support for Benin & Madagascar</p>
                 </div>
                 
-                <div class="card">
-                    <h3>🔧 System Status</h3>
-                    <div id="status">
-                        <p>Loading system status...</p>
+                <div class="chat-container">
+                    <div id="messages"></div>
+                    <div id="interrupt-zone"></div>
+                    <div class="typing-indicator" id="typing">Assistant is typing...</div>
+                    
+                    <div class="input-area">
+                        <input 
+                            type="text" 
+                            id="message-input" 
+                            placeholder="Type your legal question..."
+                            autocomplete="off"
+                        />
+                        <button id="send-btn" onclick="sendMessage()">
+                            Send
+                        </button>
+                    </div>
+                    
+                    <div class="status-bar">
+                        <span>Session: <strong id="session-id">Not started</strong></span>
+                        <span id="status">System ready</span>
                     </div>
                 </div>
-                
-                <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';
+            </div>
+
+            <script>
+                let sessionId = null;
+                let currentAssistantMessage = null;
+
+                function addMessage(content, type) {
+                    const messagesDiv = document.getElementById('messages');
+                    const messageDiv = document.createElement('div');
+                    messageDiv.className = `message ${type}-message`;
+                    messageDiv.textContent = content;
+                    messagesDiv.appendChild(messageDiv);
+                    messagesDiv.scrollTop = messagesDiv.scrollHeight;
+                    return messageDiv;
+                }
+
+                function showInterrupt(data) {
+                    const interruptZone = document.getElementById('interrupt-zone');
+                    interruptZone.innerHTML = `
+                        <div class="interrupt-panel">
+                            <h3>🔒 Human Approval Required</h3>
+                            <div class="interrupt-details">
+                                <p><strong>📧 Email:</strong> ${data.user_email || 'N/A'}</p>
+                                <p><strong>🌍 Country:</strong> ${data.country || 'N/A'}</p>
+                                <p><strong>📝 Description:</strong> ${data.description || 'N/A'}</p>
+                            </div>
+                            <div class="approval-buttons">
+                                <button class="btn btn-approve" onclick="handleApproval('approve')">
+                                    ✅ Approve Request
+                                </button>
+                                <button class="btn btn-reject" onclick="handleApproval('reject')">
+                                    ❌ Reject Request
+                                </button>
+                            </div>
+                        </div>
+                    `;
+                }
+
+                function hideInterrupt() {
+                    document.getElementById('interrupt-zone').innerHTML = '';
+                }
+
+                async function handleApproval(decision) {
+                    const reason = decision === 'approve' 
+                        ? 'Approved via web interface' 
+                        : 'Rejected via web interface';
+                    
+                    try {
+                        addMessage(`${decision === 'approve' ? '✅' : '❌'} Decision: ${decision}`, 'system');
+                        hideInterrupt();
+                        
+                        const response = await fetch(`/sessions/${sessionId}/approve`, {
+                            method: 'POST',
+                            headers: { 'Content-Type': 'application/json' },
+                            body: JSON.stringify({ decision, reason })
+                        });
+
+                        const data = await response.json();
+                        addMessage(data.response, 'assistant');
+                        
+                    } catch (error) {
+                        addMessage('Error: ' + error.message, 'system');
+                    }
+                }
+
+                async function sendMessage() {
+                    const input = document.getElementById('message-input');
+                    const message = input.value.trim();
+                    
+                    if (!message) return;
+                    
+                    // Initialize session
+                    if (!sessionId) {
+                        sessionId = 'web_' + Date.now();
+                        document.getElementById('session-id').textContent = sessionId;
+                    }
+
+                    // Add user message
+                    addMessage(message, 'user');
+                    input.value = '';
+                    
+                    // Disable input
+                    const sendBtn = document.getElementById('send-btn');
+                    sendBtn.disabled = true;
+                    input.disabled = true;
+                    
+                    // Show typing indicator
+                    document.getElementById('typing').classList.add('active');
+
+                    try {
+                        // Use streaming endpoint
+                        const response = await fetch('/chat/stream', {
+                            method: 'POST',
+                            headers: { 'Content-Type': 'application/json' },
+                            body: JSON.stringify({
+                                message: message,
+                                session_id: sessionId,
+                                stream: true
+                            })
+                        });
+
+                        const reader = response.body.getReader();
+                        const decoder = new TextDecoder();
+                        currentAssistantMessage = addMessage('', 'assistant');
+
+                        while (true) {
+                            const { value, done } = await reader.read();
+                            if (done) break;
+
+                            const chunk = decoder.decode(value);
+                            const lines = chunk.split('\\n\\n');
+
+                            for (const line of lines) {
+                                if (line.startsWith('data: ')) {
+                                    try {
+                                        const data = JSON.parse(line.slice(6));
+                                        
+                                        if (data.type === 'token') {
+                                            currentAssistantMessage.textContent += data.content;
+                                        } else if (data.type === 'interrupt') {
+                                            showInterrupt(data.data);
+                                        } else if (data.type === 'done') {
+                                            document.getElementById('status').textContent = 'Ready';
+                                        } else if (data.type === 'error') {
+                                            addMessage('Error: ' + data.message, 'system');
+                                        }
+                                    } catch (e) {
+                                        console.error('Parse error:', e);
+                                    }
+                                }
                             }
-                            
-                            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>';
                         }
+                        
+                    } catch (error) {
+                        addMessage('Error: ' + error.message, 'system');
+                    } finally {
+                        sendBtn.disabled = false;
+                        input.disabled = false;
+                        input.focus();
+                        document.getElementById('typing').classList.remove('active');
                     }
-                    
-                    updateStatus();
-                    setInterval(updateStatus, 5000);
-                </script>
-            </div>
+                }
+
+                // Enter key to send
+                document.getElementById('message-input').addEventListener('keypress', (e) => {
+                    if (e.key === 'Enter' && !e.shiftKey) {
+                        e.preventDefault();
+                        sendMessage();
+                    }
+                });
+
+                // Focus input on load
+                document.getElementById('message-input').focus();
+            </script>
         </body>
     </html>
     """
 
 @app.get("/health")
 async def health_check():
-    """Enhanced health check with your specific environment variables"""
+    """Health check endpoint"""
     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) -> str:
-    if not session_id:
-        session_id = f"api_{uuid4()}"
-
-    # CRITICAL FIX: Create input state as a dictionary first, then convert to Pydantic model
-    # This ensures proper serialization for PostgreSQL checkpointing
-    input_state_dict = {
-        "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": [],
-        "supplemental_message": "",
-        "country": None,
-        "assistance_requested": False,
-        "user_email": None,
-        "assistance_description": None,
-        "email_status": None,
-        "assistance_step": None,
-        "pending_assistance_data": {},
-        "repair_type": None,
-        "original_query": None,
-        "misunderstanding_count": 0,
-        "primary_intent": None,
-        "approval_status": None,
-        "approval_reason": None,
-        "approved_by": None,
-        "approval_timestamp": None,
-        "summary_generated": False,
-        "last_summary_timestamp": None,
-        "search_metadata": {}
-    }
-
-    # Convert to Pydantic model (this will use our custom model_dump for serialization)
-    input_state = MultiCountryLegalState(**input_state_dict)
-
-    config = {
-        "configurable": {
-            "thread_id": session_id
-        }
+        "timestamp": datetime.now().isoformat()
     }
 
-    # Stream events from the graph
-    events = graph.astream_events(
-        input_state,  # Pass the Pydantic model directly
-        version="v2",
-        config=config
-    )
-
-    current_content = ""
-    current_node = ""
-    final_state = None
-
+@app.post("/chat")
+async def chat_simple(request: ChatRequest):
+    """
+    Simple non-streaming chat endpoint.
+    Best for handling interrupts - returns complete response.
+    """
+    if not system_initialized or not chat_manager:
+        raise HTTPException(status_code=503, detail="System initializing...")
+    
     try:
-        async for event in events:
-            event_type = event["event"]
-            node_name = event.get("name", "")
+        session_id = request.session_id or f"api_{uuid4()}"
+        
+        # Check for pending interrupt
+        if chat_manager.has_pending_interrupt(session_id):
+            interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
+            interrupt_data_obj = interrupt_info.get("interrupt_data")
             
-            if node_name != current_node:
-                current_node = node_name
-                yield f"data: {safe_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: {safe_json_dumps({'type': 'content', 'content': chunk_content})}\n\n"
-                
-            elif event_type == "on_chat_model_end":
-                yield f"data: {safe_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: {safe_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: {safe_json_dumps({'type': 'search_end', 'country': country})}\n\n"
-                
-            elif event_type == "on_tool_end":
-                tool_name = event["name"]
-                yield f"data: {safe_json_dumps({'type': 'tool_complete', 'tool': tool_name})}\n\n"
-
-            elif event_type == "on_graph_end":
-                # Capture and convert the final state - WITH SAFE SERIALIZATION
-                try:
-                    state = event.get("data", {}).get("output")
-                    if state:
-                        if isinstance(state, MultiCountryLegalState):
-                            final_state = state
-                            # Use our custom model_dump method for proper serialization
-                            state_dict = state.model_dump()
-                        elif isinstance(state, dict):
-                            state_dict = state
-                        else:
-                            # Fallback: convert to string
-                            state_dict = {"state": str(state)}
-                        
-                        yield f"data: {safe_json_dumps({'type': 'state', 'content': state_dict})}\n\n"
-                except Exception as state_error:
-                    logger.warning(f"Could not serialize state: {state_error}")
-                    # Don't fail, just skip state output
-                
-                yield f"data: {safe_json_dumps({'type': 'graph_end'})}\n\n"
-
+            if hasattr(interrupt_data_obj, 'value'):
+                interrupt_value = interrupt_data_obj.value
+            else:
+                interrupt_value = {}
+            
+            return ChatResponse(
+                response="⏸️ Pending approval request. Please approve/reject first.",
+                session_id=session_id,
+                has_interrupt=True,
+                interrupt_type="human_approval",
+                interrupt_data=interrupt_value
+            )
+        
+        # Process message
+        response_text = await chat_manager.chat(
+            message=request.message,
+            session_id=session_id
+        )
+        
+        # Check if interrupt occurred
+        has_interrupt = chat_manager.has_pending_interrupt(session_id)
+        interrupt_value = None
+        
+        if has_interrupt:
+            interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
+            interrupt_data_obj = interrupt_info.get("interrupt_data")
+            
+            if hasattr(interrupt_data_obj, 'value'):
+                interrupt_value = interrupt_data_obj.value
+            elif isinstance(interrupt_data_obj, dict):
+                interrupt_value = interrupt_data_obj.get("value", {})
+        
+        return ChatResponse(
+            response=response_text,
+            session_id=session_id,
+            has_interrupt=has_interrupt,
+            interrupt_type="human_approval" if has_interrupt else None,
+            interrupt_data=interrupt_value
+        )
+        
     except Exception as e:
-        logger.error(f"Error in generate_legal_chat_responses: {e}", exc_info=True)
-        yield f"data: {safe_json_dumps({'type': 'error', 'message': str(e)})}\n\n"
+        logger.exception(f"Error in chat: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
 
-    # Yield final state if captured - WITH SAFE SERIALIZATION
-    try:
-        if final_state and isinstance(final_state, MultiCountryLegalState):
-            final_state_dict = final_state.model_dump()
-            yield f"data: {safe_json_dumps({'type': 'final_state', 'content': final_state_dict})}\n\n"
-    except Exception as final_error:
-        logger.warning(f"Could not serialize final state: {final_error}")
-        # Don't fail, just skip final state output
-
-    yield f"data: {safe_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."
-        )
+@app.post("/chat/stream")
+async def chat_stream(request: ChatRequest):
+    """
+    Streaming chat endpoint using Server-Sent Events (SSE).
+    Streams tokens in real-time and handles interrupts.
+    """
+    if not system_initialized or not chat_manager:
+        raise HTTPException(status_code=503, detail="System initializing...")
+    
+    session_id = request.session_id or f"api_{uuid4()}"
     
     return StreamingResponse(
-        generate_legal_chat_responses(message, session_id), 
+        stream_chat_response(request.message, session_id),
         media_type="text/event-stream",
         headers={
             "Cache-Control": "no-cache",
             "Connection": "keep-alive",
+            "X-Accel-Buffering": "no"
         }
     )
 
+@app.post("/sessions/{session_id}/approve")
+async def approve_assistance(session_id: str, request: ApprovalRequest):
+    """Approve or reject assistance request"""
+    if not chat_manager:
+        raise HTTPException(status_code=503, detail="System not initialized")
+    
+    if not chat_manager.has_pending_interrupt(session_id):
+        raise HTTPException(status_code=404, detail="No pending interrupt")
+    
+    try:
+        decision_text = f"{request.decision} {request.reason or ''}"
+        response_text = await chat_manager.chat(
+            message=decision_text,
+            session_id=session_id
+        )
+        
+        return {
+            "status": "success",
+            "decision": request.decision,
+            "response": response_text,
+            "session_id": session_id
+        }
+    except Exception as e:
+        logger.exception(f"Error handling approval: {e}")
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/sessions/{session_id}/status")
+async def get_session_status(session_id: str):
+    """Get session status including interrupt state"""
+    if not chat_manager:
+        raise HTTPException(status_code=503, detail="System not initialized")
+    
+    has_interrupt = chat_manager.has_pending_interrupt(session_id)
+    
+    interrupt_data = None
+    if has_interrupt:
+        interrupt_info = chat_manager.pending_interrupts.get(session_id, {})
+        interrupt_data_obj = interrupt_info.get("interrupt_data")
+        
+        if hasattr(interrupt_data_obj, 'value'):
+            interrupt_data = interrupt_data_obj.value
+    
+    return {
+        "session_id": session_id,
+        "has_pending_interrupt": has_interrupt,
+        "interrupt_data": interrupt_data
+    }
+
 @app.get("/sessions/{session_id}/history")
-async def get_conversation_history(session_id: str):
-    """Get conversation history for a session"""
+async def get_history(session_id: str):
+    """Get conversation history"""
     if not chat_manager:
-        return {"error": "System not initialized"}
+        raise HTTPException(status_code=503, detail="System not initialized")
     
     try:
         history = await chat_manager.get_conversation_history(session_id)
@@ -411,12 +694,11 @@ async def get_conversation_history(session_id: str):
             "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)
+                    "role": msg.type if hasattr(msg, 'type') else "unknown",
+                    "content": msg.content if hasattr(msg, 'content') else str(msg)
                 }
                 for msg in history
             ]
         }
     except Exception as e:
-        return {"error": str(e)}
+        raise HTTPException(status_code=500, detail=str(e))