added ambiguity node after schema explorer

agent/src/agent/config.py

@@ -43,6 +43,7 @@ class AgentSettings(BaseSettings):
         "text2sql/finalizer_sql_explanation"
     )
     LANGFUSE_PROMPT_REJECTION_ROUTER: str = "text2sql/rejection_router"
+    LANGFUSE_PROMPT_DETECT_AMBIGUITY: str = "text2sql/detect_ambiguity"
 
     MAX_REFINER_ITERATIONS: int = Field(default=3, gt=0)
     REFINER_SCHEMA_CONTEXT_TABLES: int = Field(default=4, gt=0)
@@ -54,7 +55,6 @@ class AgentSettings(BaseSettings):
     ENABLE_SEMANTIC_TYPING: bool = False   # single batched LLM call — adds id/timestamp/category labels
     ENABLE_JOIN_GRAPH: bool = False
     ENABLE_SCHEMA_SUMMARIZATION: bool = False  # generated once at profile-time, not at runtime
-    ENABLE_AMBIGUITY_DETECT: bool = False
     ENABLE_SKILL_INJECTION: bool = False
 
     # ── G2-04: Satisfaction Check ─────────────────────────────────────────────
@@ -68,6 +68,10 @@ class AgentSettings(BaseSettings):
     SATISFACTION_SEMANTIC_THRESHOLD: float = 0.75
     SATISFACTION_MAX_FAILURES: int = 2  # escalate to HITL after this many check failures
 
+    # ── Ambiguity Resolution ──────────────────────────────────────────────────
+    ENABLE_AMBIGUITY_DETECT: bool = True
+    MAX_AMBIGUITY_RETRIES: int = 2  # max times user can clarify before hard stop
+
     # ── G2-05: Redis Schema Cache ─────────────────────────────────────────────
     SCHEMA_CACHE_TTL: int = Field(default=600, gt=0)    # seconds — DDL content
     PROFILE_CACHE_TTL: int = Field(default=1800, gt=0)  # seconds — table profile statistics

agent/src/agent/graph.py

@@ -2,13 +2,21 @@
 LangGraph agent graph — Group 2 hardened topology.
 
 Node order:
-  START → validate_config → extractor → schema_explorer → ...
-          (G2-01 fail-fast)
-
-HITL escalation compiles with interrupt_before=["hitl_escalation"] so
-LangGraph pauses before executing that node.  After a human injects
-corrected state the graph resumes from hitl_escalation which immediately
-routes to extractor (full state reset path).
+  START → init_flags → validate_config
+                     → init_skills → extractor → schema_explorer
+                                                    → sql_static_validations
+                                                    → detect_ambiguity
+                                                         ├─ [clear]        → query_builder → ...
+                                                         ├─ [ambiguous]    → ambiguity_resolution  ← interrupt
+                                                         │                      └─ schema_explorer (retry)
+                                                         └─ [unanswerable] → END
+
+HITL escalation compiles with interrupt_before=["hitl_escalation", "ambiguity_resolution"] so
+LangGraph pauses before executing either node. After a human injects corrected
+state (feedback / clarification) the graph resumes:
+  • hitl_escalation   → extractor (full reset path for hard errors)
+  • ambiguity_resolution → schema_explorer (targeted retry with user’s clarification
+                            appended as feedback)
 
 Satisfaction check sits between refiner success path and finalizer (G2-04).
 """
@@ -23,12 +31,13 @@
 from langchain_core.runnables.config import RunnableConfig
 from langchain_core.prompts import ChatPromptTemplate
 from agent.state import AgentState
-from agent.utils.redis_publisher import publish_node_event_sync
+from agent.utils.redis_publisher import publish_node_event_sync, publish_node_event
 from agent.nodes.extractor import extractor_node
 from agent.nodes.init_flags import init_flags_node
 from agent.nodes.init_skills import init_skills_node
 from agent.nodes.schema_explorer import schema_explorer_node, MAX_SCHEMA_RETRIES, sql_static_validations_node
 from agent.nodes.query_builder import query_builder_node
+from agent.nodes.detect_ambiguity import detect_ambiguity_node, ambiguity_resolution_node
 from agent.nodes.refiner_graph import refiner_subgraph
 from agent.nodes.finalizer import finalizer_node
 from agent.config import settings
@@ -167,6 +176,14 @@ def route_schema_explorer(state: AgentState) -> str:
         if (state.get("schema_explorer_retry_count") or 0) >= MAX_SCHEMA_RETRIES:
             return "hitl_escalation"
         return "schema_explorer"
+
+    runtime_flags = state.get("runtime_flags") or {}
+    enable_ambiguity = runtime_flags.get("SCHEMA_AMBIGUITY_DETECT", settings.ENABLE_AMBIGUITY_DETECT)
+    if isinstance(enable_ambiguity, str):
+        enable_ambiguity = enable_ambiguity.lower() == "true"
+
+    if enable_ambiguity:
+        return "detect_ambiguity"
     return "query_builder"
 
 
@@ -184,6 +201,27 @@ def route_refiner_subagent(state: AgentState) -> str:
     return "finalizer"
 
 
+def route_detect_ambiguity(state: AgentState) -> str:
+    """
+    Route out of detect_ambiguity based on the resolved ambiguity_type.
+
+      - "clear"        → query_builder          (proceed normally)
+      - "ambiguous"    → ambiguity_resolution   (HITL: user clarifies, then retry)
+                         → END if MAX_AMBIGUITY_RETRIES exhausted
+      - "unanswerable" → END                    (data doesn’t exist; clarification won’t help)
+    """
+    t = state.get("ambiguity_type") or "clear"
+
+    if t == "clear":
+        return "query_builder"
+
+    if t == "unanswerable":
+        return END
+
+    # ambiguous — offer HITL (max retries logic is now handled inside detect_ambiguity_node)
+    return "ambiguity_resolution"
+
+
 def route_query_builder(state: AgentState) -> str:
     if state.get("feedback"):
         return "rejection_router"
@@ -207,6 +245,8 @@ def route_rejection(state: AgentState) -> str:
 workflow.add_node("extractor", extractor_node)
 workflow.add_node("schema_explorer", schema_explorer_node)
 workflow.add_node("sql_static_validations", sql_static_validations_node)
+workflow.add_node("detect_ambiguity", detect_ambiguity_node)
+workflow.add_node("ambiguity_resolution", ambiguity_resolution_node)
 workflow.add_node("query_builder", query_builder_node)
 workflow.add_node("rejection_router", rejection_router_node)
 workflow.add_node("refiner_subagent", refiner_subgraph)
@@ -226,11 +266,26 @@ def route_rejection(state: AgentState) -> str:
     route_schema_explorer,
     {
         "schema_explorer": "schema_explorer",
+        "detect_ambiguity": "detect_ambiguity",
         "query_builder": "query_builder",
         "hitl_escalation": "hitl_escalation",  # G2-02
     },
 )
 
+workflow.add_conditional_edges(
+    "detect_ambiguity",
+    route_detect_ambiguity,
+    {
+        "query_builder": "query_builder",
+        "ambiguity_resolution": "ambiguity_resolution",
+        END: END,
+    },
+)
+
+# ambiguity_resolution → schema_explorer: targeted retry (not a full extractor reset).
+# The user’s clarification is in state["feedback"] which schema_explorer already reads.
+workflow.add_edge("ambiguity_resolution", "schema_explorer")
+
 workflow.add_conditional_edges(
     "query_builder",
     route_query_builder,
@@ -263,5 +318,8 @@ def route_rejection(state: AgentState) -> str:
 memory = MemorySaver()
 agent_graph = workflow.compile(
     checkpointer=memory,
-    interrupt_before=["hitl_escalation"],  # G2-02: pause before HITL node
+    interrupt_before=[
+        "hitl_escalation",      # G2-02: hard error escalation (full reset)
+        "ambiguity_resolution", # Ambiguity HITL: user clarifies, then schema_explorer retry
+    ],
 )

agent/src/agent/mcp_server.py

@@ -64,14 +64,24 @@ async def chat_with_agent(
                 }
             )
 
+        update_dict = {
+            "last_error": None,
+            "trino_error": None,
+            "escalated": None,
+            "escalation_reason": None,
+            "refinement_count": 0,
+        }
+
+        # If paused at a node breakpoint (interrupt_before), manually map the resume
+        # value into the state, since Command(resume=) only targets dynamic interrupt()
+        if state_snapshot.next:
+            if isinstance(resume_value, str):
+                update_dict["feedback"] = resume_value
+            elif isinstance(resume_value, dict) and "feedback" in resume_value:
+                update_dict["feedback"] = resume_value["feedback"]
+
         result = await agent_graph.ainvoke(Command(
-            update={
-                "last_error": None,
-                "trino_error": None,
-                "escalated": None,
-                "escalation_reason": None,
-                "refinement_count": 0,
-            },
+            update=update_dict,
             resume=resume_value
         ), config=config)
     else:
@@ -136,6 +146,9 @@ async def chat_with_agent(
     # Get trace_id from the Langfuse handler if available
     trace_id = getattr(langfuse_handler, "last_trace_id", None) if langfuse_handler else None
 
+    if langfuse_handler and hasattr(langfuse_handler, "flush"):
+        langfuse_handler.flush()
+
     from agent.langfuse_client import langfuse_client
     langfuse_client.flush()
 
@@ -177,6 +190,10 @@ async def chat_with_agent(
             }
         )
 
+    is_unans = False
+    if final_state.values.get("ambiguity_type") == "unanswerable" or final_state.values.get("failure_reason"):
+        is_unans = True
+
     return json.dumps(
         {
             "thread_id": thread_id,
@@ -188,6 +205,7 @@ async def chat_with_agent(
             "schema_plan": result.get("schema_plan"),
             "trace_id": trace_id,
             "execution_path": result.get("execution_path", []),
+            "is_unanswerable": is_unans,
         }
     )