[Deep Agent 15/15] Serve the deep agent behind chat_with_agent: engine switch, durable threads, streaming, HITL (#43)

agent/README.md

@@ -110,3 +110,22 @@ For a Langfuse dataset, items use `input = {"question": ...}` and `expected_outp
 for every PR touching `agent/` or `core/`. It needs a repository secret **`PYTHON_CORE_UTILS_TOKEN`**: a fine-grained
 token with read access to `matzpen-agency/python-core-utils` (used instead of the `git+ssh` URL). Real-service e2e
 tests (`tests/refiner/test_refiner_e2e_real.py`) are excluded.
+
+### Serving the deep agent (rollout)
+
+The UI → backend → agent path is the `chat_with_agent` MCP tool. It runs the legacy graph or the deep agent:
+
+| Setting / argument | Effect |
+|---|---|
+| `AGENT_ENGINE=legacy` (default) / `deep` | default engine for new threads |
+| `engine` argument (backend `POST /api/agent/chat` body: `"engine": "deep"`) | per-request choice |
+| thread id `deep-…` | resumes always go back to the deep agent (and legacy threads to the legacy graph) |
+| `DEEP_AGENT_CHECKPOINT_DB_URL=postgresql://…` | durable deep-agent threads (shared by replicas, survive restarts); unset = in-memory |
+
+The deep agent is built once per process and streams the same Redis `agent_stream:{thread_id}` events as the legacy graph
+(nodes `deep:<tool>` and `deep:answer`). Responses keep the legacy contract and add `engine`, `verification`, `confidence`.
+A clarification comes back as `status: "interrupted"` with `interrupt_details.type = "hitl_escalation"`
+(`subtype: "clarification"`); the user's answer is sent as `resume_value` (string or `{"feedback": ...}`).
+
+Rollout: `AGENT_ENGINE=legacy` + per-request `engine: "deep"` for internal users → `deep-agent-eval --gate` passes on the
+expert gold set → `AGENT_ENGINE=deep` → remove the legacy graph.

agent/pyproject.toml

@@ -26,6 +26,8 @@ dependencies = [
     "urllib3>=2.7.0",
     "pyyaml>=6.0",
     "langchain-mcp-adapters>=0.2,<0.3",
+    "langgraph-checkpoint-postgres>=2.0",
+    "psycopg[binary,pool]>=3.2",
 ]
 
 [project.scripts]

agent/src/agent/config.py

@@ -72,6 +72,9 @@ class AgentSettings(BaseSettings):
     JARVIS_MCP_NAMESPACE: str = "global"            # gateway namespace (provider slug) exposing Utils-Server
     JARVIS_API_KEY: str = ""                        # sent as "Authorization: ApiKey <key>"
     UTILS_MCP_TIMEOUT_S: float = Field(default=10.0, gt=0)
+    # Serving (see deep_agent/service.py)
+    AGENT_ENGINE: Literal["legacy", "deep"] = "legacy"   # default engine for chat_with_agent (per-call `engine` overrides)
+    DEEP_AGENT_CHECKPOINT_DB_URL: str | None = None       # postgresql://... for durable deep-agent threads; unset = in-memory
 
     # ── Trino connection info for explicit catalog/schema qualification ──
     TRINO_CATALOG: str = ""

agent/src/agent/mcp_server.py

@@ -25,6 +25,8 @@ async def chat_with_agent(
     active_skills: list[str] | None = None,
     execution_mode: str | None = None,
     hitl_enabled: bool = True,
+    engine: str | None = None,
+    user_id: str | None = None,
 ) -> str:
     """Run the Text2SQL agent to answer database queries.
 
@@ -40,7 +42,14 @@ async def chat_with_agent(
                         'high_quality', 'benchmark'). Overrides flag defaults
                         for this invocation only.
         hitl_enabled:   If False, skip all human-in-the-loop interrupts.
+        engine:         'legacy' (LangGraph workflow) or 'deep' (Deep Agent). Defaults to AGENT_ENGINE.
+                        Threads started on the deep agent are routed back to it on resume.
+        user_id:        End-user id for tracing (Deep Agent).
     """
+    if _use_deep_agent(engine, thread_id):
+        return json.dumps(await _deep_chat(connection_id, query, thread_id, resume_value, hitl_enabled, user_id),
+                          default=str)
+
     thread_id = thread_id or str(uuid.uuid4())
 
     langfuse_handler = None
@@ -211,6 +220,35 @@ async def chat_with_agent(
         }
     )
 
+DEEP_THREAD_PREFIX = "deep-"
+
+
+def _use_deep_agent(engine: str | None, thread_id: str | None) -> bool:
+    if thread_id:  # a resumed thread stays on the engine that started it
+        return thread_id.startswith(DEEP_THREAD_PREFIX)
+    from agent.config import settings
+
+    return (engine or settings.AGENT_ENGINE) == "deep"
+
+
+async def _deep_chat(connection_id, query, thread_id, resume_value, hitl_enabled, user_id) -> dict:
+    from deep_agent.service import run_deep_chat
+
+    response = await run_deep_chat(
+        connection_id=connection_id,
+        query=query,
+        thread_id=thread_id or f"{DEEP_THREAD_PREFIX}{uuid.uuid4()}",
+        resume_value=resume_value,
+        user_id=user_id,
+    )
+    if not hitl_enabled and response.get("status") == "interrupted":
+        # Non-interactive callers get the clarifying question as the answer instead of a pause.
+        question = (response.get("interrupt_details") or {}).get("reason", "")
+        response = {**response, "status": "completed", "summary": question, "interrupt_details": None,
+                    "is_unanswerable": True}
+    return response
+
+
 @mcp.tool()
 async def suggest_fixes(thread_id: str, category: str) -> str:
     """Generate quick fix suggestions during an interruption."""

agent/src/deep_agent/agent.py

@@ -28,12 +28,14 @@ def _default_llm() -> BaseChatModel:
     return get_llm().model_copy(update={"request_timeout": settings.DEEP_AGENT_LLM_TIMEOUT_S, "max_retries": 0})
 
 
-async def acreate_deep_agent(model: BaseChatModel | None = None):
+async def acreate_deep_agent(model: BaseChatModel | None = None, checkpointer=None):
     """create_deep_agent plus the curated Utils-Server tools from the Jarvis Registry gateway (when configured)."""
-    return create_deep_agent(model=model, extra_tools=await load_utils_tools())
+    return create_deep_agent(model=model, extra_tools=await load_utils_tools(), checkpointer=checkpointer)
 
 
-def create_deep_agent(model: BaseChatModel | None = None, extra_tools: dict[str, list] | None = None):
+def create_deep_agent(
+    model: BaseChatModel | None = None, extra_tools: dict[str, list] | None = None, checkpointer=None
+):
     """
     Creates and returns a compiled LangGraph deep agent equipped with all exploration,
     enrichment, temporal, geospatial, and semantic validation tools.
@@ -44,10 +46,11 @@ def create_deep_agent(model: BaseChatModel | None = None, extra_tools: dict[str,
     Args:
         model: Chat model override (tests); defaults to the configured LLM.
         extra_tools: Additional tools per agent ("main" or a subagent name), e.g. from load_utils_tools().
+        checkpointer: Thread persistence (e.g. AsyncPostgresSaver); defaults to an in-memory saver.
     """
     extra_tools = extra_tools or {}
     llm = model or _default_llm()
-    memory = MemorySaver()
+    memory = checkpointer or MemorySaver()
     resources = build_resources()
     system_prompt = load_prompt("system").text
 

agent/src/deep_agent/answer.py

@@ -145,10 +145,26 @@ def verify_answer(answer: Any, executed: list[dict[str, Any]], stop_reason: str
 
 
 class AnswerVerificationMiddleware(AgentMiddleware):
-    """After the agent finishes, bind the structured answer to executed SQL and real rows."""
+    """Bind each turn's structured answer to executed SQL and real rows.
+
+    before_agent clears the previous turn's structured answer (it would otherwise be reported
+    again on a turn that ends without a new one) and marks where this turn's executions start.
+    """
+
+    def before_agent(self, state, runtime) -> dict[str, Any] | None:
+        return {"structured_response": None, "turn_execution_start": len(state.get("executed_queries") or [])}
+
+    async def abefore_agent(self, state, runtime) -> dict[str, Any] | None:
+        return self.before_agent(state, runtime)
 
     def after_agent(self, state, runtime) -> dict[str, Any] | None:
-        return {"final_result": verify_answer(state.get("structured_response"), state.get("executed_queries") or [])}
+        executed = state.get("executed_queries") or []
+        answer = state.get("structured_response")
+        if not isinstance(answer, DeepAgentAnswer):
+            # No answer this turn: fall back only to SQL that ran in this turn.
+            return {"final_result": best_effort_result(executed[state.get("turn_execution_start") or 0:],
+                                                       "no structured answer")}
+        return {"final_result": verify_answer(answer, executed)}
 
     async def aafter_agent(self, state, runtime) -> dict[str, Any] | None:
         return await asyncio.to_thread(self.after_agent, state, runtime)

agent/src/deep_agent/context.py

@@ -53,6 +53,7 @@ class DeepAgentRuntimeState(AgentState):
     jeen_catalog: NotRequired[str]  # raw catalog text, cached so tools never need the model to re-send it
     executed_queries: NotRequired[Annotated[list[dict[str, Any]], merge_executions]]  # successful executions
     final_result: NotRequired[dict[str, Any]]  # verified answer + rows (see deep_agent.answer)
+    turn_execution_start: NotRequired[int]  # index into executed_queries where the current turn began
 
 
 def connection_id(runtime: ToolRuntime | None) -> int | None:

agent/src/deep_agent/service.py

@@ -0,0 +1,185 @@
+"""
+Serving the Deep Agent behind the existing `chat_with_agent` MCP tool.
+
+- One agent per process (built lazily, shared by all requests); per-request data travels in
+  DeepAgentContext (connection, user) and the thread id.
+- Durable threads: AsyncPostgresSaver when DEEP_AGENT_CHECKPOINT_DB_URL is set (shared by all
+  replicas, survives restarts), otherwise an in-memory saver.
+- Streaming: the same Redis events the legacy graph publishes (`agent_stream:{thread_id}`,
+  {"thread_id", "node", "status"}), so the backend SSE relay and UI work unchanged. Nodes are
+  named `deep:<tool>` (e.g. deep:execute_trino_query) plus `deep:answer`.
+- Response: the legacy JSON contract (status / summary / sql_query / sql_explanation /
+  interrupt_details / trace_id / execution_path / is_unanswerable) plus `engine`,
+  `verification` and `confidence`.
+- Clarification: a DeepAgentAnswer with needs_clarification returns status "interrupted" with
+  interrupt_details {"type": "hitl_escalation", "subtype": "clarification", "reason": <question>},
+  which the UI already renders as a free-text prompt. The reply (resume_value) is sent as the
+  next user message on the same thread.
+"""
+
+import asyncio
+import logging
+import time
+from typing import Any
+from uuid import UUID
+
+from langchain_core.callbacks import AsyncCallbackHandler
+from langgraph.checkpoint.memory import MemorySaver
+
+from agent.config import settings
+from agent.utils.redis_publisher import publish_node_event
+from deep_agent import observability
+from deep_agent.answer import best_effort_result, render_markdown
+from deep_agent.context import DeepAgentContext
+
+logger = logging.getLogger(__name__)
+
+_agent = None
+_agent_lock: asyncio.Lock | None = None
+_pg_pool = None
+
+UNANSWERED = {"no_answer", "rejected", "re_execution_failed"}
+
+
+async def build_checkpointer():
+    """AsyncPostgresSaver over a process-wide pool when configured, else MemorySaver."""
+    global _pg_pool
+    if not settings.DEEP_AGENT_CHECKPOINT_DB_URL:
+        logger.warning("DEEP_AGENT_CHECKPOINT_DB_URL not set — deep-agent threads are in-memory (lost on restart)")
+        return MemorySaver()
+    from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
+    from psycopg.rows import dict_row
+    from psycopg_pool import AsyncConnectionPool
+
+    _pg_pool = AsyncConnectionPool(
+        settings.DEEP_AGENT_CHECKPOINT_DB_URL, open=False, max_size=10,
+        kwargs={"autocommit": True, "prepare_threshold": 0, "row_factory": dict_row},
+    )
+    await _pg_pool.open()
+    saver = AsyncPostgresSaver(_pg_pool)
+    await saver.setup()  # idempotent: creates/migrates the checkpoint tables
+    return saver
+
+
+async def get_deep_agent():
+    """The process-wide Deep Agent (built once, on first use)."""
+    global _agent, _agent_lock
+    if _agent is not None:
+        return _agent
+    if _agent_lock is None:
+        _agent_lock = asyncio.Lock()
+    async with _agent_lock:
+        if _agent is None:
+            from deep_agent.agent import acreate_deep_agent
+
+            _agent = await acreate_deep_agent(checkpointer=await build_checkpointer())
+    return _agent
+
+
+def reset_for_tests(agent=None) -> None:
+    global _agent, _agent_lock
+    _agent, _agent_lock = agent, None
+
+
+class StreamPublisher(AsyncCallbackHandler):
+    """Publishes tool start/end to the legacy Redis stream and records the execution path."""
+
+    def __init__(self, thread_id: str) -> None:
+        self.thread_id = thread_id
+        self.execution_path: list[str] = []
+        self._names: dict[UUID, str] = {}
+
+    async def _publish(self, node: str, status: str) -> None:
+        await publish_node_event(self.thread_id, node, status)
+
+    async def on_tool_start(self, serialized, input_str, *, run_id, **kwargs) -> None:
+        name = f"deep:{(serialized or {}).get('name') or kwargs.get('name') or 'tool'}"
+        self._names[run_id] = name
+        self.execution_path.append(name)
+        await self._publish(name, "active")
+
+    async def on_tool_end(self, output, *, run_id, **kwargs) -> None:
+        await self._publish(self._names.pop(run_id, "deep:tool"), "completed")
+
+    async def on_tool_error(self, error, *, run_id, **kwargs) -> None:
+        await self._publish(self._names.pop(run_id, "deep:tool"), "error")
+
+
+def to_response(thread_id: str, final_result: dict[str, Any], trace_id: str | None,
+                execution_path: list[str]) -> dict[str, Any]:
+    answer = final_result.get("answer") or {}
+    status = final_result.get("status", "no_answer")
+    base = {
+        "thread_id": thread_id,
+        "engine": "deep",
+        "verification": status,
+        "confidence": answer.get("confidence"),
+        "trace_id": trace_id,
+        "execution_path": execution_path,
+    }
+    if status == "clarification":
+        return {
+            **base,
+            "status": "interrupted",
+            "interrupt_details": {
+                "type": "hitl_escalation",
+                "subtype": "clarification",
+                "reason": answer.get("needs_clarification") or answer.get("explanation_he") or "",
+            },
+            "sql_query": None,
+        }
+    return {
+        **base,
+        "status": "completed",
+        "summary": render_markdown({**final_result, "sql": None}),  # SQL is returned separately
+        "raw_data_ref": None,
+        "sql_query": final_result.get("sql"),
+        "sql_explanation": answer.get("explanation_he"),
+        "is_unanswerable": status in UNANSWERED,
+    }
+
+
+def _resume_text(resume_value: str | dict | None) -> str | None:
+    if isinstance(resume_value, str):
+        return resume_value
+    if isinstance(resume_value, dict):
+        return resume_value.get("feedback") or resume_value.get("answer")
+    return None
+
+
+async def run_deep_chat(
+    *,
+    connection_id: int | None,
+    query: str | None = None,
+    thread_id: str,
+    resume_value: str | dict | None = None,
+    user_id: str | None = None,
+    agent=None,
+) -> dict[str, Any]:
+    """One chat turn on the Deep Agent, in the legacy chat_with_agent response shape."""
+    message = query if resume_value is None else _resume_text(resume_value)
+    if not message:
+        return {"error": "Query is required for new chat session." if resume_value is None
+                else "resume_value must be the user's answer (a string or {'feedback': ...})."}
+
+    agent = agent or await get_deep_agent()
+    context = DeepAgentContext(connection_id=connection_id, user_id=user_id)
+    publisher = StreamPublisher(thread_id)
+    config = observability.build_run_config(thread_id, user_id=user_id, connection_id=connection_id,
+                                            callbacks=[publisher])
+    start = time.perf_counter()
+    try:
+        async with asyncio.timeout(settings.DEEP_AGENT_RUN_TIMEOUT_S):
+            state = await agent.ainvoke({"messages": [("user", message)]}, config, context=context)
+        final_result = state.get("final_result") or {"status": "no_answer", "verified": False}
+    except TimeoutError:
+        state = (await agent.aget_state(config)).values
+        final_result = best_effort_result(state.get("executed_queries") or [],
+                                          f"timeout after {settings.DEEP_AGENT_RUN_TIMEOUT_S:.0f}s")
+
+    await publish_node_event(thread_id, "deep:answer", "completed")
+    trace_id = observability.langfuse_trace_id(config)
+    observability.score_run(trace_id, observability.run_metrics(state.get("messages") or [],
+                                                               time.perf_counter() - start, final_result))
+    observability.flush()
+    return to_response(thread_id, final_result, trace_id, publisher.execution_path)