""" LangGraph Multi-Agent RAG Pipeline for Healthcare FAQ. Agent Flow: User Query └─► ROUTER AGENT — Classifies intent, detects emergencies, routes query ├─► RETRIEVER AGENT — Fetches relevant chunks from FAISS │ └─► RESPONDER AGENT — Generates grounded, safe medical response │ └─► EVALUATOR AGENT — Scores response quality & flags issues └─► (Emergency) ── Direct safety response bypassing retrieval """ import sys from pathlib import Path from typing import TypedDict, Annotated, List, Dict, Any from enum import Enum from langchain_openai import ChatOpenAI from langchain_core.messages import SystemMessage, HumanMessage, AIMessage from langgraph.graph import StateGraph, END # langchain-nvidia-ai-endpoints and tavily are optional (not installed on Render). # They are imported lazily inside the functions that need them. import operator from loguru import logger sys.path.append(str(Path(__file__).parent.parent)) from utils.config import config from vectorstore.retriever import HybridRetriever # ─── Enums ──────────────────────────────────────────────────────────────────── class QueryIntent(str, Enum): MEDICAL_FAQ = "medical_faq" EMERGENCY = "emergency" GENERAL_GREETING = "greeting" OUT_OF_SCOPE = "out_of_scope" WEB_SEARCH = "web_search" FOLLOW_UP = "follow_up" # ─── Graph State ────────────────────────────────────────────────────────────── class AgentState(TypedDict): """The state of the RAG pipeline.""" # Input user_query: str conversation_history: Annotated[list, operator.add] # Router outputs intent: QueryIntent is_emergency: bool reformulated_query: str decision: str # Added for routing decisions # Retriever outputs retrieved_chunks: List[Dict[str, Any]] # More specific type context: str retrieval_confidence: float # Responder outputs response: str disclaimer_added: bool # Evaluator outputs quality_score: float hallucination_risk: str evaluation_notes: str # Metadata agent_trace: Annotated[list, operator.add] error: str retry_count: int # ─── LLM Setup ──────────────────────────────────────────────────────────────── def get_llm(temperature: float = 0.0, streaming: bool = False): """ LLM factory — 3 modes, resolved in priority order: 1. LOCAL_MODE=true → AirLLM (Llama 3 8B on-device, no API key, private) 2. NVIDIA_API_KEY → NVIDIA NIM (cloud, fast, powerful) 3. default → OpenAI GPT-4o-mini (cloud, fast, cheap) """ # ── Mode 1: Local privacy mode via AirLLM ──────────────────────────────── if config.LOCAL_MODE: try: from utils.local_llm import LocalLLM, is_apple_silicon if is_apple_silicon(): logger.info("[LLM] Using LOCAL mode (AirLLM + Llama 3 8B on Apple MLX)") return LocalLLM(model_id=config.LOCAL_MODEL_ID) else: logger.warning("[LLM] LOCAL_MODE=true but not Apple Silicon — falling back to OpenAI") except ImportError: logger.warning("[LLM] AirLLM not installed — falling back to OpenAI. Run: pip install airllm mlx mlx-lm") # ── Mode 2: NVIDIA NIM cloud ───────────────────────────────────────────── if config.NVIDIA_API_KEY and "your-" not in config.NVIDIA_API_KEY: try: from langchain_nvidia_ai_endpoints import ChatNVIDIA logger.debug(f"[LLM] Using NVIDIA NIM: {config.NVIDIA_MODEL}") return ChatNVIDIA( model=config.NVIDIA_MODEL, api_key=config.NVIDIA_API_KEY, temperature=temperature, streaming=streaming, ) except ImportError: logger.warning("[LLM] langchain_nvidia_ai_endpoints not installed — falling back to OpenAI") # ── Mode 3: OpenAI (default) ───────────────────────────────────────────── logger.debug(f"[LLM] Using OpenAI: {config.OPENAI_MODEL}") return ChatOpenAI( api_key=config.OPENAI_API_KEY, model=config.OPENAI_MODEL, temperature=temperature, streaming=streaming, ) # ─── Agent 1: Router ────────────────────────────────────────────────────────── ROUTER_SYSTEM_PROMPT = """You are a medical query router for a Healthcare FAQ assistant. Analyze the user's query and respond with ONLY a JSON object containing: { "intent": "", "is_emergency": , "reformulated_query": "", "reasoning": "" } Intent definitions: - medical_faq: General non-emergency medical questions. - emergency: Life-threatening situations (heart attack, stroke, etc.) - web_search: Queries about RECENT news, recalls (2024-2025), or current health outbreaks. - greeting: Simple hello/hi. - out_of_scope: Non-medical topics (sports, cooking, finance, etc.) - follow_up: Continuation of previous medical discussion Emergency examples: "I have crushing chest pain", "Someone is having a seizure", "I think I overdosed" IMPORTANT: When in doubt about emergency status, set is_emergency=true. Safety first. Respond ONLY with valid JSON. No markdown, no explanation outside JSON.""" async def router_agent(state: AgentState) -> AgentState: """ Agent 1: Routes the query, detects emergencies, reformulates for retrieval. """ logger.info(f"[ROUTER] Processing: '{state['user_query'][:60]}...'") llm = get_llm(temperature=0.0) messages = [ SystemMessage(content=ROUTER_SYSTEM_PROMPT), HumanMessage(content=f"User query: {state['user_query']}"), ] try: import json response = await llm.ainvoke(messages) result = json.loads(response.content.strip()) logger.info( f"[ROUTER] Intent: {result['intent']} | " f"Emergency: {result['is_emergency']} | " f"Reformulated: {result['reformulated_query'][:50]}" ) return { **state, "intent": result["intent"], "is_emergency": result["is_emergency"], "reformulated_query": result.get("reformulated_query", state["user_query"]), "agent_trace": [f"ROUTER: {result['intent']} | emergency={result['is_emergency']}"], } except Exception as e: logger.error(f"[ROUTER] Failed: {e}") return { **state, "intent": QueryIntent.MEDICAL_FAQ, "is_emergency": False, "reformulated_query": state["user_query"], "agent_trace": [f"ROUTER: fallback due to error: {e}"], "error": str(e), } # ─── Agent 2: Retriever ─────────────────────────────────────────────────────── _retriever_instance = None def get_retriever() -> HybridRetriever: global _retriever_instance if _retriever_instance is None: _retriever_instance = HybridRetriever() return _retriever_instance async def retriever_agent(state: AgentState) -> AgentState: """ Agent 2: Fetches relevant chunks from FAISS + reranks them. """ query = state.get("reformulated_query") or state["user_query"] logger.info(f"[RETRIEVER] Searching: '{query[:60]}...'") try: retriever = get_retriever() # In a real async system, we'd use an async retriever, but for now we'll run in a thread import asyncio chunks = await asyncio.to_thread(retriever.retrieve, query) context = retriever.format_context(chunks) # Confidence = average of top rerank scores # Use rerank_score if available (cross-encoder), else fall back to RRF score confidence = ( sum(c.rerank_score if c.rerank_score > 0 else c.score for c in chunks) / len(chunks) if chunks else 0.0 ) logger.info(f"[RETRIEVER] Found {len(chunks)} chunks | Confidence: {confidence:.3f}") return { **state, "retrieved_chunks": [ {"text": c.text, "metadata": c.metadata, "score": c.rerank_score if c.rerank_score > 0 else c.score} for c in chunks ], "context": context, "retrieval_confidence": confidence, "agent_trace": [f"RETRIEVER: {len(chunks)} chunks | confidence={confidence:.3f}"], } except Exception as e: logger.error(f"[RETRIEVER] Failed: {e}") return { **state, "retrieved_chunks": [], "context": "Knowledge base unavailable.", "retrieval_confidence": 0.0, "agent_trace": [f"RETRIEVER: error - {e}"], "error": str(e), } # ─── Agent 3: Web Search ────────────────────────────────────────────────────── async def web_search_agent(state: AgentState) -> AgentState: """Agent 3b: Real-time search fallback for recent events.""" query = state.get("reformulated_query") or state["user_query"] logger.info(f"[WEB_SEARCH] Searching for: {query}") if not config.TAVILY_API_KEY or config.TAVILY_API_KEY == "tvly-your-key-here": logger.warning("Tavily API key not found. Skipping web search.") return {**state, "context": "Error: Web search required but API key missing.", "agent_trace": ["WEB_SEARCH: error - missing key"]} try: from tavily import TavilyClient tavily = TavilyClient(api_key=config.TAVILY_API_KEY) # Search for medical news/recalls response = tavily.search(query=query, search_depth="advanced", max_results=5) context_parts = [] sources = [] for result in response.get("results", []): context_parts.append(f"Source: {result['url']}\nContent: {result['content']}") sources.append({"source": result['url'], "title": result.get("title", "Web Result")}) full_context = "\n\n---\n\n".join(context_parts) return { **state, "context": full_context, "retrieved_chunks": sources, # Mocking structure for SSE "retrieval_confidence": 0.8, # Web search is usually relevant "agent_trace": [f"WEB_SEARCH: {len(sources)} results found"], } except Exception as e: logger.error(f"[WEB_SEARCH] Failed: {e}") return {**state, "error": str(e), "agent_trace": [f"WEB_SEARCH: error - {e}"]} # ─── Agent 3: Responder ─────────────────────────────────────────────────────── RESPONDER_SYSTEM_PROMPT = """You are a knowledgeable, empathetic Healthcare FAQ Assistant. Your role: - Answer medical questions accurately based ONLY on the provided context - Be clear, concise, and compassionate - Use plain language (avoid excessive jargon) - Structure answers with clear sections when appropriate STRICT RULES: 1. ONLY use information from the provided context. Do NOT use outside knowledge. 2. If context is insufficient, honestly say you don't have enough information. 3. NEVER diagnose conditions or prescribe treatments. 4. ALWAYS recommend consulting a healthcare provider for personal medical decisions. 5. For any life-threatening symptoms mentioned, direct to 911 or emergency services immediately. 6. Add the medical disclaimer at the end of EVERY response. Response format: - Direct answer to the question - Key points (if applicable) - When to seek immediate care (if relevant) - Medical disclaimer (ALWAYS include) Disclaimer to include: "⚕️ Medical Disclaimer: This information is for educational purposes only and does not constitute medical advice. Always consult a qualified healthcare provider for personal medical decisions." """ EMERGENCY_RESPONSE = """🚨 **MEDICAL EMERGENCY — CALL 911 IMMEDIATELY** Based on your description, you may be experiencing a **medical emergency** that requires immediate professional attention. **Please take these steps RIGHT NOW:** 1. **Call 911** (or have someone call for you) 2. **Do not drive yourself** to the hospital 3. **Stay calm** and follow the dispatcher's instructions 4. **Unlock your door** so emergency responders can enter **Emergency Resources:** - 🚑 **Emergency:** 911 - 🧠 **Stroke:** 911 (note the time symptoms started) - ❤️ **Cardiac Emergency:** 911 (chew aspirin 325mg if not allergic) - 🆘 **Suicide/Crisis:** 988 Suicide & Crisis Lifeline - ☠️ **Poison Control:** 1-800-222-1222 ⚕️ *I am an AI assistant and cannot provide emergency medical care. Please contact emergency services immediately.*""" GREETING_RESPONSE = """👋 Hello! I'm your **Healthcare FAQ Assistant**, powered by AI. I can help you with: - 🩺 **Symptoms & Conditions** — Understanding common medical symptoms - 💊 **Medications** — General information about common drugs - 🛡️ **Preventive Care** — Vaccines, screenings, wellness tips - ❤️ **Heart Health** — Cardiac symptoms and risk factors - 🧠 **Mental Health** — Depression, anxiety, and mental wellness - 👩 **Women's Health** — PCOS, hormonal health, and more - 🚨 **Emergency Guidance** — When to call 911 **How to use:** Just ask your health question in plain English! *Example: "What are symptoms of high blood pressure?" or "Can I take ibuprofen with acetaminophen?"* ⚕️ *Remember: I provide general health information only. For personal medical advice, always consult your healthcare provider.*""" OUT_OF_SCOPE_RESPONSE = """I'm sorry, but I'm specifically designed for **healthcare and medical FAQ** questions only. I'm not able to help with that topic. However, I'd be happy to answer questions about: - Medical symptoms and conditions - Medications and drug interactions - Preventive care and screenings - Mental health information - Emergency guidance Is there a health-related question I can help you with? 😊""" async def responder_agent(state: AgentState) -> AgentState: """ Agent 3: Generates the final response grounded in retrieved context. Handles special intents (emergency, greeting, out-of-scope) directly. """ intent = state.get("intent", QueryIntent.MEDICAL_FAQ) logger.info(f"[RESPONDER] Generating response for intent: {intent}") # Handle non-retrieval intents directly if state.get("is_emergency"): return {**state, "response": EMERGENCY_RESPONSE, "disclaimer_added": True, "agent_trace": ["RESPONDER: emergency direct response"]} if intent == QueryIntent.GENERAL_GREETING: return {**state, "response": GREETING_RESPONSE, "disclaimer_added": True, "agent_trace": ["RESPONDER: greeting direct response"]} if intent == QueryIntent.OUT_OF_SCOPE: return {**state, "response": OUT_OF_SCOPE_RESPONSE, "disclaimer_added": False, "agent_trace": ["RESPONDER: out-of-scope response"]} # Generate LLM response grounded in context llm = get_llm(temperature=0.1) context = state.get("context", "No relevant context found.") confidence = state.get("retrieval_confidence", 0.0) low_confidence_note = ( "\n\n*Note: I couldn't find highly specific information about this query in my knowledge base. " "The response below is based on the closest available information.*" if confidence < config.CONFIDENCE_THRESHOLD else "" ) # Build conversation history for context history_messages = [] for msg in state.get("conversation_history", [])[-4:]: # Last 2 turns if msg["role"] == "user": history_messages.append(HumanMessage(content=msg["content"])) elif msg["role"] == "assistant": history_messages.append(AIMessage(content=msg["content"])) # Build system prompt based on retry status system_prompt = RESPONDER_SYSTEM_PROMPT if state.get("retry_count", 0) > 0: system_prompt += ( "\n\n**IMPORTANT: SELF-CORRECTION MODE**\n" "Your previous response was rated low quality. Please focus more on accuracy, groundedness, " "and addressing the user's question directly from the provided context." ) messages = [ SystemMessage(content=system_prompt), *history_messages, HumanMessage(content=( f"RETRIEVED CONTEXT:\n{context}\n\n" f"USER QUESTION: {state['user_query']}\n\n" f"Please provide a helpful, accurate answer based ONLY on the context above." )), ] try: response = await llm.ainvoke(messages) final_response = response.content + low_confidence_note logger.info(f"[RESPONDER] Generated {len(final_response)} char response") return { **state, "response": final_response, "disclaimer_added": "Medical Disclaimer" in final_response, "agent_trace": [f"RESPONDER: {len(final_response)} chars | confidence={confidence:.2f}"], } except Exception as e: logger.error(f"[RESPONDER] Failed: {e}") return { **state, "response": "I apologize, but I encountered an error generating a response. Please try again.", "disclaimer_added": False, "agent_trace": [f"RESPONDER: error - {e}"], "error": str(e), } # ─── Agent 4: Evaluator ─────────────────────────────────────────────────────── EVALUATOR_SYSTEM_PROMPT = """You are a medical AI response quality evaluator. Evaluate the assistant's response and return ONLY a JSON object: { "quality_score": <0.0-1.0>, "hallucination_risk": "", "has_disclaimer": , "is_grounded": , "notes": "", "decision": "" } Note: If quality_score < 0.7 AND retry_count < 1, set decision to 'RETRY'. Otherwise set to 'FINISH'. Evaluation criteria: - quality_score: 0.0 (poor) to 1.0 (excellent) - Accuracy: Is the response consistent with the context? - Completeness: Does it address the question fully? - Safety: Does it recommend professional care appropriately? - Clarity: Is it easy to understand? - hallucination_risk: Does the response contain claims NOT supported by the context? - is_grounded: Is the response based on the provided context only? Respond ONLY with valid JSON.""" async def evaluator_agent(state: AgentState) -> AgentState: """ Agent 4: Quality control — scores response, flags hallucination risks. """ logger.info("[EVALUATOR] Evaluating response quality...") # Skip evaluation for direct responses if state.get("is_emergency") or state.get("intent") in [ QueryIntent.GENERAL_GREETING, QueryIntent.OUT_OF_SCOPE ]: return { **state, "quality_score": 1.0, "hallucination_risk": "low", "evaluation_notes": "Direct response — evaluation skipped.", "agent_trace": ["EVALUATOR: skipped (direct response)"], } llm = get_llm(temperature=0.0) messages = [ SystemMessage(content=EVALUATOR_SYSTEM_PROMPT), HumanMessage(content=( f"CONTEXT PROVIDED TO ASSISTANT:\n{state.get('context', '')[:1000]}\n\n" f"USER QUESTION: {state['user_query']}\n\n" f"ASSISTANT RESPONSE:\n{state.get('response', '')[:1500]}" )), ] try: import json eval_response = await llm.ainvoke(messages) cleaned_content = eval_response.content.strip() eval_data = json.loads(cleaned_content) quality_score = eval_data.get("quality_score", 0.0) decision = eval_data.get("decision", "FINISH") # Enforce max 1 retry via logic if state.get("retry_count", 0) >= 1: decision = "FINISH" logger.info(f"[EVALUATOR] Score: {quality_score:.2f} | Decision: {decision}") return { **state, "quality_score": quality_score, "hallucination_risk": eval_data.get("hallucination_risk", "high"), "evaluation_notes": eval_data.get("notes", ""), "agent_trace": [f"EVALUATOR: score={quality_score:.2f} | decision={decision}"], "retry_count": state.get("retry_count", 0) + (1 if decision == "RETRY" else 0) } except Exception as e: logger.error(f"[EVALUATOR] Failed: {e}") return { **state, "quality_score": 0.5, "hallucination_risk": "unknown", "evaluation_notes": f"Evaluation error: {e}", "agent_trace": [f"EVALUATOR: error - {e}"], "error": str(e), } # ─── Routing Logic ──────────────────────────────────────────────────────────── def route_after_router(state: AgentState) -> str: """Conditional edge: decide which agent to call after the router.""" if state.get("is_emergency"): return "responder" # Skip retrieval for emergencies intent = state.get("intent", QueryIntent.MEDICAL_FAQ) if intent == QueryIntent.WEB_SEARCH: return "web_search" if intent in [QueryIntent.GENERAL_GREETING, QueryIntent.OUT_OF_SCOPE]: return "responder" # Skip retrieval for non-medical queries return "retriever" def route_after_evaluator(state: AgentState) -> str: """Reflective loop: decide whether to finish or retry responder.""" # We use agent_trace to check for RETRY decision parsed in evaluator_agent trace = state.get("agent_trace", []) if any("decision=RETRY" in t for t in trace): logger.warning(f"[GRAHP] Triggering self-correction (retry {state.get('retry_count', 0)})") return "retry" return "finish" # ─── Graph Builder ──────────────────────────────────────────────────────────── def build_rag_graph() -> StateGraph: """Build and compile the LangGraph multi-agent pipeline.""" workflow = StateGraph(AgentState) # Add nodes workflow.add_node("router", router_agent) workflow.add_node("retriever", retriever_agent) workflow.add_node("web_search", web_search_agent) workflow.add_node("responder", responder_agent) workflow.add_node("evaluator", evaluator_agent) # Define edges workflow.set_entry_point("router") workflow.add_conditional_edges( "router", route_after_router, { "retriever": "retriever", "web_search": "web_search", "responder": "responder", } ) workflow.add_edge("retriever", "responder") workflow.add_edge("web_search", "responder") # Conditional edge from Evaluator (Self-Correction Loop) workflow.add_conditional_edges( "evaluator", route_after_evaluator, { "retry": "responder", "finish": END, } ) workflow.add_edge("evaluator", END) return workflow.compile() # ─── Main Pipeline Interface ────────────────────────────────────────────────── class HealthcareRAGPipeline: """High-level interface for running the multi-agent RAG pipeline.""" def __init__(self): logger.info("Initializing HealthcareRAGPipeline...") config.validate() self.graph = build_rag_graph() self.conversation_history = [] logger.success("Pipeline ready.") def _get_initial_state(self, query: str, history: List[dict] = None) -> AgentState: """Helper to create the standardized initial state.""" return { "user_query": query, "conversation_history": history or [], "agent_trace": [], "retry_count": 0, "error": "", "intent": "", "is_emergency": False, "reformulated_query": "", "retrieved_chunks": [], "context": "", "retrieval_confidence": 0.0, "response": "", "disclaimer_added": False, "quality_score": 0.0, "hallucination_risk": "low", "evaluation_notes": "" } async def run(self, user_query: str) -> dict: """Run a query through the full multi-agent pipeline.""" initial_state = self._get_initial_state(user_query) result = await self.graph.ainvoke(initial_state) # Update conversation history self.conversation_history.append({"role": "user", "content": user_query}) self.conversation_history.append({"role": "assistant", "content": result["response"]}) if len(self.conversation_history) > 20: self.conversation_history = self.conversation_history[-20:] return { "response": result["response"], "intent": result.get("intent"), "is_emergency": result.get("is_emergency", False), "retrieval_confidence": result.get("retrieval_confidence", 0.0), "quality_score": result.get("quality_score", 0.0), "hallucination_risk": result.get("hallucination_risk", "unknown"), "evaluation_notes": result.get("evaluation_notes", ""), "agent_trace": result.get("agent_trace", []), "sources": result.get("retrieved_chunks", []), } async def astream(self, user_query: str): """ Stream the response token-by-token. Yields tokens for the response, and finally a metadata object for UI sync. """ initial_state = self._get_initial_state(user_query) final_state = initial_state.copy() full_response = "" async for event in self.graph.astream_events(initial_state, version="v2"): kind = event["event"] # Update state tracker as nodes finish (important for capturing final metadata) if kind == "on_chain_end": if "output" in event["data"] and isinstance(event["data"]["output"], dict): final_state.update(event["data"]["output"]) # Stream tokens if kind == "on_chat_model_stream": content = event["data"]["chunk"].content if content: full_response += content yield content # Update conversation history with the full gathered response self.conversation_history.append({"role": "user", "content": user_query}) self.conversation_history.append({"role": "assistant", "content": full_response}) if len(self.conversation_history) > 20: self.conversation_history = self.conversation_history[-20:] # Yield final metadata for UI sidebars yield { "type": "metadata", "content": { "intent": final_state.get("intent", "medical_faq"), "is_emergency": final_state.get("is_emergency", False), "retrieved_chunks": [ {"text": c.text, "metadata": c.metadata, "score": c.score} for c in final_state.get("retrieved_chunks", []) if hasattr(c, "text") ], "quality_score": final_state.get("quality_score", 0.0), "hallucination_risk": final_state.get("hallucination_risk", "unknown"), "agent_trace": final_state.get("agent_trace", []), } } def reset_conversation(self): """Clear conversation history.""" self.conversation_history = [] logger.info("Conversation history cleared.")