import os import sys import time import json import asyncio from datetime import datetime from typing import List from fastapi import FastAPI, HTTPException, BackgroundTasks from fastapi.middleware.cors import CORSMiddleware from fastapi.staticfiles import StaticFiles from fastapi.responses import FileResponse from loguru import logger import subprocess import httpx from starlette.middleware.base import BaseHTTPMiddleware from starlette.responses import StreamingResponse as StarletteStreamingResponse # Ensure project root is on sys.path sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from backend.models import ( ChatRequest, ChatResponse, HealthResponse, SessionHistoryResponse, ChatMessage, ErrorResponse, WorkflowTraceResponse, WorkflowStep, WorkflowDiagramResponse, MemoryFact, MemoryResponse, ) from orchestration.tools import preload_retrievers from orchestration.memory import get_all_memories from orchestration.semantic_cache import cache_manager from backend.session_manager import manager from fastapi.responses import StreamingResponse import io from graphviz import Digraph # Setup logging logger.add("logs/backend.log", rotation="500 MB", level="INFO") # Detect if running in monolith mode (e.g. on Hugging Face Spaces) IS_MONOLITH = os.getenv("RUN_MONOLITH", "false").lower() == "true" or os.getenv("PORT") == "7860" app = FastAPI( title="Health Insurance AI Copilot API", description="Backend API for the RAG-based Health Insurance Assistant", version="1.0.0" ) # Enable CORS for Streamlit app.add_middleware( CORSMiddleware, allow_origins=["*"], # In production, specify actual origins allow_credentials=True, allow_methods=["*"], allow_headers=["*"], ) # ── Proxy Middleware for Single-Port Deployment (Hugging Face) ── if IS_MONOLITH and os.getenv("USE_NGINX", "false").lower() != "true": class StreamlitProxyMiddleware(BaseHTTPMiddleware): async def dispatch(self, request, call_next): # Let FastAPI handle these routes directly if (request.url.path.startswith("/chat") or request.url.path.startswith("/session") or request.url.path.startswith("/health") or request.url.path.startswith("/dev-console") or request.url.path.startswith("/static")): return await call_next(request) # Everything else goes to Streamlit (running on port 8501 internally) target_url = f"http://localhost:8501{request.url.path}" if request.url.query: target_url += f"?{request.url.query}" async with httpx.AsyncClient() as client: try: proxy_req = client.build_request( request.method, target_url, headers=request.headers.raw, content=await request.body() ) proxy_res = await client.send(proxy_req, stream=True) return StarletteStreamingResponse( proxy_res.aiter_raw(), status_code=proxy_res.status_code, headers=proxy_res.headers ) except Exception as e: logger.error(f"Proxy error: {str(e)}") return await call_next(request) app.add_middleware(StreamlitProxyMiddleware) @app.on_event("startup") async def startup_event(): logger.info("Initializing retrieval pipeline preload synchronously...") preload_retrievers() if IS_MONOLITH: logger.info("Starting Streamlit frontend sidecar on port 8501...") subprocess.Popen([ "streamlit", "run", "frontend/app.py", "--server.port", "8501", "--server.address", "0.0.0.0", "--server.headless", "true" ]) logger.info("Backend and Frontend sidecar are fully ready.") else: logger.info("Backend is fully ready.") # Mount static frontend files _frontend_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "frontend") app.mount("/static", StaticFiles(directory=_frontend_dir), name="static") @app.get("/dev-console") async def dev_console(): """Serve the beautiful live Developer Console.""" html_path = os.path.join(_frontend_dir, "dev_console.html") return FileResponse(html_path, media_type="text/html") @app.get("/health", response_model=HealthResponse) async def health_check(): """Check the health of the backend and its components.""" # Simplified health check components = { "openai_api": "ok" if os.getenv("OPENAI_API_KEY") else "error", "vector_store": "ok", # In a real app, check Chroma connectivity "knowledge_graph": "ok" } status = "healthy" if any(v == "error" for v in components.values()): status = "degraded" return HealthResponse( status=status, timestamp=datetime.now().isoformat(), components=components, active_sessions=manager.get_active_session_count() ) @app.post("/chat", response_model=ChatResponse, responses={400: {"model": ErrorResponse}, 500: {"model": ErrorResponse}}) async def chat(request: ChatRequest): """Process a chat query for a specific session.""" try: logger.info(f"Received query for session {request.session_id}: {request.query[:50]}...") # Get or create orchestrator for this session orch = manager.get_orchestrator(request.session_id) # Execute query using the detailed method we added start_time = time.time() final_query = request.query if request.plan_tier and request.plan_tier.lower() != "unknown": final_query = f"I am on the {request.plan_tier} plan. {request.query}" # Check Semantic Cache cached_result = cache_manager.check(final_query, plan_tier=request.plan_tier or "Unknown") if cached_result: if len(orch.chat_history) == 0 or orch.chat_history[-1] != ("ai", cached_result["answer"]): orch.chat_history.append(("human", request.query)) orch.chat_history.append(("ai", cached_result["answer"])) # Sync orchestrator trace for downstream requests (diagram, trace endpoints) orch.last_detailed_result = cached_result.copy() # Prepend a step identifying the cache hit display_steps = [ "⚡ Semantic Cache HIT!", f"Matched with cached query: '{cached_result.get('matched_query', '')}'", f"Similarity Score: {cached_result.get('cache_similarity', 0)}%" ] + cached_result.get("steps_log", []) return ChatResponse( session_id=request.session_id, query=request.query, answer=cached_result["answer"], intent=cached_result["intent"], steps_log=display_steps, memories_used=cached_result.get("memories_used", []), timestamp=datetime.now().isoformat(), run_id=cached_result.get("run_id"), confidence=cached_result.get("confidence", "HIGH"), confidence_reason=cached_result.get("confidence_reason", "") + " (Cached)", blocked=cached_result.get("blocked", False), sub_questions=cached_result.get("sub_questions", []), ) result = orch.ask_detailed(final_query) duration = time.time() - start_time logger.info(f"Query processed in {duration:.2f}s for session {request.session_id}") # Save to semantic cache cache_manager.store(final_query, result, plan_tier=request.plan_tier or "Unknown") return ChatResponse( session_id=request.session_id, query=request.query, answer=result["answer"], intent=result["intent"], steps_log=result["steps_log"], memories_used=result.get("memories_used", []), timestamp=datetime.now().isoformat(), run_id=result.get("run_id"), confidence=result.get("confidence", ""), confidence_reason=result.get("confidence_reason", ""), blocked=result.get("blocked", False), sub_questions=result.get("sub_questions", []), ) except Exception as e: logger.error(f"Error processing chat: {str(e)}") raise HTTPException( status_code=500, detail=ErrorResponse(error="Internal Server Error", detail=str(e), session_id=request.session_id).dict() ) @app.get("/chat/stream") async def chat_stream(session_id: str, query: str, plan_tier: str = "Unknown"): """ Streamed version of the chat process for the Developer Console. Emits granular JSON chunks: startup nodes, then per-step events within each LangGraph node. """ import threading orch = manager.get_orchestrator(session_id) final_query = query if plan_tier and plan_tier.lower() != "unknown": final_query = f"I am on the {plan_tier} plan. {query}" # Check Semantic Cache norm_query = cache_manager.normalize_query(final_query) cached_result = cache_manager.check(final_query, plan_tier=plan_tier, normalized_query=norm_query) if cached_result: # Sync orchestrator history & trace if len(orch.chat_history) == 0 or orch.chat_history[-1] != ("ai", cached_result["answer"]): orch.chat_history.append(("human", query)) orch.chat_history.append(("ai", cached_result["answer"])) orch.last_detailed_result = cached_result.copy() async def cache_hit_event_generator(): def emit(data: dict) -> str: return f"data: {json.dumps(data)}\n\n" # ── Startup handshake events ────────────────────────────── yield emit({"type": "node_start", "node": "user", "msg": query}) await asyncio.sleep(0.15) yield emit({"type": "node_start", "node": "fastapi", "msg": "POST /chat/stream received"}) await asyncio.sleep(0.15) # Query Analyzer yield emit({"type": "node_start", "node": "query_analyzer", "msg": "Analyzing search intent & phrasing..."}) await asyncio.sleep(0.2) yield emit({ "type": "substep", "node": "query_analyzer", "step": f"Normalized query: '{final_query[:40]}...' -> '{norm_query}'" }) await asyncio.sleep(0.2) yield emit({"type": "node_done", "node": "query_analyzer", "normalized_query": norm_query}) await asyncio.sleep(0.15) # Redis Cache Check yield emit({"type": "node_start", "node": "semantic_cache", "msg": "Checking semantic cache…"}) await asyncio.sleep(0.2) sim_score = cached_result.get('cache_similarity', 0.0) matched_q = cached_result.get('matched_query', '') steps = [ "⚡ Semantic Cache HIT!", f"Matched cached query: '{matched_q}'", f"Similarity Score: {sim_score}%", "Bypassing intent classifier, hybrid retrievers, and synthesis LLM.", "Retrieving cached answer payload from Redis." ] for step in steps: yield emit({ "type": "substep", "node": "semantic_cache", "intent": cached_result.get("intent", "POLICY_QUESTION"), "step": step, "all_steps": steps }) await asyncio.sleep(0.08) # Signal cache node completion yield emit({ "type": "node_done", "node": "semantic_cache", "intent": cached_result.get("intent", "POLICY_QUESTION"), "steps": steps + cached_result.get("steps_log", []), "answer": cached_result["answer"], "confidence": cached_result.get("confidence", "HIGH"), "confidence_reason": cached_result.get("confidence_reason", "") + " (Cached)", "blocked": cached_result.get("blocked", False), "sub_questions": cached_result.get("sub_questions", []), }) await asyncio.sleep(0.1) yield "data: [DONE]\n\n" return StreamingResponse(cache_hit_event_generator(), media_type="text/event-stream") # Cache Miss: run full LangGraph pipeline queue = asyncio.Queue() loop = asyncio.get_running_loop() def run_stream(): try: for event in orch.stream_detailed(final_query): loop.call_soon_threadsafe(queue.put_nowait, {"type": "event", "data": event}) loop.call_soon_threadsafe(queue.put_nowait, {"type": "done"}) except Exception as e: loop.call_soon_threadsafe(queue.put_nowait, {"type": "error", "error": e}) thread = threading.Thread(target=run_stream, daemon=True) thread.start() async def event_generator(): def emit(data: dict) -> str: return f"data: {json.dumps(data)}\n\n" # ── Startup handshake events ────────────────────────────── yield emit({"type": "node_start", "node": "user", "msg": query}) await asyncio.sleep(0.15) yield emit({"type": "node_start", "node": "fastapi", "msg": "POST /chat/stream received"}) await asyncio.sleep(0.15) # Query Analyzer yield emit({"type": "node_start", "node": "query_analyzer", "msg": "Analyzing search intent & phrasing..."}) await asyncio.sleep(0.2) yield emit({ "type": "substep", "node": "query_analyzer", "step": f"Normalized query: '{final_query[:40]}...' -> '{norm_query}'" }) await asyncio.sleep(0.2) yield emit({"type": "node_done", "node": "query_analyzer", "normalized_query": norm_query}) await asyncio.sleep(0.15) # Redis Cache Check yield emit({"type": "node_start", "node": "semantic_cache", "msg": "Checking semantic cache…"}) await asyncio.sleep(0.2) yield emit({ "type": "substep", "node": "semantic_cache", "step": "Cache MISS. No matching normalized query found." }) await asyncio.sleep(0.15) yield emit({"type": "node_done", "node": "semantic_cache", "msg": "Proceeding to full RAG pipeline"}) await asyncio.sleep(0.15) yield emit({"type": "node_start", "node": "orchestrator", "msg": "Building LangGraph state…"}) await asyncio.sleep(0.2) prev_steps: list[str] = [] final_answer = "" last_intent = "" last_state = None # ── Stream LangGraph events ─────────────────────────────── try: while True: try: item = await asyncio.wait_for(queue.get(), timeout=2.0) except asyncio.TimeoutError: yield ": keepalive\n\n" continue if item["type"] == "done": break elif item["type"] == "error": raise item["error"] event = item["data"] node = event["node"] state = event["state"] last_state = state current_steps = state.get("steps_log", []) new_steps = current_steps[len(prev_steps):] intent = state.get("intent", last_intent) last_intent = intent yield emit({"type": "node_start", "node": node, "intent": intent, "msg": f"Node '{node}' executing…"}) await asyncio.sleep(0.15) for step in new_steps: yield emit({"type": "substep", "node": node, "intent": intent, "step": step, "all_steps": current_steps}) # Skip artificial delay for raw graph DB edge logs and entity listings to prevent stream lagging if not step.startswith(("[GraphDB-Edge]", "[GraphDB]")): await asyncio.sleep(0.05) final_answer = state.get("answer", "") yield emit({ "type": "node_done", "node": node, "intent": intent, "steps": current_steps, "answer": final_answer, "confidence": state.get("confidence", ""), "confidence_reason": state.get("confidence_reason", ""), "blocked": state.get("blocked", False), "sub_questions": state.get("sub_questions", []), }) if node == "synthesize": if len(orch.chat_history) == 0 or orch.chat_history[-1] != ("ai", final_answer): orch.chat_history.append(("human", query)) orch.chat_history.append(("ai", final_answer)) orch.last_detailed_result = { "query": query, "answer": final_answer, "intent": intent, "steps_log": current_steps, "retrieved_context": state.get("retrieved_context", ""), } prev_steps = current_steps await asyncio.sleep(0.1) # Stream finished successfully, store result in cache if last_state: cache_payload = { "query": query, "answer": last_state.get("answer", ""), "intent": last_state.get("intent", "POLICY_QUESTION"), "steps_log": last_state.get("steps_log", []), "retrieved_context": last_state.get("retrieved_context", ""), "memories_used": last_state.get("memories_used", []), "run_id": orch.last_detailed_result.get("run_id", "") if orch.last_detailed_result else "", "blocked": last_state.get("blocked", False), "confidence": last_state.get("confidence", ""), "confidence_reason": last_state.get("confidence_reason", ""), "sub_questions": last_state.get("sub_questions", []), } cache_manager.store(final_query, cache_payload, plan_tier=plan_tier, normalized_query=norm_query) except Exception as e: logger.error(f"Error in chat_stream: {str(e)}") yield emit({"type": "error", "msg": f"Backend Error: {str(e)}"}) yield "data: [DONE]\n\n" return StreamingResponse(event_generator(), media_type="text/event-stream") @app.get("/session/{session_id}/history", response_model=SessionHistoryResponse) async def get_history(session_id: str): """Retrieve the chat history for a specific session.""" try: orch = manager.get_orchestrator(session_id) history = [] for role, content in orch.chat_history: history.append(ChatMessage(role=role, content=content)) return SessionHistoryResponse( session_id=session_id, history=history, message_count=len(history) ) except Exception as e: raise HTTPException(status_code=404, detail=f"Session {session_id} not found or error retrieving history") @app.delete("/session/{session_id}") async def clear_session(session_id: str): """Clear a session's history, state, and in-memory Mem0 facts.""" manager.clear_session(session_id) return {"status": "success", "message": f"Session {session_id} cleared (including memory)"} @app.get("/memory/{session_id}", response_model=MemoryResponse) async def get_memory(session_id: str): """ Inspect all facts Mem0 has stored for this session. Useful for debugging — shows exactly what the agent remembers. Note: ephemeral — resets when the server restarts. """ try: mem = manager.get_memory(session_id) if mem is None: return MemoryResponse( session_id=session_id, memory_count=0, facts=[], note="Session not found or memory disabled" ) raw_facts = get_all_memories(mem) facts = [ MemoryFact( id=str(f.get("id", "")), memory=f.get("memory", ""), created_at=str(f.get("created_at", "")) ) for f in raw_facts ] return MemoryResponse( session_id=session_id, memory_count=len(facts), facts=facts ) except Exception as e: logger.error(f"Error fetching memory for {session_id}: {e}") raise HTTPException(status_code=500, detail=str(e)) @app.delete("/memory/{session_id}") async def clear_memory(session_id: str): """ Wipe all Mem0 facts for this session without clearing the full session. The session and chat history remain intact. """ try: mem = manager.get_memory(session_id) if mem is None: raise HTTPException(status_code=404, detail=f"Session {session_id} not found") # Re-create a fresh memory instance for the session from orchestration.memory import create_session_memory orch = manager.get_orchestrator(session_id) new_mem = create_session_memory() orch._mem = new_mem orch._search_memories = lambda query: __import__('orchestration.memory', fromlist=['search_memories']).search_memories(new_mem, query) orch._add_memory = lambda q, a: __import__('orchestration.memory', fromlist=['add_memory']).add_memory(new_mem, q, a) return {"status": "success", "message": f"Memory cleared for session {session_id}"} except HTTPException: raise except Exception as e: logger.error(f"Error clearing memory for {session_id}: {e}") raise HTTPException(status_code=500, detail=str(e)) # Helper to turn steps into a Graphviz diagram (PNG bytes) def _steps_to_graphviz(query: str, intent: str, steps: list[str]) -> io.BytesIO: """Generate a PNG image of the workflow using Graphviz. Nodes: FastAPI, Orchestrator, Intent, Retrieval, Synthesis, Answer. If cached, bypasses all retrieval nodes and shows a direct link via Redis Cache. """ is_cached = any("Cache HIT" in s for s in steps) dot = Digraph(comment='Workflow') dot.attr(rankdir='LR') if is_cached: # Pruned cache hit nodes dot.node('A', 'FastAPI Endpoint') dot.node('R', 'Redis Semantic Cache') dot.node('F', 'Answer') # Edges dot.edge('A', 'R') dot.edge('R', 'F') # Detailed steps subgraph if steps: with dot.subgraph(name='cluster_details') as c: c.attr(label='Cache Hit Details') prev = None for i, s in enumerate(steps, start=1): node_id = f'S{i}' safe = s.replace('"', '\\"') c.node(node_id, safe, shape='box') if i == 1: dot.edge('R', node_id) else: c.edge(prev, node_id) prev = node_id else: # Core nodes dot.node('A', 'FastAPI Endpoint') dot.node('B', 'Orchestrator') dot.node('C', 'Intent Classification') dot.node('D', 'Retrieval Pipeline') dot.node('E', 'Synthesis Agent') dot.node('F', 'Answer') # Edges dot.edge('A', 'B') dot.edge('B', 'C') dot.edge('C', 'D') dot.edge('D', 'E') dot.edge('E', 'F') # Detailed steps subgraph if steps: with dot.subgraph(name='cluster_details') as c: c.attr(label='Retrieval Details') prev = None for i, s in enumerate(steps, start=1): node_id = f'S{i}' safe = s.replace('"', '\\"') c.node(node_id, safe, shape='box') if i == 1: dot.edge('D', node_id) else: c.edge(prev, node_id) prev = node_id # Render to PNG in memory png_bytes = dot.pipe(format='png') return io.BytesIO(png_bytes) def _steps_to_mermaid(query: str, intent: str, steps: list[str]) -> str: """Generate a simple mermaid diagram describing the workflow. If cached, bypasses all retrieval nodes and shows a direct link via Redis Cache. """ is_cached = any("Cache HIT" in s for s in steps) lines = ["graph LR"] if is_cached: lines.append(' A[FastAPI Endpoint] --> R[Redis Semantic Cache]') lines.append(' R --> F[Answer]') if steps: lines.append(' subgraph Details [Cache Hit Details]') for i, s in enumerate(steps, start=1): safe = s.replace('"', '\\"') node_id = f"S{i}" lines.append(f' {node_id}["{safe}"]') if i == 1: lines.append(f' R --> {node_id}') else: lines.append(f' S{i-1} --> {node_id}') lines.append(' end') else: lines.append(' A[FastAPI Endpoint] --> B[Orchestrator]') lines.append(' B --> C[Intent Classification]') lines.append(' C --> D[Retrieval Pipeline]') lines.append(' D --> E[Synthesis Agent]') lines.append(' E --> F[Answer]') if steps: lines.append(' subgraph Details [Retrieval Details]') for i, s in enumerate(steps, start=1): safe = s.replace('"', '\\"') node_id = f"S{i}" lines.append(f' {node_id}["{safe}"]') if i == 1: lines.append(f' D --> {node_id}') else: lines.append(f' S{i-1} --> {node_id}') lines.append(' end') return "\n".join(lines) @app.get("/session/{session_id}/diagram", response_model=WorkflowDiagramResponse) async def get_diagram(session_id: str): """Return a Mermaid diagram visualising the most recent query workflow for a session.""" try: orch = manager.get_orchestrator(session_id) result = orch.last_detailed_result if not result: raise HTTPException(status_code=404, detail="No trace found for this session. Ask a question first.") diagram = _steps_to_mermaid( query=result.get("query", ""), intent=result.get("intent", ""), steps=result.get("steps_log", []), ) return WorkflowDiagramResponse( session_id=session_id, query=result.get("query", ""), intent=result.get("intent", ""), diagram=diagram, steps_log=result.get("steps_log", []), ) except HTTPException: raise except Exception as e: logger.error(f"Error generating diagram: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}") @app.get("/session/{session_id}/graph", response_model=None) async def get_graph(session_id: str): """Return a PNG image visualising the latest query workflow for a session.""" try: orch = manager.get_orchestrator(session_id) result = orch.last_detailed_result if not result: raise HTTPException(status_code=404, detail="No trace found for this session. Ask a question first.") img_io = _steps_to_graphviz( query=result.get("query", ""), intent=result.get("intent", ""), steps=result.get("steps_log", []), ) return StreamingResponse(img_io, media_type="image/png") except HTTPException: raise except Exception as e: logger.error(f"Error generating graph: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}") @app.get("/session/{session_id}/trace", response_model=WorkflowTraceResponse) async def get_trace(session_id: str): """Retrieve a detailed workflow trace for the last query in a session.""" try: orch = manager.get_orchestrator(session_id) if not orch.last_detailed_result: raise HTTPException(status_code=404, detail="No trace found for this session. Ask a question first.") res = orch.last_detailed_result # Map the flat steps_log to a more structured workflow if possible # This is a bit heuristic but provides the requested "agent1 -> agent2" feel steps = [] for i, s in enumerate(res["steps_log"]): name = "Process Step" if "Intent" in s: name = "Intent Classification" elif "Graph" in s: name = "Knowledge Graph Retrieval" elif "Hybrid" in s or "document" in s: name = "Policy Retrieval" elif "synthesized" in s: name = "LLM Synthesis" steps.append(WorkflowStep( step_id=f"step_{i+1}", name=name, description=s )) return WorkflowTraceResponse( session_id=session_id, query=res["query"], intent=res["intent"], steps=steps, full_trace=res["steps_log"], retrieved_context_preview=res["retrieved_context"][:500] + "..." if len(res["retrieved_context"]) > 500 else res["retrieved_context"] ) except HTTPException: raise except Exception as e: logger.error(f"Error retrieving trace: {str(e)}") raise HTTPException(status_code=500, detail=f"Internal error: {str(e)}") if __name__ == "__main__": import uvicorn uvicorn.run(app, host="0.0.0.0", port=8000)