Implement Redis Semantic Caching, plan-tier scoping, cognitive query normalization, and Dev Console Redis visual flow integration

Dockerfile

@@ -11,13 +11,14 @@ ENV USE_NGINX=true
 # Disable Mem0 telemetry
 ENV MEM0_TELEMETRY=false
 
-# Install system dependencies (Graphviz for diagrams, Nginx for reverse proxy)
+# Install system dependencies (Graphviz for diagrams, Nginx for reverse proxy, Redis for caching)
 RUN apt-get update && apt-get install -y \
     graphviz \
     libgraphviz-dev \
     pkg-config \
     build-essential \
     nginx \
+    redis-server \
     && rm -rf /var/lib/apt/lists/*
 
 # Set working directory
@@ -36,9 +37,10 @@ RUN mkdir -p logs storage && chmod -R 777 /app
 # Expose the port Hugging Face expects (7860)
 EXPOSE 7860
 
-# Run uvicorn backend + nginx reverse proxy
+# Run uvicorn backend + nginx reverse proxy + redis server
 # Streamlit frontend is launched as a sidecar by FastAPI startup event
 CMD ["sh", "-c", "\
   mkdir -p /tmp/nginx_client_body /tmp/nginx_proxy /tmp/nginx_fastcgi /tmp/nginx_uwsgi /tmp/nginx_scgi && \
+  redis-server --port 6379 --daemonize yes --protected-mode no --dir /app/storage & \
   python3 -m uvicorn backend.main:app --host 127.0.0.1 --port 8000 & \
   nginx -c /app/nginx.conf -g 'daemon off;'"]

backend/main.py

@@ -28,6 +28,7 @@ from backend.models import (
 )
 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
@@ -157,11 +158,46 @@ async def chat(request: ChatRequest):
         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,
@@ -170,7 +206,11 @@ async def chat(request: ChatRequest):
             steps_log=result["steps_log"],
             memories_used=result.get("memories_used", []),
             timestamp=datetime.now().isoformat(),
-            run_id=result.get("run_id"),   # LangSmith trace UUID (None if tracing off)
+            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)}")
@@ -187,13 +227,74 @@ async def chat_stream(session_id: str, query: str, plan_tier: str = "Unknown"):
     """
     import threading
     orch = manager.get_orchestrator(session_id)
-    queue = asyncio.Queue()
-    loop = asyncio.get_running_loop()
     
     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
+    cached_result = cache_manager.check(final_query, plan_tier=plan_tier)
+    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)
+            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):
@@ -202,7 +303,6 @@ async def chat_stream(session_id: str, query: str, plan_tier: str = "Unknown"):
         except Exception as e:
             loop.call_soon_threadsafe(queue.put_nowait, {"type": "error", "error": e})
 
-    # Start the LangGraph execution in a background thread so it doesn't block the main event loop
     thread = threading.Thread(target=run_stream, daemon=True)
     thread.start()
 
@@ -221,15 +321,14 @@ async def chat_stream(session_id: str, query: str, plan_tier: str = "Unknown"):
         prev_steps: list[str] = []
         final_answer = ""
         last_intent  = ""
+        last_state = None
 
         # ── Stream LangGraph events ───────────────────────────────
         try:
             while True:
                 try:
-                    # Wait for 2.0s for the next event, if none, send keepalive
                     item = await asyncio.wait_for(queue.get(), timeout=2.0)
                 except asyncio.TimeoutError:
-                    # SSE keepalive comment to keep the connection alive through reverse proxies
                     yield ": keepalive\n\n"
                     continue
 
@@ -241,31 +340,32 @@ async def chat_stream(session_id: str, query: str, plan_tier: str = "Unknown"):
                 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):]   # only new lines
+                new_steps     = current_steps[len(prev_steps):]
                 intent        = state.get("intent", last_intent)
                 last_intent   = intent
 
-                # Tell the frontend the LangGraph node just started
                 yield emit({"type": "node_start", "node": node, "intent": intent, "msg": f"Node '{node}' executing…"})
                 await asyncio.sleep(0.15)
 
-                # Emit individual sub-steps (each line from steps_log) with a small delay
                 for step in new_steps:
                     yield emit({"type": "substep", "node": node, "intent": intent, "step": step, "all_steps": current_steps})
-                    await asyncio.sleep(0.05)   # minimal pause for smooth UI without blocking stream
+                    await asyncio.sleep(0.05)
 
-                # Signal node completion
                 final_answer = state.get("answer", "")
                 yield emit({
-                    "type":   "node_done",
-                    "node":   node,
-                    "intent": intent,
-                    "steps":  current_steps,
-                    "answer": final_answer,
+                    "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", []),
                 })
 
-                # Persist completed state for trace/graph endpoints
                 if node == "synthesize":
                     if len(orch.chat_history) == 0 or orch.chat_history[-1] != ("ai", final_answer):
                         orch.chat_history.append(("human", query))
@@ -280,6 +380,24 @@ async def chat_stream(session_id: str, query: str, plan_tier: str = "Unknown"):
 
                 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)
+                
         except Exception as e:
             logger.error(f"Error in chat_stream: {str(e)}")
             yield emit({"type": "error", "msg": f"Backend Error: {str(e)}"})
@@ -376,66 +494,108 @@ async def clear_memory(session_id: str):
 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.
-    Detailed steps are added as chained nodes attached to Retrieval.
+    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')
-    # 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
+    
+    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.
-    The diagram will have the form:
-        FastAPI --> Orchestrator --> Intent --> Retrieval --> Synthesis
-    and then each step from steps will be added as a sub‑node.
+    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"]
-    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]')
-    # Add detailed steps as a vertical sub‑graph attached to D
-    if steps:
-        lines.append('    subgraph Details [Retrieval Details]')
-        for i, s in enumerate(steps, start=1):
-            # sanitize characters that break mermaid syntax
-            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')
+    
+    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)

backend/models.py

@@ -31,11 +31,15 @@ class ChatResponse(BaseModel):
     session_id: str
     query: str
     answer: str
-    intent: str         # SIMPLE_LOOKUP | POLICY_QUESTION | MULTI_HOP | COMPARISON
+    intent: str              # SIMPLE_LOOKUP | POLICY_QUESTION | MULTI_HOP | COMPARISON
     steps_log: List[str]
     memories_used: List[str] = []   # Facts Mem0 extracted and stored from this turn
     timestamp: str
-    run_id: Optional[str] = None    # LangSmith trace run ID (None if tracing disabled)
+    run_id: Optional[str] = None         # LangSmith trace run ID (None if tracing disabled)
+    confidence: Optional[str] = None    # HIGH | MEDIUM | LOW | BLOCKED (new confidence scorer)
+    confidence_reason: Optional[str] = None  # Human-readable confidence explanation
+    blocked: bool = False                # True if query_guard blocked the query
+    sub_questions: List[str] = []       # Sub-questions from query_decomposer (MULTI_HOP)
 
 
 class SessionHistoryResponse(BaseModel):

config.py

@@ -111,3 +111,12 @@ GUIDELINES:
    - Use 'policy_search' for general coverage rules, FAQs, and procedures.
    - Use 'relational_search' for specific data like copays for a drug, provider lookups, or plan-specific relational details.
 """
+
+# ──────────────────────────────────────────────
+# Redis Semantic Cache Settings
+# ──────────────────────────────────────────────
+REDIS_URL = os.getenv("REDIS_URL", "redis://localhost:6379")
+# Cosine similarity threshold for cache hits (1 - cosine_distance).
+# 0.85 is a standard threshold for text-embedding-3-small semantic similarity.
+SEMANTIC_CACHE_THRESHOLD = float(os.getenv("SEMANTIC_CACHE_THRESHOLD", "0.85"))
+SEMANTIC_CACHE_COLLECTION = "semantic_cache"

orchestration/orchestrator.py

@@ -1,49 +1,58 @@
 """
 Health Insurance AI Copilot — LangGraph Sequential Chain Orchestrator.
 
-Architecture (5 nodes in a directed StateGraph):
+Architecture (9 nodes in a directed StateGraph):
 
     START
       │
       ▼
-  [1. memory_search]  ← NEW: Retrieve relevant past facts from Mem0
+  [1. query_guard]     ← NEW: PII detection, off-topic guard, query sanitization
+      │ (blocked → jump to memory_add with pre-set answer)
+      ▼
+  [2. memory_search]   ← Retrieve relevant past facts from Mem0
       │
       ▼
-  [2. classify_intent]
+  [3. classify_intent]
       Uses GPT-4o-mini to classify the query into one of four intents:
       SIMPLE_LOOKUP | POLICY_QUESTION | MULTI_HOP | COMPARISON
       │
       ▼
-  [3. retrieve]
-      Routes internally by intent to the right retrieval strategy:
-      ┌─ SIMPLE_LOOKUP   → Graph entity lookup → Hybrid retrieval
-      ├─ POLICY_QUESTION → Hybrid retrieval only
-      ├─ MULTI_HOP       → Graph → Policy → Prior-Auth (3 sequential steps)
-      └─ COMPARISON      → Hybrid retrieval per plan tier (Bronze / Silver / Gold)
+  [4. query_decomposer] ← NEW: Splits MULTI_HOP into atomic sub-questions
+      │
+      ▼
+  [5. retrieve]
+      Routes internally by intent to the right retrieval strategy.
+      If sub_questions exist (MULTI_HOP), retrieves per sub-question.
       │
       ▼
-  [4. synthesize]
+  [6. context_quality_check] ← NEW: Validates context richness before synthesis
+      │ (skip_synthesis → jump to confidence_scorer with fallback answer)
+      ▼
+  [7. synthesize]
       Uses GPT-4o to generate a cited, safety-compliant final answer
       │
       ▼
-  [5. memory_add]     ← NEW: Persist Q&A facts to Mem0 for future sessions
+  [8. self_critique]   ← NEW: Verifies answer quality (Self-RAG style). Max 1 retry.
+      │ (needs_retry → loop back to retrieve)
+      ▼
+  [9. confidence_scorer] ← NEW: Heuristic HIGH/MEDIUM/LOW confidence rating
+      │
+      ▼
+  [10. memory_add]     ← Persist Q&A facts to Mem0 for future sessions
       │
       ▼
     END
 """
 
 import os
+import re
 import sys
-from typing import TypedDict, List
+from typing import TypedDict, List, Optional
 import concurrent.futures
 import contextvars
 
-# Ensure project root is on sys.path so `config` and sibling packages resolve
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
-# ── Load .env FIRST — before any LangSmith/LangChain imports ─────────────────
-# LANGCHAIN_TRACING_V2 must be in os.environ before langsmith is imported,
-# otherwise the SDK and our _TRACING_ENABLED flag will always read False.
 from dotenv import load_dotenv
 load_dotenv()
 
@@ -52,12 +61,10 @@ from langchain_core.messages import HumanMessage, SystemMessage
 from langgraph.graph import StateGraph, END
 from rich.console import Console
 
-# LangSmith observability — graceful no-ops if LANGCHAIN_TRACING_V2 is not set
 try:
     from langsmith import traceable as _langsmith_traceable
     _LANGSMITH_AVAILABLE = True
 except ImportError:
-    # If langsmith is not installed, create a passthrough decorator
     def _langsmith_traceable(*args, **kwargs):  # type: ignore[misc]
         def _decorator(fn):
             return fn
@@ -65,7 +72,6 @@ except ImportError:
     _LANGSMITH_AVAILABLE = False
 
 from orchestration.langsmith_tracing import tag_current_run, get_run_id
-
 from config import (
     LLM_MODEL,
     CLASSIFIER_LLM_MODEL,
@@ -80,23 +86,29 @@ console = Console()
 
 
 # ══════════════════════════════════════════════════════════════
-# 1.  STATE  — the shared data bag that flows through all nodes
+# 1.  STATE — the shared data bag that flows through all nodes
 # ══════════════════════════════════════════════════════════════
 
 class AgentState(TypedDict):
-    """
-    What gets passed between every node in the graph.
-    Each node receives this dict and returns an updated copy.
-    """
-    query:             str          # Original user question
-    user_id:           str          # Session identifier (for logging)
-    intent:            str          # Classified intent (set by classify_intent)
-    retrieved_context: str          # All retrieved text (set by retrieve)
-    past_memories:     str          # Relevant session memories (set by memory_search)
-    memories_used:     List[str]    # Facts stored to Mem0 (set by memory_add)
-    answer:            str          # Final answer (set by synthesize)
-    chat_history:      List[tuple]  # [(role, message), ...] — last 5 turns
-    steps_log:         List[str]    # Human-readable trace of what happened
+    """What gets passed between every node in the graph."""
+    query:              str           # Original user question
+    user_id:            str           # Session identifier
+    intent:             str           # Classified intent
+    retrieved_context:  str           # All retrieved text
+    past_memories:      str           # Relevant session memories
+    memories_used:      List[str]     # Facts stored to Mem0
+    answer:             str           # Final answer
+    chat_history:       List[tuple]   # [(role, message), ...]
+    steps_log:          List[str]     # Human-readable trace
+    # ── NEW fields ────────────────────────────────────────────
+    blocked:            bool          # query_guard: was query blocked?
+    block_reason:       str           # query_guard: reason for blocking
+    sub_questions:      List[str]     # query_decomposer: MULTI_HOP sub-questions
+    skip_synthesis:     bool          # context_quality_check: skip synthesis?
+    needs_retry:        bool          # self_critique: trigger retrieve retry?
+    retry_count:        int           # self_critique: how many retries done
+    confidence:         str           # confidence_scorer: HIGH / MEDIUM / LOW
+    confidence_reason:  str           # confidence_scorer: explanation
 
 
 # ══════════════════════════════════════════════════════════════
@@ -104,14 +116,13 @@ class AgentState(TypedDict):
 # ══════════════════════════════════════════════════════════════
 
 def _classifier_llm() -> ChatOpenAI:
-    """Fast, cheap model for intent classification."""
+    """Fast, cheap model for intent classification and auxiliary LLM calls."""
     return ChatOpenAI(
         model=CLASSIFIER_LLM_MODEL,
         temperature=LLM_TEMPERATURE,
         openai_api_key=os.getenv("OPENAI_API_KEY"),
     )
 
-
 def _synthesis_llm() -> ChatOpenAI:
     """High-accuracy model for final answer synthesis."""
     return ChatOpenAI(
@@ -122,30 +133,31 @@ def _synthesis_llm() -> ChatOpenAI:
 
 
 # ══════════════════════════════════════════════════════════════
-# 3.  NODE 1 — memory_search  (NEW)
+# 3.  NODE FUNCTIONS (stateless — defined at module level)
 # ══════════════════════════════════════════════════════════════
 
-def memory_search(state: AgentState) -> AgentState:
-    """NODE 1: Retrieve relevant past memories from Mem0 before classification."""
-    query   = state["query"]
-    user_id = state.get("user_id", "default")
-    log     = list(state.get("steps_log", []))
-
-    memories = search_memories(user_id=user_id, query=query)
-    if memories:
-        log.append(f"🧠 Mem0 recalled {memories.count(chr(10) + '  ')+1} memory fact(s) for user")
-    else:
-        log.append("🧠 Mem0: no relevant past memories found")
+# ── PII patterns to detect and redact ─────────────────────────
+_PII_PATTERNS = [
+    (r'\b\d{3}-\d{2}-\d{4}\b',                                'SSN pattern'),
+    (r'\b(0[1-9]|1[0-2])[-/](0[1-9]|[12]\d|3[01])[-/](19|20)\d{2}\b', 'date-of-birth pattern'),
+    (r'\b[2-9]\d{3}[\s-]?\d{4}[\s-]?\d{4}[\s-]?\d{4}\b',   'credit/debit card pattern'),
+    (r'\bMRN[\s:#]*\d{6,}\b',                                 'medical record number'),
+]
+
+# ── Off-topic topic signal words ───────────────────────────────
+_OFF_TOPIC_SIGNALS = [
+    'weather', 'forecast', 'cryptocurrency', 'bitcoin', 'stock price',
+    'sports score', 'movie review', 'recipe', 'video game', 'dating',
+    'political party', 'election result',
+]
 
-    return {**state, "past_memories": memories, "steps_log": log}
-
-
-
-# ══════════════════════════════════════════════════════════════
-# 4.  NODE 2 — classify_intent
-# ══════════════════════════════════════════════════════════════
+def classify_intent(state: AgentState) -> AgentState:
+    """NODE 3: Intent Classification."""
+    llm   = _classifier_llm()
+    query = state["query"]
+    log   = list(state.get("steps_log", []))
 
-_INTENT_SYSTEM = """You are a query intent classifier for a Health Insurance AI assistant.
+    _INTENT_SYSTEM = """You are a query intent classifier for a Health Insurance AI assistant.
 Classify the user query into EXACTLY ONE of these four categories:
 
   SIMPLE_LOOKUP    — Quick single-fact lookups.
@@ -169,13 +181,6 @@ Classify the user query into EXACTLY ONE of these four categories:
 
 Respond with ONLY the category name — no explanation, no punctuation, just the word."""
 
-
-def classify_intent(state: AgentState) -> AgentState:
-    """NODE 1: Intent Classification."""
-    llm   = _classifier_llm()
-    query = state["query"]
-    log   = list(state.get("steps_log", []))
-
     response = llm.invoke([
         SystemMessage(content=_INTENT_SYSTEM),
         HumanMessage(content=query),
@@ -189,38 +194,38 @@ def classify_intent(state: AgentState) -> AgentState:
     return {**state, "intent": intent, "steps_log": log}
 
 
-# ══════════════════════════════════════════════════════════════
-# 4.  NODE 2 — retrieve
-# ══════════════════════════════════════════════════════════════
-
 def retrieve(state: AgentState) -> AgentState:
-    """NODE 2: Smart Retrieval."""
-    query  = state["query"]
-    intent = state["intent"]
-    log    = list(state.get("steps_log", []))
+    """NODE 5: Smart Retrieval — routes by intent, uses sub-questions when available."""
+    query         = state["query"]
+    intent        = state["intent"]
+    sub_questions = state.get("sub_questions", [])
+    retry_count   = state.get("retry_count", 0)
+    log           = list(state.get("steps_log", []))
     parts: list[str] = []
 
-    # Initialize context-aware trace log for this node execution
     internal_logs = []
     token = trace_log.set(internal_logs)
 
     def _run_tool(tool_func, kwargs):
         return contextvars.copy_context().run(tool_func, kwargs)
 
-    # ── SIMPLE_LOOKUP ────────────────────────────────────────────
+    if retry_count > 0:
+        log.append(f"🔄 Retry #{retry_count}: re-running retrieval with broader search")
+
+    # ── SIMPLE_LOOKUP ──────────────────────────────────────────
     if intent == "SIMPLE_LOOKUP":
         log.append("📊 [SIMPLE_LOOKUP] Executing Graph & Hybrid retrieval concurrently")
         with concurrent.futures.ThreadPoolExecutor(max_workers=2) as executor:
             fut_g = executor.submit(_run_tool, relational_search.invoke, {"query": query})
             fut_p = executor.submit(_run_tool, policy_search.invoke, {"query": query})
-            
+
             graph_ctx = fut_g.result()
             if graph_ctx and "No structured" not in graph_ctx:
                 parts.append(f"[STRUCTURED GRAPH FACTS]\n{graph_ctx}")
                 log.append("🕸️ Graph entity lookup completed")
             else:
                 log.append("🕸️ Graph entity lookup — no structured results")
-                
+
             policy_ctx = fut_p.result()
             if policy_ctx and "No relevant" not in policy_ctx:
                 parts.append(f"[POLICY DOCUMENTS]\n{policy_ctx}")
@@ -228,35 +233,39 @@ def retrieve(state: AgentState) -> AgentState:
             else:
                 log.append("📄 Hybrid policy retrieval — no relevant results")
 
-    # ── POLICY_QUESTION ──────────────────────────────────────────
+    # ── POLICY_QUESTION ────────────────────────────────────────
     elif intent == "POLICY_QUESTION":
         log.append("📄 [POLICY_QUESTION] Full hybrid document retrieval")
         policy_ctx = policy_search.invoke({"query": query})
         if policy_ctx and "No relevant" not in policy_ctx:
             parts.append(f"[POLICY DOCUMENTS]\n{policy_ctx}")
 
-    # ── MULTI_HOP ────────────────────────────────────────────────
+    # ── MULTI_HOP ──────────────────────────────────────────────
     elif intent == "MULTI_HOP":
-        log.append("🔗 [MULTI_HOP] Executing Graph, Policy, and Prior-Auth retrievals concurrently")
+        # Use decomposed sub-questions for targeted retrieval per hop
+        queries_to_run = sub_questions if sub_questions else [query]
+        log.append(f"🔗 [MULTI_HOP] Executing {len(queries_to_run)} targeted retrieval(s)")
+
         with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
+            # Always run graph, policy, and prior-auth in parallel on main query
             fut_g = executor.submit(_run_tool, relational_search.invoke, {"query": query})
             fut_p = executor.submit(_run_tool, policy_search.invoke, {"query": query})
             fut_a = executor.submit(_run_tool, prior_auth_search.invoke, {"query": query})
-            
+
             graph_ctx = fut_g.result()
             if graph_ctx and "No structured" not in graph_ctx:
                 parts.append(f"[STRUCTURED GRAPH FACTS]\n{graph_ctx}")
                 log.append("🕸️ Graph entity lookup completed")
             else:
                 log.append("🕸️ Graph entity lookup — no structured results")
-                
+
             policy_ctx = fut_p.result()
             if policy_ctx and "No relevant" not in policy_ctx:
                 parts.append(f"[POLICY DOCUMENTS]\n{policy_ctx}")
                 log.append("📄 Hybrid policy retrieval completed")
             else:
                 log.append("📄 Hybrid policy retrieval — no relevant results")
-                
+
             auth_ctx = fut_a.result()
             if auth_ctx and "No prior authorization" not in auth_ctx:
                 parts.append(f"[PRIOR AUTHORIZATION RULES]\n{auth_ctx}")
@@ -264,9 +273,18 @@ def retrieve(state: AgentState) -> AgentState:
             else:
                 log.append("📋 Prior authorization — no relevant rules found")
 
-    # ── COMPARISON ───────────────────────────────────────────────
+        # Additionally run per-sub-question retrieval for targeted coverage
+        if len(sub_questions) > 1:
+            log.append(f"✂️ Running targeted retrieval for {len(sub_questions)} sub-question(s)")
+            for i, sq in enumerate(sub_questions):
+                sq_ctx = policy_search.invoke({"query": sq})
+                if sq_ctx and "No relevant" not in sq_ctx:
+                    parts.append(f"[SUB-QUESTION {i+1}: {sq}]\n{sq_ctx}")
+                    log.append(f"  ↳ Sub-Q {i+1} retrieved successfully")
+
+    # ── COMPARISON ─────────────────────────────────────────────
     elif intent == "COMPARISON":
-        log.append(f"⚖️  [COMPARISON] Retrieving Bronze, Silver, and Gold plan contexts concurrently")
+        log.append("⚖️  [COMPARISON] Retrieving Bronze, Silver, and Gold plan contexts concurrently")
         with concurrent.futures.ThreadPoolExecutor(max_workers=3) as executor:
             futures = {
                 tier: executor.submit(_run_tool, plan_comparison_search.invoke, {"query": query, "tier": tier})
@@ -280,25 +298,18 @@ def retrieve(state: AgentState) -> AgentState:
                 else:
                     log.append(f"⚠️ [{tier} Tier] No relevant context found")
 
-    separator   = "\n\n" + "─" * 60 + "\n\n"
+    separator    = "\n\n" + "─" * 60 + "\n\n"
     full_context = separator.join(parts) if parts else "No relevant context found."
-    
-    # Capture any internal logs (like Multi-Query variants)
+
     for l in internal_logs:
         log.append(l)
-    
+
     log.append(f"✅ Retrieved {len(parts)} context section(s)")
-    
-    # Cleanup context
     trace_log.reset(token)
 
     return {**state, "retrieved_context": full_context, "steps_log": log}
 
 
-# ══════════════════════════════════════════════════════════════
-# 6.  NODE 4 — synthesize
-# ══════════════════════════════════════════════════════════════
-
 _SYNTHESIS_TEMPLATE = """{system_prompt}
 
 ─── PAST USER MEMORIES (from previous sessions) ──────────────
@@ -314,14 +325,14 @@ _SYNTHESIS_TEMPLATE = """{system_prompt}
 ──────────────────────────────────────────────────────────────
 
 Using ONLY the retrieved context above, answer the user's question.
-IMPORTANT: You MUST also follow any user preferences (e.g., formatting, language, or specific plan details) found in the PAST USER MEMORIES section.
-If the user asks for entities with multiple criteria (e.g., "Specialist X in City Y"), you MUST verify that the same entity satisfies ALL criteria in the context. Do NOT assume an entity has a property just because it is listed in the context alongside another entity.
-Always cite the source file and page number for every fact you state EXACTLY as it appears in the context (e.g., "(Source: filename.pdf, Page: 2)"). Failure to include the page number is a strict violation.
+IMPORTANT: You MUST also follow any user preferences found in the PAST USER MEMORIES section.
+If the user asks for entities with multiple criteria, verify ALL criteria match in the context.
+Always cite the source file and page number for every fact EXACTLY as it appears (e.g., "(Source: filename.pdf, Page: 2)").
 If the context does not contain enough information, say so explicitly — do NOT guess."""
 
 
 def synthesize(state: AgentState) -> AgentState:
-    """NODE 3: Answer Synthesis."""
+    """NODE 7: Answer Synthesis."""
     llm   = _synthesis_llm()
     query = state["query"]
     log   = list(state.get("steps_log", []))
@@ -350,42 +361,110 @@ def synthesize(state: AgentState) -> AgentState:
 
 
 # ══════════════════════════════════════════════════════════════
-# 7.  PUBLIC ORCHESTRATOR CLASS
+# 4.  ORCHESTRATOR CLASS — builds the LangGraph with all 9 nodes
 # ══════════════════════════════════════════════════════════════
 
 class Orchestrator:
     def __init__(self, user_id: str = "default", mem=None):
-        """
-        Args:
-            user_id: The session ID — used for logging/tracing.
-            mem:     A pre-created Mem0 Memory instance (in-memory, per-session).
-                     Created by SessionManager alongside this Orchestrator.
-                     Pass None to disable memory features.
-        """
         if not os.getenv("OPENAI_API_KEY"):
             raise ValueError("OPENAI_API_KEY not found in environment.")
-        self._mem = mem                      # Mem0 in-memory instance for this session
+        self._mem = mem
         self.user_id: str = user_id
         self.chat_history: List[tuple] = []
         self.last_detailed_result: dict = {}
 
-        # Bind memory helpers to this session's Memory instance so nodes can call them
         self._search_memories = lambda query: search_memories(self._mem, query)
         self._add_memory       = lambda q, a:  add_memory(self._mem, q, a)
 
         self.graph: any = self._build_graph()
 
-    # ── Build graph with closures over self._mem ──────────────────────────────
-
     def _build_graph(self):
-        """Build the LangGraph StateGraph with memory nodes bound to this session."""
+        """Build the LangGraph with 9 nodes, conditional edges, and retry loop."""
         _search = self._search_memories
         _add    = self._add_memory
 
-        def memory_search_node(state: AgentState) -> AgentState:
-            """NODE 1: Retrieve relevant facts from this session's Mem0."""
+        # ── Node 1: Query Guard ─────────────────────────────────
+        def query_guard_node(state: AgentState) -> AgentState:
+            """PII detection, off-topic guard, query sanitization."""
             query = state["query"]
             log   = list(state.get("steps_log", []))
+
+            # 1a. PII detection — redact but don't block (still useful info to answer)
+            sanitized = query
+            pii_found = []
+            for pattern, label in _PII_PATTERNS:
+                if re.search(pattern, query, re.IGNORECASE):
+                    sanitized = re.sub(pattern, '[REDACTED]', sanitized, flags=re.IGNORECASE)
+                    pii_found.append(label)
+
+            if pii_found:
+                log.append(f"🛡️ Guard: PII detected ({', '.join(pii_found)}) — sensitive data redacted")
+                # Don't block, just sanitize
+                return {**state, "query": sanitized, "blocked": False, "block_reason": "",
+                        "sub_questions": [], "skip_synthesis": False, "needs_retry": False,
+                        "retry_count": 0, "confidence": "", "confidence_reason": "",
+                        "steps_log": log}
+
+            # 1b. Off-topic detection
+            q_lower = query.lower()
+            for signal in _OFF_TOPIC_SIGNALS:
+                if signal in q_lower:
+                    log.append(f"🛡️ Guard BLOCKED: off-topic signal '{signal}' detected")
+                    return {
+                        **state,
+                        "blocked":     True,
+                        "block_reason": "OFF_TOPIC",
+                        "answer":      (
+                            f"I'm a Health Insurance AI assistant and can only help with health "
+                            f"insurance questions (coverage, claims, providers, plan comparisons, "
+                            f"drug formularies, etc.). Your question appears to be about '{signal}' "
+                            f"which is outside my scope. Please ask something related to your health plan."
+                        ),
+                        "sub_questions":   [],
+                        "skip_synthesis":  True,
+                        "needs_retry":     False,
+                        "retry_count":     0,
+                        "confidence":      "BLOCKED",
+                        "confidence_reason": "Off-topic query",
+                        "steps_log": log,
+                    }
+
+            # 1c. Empty / trivially short query
+            if len(query.strip()) < 4:
+                log.append("🛡️ Guard BLOCKED: query too short or empty")
+                return {
+                    **state,
+                    "blocked":     True,
+                    "block_reason": "INVALID_QUERY",
+                    "answer":      "Please provide a more detailed health insurance question.",
+                    "sub_questions":   [],
+                    "skip_synthesis":  True,
+                    "needs_retry":     False,
+                    "retry_count":     0,
+                    "confidence":      "BLOCKED",
+                    "confidence_reason": "Query too short",
+                    "steps_log": log,
+                }
+
+            log.append("✅ Guard: query validated — all checks passed")
+            return {
+                **state,
+                "blocked":     False,
+                "block_reason": "",
+                "sub_questions":   [],
+                "skip_synthesis":  False,
+                "needs_retry":     False,
+                "retry_count":     state.get("retry_count", 0),
+                "confidence":      "",
+                "confidence_reason": "",
+                "steps_log": log,
+            }
+
+        # ── Node 2: Memory Search ───────────────────────────────
+        def memory_search_node(state: AgentState) -> AgentState:
+            """Retrieve relevant facts from this session's Mem0."""
+            query    = state["query"]
+            log      = list(state.get("steps_log", []))
             memories = _search(query)
             if memories:
                 count = memories.count("\n  ") + 1
@@ -394,31 +473,264 @@ class Orchestrator:
                 log.append("🧠 Mem0: no relevant facts recalled yet")
             return {**state, "past_memories": memories, "steps_log": log}
 
+        # ── Node 4: Query Decomposer ────────────────────────────
+        def query_decomposer_node(state: AgentState) -> AgentState:
+            """Split MULTI_HOP queries into atomic sub-questions for targeted retrieval."""
+            intent = state.get("intent", "")
+            log    = list(state.get("steps_log", []))
+
+            # Only decompose MULTI_HOP — pass-through for everything else
+            if intent != "MULTI_HOP":
+                log.append(f"✂️ Decomposer: skipped (intent={intent}, no decomposition needed)")
+                return {**state, "sub_questions": [], "steps_log": log}
+
+            query = state["query"]
+            llm   = _classifier_llm()
+
+            prompt = (
+                f"Break this health insurance question into 2-4 atomic sub-questions "
+                f"that can each be answered independently from separate document lookups.\n\n"
+                f"Question: {query}\n\n"
+                f"Rules:\n"
+                f"- Each sub-question must be self-contained (no pronouns referencing other sub-questions)\n"
+                f"- Focus on ONE entity or fact per sub-question\n"
+                f"- Return ONLY the sub-questions, one per line, no numbering or bullet points"
+            )
+
+            try:
+                response = llm.invoke(prompt)
+                sub_questions = [q.strip() for q in response.content.splitlines() if q.strip()]
+                # Fallback: if decomposition fails or returns nothing useful
+                if not sub_questions or len(sub_questions) == 1:
+                    sub_questions = [query]
+                    log.append(f"✂️ Decomposer: could not split — using original query")
+                else:
+                    log.append(f"✂️ Decomposer: split into {len(sub_questions)} sub-question(s)")
+                    for i, sq in enumerate(sub_questions):
+                        log.append(f"  ↳ Sub-Q {i+1}: {sq}")
+            except Exception as e:
+                sub_questions = [query]
+                log.append(f"✂️ Decomposer: error ({e}) — using original query")
+
+            return {**state, "sub_questions": sub_questions, "steps_log": log}
+
+        # ── Node 6: Context Quality Check ───────────────────────
+        def context_quality_check_node(state: AgentState) -> AgentState:
+            """Validate retrieved context richness before synthesis."""
+            log = list(state.get("steps_log", []))
+            ctx = state.get("retrieved_context", "")
+
+            # Count meaningful sections
+            empty_phrases = [
+                "No relevant context found.",
+                "No structured relational data found",
+                "No relevant policy information found",
+            ]
+            is_empty = (
+                not ctx or
+                ctx.strip() in empty_phrases or
+                all(ep in ctx for ep in ["No relevant"]) or
+                len(ctx.strip()) < 80
+            )
+
+            if is_empty:
+                log.append("⚠️ Context Quality: insufficient context retrieved — synthesis skipped")
+                return {
+                    **state,
+                    "skip_synthesis": True,
+                    "answer": (
+                        "I don't have sufficient information in my knowledge base to answer this "
+                        "question accurately. This could be because:\n"
+                        "• The specific plan detail isn't in our documents\n"
+                        "• The question may need more specific terms\n\n"
+                        "Please contact your insurance provider directly or refer to your plan documents."
+                    ),
+                    "confidence":        "LOW",
+                    "confidence_reason": "No relevant context retrieved",
+                    "steps_log": log,
+                }
+
+            # Count context sections for logging
+            section_count = ctx.count("─" * 20) + 1
+            char_count    = len(ctx)
+            log.append(
+                f"✅ Context Quality: {section_count} section(s) | {char_count:,} chars — synthesis proceeding"
+            )
+            return {**state, "skip_synthesis": False, "steps_log": log}
+
+        # ── Node 8: Self-Critique ───────────────────────────────
+        def self_critique_node(state: AgentState) -> AgentState:
+            """Self-RAG style reflection: verify answer quality, trigger retry if poor."""
+            log         = list(state.get("steps_log", []))
+            retry_count = state.get("retry_count", 0)
+
+            # Skip if already retried, blocked, or synthesis was skipped
+            if retry_count >= 1 or state.get("skip_synthesis"):
+                log.append("🔄 Self-Critique: skipped (max retries reached or synthesis was skipped)")
+                return {**state, "needs_retry": False, "steps_log": log}
+
+            answer = state.get("answer", "")
+            query  = state["query"]
+
+            if not answer:
+                log.append("🔄 Self-Critique: no answer to evaluate")
+                return {**state, "needs_retry": False, "steps_log": log}
+
+            llm = _classifier_llm()
+            prompt = (
+                f"You are evaluating a health insurance AI assistant's answer quality.\n\n"
+                f"Question: {query}\n\n"
+                f"Answer (first 600 chars): {answer[:600]}\n\n"
+                f"Evaluate:\n"
+                f"1. Does the answer directly address the specific question?\n"
+                f"2. Does it contain concrete details (numbers, plan names, coverage specifics)?\n"
+                f"3. Does it cite sources?\n\n"
+                f"Respond with ONLY one word:\n"
+                f"- GOOD: answer is adequate, specific, and addresses the question\n"
+                f"- RETRY: answer is vague, generic, or says 'I don't have info' when a better retrieval might help"
+            )
+
+            try:
+                response = llm.invoke(prompt)
+                verdict  = response.content.strip().upper()
+
+                if "RETRY" in verdict:
+                    log.append(f"🔄 Self-Critique: answer quality INSUFFICIENT — triggering retry #{retry_count + 1}")
+                    return {**state, "needs_retry": True, "retry_count": retry_count + 1, "steps_log": log}
+                else:
+                    log.append("✅ Self-Critique: answer quality VERIFIED — proceeding")
+                    return {**state, "needs_retry": False, "steps_log": log}
+            except Exception as e:
+                log.append(f"🔄 Self-Critique: evaluation error ({e}) — skipping retry")
+                return {**state, "needs_retry": False, "steps_log": log}
+
+        # ── Node 9: Confidence Scorer ───────────────────────────
+        def confidence_scorer_node(state: AgentState) -> AgentState:
+            """Heuristic HIGH/MEDIUM/LOW confidence rating based on context and answer quality."""
+            log = list(state.get("steps_log", []))
+
+            # Already set during guard or quality check
+            if state.get("confidence") in ("BLOCKED", "LOW"):
+                log.append(f"📊 Confidence: {state['confidence']} (pre-set by earlier node)")
+                return {**state, "steps_log": log}
+
+            answer  = state.get("answer", "")
+            ctx     = state.get("retrieved_context", "")
+            reasons = []
+            score   = 0
+
+            # Citations in answer
+            citation_count = answer.count("(Source:")
+            if citation_count >= 3:
+                score += 3;  reasons.append(f"{citation_count} citations")
+            elif citation_count >= 1:
+                score += 2;  reasons.append(f"{citation_count} citation(s)")
+            else:
+                score -= 1;  reasons.append("no citations found")
+
+            # Context richness
+            ctx_len = len(ctx)
+            if ctx_len > 2000:
+                score += 2;  reasons.append("rich context")
+            elif ctx_len > 500:
+                score += 1;  reasons.append("moderate context")
+            else:
+                reasons.append("sparse context")
+
+            # Hedging language in answer (uncertainty signals)
+            hedging_phrases = [
+                "i don't have", "cannot find", "not available in",
+                "no information", "unable to", "not in my knowledge"
+            ]
+            if any(h in answer.lower() for h in hedging_phrases):
+                score -= 1;  reasons.append("hedged answer")
+
+            # Retry needed (lower confidence if answer required retry)
+            if state.get("retry_count", 0) > 0:
+                score -= 1;  reasons.append("required retry")
+
+            # Final rating
+            if score >= 4:
+                confidence = "HIGH"
+            elif score >= 2:
+                confidence = "MEDIUM"
+            else:
+                confidence = "LOW"
+
+            reason_str = " | ".join(reasons)
+            log.append(f"📊 Confidence: {confidence} (score={score} — {reason_str})")
+
+            return {**state, "confidence": confidence, "confidence_reason": reason_str, "steps_log": log}
+
+        # ── Node 10: Memory Add ─────────────────────────────────
         def memory_add_node(state: AgentState) -> AgentState:
-            """NODE 5: Persist Q&A facts back to this session's Mem0."""
+            """Persist Q&A facts back to this session's Mem0."""
+            # Skip memory storage if query was blocked
+            if state.get("blocked"):
+                log = list(state.get("steps_log", []))
+                log.append("💾 Mem0: skipped (query was blocked)")
+                return {**state, "memories_used": [], "steps_log": log}
+
             stored = _add(state["query"], state["answer"])
             log    = list(state.get("steps_log", []))
-            log.append(f"💾 Mem0 stored {len(stored)} fact(s) from this turn" if stored else "💾 Mem0: no new facts extracted")
+            if stored:
+                log.append(f"💾 Mem0 stored {len(stored)} fact(s) from this turn")
+            else:
+                log.append("💾 Mem0: no new facts extracted")
             return {**state, "memories_used": stored, "steps_log": log}
 
+        # ── Build the StateGraph ────────────────────────────────
         builder = StateGraph(AgentState)
-        builder.add_node("memory_search",   memory_search_node)
-        builder.add_node("classify_intent", classify_intent)
-        builder.add_node("retrieve",        retrieve)
-        builder.add_node("synthesize",      synthesize)
-        builder.add_node("memory_add",      memory_add_node)
 
-        builder.set_entry_point("memory_search")
+        builder.add_node("query_guard",           query_guard_node)
+        builder.add_node("memory_search",          memory_search_node)
+        builder.add_node("classify_intent",        classify_intent)
+        builder.add_node("query_decomposer",       query_decomposer_node)
+        builder.add_node("retrieve",               retrieve)
+        builder.add_node("context_quality_check",  context_quality_check_node)
+        builder.add_node("synthesize",             synthesize)
+        builder.add_node("self_critique",          self_critique_node)
+        builder.add_node("confidence_scorer",      confidence_scorer_node)
+        builder.add_node("memory_add",             memory_add_node)
+
+        builder.set_entry_point("query_guard")
+
+        # Conditional: blocked queries skip straight to memory_add
+        builder.add_conditional_edges(
+            "query_guard",
+            lambda s: "blocked" if s.get("blocked") else "ok",
+            {"blocked": "memory_add", "ok": "memory_search"},
+        )
+
         builder.add_edge("memory_search",   "classify_intent")
-        builder.add_edge("classify_intent", "retrieve")
-        builder.add_edge("retrieve",        "synthesize")
-        builder.add_edge("synthesize",      "memory_add")
-        builder.add_edge("memory_add",      END)
+        builder.add_edge("classify_intent", "query_decomposer")
+        builder.add_edge("query_decomposer","retrieve")
+        builder.add_edge("retrieve",        "context_quality_check")
+
+        # Conditional: no context → skip synthesis, go straight to confidence_scorer
+        builder.add_conditional_edges(
+            "context_quality_check",
+            lambda s: "skip" if s.get("skip_synthesis") else "ok",
+            {"skip": "confidence_scorer", "ok": "synthesize"},
+        )
+
+        builder.add_edge("synthesize", "self_critique")
+
+        # Conditional: poor answer → retry retrieve (max once)
+        builder.add_conditional_edges(
+            "self_critique",
+            lambda s: "retry" if s.get("needs_retry") else "ok",
+            {"retry": "retrieve", "ok": "confidence_scorer"},
+        )
+
+        builder.add_edge("confidence_scorer", "memory_add")
+        builder.add_edge("memory_add",         END)
 
         return builder.compile()
 
+    # ── Base state builder ─────────���──────────────────────────────────────────
+
     def _base_state(self, query: str) -> AgentState:
-        """Build the initial AgentState for a new query."""
         return {
             "query":             query,
             "user_id":           self.user_id,
@@ -429,8 +741,19 @@ class Orchestrator:
             "answer":            "",
             "chat_history":      self.chat_history.copy(),
             "steps_log":         [],
+            # New fields — always initialized
+            "blocked":           False,
+            "block_reason":      "",
+            "sub_questions":     [],
+            "skip_synthesis":    False,
+            "needs_retry":       False,
+            "retry_count":       0,
+            "confidence":        "",
+            "confidence_reason": "",
         }
 
+    # ── Public API ────────────────────────────────────────────────────────────
+
     def ask(self, query: str, verbose: bool = False) -> str:
         result = self.graph.invoke(self._base_state(query))
 
@@ -440,7 +763,6 @@ class Orchestrator:
                 console.print(f"  [dim]{step}[/dim]")
 
         answer = result["answer"]
-
         self.chat_history.append(("human", query))
         self.chat_history.append(("ai",    answer))
         if len(self.chat_history) > 10:
@@ -450,32 +772,19 @@ class Orchestrator:
 
     @_langsmith_traceable(name="health-insurance-rag-query", run_type="chain")  # type: ignore[misc]
     def ask_detailed(self, query: str) -> dict:
-        """
-        Like ask(), but returns the full result dict for the API layer.
-        Decorated with @traceable so every call creates a named top-level
-        LangSmith parent span that groups all 6 LangGraph node child spans.
-
-        Returns:
-            {
-                "answer":            str,
-                "intent":            str,
-                "steps_log":         list[str],
-                "retrieved_context": str,
-                "memories_used":     list[str],
-                "run_id":            str | None,  # LangSmith trace UUID
-            }
-        """
+        """Like ask(), but returns the full result dict for the API layer."""
         result = self.graph.invoke(self._base_state(query))
         answer = result["answer"]
 
-        # ── Tag the LangSmith trace with structured metadata ──────────────────
-        # Called after invoke so intent & language are already resolved.
         tag_current_run(
             session_id=self.user_id,
             intent=result.get("intent", ""),
-            extra_metadata={"query_length": len(query)},
+            extra_metadata={
+                "query_length": len(query),
+                "blocked":      result.get("blocked", False),
+                "confidence":   result.get("confidence", ""),
+            },
         )
-        # Capture run_id while still inside the @traceable scope
         run_id = get_run_id()
 
         self.chat_history.append(("human", query))
@@ -490,28 +799,19 @@ class Orchestrator:
             "steps_log":         result.get("steps_log", []),
             "retrieved_context": result.get("retrieved_context", ""),
             "memories_used":     result.get("memories_used", []),
-            "run_id":            run_id,   # None when tracing is disabled
+            "run_id":            run_id,
+            "blocked":           result.get("blocked", False),
+            "confidence":        result.get("confidence", ""),
+            "confidence_reason": result.get("confidence_reason", ""),
+            "sub_questions":     result.get("sub_questions", []),
         }
         return self.last_detailed_result
 
     def stream_detailed(self, query: str):
-        """
-        Generator that yields intermediate AgentState updates as they happen.
-        Useful for "Live" Developer Console views.
-        """
-        # Use LangGraph's streaming mode
+        """Generator that yields intermediate AgentState updates for the Developer Console."""
         for event in self.graph.stream(self._base_state(query)):
-            # event is a dict like {"node_name": state_update}
             for node, state in event.items():
-                # We yield the current state and which node just finished
-                yield {
-                    "node": node,
-                    "state": state
-                }
-        
-        # After completion, update internal history (last yield will have the full state)
-        # Note: In a production stream, you might want to handle this differently
-        # but for this POC, the last event from 'synthesize' has the answer.
+                yield {"node": node, "state": state}
 
 
 if __name__ == "__main__":

orchestration/semantic_cache.py

@@ -0,0 +1,382 @@
+import os
+import json
+import hashlib
+import time
+from typing import Optional, Dict, List, Tuple
+from loguru import logger
+import redis
+
+from langchain_openai import OpenAIEmbeddings, ChatOpenAI
+from langchain_core.messages import SystemMessage, HumanMessage
+from langchain_chroma import Chroma
+from langchain_core.documents import Document
+
+from config import (
+    REDIS_URL,
+    SEMANTIC_CACHE_THRESHOLD,
+    SEMANTIC_CACHE_COLLECTION,
+    EMBEDDING_MODEL,
+    CHROMA_PERSIST_DIR,
+)
+
+class SemanticCache:
+    """
+    Cognitive Semantic Cache for LLM responses.
+    
+    1. Query Normalization:
+       - Uses a lightweight ChatOpenAI call (gpt-4o-mini) to translate conversational, 
+         first-person user queries into standard, formal third-person policy search statements
+         before vector search. This bridges the semantic gap (e.g. "asthma inhalers" -> "chronic pre-existing conditions and maintenance medications").
+         
+    2. Redis Mode:
+       - Persists query text, embedding vector of normalized query, and response JSON in Redis.
+       - Caches vectors and queries in Python memory on startup for sub-millisecond similarity search.
+       - Uses plain cosine similarity in Python.
+       - Does NOT require Redis Stack/RediSearch, making it 100% compatible with standard Redis (e.g. local redis-server on HF, Upstash, AWS).
+       
+    3. Local Fallback Mode:
+       - Used if Redis is unavailable.
+       - Stores vectors and response JSON in a local ChromaDB collection (`semantic_cache`).
+       - Matches using ChromaDB's native vector similarity search on normalized query vectors.
+    """
+    def __init__(self):
+        # Initialize Embeddings
+        logger.info(f"Initializing SemanticCache embeddings with model {EMBEDDING_MODEL}...")
+        self.embeddings = OpenAIEmbeddings(model=EMBEDDING_MODEL)
+        
+        # Initialize query normalizer LLM
+        logger.info("Initializing SemanticCache normalizer LLM (gpt-4o-mini)...")
+        self.normalizer_llm = ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
+        
+        # Redis configuration
+        self.redis_url = REDIS_URL
+        self.redis_client: Optional[redis.Redis] = None
+        self.redis_available = False
+        
+        # In-memory index of vectors for fast search in Redis Mode
+        # Format: {cache_id: {"query": query_text, "vector": list[float], "plan_tier": plan_tier}}
+        self._cache_memory: Dict[str, Dict] = {}
+        
+        # Local Fallback Store (Chroma)
+        self.vector_store: Optional[Chroma] = None
+        
+        # Try to connect to Redis
+        try:
+            logger.info(f"Connecting to Redis at {self.redis_url}...")
+            self.redis_client = redis.Redis.from_url(
+                self.redis_url, 
+                socket_timeout=1.5,
+                socket_connect_timeout=1.5,
+                decode_responses=True  # Decode hash values as string
+            )
+            self.redis_client.ping()
+            self.redis_available = True
+            logger.info("Successfully connected to Redis. Running in Redis Cache Mode.")
+            
+            # Load existing cache items into memory
+            self._load_cache_into_memory()
+            
+        except Exception as e:
+            logger.warning(f"Redis not available: {e}. Falling back to ChromaDB Local Cache Mode.")
+            self.redis_available = False
+            self._init_chroma_store()
+
+    def _init_chroma_store(self):
+        """Initialize Chroma collection for fallback caching."""
+        try:
+            self.vector_store = Chroma(
+                persist_directory=CHROMA_PERSIST_DIR,
+                embedding_function=self.embeddings,
+                collection_name=SEMANTIC_CACHE_COLLECTION
+            )
+            logger.info(f"Initialized fallback ChromaDB cache collection: '{SEMANTIC_CACHE_COLLECTION}'")
+        except Exception as e:
+            logger.error(f"Failed to initialize fallback ChromaDB: {e}")
+
+    def _load_cache_into_memory(self):
+        """Pre-load all cached vectors from Redis into Python memory for fast scanning."""
+        if not self.redis_client:
+            return
+            
+        start_time = time.time()
+        try:
+            # Retrieve all cache IDs
+            cache_ids = self.redis_client.smembers("cache:ids")
+            if not cache_ids:
+                logger.info("Redis cache is empty. In-memory index initialized empty.")
+                return
+                
+            logger.info(f"Pre-loading {len(cache_ids)} cached queries from Redis...")
+            
+            # Fetch the query and vector hashes in pipelined execution
+            pipeline = self.redis_client.pipeline()
+            for cid in cache_ids:
+                pipeline.hmget(f"cache:data:{cid}", ["query", "vector", "plan_tier"])
+            
+            results = pipeline.execute()
+            
+            loaded_count = 0
+            for cid, (query, vector_str, plan_tier) in zip(cache_ids, results):
+                if query and vector_str:
+                    try:
+                        vector = json.loads(vector_str)
+                        self._cache_memory[cid] = {
+                            "query": query,
+                            "vector": vector,
+                            "plan_tier": plan_tier or "Unknown"
+                        }
+                        loaded_count += 1
+                    except Exception as ve:
+                        logger.error(f"Error parsing vector for key {cid}: {ve}")
+            
+            logger.info(f"Loaded {loaded_count} cache keys into Python memory in {time.time() - start_time:.3f}s")
+            
+        except Exception as e:
+            logger.error(f"Error pre-loading cache from Redis: {e}")
+
+    def _cosine_similarity(self, vec_a: List[float], vec_b: List[float]) -> float:
+        """
+        Calculate cosine similarity between two vectors.
+        OpenAI text-embedding-3-small vectors are already normalized (L2 norm = 1.0).
+        Thus, cosine similarity is exactly the dot product.
+        """
+        if len(vec_a) != len(vec_b):
+            return 0.0
+        return sum(a * b for a, b in zip(vec_a, vec_b))
+
+    def normalize_query(self, query: str) -> str:
+        """
+        Normalize conversational, user-specific queries into formal, third-person health insurance search queries.
+        e.g., "asthma inhaler wait" -> "waiting periods for pre-existing conditions and maintenance medications"
+        """
+        if not query or len(query.strip()) < 4:
+            return query
+            
+        try:
+            system_prompt = (
+                "You are a health insurance query normalizer. Convert the user's conversational, first-person query "
+                "into a standard, formal, third-person search query about health insurance policies, deductibles, "
+                "coverage rules, or providers. Strip away personal details (like names, specific diagnosis dates, "
+                "pronouns, conversational filler). Keep plan names and general terms. Keep it concise.\n\n"
+                "Examples:\n"
+                "- 'What is the policy regarding waiting periods for chronic, pre-existing conditions before the plan starts paying for specialist visits and maintenance medications?' "
+                "-> 'Waiting periods for pre-existing conditions, specialist visits, and maintenance medications'\n"
+                "- 'I was diagnosed with asthma last year and just signed up for this insurance. Do I have to wait a certain number of months before you guys will cover my inhalers and pulmonologist appointments, or am I good to go right now?' "
+                "-> 'Waiting periods for pre-existing conditions, specialist visits, and maintenance medications'\n"
+                "- 'Who are the skin doctors in Aurora?' -> 'In-network dermatologists in Aurora'\n"
+                "- 'How much do I pay for Metformin on Silver?' -> 'Copay for Metformin on Silver plan'"
+            )
+            
+            response = self.normalizer_llm.invoke([
+                SystemMessage(content=system_prompt),
+                HumanMessage(content=query)
+            ])
+            
+            normalized = response.content.strip()
+            logger.info(f"Normalized query: '{query[:50]}...' -> '{normalized}'")
+            return normalized
+        except Exception as e:
+            logger.warning(f"Failed to normalize query: {e}. Using original query.")
+            return query
+
+    def check(self, query: str, plan_tier: str = "Unknown") -> Optional[dict]:
+        """
+        Check the semantic cache for a match.
+        
+        Returns:
+            dict containing the response data and hit metadata if found, else None.
+        """
+        if not query or len(query.strip()) < 4:
+            return None
+            
+        start_time = time.time()
+        try:
+            # 1. Normalize query to standard terminology
+            norm_query = self.normalize_query(query)
+            
+            # 2. Embed the normalized query
+            query_vector = self.embeddings.embed_query(norm_query)
+            
+            if self.redis_available and self.redis_client:
+                # Redis Mode: Scan memory cache for similarity
+                if not self._cache_memory:
+                    return None
+                    
+                best_id = None
+                best_score = -1.0
+                
+                # Compute similarities in Python memory, filtering by plan_tier
+                for cid, item in self._cache_memory.items():
+                    # Filter by plan_tier case-insensitively
+                    if item.get("plan_tier", "Unknown").lower() != plan_tier.lower():
+                        continue
+                        
+                    sim = self._cosine_similarity(query_vector, item["vector"])
+                    if sim > best_score:
+                        best_score = sim
+                        best_id = cid
+                
+                # Check threshold
+                if best_id and best_score >= SEMANTIC_CACHE_THRESHOLD:
+                    # Cache Hit! Retrieve full response from Redis
+                    cached_json = self.redis_client.hget(f"cache:data:{best_id}", "response")
+                    if cached_json:
+                        response = json.loads(cached_json)
+                        # Inject cache metadata
+                        response["cached"] = True
+                        response["cache_similarity"] = round(best_score * 100, 1)
+                        response["matched_query"] = self._cache_memory[best_id]["query"]
+                        
+                        logger.info(f"⚡ Redis Semantic Cache HIT (Plan: {plan_tier}, Similarity: {response['cache_similarity']}%) in {time.time() - start_time:.3f}s")
+                        return response
+                        
+            else:
+                # Fallback ChromaDB Mode
+                if not self.vector_store:
+                    self._init_chroma_store()
+                if not self.vector_store:
+                    return None
+                    
+                # Search ChromaDB with plan_tier filter using the normalized query
+                results = self.vector_store.similarity_search_with_score(
+                    norm_query, 
+                    k=1, 
+                    filter={"plan_tier": plan_tier}
+                )
+                if results:
+                    doc, distance = results[0]
+                    # Convert distance to similarity score
+                    similarity = 1.0 - (distance / 2.0)
+                    
+                    if similarity >= SEMANTIC_CACHE_THRESHOLD:
+                        response_json = doc.metadata.get("response_json")
+                        if response_json:
+                            response = json.loads(response_json)
+                            response["cached"] = True
+                            response["cache_similarity"] = round(similarity * 100, 1)
+                            response["matched_query"] = doc.metadata.get("original_query", doc.page_content)
+                            
+                            logger.info(f"⚡ Local ChromaDB Cache HIT (Plan: {plan_tier}, Similarity: {response['cache_similarity']}%) in {time.time() - start_time:.3f}s")
+                            return response
+                            
+        except Exception as e:
+            logger.error(f"Error checking semantic cache: {e}")
+            
+        return None
+
+    def store(self, query: str, response: dict, plan_tier: str = "Unknown") -> None:
+        """
+        Store query response in the semantic cache.
+        """
+        # Skip caching error responses or empty results
+        if not query or not response or "error" in response:
+            return
+            
+        try:
+            # 1. Normalize query to standard terminology
+            norm_query = self.normalize_query(query)
+            
+            # Generate deterministic cache ID based on normalized query
+            cache_id = hashlib.sha256(norm_query.encode("utf-8")).hexdigest()
+            query_vector = self.embeddings.embed_query(norm_query)
+            
+            # Remove any existing cache flags from the saved dictionary
+            clean_response = response.copy()
+            clean_response.pop("cached", None)
+            clean_response.pop("cache_similarity", None)
+            clean_response.pop("matched_query", None)
+            
+            response_json = json.dumps(clean_response)
+            
+            if self.redis_available and self.redis_client:
+                # Store in Redis
+                pipeline = self.redis_client.pipeline()
+                # Hash contents
+                pipeline.hset(f"cache:data:{cache_id}", mapping={
+                    "query": query,  # original user query for UI
+                    "vector": json.dumps(query_vector),
+                    "response": response_json,
+                    "plan_tier": plan_tier,
+                    "timestamp": time.time()
+                })
+                # Add to ID set
+                pipeline.sadd("cache:ids", cache_id)
+                # Set TTL on the hash (e.g. 7 days = 604800 seconds)
+                pipeline.expire(f"cache:data:{cache_id}", 604800)
+                pipeline.execute()
+                
+                # Update Python memory index (storing normalized vector)
+                self._cache_memory[cache_id] = {
+                    "query": query,
+                    "vector": query_vector,
+                    "plan_tier": plan_tier
+                }
+                logger.info(f"💾 Query cached in Redis under ID cache:data:{cache_id[:8]}... (Plan: {plan_tier})")
+                
+            else:
+                # Store in fallback ChromaDB
+                if not self.vector_store:
+                    self._init_chroma_store()
+                if not self.vector_store:
+                    return
+                    
+                # To prevent bloating, check if this query ID is already in Chroma
+                existing = self.vector_store.get(ids=[cache_id])
+                if existing and existing["ids"]:
+                    # Update metadata
+                    self.vector_store.update_document(
+                        document_id=cache_id,
+                        document=Document(
+                            page_content=norm_query,
+                            metadata={
+                                "cache_id": cache_id, 
+                                "original_query": query, 
+                                "response_json": response_json, 
+                                "plan_tier": plan_tier, 
+                                "timestamp": time.time()
+                            }
+                        )
+                    )
+                    logger.info(f"💾 Updated query cache in ChromaDB under ID {cache_id[:8]}... (Plan: {plan_tier})")
+                else:
+                    # Add new
+                    self.vector_store.add_texts(
+                        texts=[norm_query],
+                        metadatas=[{
+                            "cache_id": cache_id, 
+                            "original_query": query, 
+                            "response_json": response_json, 
+                            "plan_tier": plan_tier, 
+                            "timestamp": time.time()
+                        }],
+                        ids=[cache_id]
+                    )
+                    logger.info(f"💾 Created query cache in ChromaDB under ID {cache_id[:8]}... (Plan: {plan_tier})")
+                    
+        except Exception as e:
+            logger.error(f"Error storing query in semantic cache: {e}")
+
+    def clear(self) -> None:
+        """Clear all entries in the cache."""
+        try:
+            if self.redis_available and self.redis_client:
+                cache_ids = self.redis_client.smembers("cache:ids")
+                if cache_ids:
+                    pipeline = self.redis_client.pipeline()
+                    for cid in cache_ids:
+                        pipeline.delete(f"cache:data:{cid}")
+                    pipeline.delete("cache:ids")
+                    pipeline.execute()
+                self._cache_memory.clear()
+                logger.info("Cleared Redis semantic cache.")
+            else:
+                if self.vector_store:
+                    # Recreate collection to wipe it
+                    self.vector_store.delete_collection()
+                    self._init_chroma_store()
+                    logger.info("Cleared local ChromaDB semantic cache.")
+        except Exception as e:
+            logger.error(f"Error clearing cache: {e}")
+
+# Global Cache Manager Singleton
+cache_manager = SemanticCache()

orchestration/tools.py

@@ -7,14 +7,16 @@ Defines 4 retrieval tools the LangGraph orchestrator can invoke:
   - plan_comparison_search : Hybrid retrieval scoped to a specific plan tier
   - prior_auth_search      : Hybrid retrieval filtered to prior-authorization docs
 
-The underlying retrieval pipeline (BM25 + Vector + MultiQuery + CrossEncoder + Graph)
-is UNCHANGED — these tools are thin wrappers that control WHAT to query and HOW to filter.
+ENHANCEMENTS:
+  - Doc-type routing: pre-routes queries to the most relevant document collection
+  - Metadata pre-filtering: scopes vector searches to relevant doc types
+  - Per-tool query augmentation for higher precision
 """
 
 import sys
 import os
+from typing import Optional
 
-# Ensure project root is on sys.path so `config` and sibling packages resolve
 sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
 
 from langchain_core.tools import tool
@@ -23,7 +25,7 @@ from langchain_core.documents import Document
 from retrieval.retriever import get_hybrid_retriever
 from retrieval.graph_retriever import GraphRetriever
 
-# ── Lazy singletons (built once, reused across all tool calls) ──
+# ── Lazy singletons ────────────────────────────────────────────────────────────
 _hybrid_retriever = None
 _graph_retriever  = None
 
@@ -52,6 +54,62 @@ def _format_docs(docs: list[Document], max_per_tool: int = 5) -> str:
     return "\n\n---\n\n".join(formatted) if formatted else ""
 
 
+# ── Doc-Type Router ────────────────────────────────────────────────────────────
+# Maps keyword signals → document collection types for pre-filtering
+_DOC_TYPE_KEYWORDS: dict[str, list[str]] = {
+    "drug_formulary": [
+        "drug", "formulary", "medication", "prescription", "generic", "brand",
+        "metformin", "lisinopril", "atorvastatin", "tier 1", "tier 2", "tier 3",
+        "covered drug", "formulary list", "pharmacy benefit",
+    ],
+    "claim_submission_guidelines": [
+        "claim", "submit a claim", "reimburs", "billing", "appeal", "dispute",
+        "out-of-pocket claim", "claim form", "eoob", "explanation of benefits",
+    ],
+    "preventive_care_schedule": [
+        "preventive", "annual exam", "wellness visit", "vaccine", "immunization",
+        "mammogram", "colonoscopy", "screening", "checkup", "physical exam",
+    ],
+    "prior_authorization": [
+        "prior auth", "prior authorization", "step therapy", "pre-approval",
+        "preauthorization", "pa required", "approval required", "step-therapy",
+    ],
+    "provider_directory": [
+        "provider", "doctor", "specialist", "in-network", "out-of-network",
+        "hospital", "physician", "clinic", "dermatologist", "cardiologist",
+        "network provider", "find a doctor",
+    ],
+    "evidence_of_coverage": [
+        "coverage", "covered service", "benefit", "exclusion", "limitation",
+        "deductible", "copay", "coinsurance", "out-of-pocket maximum",
+        "emergency", "urgent care", "what is covered",
+    ],
+}
+
+def _infer_doc_type(query: str) -> Optional[str]:
+    """Infer the most relevant document type from the query text."""
+    q_lower = query.lower()
+    best_doc_type = None
+    best_match_count = 0
+
+    for doc_type, keywords in _DOC_TYPE_KEYWORDS.items():
+        match_count = sum(1 for kw in keywords if kw in q_lower)
+        if match_count > best_match_count:
+            best_match_count = match_count
+            best_doc_type = doc_type
+
+    # Only return a type if we have at least one strong signal
+    return best_doc_type if best_match_count >= 1 else None
+
+
+def _get_doc_type_filtered_docs(vs, query: str, doc_type: str, k: int = 3) -> list[Document]:
+    """Fetch docs pre-filtered by doc_type metadata from ChromaDB."""
+    try:
+        return vs.similarity_search(query, k=k, filter={"doc_type": doc_type})
+    except Exception:
+        return []
+
+
 # ── Tool 1: General policy search ─────────────────────────────────────────────
 @tool
 def policy_search(query: str) -> str:
@@ -60,14 +118,32 @@ def policy_search(query: str) -> str:
     Use this for questions about coverage rules, claim procedures, benefit summaries,
     and general insurance terms.
     Uses the full hybrid pipeline: BM25 + Vector + MultiQuery + CrossEncoder Reranker + Graph.
+    Enhanced with doc-type pre-filtering for higher precision.
     """
     hybrid, _ = _get_retrievers()
-    
-    if "claim" in query.lower() or "diagnosis" in query.lower() or "eligib" in query.lower():
-        docs = hybrid.invoke(f"claim submission eligibility diagnosis {query}")
-    else:
-        docs = hybrid.invoke(query)
-        
+
+    # Augment query for claim/eligibility-related searches
+    search_query = query
+    if any(kw in query.lower() for kw in ["claim", "diagnosis", "eligib"]):
+        search_query = f"claim submission eligibility diagnosis {query}"
+
+    docs = hybrid.invoke(search_query)
+
+    # Boost with doc-type pre-filtered results if a strong signal is found
+    inferred_type = _infer_doc_type(query)
+    if inferred_type:
+        try:
+            from retrieval.retriever import _load_vectorstore
+            vs = _load_vectorstore()
+            type_docs = _get_doc_type_filtered_docs(vs, query, inferred_type, k=3)
+            if type_docs:
+                # Prepend targeted docs — they're likely more relevant
+                seen = {d.page_content[:50] for d in docs}
+                new_type_docs = [d for d in type_docs if d.page_content[:50] not in seen]
+                docs = new_type_docs + docs
+        except Exception:
+            pass  # Graceful fallback to standard retrieval
+
     result = _format_docs(docs)
     return result if result else "No relevant policy information found."
 
@@ -95,30 +171,34 @@ def plan_comparison_search(query: str, tier: str) -> str:
     Use this when the user wants to compare plans or asks about a specific plan tier.
     Runs the full hybrid retrieval pipeline with a tier-augmented query, then
     prioritises documents whose metadata or content mentions that tier.
+    Enhanced: uses ChromaDB metadata pre-filter for plan_tier before hybrid search.
     """
     hybrid, _ = _get_retrievers()
 
-    # Augment the query with the specific tier to force relevant documents up the ranking
+    # Tier-augmented query forces relevant docs up the ranking
     tier_query = f"{tier} plan {query}"
     docs = hybrid.invoke(tier_query)
 
-    # Guarantee tier-specific documents (like the SBC deductibles table) are included
+    # Guarantee tier-specific documents via metadata pre-filter (high precision boost)
     try:
         from retrieval.retriever import _load_vectorstore
         vs = _load_vectorstore()
-        tier_specific_docs = vs.similarity_search(query, k=3, filter={"plan_tier": tier.capitalize()})
+
+        # Pre-filter by both plan_tier AND inferred doc_type for maximum precision
+        tier_specific_docs = vs.similarity_search(query, k=4, filter={"plan_tier": tier.capitalize()})
         docs = tier_specific_docs + docs
-    except Exception as e:
+    except Exception:
         pass
 
+    # Filter to tier-relevant docs
     tier_docs = [
         d for d in docs
         if d.metadata.get("plan_tier", "").lower() in [tier.lower(), "all"]
         or tier.lower() in d.page_content.lower()
-        or d.metadata.get("doc_type") == "sbc"
+        or d.metadata.get("doc_type") in ("sbc", f"summary_of_benefits_{tier.lower()}")
     ]
     if not tier_docs:
-        tier_docs = docs
+        tier_docs = docs  # Fallback: use all results
 
     result = _format_docs(tier_docs, max_per_tool=5)
     return result if result else f"No {tier} plan information found for this query."
@@ -133,25 +213,41 @@ def prior_auth_search(query: str) -> str:
     requirements for drugs or procedures, or PA criteria.
     Runs the full hybrid pipeline with a PA-focused query and filters to
     documents tagged as prior_authorization type or containing PA keywords.
+    Enhanced: uses ChromaDB metadata pre-filter for prior_authorization docs.
     """
     hybrid, _ = _get_retrievers()
 
-    auth_query = f"prior authorization {query}"
+    auth_query = f"prior authorization requirements {query}"
     docs = hybrid.invoke(auth_query)
 
-    pa_keywords = {"prior auth", "authorization required", "step therapy",
-                   "pre-approval", "pre-authorization", "pa required", "preauth"}
+    # Boost with metadata-filtered PA docs
+    try:
+        from retrieval.retriever import _load_vectorstore
+        vs = _load_vectorstore()
+        pa_docs = _get_doc_type_filtered_docs(vs, query, "prior_authorization", k=3)
+        if pa_docs:
+            seen = {d.page_content[:50] for d in docs}
+            new_pa = [d for d in pa_docs if d.page_content[:50] not in seen]
+            docs = new_pa + docs
+    except Exception:
+        pass
+
+    pa_keywords = {
+        "prior auth", "authorization required", "step therapy",
+        "pre-approval", "pre-authorization", "pa required", "preauth",
+        "prior authorization", "requires authorization",
+    }
 
     auth_docs = [
         d for d in docs
-        if d.metadata.get("doc_type") == "prior_authorization"
+        if d.metadata.get("doc_type") in ("prior_authorization",)
         or any(kw in d.page_content.lower() for kw in pa_keywords)
     ]
 
     if not auth_docs:
         return "No prior authorization information found for this query."
 
-    result = _format_docs(auth_docs, max_per_tool=3)
+    result = _format_docs(auth_docs, max_per_tool=4)
     return result if result else "No prior authorization information found for this query."
 
 

orchestration/tracing.py

@@ -1,4 +1,5 @@
 import contextvars
+import concurrent.futures
 
 # ContextVar to store a list of trace logs for the current request thread
 trace_log = contextvars.ContextVar("trace_log", default=None)
@@ -8,3 +9,17 @@ def log_event(msg: str):
     log = trace_log.get()
     if log is not None:
         log.append(msg)
+    else:
+        print(f"TRACE_LOG IS NONE FOR MSG: {msg}", flush=True)
+
+# Global Monkey Patch: Ensure all ThreadPoolExecutor threads inherit contextvars
+_original_submit = concurrent.futures.ThreadPoolExecutor.submit
+
+def _patched_submit(self, fn, *args, **kwargs):
+    ctx = contextvars.copy_context()
+    def _wrapper(*wargs, **wkwargs):
+        return ctx.run(fn, *wargs, **wkwargs)
+    return _original_submit(self, _wrapper, *args, **kwargs)
+
+concurrent.futures.ThreadPoolExecutor.submit = _patched_submit
+

requirements.txt

@@ -41,6 +41,7 @@ fastapi>=0.115.0
 uvicorn[standard]>=0.30.0
 pydantic>=2.0.0
 loguru>=0.7.0
+redis>=5.0.0
 
 # Web Frontend (Streamlit)
 streamlit>=1.40.0

retrieval/graph_retriever.py

@@ -17,6 +17,7 @@ from langchain_core.documents import Document
 from rich.console import Console
 
 from config import GRAPH_DATA_PATH
+from orchestration.tracing import log_event
 
 # ── Medical synonym expansion map ────────────────────────────
 # Maps colloquial / lay terms → canonical specialty names in the graph
@@ -114,6 +115,7 @@ class GraphRetriever:
             return ""
         data = self.G.nodes[node]
         node_type = data.get("type", "Unknown")
+        node_display = data.get("name", str(node))
 
         context = f"### GRAPH ENTITY: {node} ({node_type})\n"
         props = [f"{k}: {v}" for k, v in data.items() if k != "type"]
@@ -130,6 +132,9 @@ class GraphRetriever:
                 target_type = target_data.get("type", "")
                 target_name = target_data.get("name", "")
                 
+                target_display = target_name if target_name else str(target)
+                log_event(f"[GraphDB-Edge] {node_display}|||{rel}|||{target_display}")
+                
                 display_target = f"{target_name} ({target})" if target_name else str(target)
                 context += f"  - [{rel}] -> {display_target} ({target_type})"
                 
@@ -165,6 +170,9 @@ class GraphRetriever:
                 source_type = source_data.get("type", "")
                 source_name = source_data.get("name", "")
                 
+                source_display = source_name if source_name else str(source)
+                log_event(f"[GraphDB-Edge] {source_display}|||{rel}|||{node_display}")
+                
                 display_source = f"{source_name} (NPI: {source})" if source_name and source_type == "Provider" else str(source)
                 context += f"  - {display_source} ({source_type}) -> [{rel}] -> [THIS ENTITY]"
                 
@@ -249,7 +257,10 @@ class GraphRetriever:
                 context += f"Note: No in-network {spec_label} providers found in {city}. "
                 context += f"Showing the closest in-network options in {state}:\n\n"
                 for npi, d in same_state[:8]:
-                    context += (f"- {d.get('name','Unknown')} | {d.get('specialty','')} | "
+                    provider_name = d.get('name', str(npi))
+                    log_event(f"[GraphDB-Edge] {provider_name}|||located in|||{d.get('state','')}")
+                    log_event(f"[GraphDB-Edge] {provider_name}|||specializes in|||{d.get('specialty','')}")
+                    context += (f"- {provider_name} | {d.get('specialty','')} | "
                                f"{d.get('city','')}, {d.get('state','')} (NPI: {npi})\n")
                 return context
 
@@ -269,7 +280,10 @@ class GraphRetriever:
             for st, providers in list(by_state.items())[:5]:
                 context += f"**{st}:**\n"
                 for npi, d in providers[:4]:
-                    context += (f"  - {d.get('name','Unknown')} | {d.get('specialty','')} | "
+                    provider_name = d.get('name', str(npi))
+                    log_event(f"[GraphDB-Edge] {provider_name}|||located in|||{st}")
+                    log_event(f"[GraphDB-Edge] {provider_name}|||specializes in|||{d.get('specialty','')}")
+                    context += (f"  - {provider_name} | {d.get('specialty','')} | "
                                f"{d.get('city','')}, {st} (NPI: {npi})\n")
             return context
 
@@ -279,8 +293,10 @@ class GraphRetriever:
         """Perform graph retrieval and return LangChain Documents."""
         entities = self._extract_entities(query)
         if not entities:
+            log_event("[GraphDB] No entities matched.")
             return []
 
+        log_event(f"[GraphDB] Retrieved entities: {', '.join(entities)}")
         docs = []
         
         # 1. Add intersection context if multiple entities found

retrieval/retriever.py

@@ -180,6 +180,43 @@ def get_hybrid_retriever(
 
     console.print("\n⚙️  Building retrieval pipeline...", style="bold cyan")
 
+    class LoggingRetriever(Runnable):
+        def __init__(self, base_retriever, name):
+            self.base_retriever = base_retriever
+            self.name = name
+
+        def invoke(self, query: str, *args, **kwargs) -> list[Document]:
+            # Custom logic to extract scores for visualization
+            if self.name == "VectorDB":
+                # base_retriever is VectorStoreRetriever
+                docs_and_scores = self.base_retriever.vectorstore.similarity_search_with_score(query, **self.base_retriever.search_kwargs)
+                docs = []
+                for d, score in docs_and_scores:
+                    d.metadata["score"] = score
+                    docs.append(d)
+            elif self.name == "BM25":
+                # base_retriever is BM25Retriever
+                docs = self.base_retriever.invoke(query, *args, **kwargs)
+                try:
+                    q_tokens = self.base_retriever.preprocess_func(query)
+                    scores = self.base_retriever.vectorizer.get_scores(q_tokens)
+                    # Assign the corresponding score based on the original document index
+                    for d in docs:
+                        idx = self.base_retriever.docs.index(d)
+                        d.metadata["score"] = scores[idx]
+                except Exception:
+                    pass
+            else:
+                docs = self.base_retriever.invoke(query, *args, **kwargs)
+
+            for i, d in enumerate(docs[:4]):
+                source = d.metadata.get("source_file", "unknown")
+                page = d.metadata.get("page", "?")
+                score = d.metadata.get("score", 0.0)
+                # Format: [Name] Retrieved: source | Score: 1.23 | Content: ...
+                log_event(f"[{self.name}] Retrieved: {source} (Page: {page})|{score:.4f}|{d.page_content[:40]}...")
+            return docs
+
     # ── Stage 1: Vector Retriever ───────────────────────────────
     console.print("  📦 Loading ChromaDB vector store...")
     vectorstore = _load_vectorstore()
@@ -187,17 +224,19 @@ def get_hybrid_retriever(
         search_type="similarity",
         search_kwargs={"k": k},
     )
+    logging_vector_retriever = LoggingRetriever(vector_retriever, "VectorDB")
     console.print(f"    ✅ Vector retriever ready (k={k})")
 
     # ── Stage 2: BM25 Retriever ─────────────────────────────────
     console.print("  📝 Building BM25 index from stored documents...")
     all_docs = _get_all_documents(vectorstore)
     bm25_retriever = BM25Retriever.from_documents(all_docs, k=k)
+    logging_bm25_retriever = LoggingRetriever(bm25_retriever, "BM25")
     console.print(f"    ✅ BM25 retriever ready (k={k}, {len(all_docs)} documents indexed)")
 
     # ── Stage 3: Ensemble Retriever ─────────────────────────────
     ensemble_retriever = EnsembleRetriever(
-        retrievers=[bm25_retriever, vector_retriever],
+        retrievers=[logging_bm25_retriever, logging_vector_retriever],
         weights=weights,
     )
     console.print(f"    ✅ Ensemble retriever ready (weights: BM25={weights[0]}, Vector={weights[1]})")