From eb077577a764409a2c9e71982c96c0b7fcc16b83 Mon Sep 17 00:00:00 2001 From: Stav Ponte Date: Tue, 20 Oct 2026 17:44:33 +0300 Subject: [PATCH] feat(deep-agent): serve the deep agent behind chat_with_agent with durable threads - deep_agent/service.py: process-wide agent, AsyncPostgresSaver when DEEP_AGENT_CHECKPOINT_DB_URL is set (else in-memory), legacy Redis stream events (deep:, deep:answer), legacy JSON response contract + engine/verification/confidence, clarification as free-text HITL interrupt resumed with the user's answer on the same thread - chat_with_agent: engine argument / AGENT_ENGINE default; deep-* thread ids are routed back to the deep agent on resume; non-interactive calls get the clarifying question as the answer - answer middleware resets the structured answer each turn (a turn without a new answer no longer re-reports the previous one) - backend: optional engine passthrough and engine/verification/confidence response fields - create_deep_agent(checkpointer=...) Closes #43 Co-Authored-By: Claude Opus 5.5 --- agent/README.md | 19 +++ agent/pyproject.toml | 2 + agent/src/agent/config.py | 3 + agent/src/agent/mcp_server.py | 38 +++++ agent/src/deep_agent/agent.py | 11 +- agent/src/deep_agent/answer.py | 20 ++- agent/src/deep_agent/context.py | 1 + agent/src/deep_agent/service.py | 185 +++++++++++++++++++++++++ agent/tests/deep_agent/test_service.py | 161 +++++++++++++++++++++ agent/uv.lock | 76 +++++++++- backend/app/routers/agent.py | 8 ++ backend/tests/test_agent_engine.py | 30 ++++ 12 files changed, 545 insertions(+), 9 deletions(-) create mode 100644 agent/src/deep_agent/service.py create mode 100644 agent/tests/deep_agent/test_service.py create mode 100644 backend/tests/test_agent_engine.py diff --git a/agent/README.md b/agent/README.md index 4e2075f..31e4743 100644 --- a/agent/README.md +++ b/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:` 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. diff --git a/agent/pyproject.toml b/agent/pyproject.toml index 18591ea..ccb6c3f 100644 --- a/agent/pyproject.toml +++ b/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] diff --git a/agent/src/agent/config.py b/agent/src/agent/config.py index ab0608f..c825871 100644 --- a/agent/src/agent/config.py +++ b/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 " 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 = "" diff --git a/agent/src/agent/mcp_server.py b/agent/src/agent/mcp_server.py index 331b014..6c88ae6 100644 --- a/agent/src/agent/mcp_server.py +++ b/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.""" diff --git a/agent/src/deep_agent/agent.py b/agent/src/deep_agent/agent.py index 4cd775f..4783f07 100644 --- a/agent/src/deep_agent/agent.py +++ b/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 diff --git a/agent/src/deep_agent/answer.py b/agent/src/deep_agent/answer.py index c5df58c..7716837 100644 --- a/agent/src/deep_agent/answer.py +++ b/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) diff --git a/agent/src/deep_agent/context.py b/agent/src/deep_agent/context.py index 143d241..0630770 100644 --- a/agent/src/deep_agent/context.py +++ b/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: diff --git a/agent/src/deep_agent/service.py b/agent/src/deep_agent/service.py new file mode 100644 index 0000000..a081507 --- /dev/null +++ b/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:` (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": }, + 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) diff --git a/agent/tests/deep_agent/test_service.py b/agent/tests/deep_agent/test_service.py new file mode 100644 index 0000000..261ad66 --- /dev/null +++ b/agent/tests/deep_agent/test_service.py @@ -0,0 +1,161 @@ +import json +import os +import uuid +from unittest.mock import AsyncMock, patch + +import pytest +from langchain_core.messages import AIMessage + +from core.trino import TrinoExecutionResult +from deep_agent import service, tools +from deep_agent.agent import create_deep_agent +from deep_agent.service import run_deep_chat, to_response + +SQL = "SELECT COUNT(DISTINCT flight_id) AS flights FROM flights_table WHERE launch_country = 'France'" +OK = TrinoExecutionResult(success=True, rows=[[466]], columns=["flights"], row_count=1, query_id="q-1") + + +def _answer(sql=SQL, **extra): + return AIMessage("", tool_calls=[{"name": "DeepAgentAnswer", "id": f"f{uuid.uuid4().hex[:4]}", "args": { + "sql": sql, "explanation_he": "סופר טיסות מצרפת", "confidence": "high", **extra}}]) + + +def _policy(messages): + """Turn 1 asks for clarification; after the user's reply, executes and answers.""" + humans = [m for m in messages if m.type == "human"] + if len(humans) == 1: + return _answer(sql=None, needs_clarification="האם הכוונה להמראות או לנחיתות?") + last_tool = next((m for m in reversed(messages) if m.type == "tool"), None) + if last_tool is None or last_tool.name == "DeepAgentAnswer": + return AIMessage("", tool_calls=[{"name": "execute_trino_query", "args": {"sql": SQL}, "id": f"e{len(messages)}"}]) + return _answer() + + +@pytest.fixture +def published(): + with patch.object(service, "publish_node_event", AsyncMock()) as pub: + yield pub + + +# ── response mapping ───────────────────────────────────────────────────────── + +def test_completed_response_matches_legacy_contract(): + final = {"status": "matched", "verified": True, "sql": SQL, "error": None, + "answer": {"explanation_he": "הסבר", "confidence": "high", "assumptions": []}, + "result": {"columns": ["flights"], "rows": [[466]], "row_count": 1, "truncated": False, "query_id": "q"}} + r = to_response("deep-1", final, "trace-1", ["deep:execute_trino_query"]) + assert r["status"] == "completed" and r["sql_query"] == SQL and r["sql_explanation"] == "הסבר" + assert "| 466 |" in r["summary"] and "```sql" not in r["summary"] # SQL is returned separately + assert r["is_unanswerable"] is False and r["engine"] == "deep" and r["verification"] == "matched" + assert {"thread_id", "status", "summary", "raw_data_ref", "sql_query", "sql_explanation", "trace_id", + "execution_path", "is_unanswerable"} <= set(r) + + +@pytest.mark.parametrize("status", ["no_answer", "rejected", "re_execution_failed"]) +def test_unanswered_statuses(status): + assert to_response("t", {"status": status, "answer": None}, None, [])["is_unanswerable"] is True + + +def test_clarification_maps_to_free_text_hitl(): + r = to_response("t", {"status": "clarification", "answer": {"needs_clarification": "איזו שנה?"}}, None, []) + assert r["status"] == "interrupted" + assert r["interrupt_details"] == {"type": "hitl_escalation", "subtype": "clarification", "reason": "איזו שנה?"} + + +# ── full turns ─────────────────────────────────────────────────────────────── + +@pytest.mark.asyncio +async def test_clarify_then_resume_on_the_same_thread(policy_model, published): + agent = create_deep_agent(model=policy_model(_policy)) + with patch.object(tools, "execute_query_sync", return_value=OK): + first = await run_deep_chat(connection_id=89, query="כמה טיסות מצרפת", thread_id="deep-t1", agent=agent) + second = await run_deep_chat(connection_id=89, thread_id="deep-t1", agent=agent, + resume_value={"approved": False, "feedback": "המראות"}) + + assert first["status"] == "interrupted" and "המראות" in first["interrupt_details"]["reason"] + assert second["status"] == "completed" and second["sql_query"] == SQL + assert second["verification"] == "matched" and "466" in second["summary"] + assert second["execution_path"] == ["deep:execute_trino_query"] + + events = [(c.args[1], c.args[2]) for c in published.await_args_list] + assert ("deep:execute_trino_query", "active") in events and ("deep:execute_trino_query", "completed") in events + assert events[-1] == ("deep:answer", "completed") + assert all(c.args[0] == "deep-t1" for c in published.await_args_list) + + +@pytest.mark.asyncio +async def test_turn_without_new_answer_does_not_repeat_the_previous_one(scripted_model, published): + model = scripted_model( + AIMessage("", tool_calls=[{"name": "execute_trino_query", "args": {"sql": SQL}, "id": "e1"}]), + _answer(), + AIMessage("small talk, no answer"), + ) + agent = create_deep_agent(model=model) + with patch.object(tools, "execute_query_sync", return_value=OK): + first = await run_deep_chat(connection_id=89, query="q1", thread_id="deep-t2", agent=agent) + second = await run_deep_chat(connection_id=89, query="ותודה", thread_id="deep-t2", agent=agent) + assert first["verification"] == "matched" + assert second["verification"] == "no_answer" and second["sql_query"] is None + + +@pytest.mark.asyncio +async def test_missing_query_is_an_error(): + assert "error" in await run_deep_chat(connection_id=1, query=None, thread_id="deep-x", agent=object()) + assert "error" in await run_deep_chat(connection_id=1, thread_id="deep-x", resume_value={"approved": True}, + agent=object()) + + +# ── MCP tool routing ───────────────────────────────────────────────────────── + +@pytest.mark.parametrize("engine, thread_id, default, deep", [ + ("deep", None, "legacy", True), + (None, None, "deep", True), + (None, None, "legacy", False), + ("legacy", "deep-abc", "legacy", True), # a resumed deep thread stays on the deep agent + ("deep", "legacy-thread", "deep", False), # ...and a legacy thread stays legacy +]) +def test_engine_routing(engine, thread_id, default, deep, monkeypatch): + from agent import mcp_server + from agent.config import settings + + monkeypatch.setattr(settings, "AGENT_ENGINE", default) + assert mcp_server._use_deep_agent(engine, thread_id) is deep + + +@pytest.mark.asyncio +@pytest.mark.parametrize("hitl, expected_status", [(True, "interrupted"), (False, "completed")]) +async def test_chat_with_agent_deep_path(hitl, expected_status): + from agent import mcp_server + + clarification = to_response("deep-z", {"status": "clarification", "answer": {"needs_clarification": "?"}}, None, []) + with patch("deep_agent.service.run_deep_chat", AsyncMock(return_value=clarification)) as run, \ + patch.object(mcp_server.agent_graph, "ainvoke") as legacy: + raw = await mcp_server.chat_with_agent(connection_id=89, query="q", engine="deep", hitl_enabled=hitl) + out = json.loads(raw) + legacy.assert_not_called() + assert run.await_args.kwargs["thread_id"].startswith("deep-") + assert out["status"] == expected_status + + +# ── durable threads (real Postgres) ────────────────────────────────────────── + +PG_URL = os.environ.get("DEEP_AGENT_TEST_PG_URL") + + +@pytest.mark.asyncio +@pytest.mark.skipif(not PG_URL, reason="set DEEP_AGENT_TEST_PG_URL to a throwaway Postgres database") +async def test_threads_survive_a_restart_with_postgres(scripted_model, monkeypatch): + from agent.config import settings + + monkeypatch.setattr(settings, "DEEP_AGENT_CHECKPOINT_DB_URL", PG_URL) + thread = f"deep-durable-{uuid.uuid4().hex[:8]}" + + first_process = create_deep_agent(model=scripted_model(AIMessage("first")), checkpointer=await service.build_checkpointer()) + await first_process.ainvoke({"messages": [("user", "שאלה ראשונה")]}, {"configurable": {"thread_id": thread}}) + await service._pg_pool.close() + + # "Restart": a new agent with a new pool sees the same thread. + second_process = create_deep_agent(model=scripted_model(AIMessage("second")), checkpointer=await service.build_checkpointer()) + state = await second_process.aget_state({"configurable": {"thread_id": thread}}) + assert [m.text for m in state.values["messages"] if m.type == "human"] == ["שאלה ראשונה"] + await service._pg_pool.close() diff --git a/agent/uv.lock b/agent/uv.lock index d05d65b..587aef9 100644 --- a/agent/uv.lock +++ b/agent/uv.lock @@ -19,8 +19,10 @@ dependencies = [ { name = "langchain-openai" }, { name = "langfuse" }, { name = "langgraph" }, + { name = "langgraph-checkpoint-postgres" }, { name = "mcp" }, { name = "networkx" }, + { name = "psycopg", extra = ["binary", "pool"] }, { name = "pydantic-settings" }, { name = "pyyaml" }, { name = "shapely" }, @@ -53,8 +55,10 @@ requires-dist = [ { name = "langchain-openai", specifier = "==1.3.5" }, { name = "langfuse", specifier = ">=2.0.0" }, { name = "langgraph" }, + { name = "langgraph-checkpoint-postgres", specifier = ">=2.0" }, { name = "mcp", specifier = ">=1.12.4" }, { name = "networkx", specifier = ">=3.3" }, + { name = "psycopg", extras = ["binary", "pool"], specifier = ">=3.2" }, { name = "pydantic-settings", specifier = ">=2.7.0" }, { name = "pyyaml", specifier = ">=6.0" }, { name = "shapely", specifier = ">=2.0.0" }, @@ -834,15 +838,30 @@ wheels = [ [[package]] name = "langgraph-checkpoint" -version = "4.1.1" +version = "4.3.0" source = { registry = "https://pypi.org/simple" } dependencies = [ { name = "langchain-core" }, { name = "ormsgpack" }, ] -sdist = { url = "https://files.pythonhosted.org/packages/83/47/886af6f886f0bff2273164a45f008694e48a96ff3cd25ff0228f2aa9480e/langgraph_checkpoint-4.1.1.tar.gz", hash = "sha256:6c2bdb530c91f91d7d9c1bd100925d0fc4f498d418c17f3587d1526279482a25", size = 184020, upload-time = "2026-05-22T16:57:38.503Z" } +sdist = { url = "https://files.pythonhosted.org/packages/0f/69/31fdbdc65a85bbd6178afa193c772bb926620f47b4869638bc2bc80afaaa/langgraph_checkpoint-4.3.0.tar.gz", hash = "sha256:c75965d84cc2c1d549163e910a15bcb577758001b141619d05297c463280b018", size = 182652, upload-time = "2026-10-12T22:26:31.478Z" } wheels = [ - { url = "https://files.pythonhosted.org/packages/bd/b4/71425e3e38be92611300b9cc5e46a5bf98ab23f5ea8a75b73d02a2f1413c/langgraph_checkpoint-4.1.1-py3-none-any.whl", hash = "sha256:25d29144b082827218e7bc3f1e9b0566a4bb007895cd6cc26f66a8428739f56e", size = 56212, upload-time = "2026-05-22T16:57:37.203Z" }, + { url = "https://files.pythonhosted.org/packages/1f/0c/84747e340bf4f29291c84cdd5733fc8d0a822f3d33bb24e664a18afa4a7c/langgraph_checkpoint-4.3.0-py3-none-any.whl", hash = "sha256:bedfafe2f997ded60e4fa593e79f56f436a6e45586392dc382aa810d0c751c64", size = 58063, upload-time = "2026-10-12T22:26:30.429Z" }, +] + +[[package]] +name = "langgraph-checkpoint-postgres" +version = "3.2.0" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "langgraph-checkpoint" }, + { name = "orjson" }, + { name = "psycopg" }, + { name = "psycopg-pool" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/45/28/bc0927c2770ab713edc33c4c2f87d1f7b51c344b401d7dda5c8584bd8bf8/langgraph_checkpoint_postgres-3.2.0.tar.gz", hash = "sha256:dffef0e6822d7c614019f7f2c4bdbef42859746aee2d3e6d9d9747cf3744ca90", size = 165780, upload-time = "2026-10-14T19:54:41.116Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e7/af/07090322c0ac00429ff7f4789b0855e72cbd3c47026ce6b8575f39cd0f9a/langgraph_checkpoint_postgres-3.2.0-py3-none-any.whl", hash = "sha256:4d89526ab3dff0c71d575233e4132327f95812deb07b6c2e75fc473525b5b8fc", size = 55937, upload-time = "2026-10-14T19:54:40.048Z" }, ] [[package]] @@ -1190,6 +1209,57 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/c4/72/02445137af02769918a93807b2b7890047c32bfb9f90371cbc12688819eb/protobuf-6.33.6-py3-none-any.whl", hash = "sha256:77179e006c476e69bf8e8ce866640091ec42e1beb80b213c3900006ecfba6901", size = 170656, upload-time = "2026-03-18T19:04:59.826Z" }, ] +[[package]] +name = "psycopg" +version = "3.3.6" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, + { name = "tzdata", marker = "sys_platform == 'win32'" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/76/26/3ea4ca5eaea1c0debcdf7ee7c1613fbe721dc27a03c461c0817ffd8a0601/psycopg-3.3.6.tar.gz", hash = "sha256:c081f2250df751a943036e42db6df4571c66cd0aabe8291a7a506512b12007d2", size = 168171, upload-time = "2026-09-18T13:22:55.152Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/4e/de/748bd7609c71cae5d737f0ba9192f19329f70180ecda8fff3cac02c5abe3/psycopg-3.3.6-py3-none-any.whl", hash = "sha256:a1db9f7148b06a28606767efaca51fa6f9398c5c0a3810519be69d7000bdb631", size = 215490, upload-time = "2026-09-18T13:15:29.374Z" }, +] + +[package.optional-dependencies] +binary = [ + { name = "psycopg-binary", marker = "implementation_name != 'pypy'" }, +] +pool = [ + { name = "psycopg-pool" }, +] + +[[package]] +name = "psycopg-binary" +version = "3.3.6" +source = { registry = "https://pypi.org/simple" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/e6/01/2cdd1824e58b4467ee0b9498664cd28c42d8794db6b1e35b6bcb834f0044/psycopg_binary-3.3.6-cp312-cp312-macosx_10_13_x86_64.whl", hash = "sha256:3f84dab25e0385692ee13274c68678377e0b1a70ab9d14e56264cbf61f60c62d", size = 4707086, upload-time = "2026-09-18T13:18:05.138Z" }, + { url = "https://files.pythonhosted.org/packages/f6/76/de9948ac06895261c84d5b9fbe283d8f3c5bc9f070691b8d9eaa1b51e322/psycopg_binary-3.3.6-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:612382ac3ed13651c7fa44b5fee9fbf7baaa2ddbc6f500391672682c5f1df9e0", size = 4769607, upload-time = "2026-09-18T13:18:12.83Z" }, + { url = "https://files.pythonhosted.org/packages/76/a9/72436c9915ee4905964689e7f0e182ce7767cc0a0390b3ce703be8177625/psycopg_binary-3.3.6-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.whl", hash = "sha256:366db6e97e66b37211475f20c4c1324a2dc0dd825e46d4e87f9d599304d276f9", size = 5554134, upload-time = "2026-09-18T13:18:21.175Z" }, + { url = "https://files.pythonhosted.org/packages/0a/42/948bb3d2617795093512613fd96ba380e922992c7908fbc073858147d196/psycopg_binary-3.3.6-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:1679a1cb93fbe5a6d1fd58d82cbddcc6fcb8c61446ba7cae6eb2a7b19bc585de", size = 5235723, upload-time = "2026-09-18T13:18:27.071Z" }, + { url = "https://files.pythonhosted.org/packages/99/47/93e823ff1b0088400703410939c9bda3e63ed9c850b3ee088e8769f4c10b/psycopg_binary-3.3.6-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:37d40450659401600e6d043ff586c89a71a69f33cbb8bcdba6cdb2569beecdbe", size = 6833587, upload-time = "2026-09-18T13:18:33.794Z" }, + { url = "https://files.pythonhosted.org/packages/5e/2d/ecc69c847795aa704041a9f5667a6b0938a088cf1853636d762a6938e493/psycopg_binary-3.3.6-cp312-cp312-manylinux_2_38_riscv64.manylinux_2_39_riscv64.whl", hash = "sha256:a5165300324efd5a772c48a88ab3a928513ab3979fca76553e62ee815f7b2b9c", size = 5070013, upload-time = "2026-09-18T13:18:39.628Z" }, + { url = "https://files.pythonhosted.org/packages/92/36/6126f0dac21713dcae91404f2a76da18598a6252339a8c669c46370d43b2/psycopg_binary-3.3.6-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:d636338c8f21b0df2f84657b00bc34f9313f826ef93f1155bc743607e4a0c5eb", size = 4597367, upload-time = "2026-09-18T13:18:45.023Z" }, + { url = "https://files.pythonhosted.org/packages/4d/29/7ecfc04243b46c89ffd49924e9c5634ea904ef96c7d0f37e4073623584c1/psycopg_binary-3.3.6-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:a4ee3bdd5468a725f2a4d9aab8a74b6d0279f768c8b5d3aeb102c5307ff3d59c", size = 4275419, upload-time = "2026-09-18T13:18:49.299Z" }, + { url = "https://files.pythonhosted.org/packages/6e/90/2f46d2e0de79706ac170df0a3637fe63c4498fc04f131f6049520b78b806/psycopg_binary-3.3.6-cp312-cp312-musllinux_1_2_riscv64.whl", hash = "sha256:289aadd6a00e151203c081f708348ec89f1e483c9b510ef4ac3981f847f01f79", size = 4007358, upload-time = "2026-09-18T13:18:53.944Z" }, + { url = "https://files.pythonhosted.org/packages/03/48/6744e91291b751a8cf12d63d719977974bb94c84ceba913e7ddb2e478e51/psycopg_binary-3.3.6-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:f21d057f3e5f5491067e5b292498073b73847d48799b099803fef100775fcc52", size = 4320156, upload-time = "2026-09-18T13:18:59.258Z" }, + { url = "https://files.pythonhosted.org/packages/1a/9b/94ff7fce53a64d5b286e2ec454e0a025cf3d6e6b4a9189bef16aa5de98b2/psycopg_binary-3.3.6-cp312-cp312-win_amd64.whl", hash = "sha256:e23a66a763fbe83fcc210bc77c27e5a5ea380ebf091c06f34d8561b695e5a40f", size = 3658864, upload-time = "2026-09-18T13:19:06.503Z" }, +] + +[[package]] +name = "psycopg-pool" +version = "3.3.3" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "typing-extensions" }, +] +sdist = { url = "https://files.pythonhosted.org/packages/74/5e/c0664b968b102ff68b811d999c728546c48d5c1eec03e3bbaf88c0cb4472/psycopg_pool-3.3.3.tar.gz", hash = "sha256:df87b5d9d0ad7db37f6cdad4fa8ce113d250f5997f6db38e9a99192fb67f9e1d", size = 32006, upload-time = "2026-09-22T15:53:24.947Z" } +wheels = [ + { url = "https://files.pythonhosted.org/packages/5d/b4/452c6607a0f479465cd8a9b0d9956919fcb150050c1f83f9f11e6b8ee8dc/psycopg_pool-3.3.3-py3-none-any.whl", hash = "sha256:9b9cd6a4fcec47a410f7e82d408540e7f77b478509e91b44c1a5457a13e5ff37", size = 40304, upload-time = "2026-09-22T15:53:23.712Z" }, +] + [[package]] name = "psycopg2-binary" version = "2.9.9" diff --git a/backend/app/routers/agent.py b/backend/app/routers/agent.py index 94a93b7..c059ef2 100644 --- a/backend/app/routers/agent.py +++ b/backend/app/routers/agent.py @@ -20,6 +20,8 @@ from fastapi.responses import StreamingResponse from mcp.client.session import ClientSession from mcp.client.streamable_http import streamablehttp_client +from typing import Literal + from pydantic import BaseModel from app.config import settings @@ -49,6 +51,7 @@ class ChatRequest(BaseModel): execution_mode: str | None = None hitl_enabled: bool = True connection_id: int + engine: Literal["legacy", "deep"] | None = None # None -> the agent service's AGENT_ENGINE class SuggestFixesRequest(BaseModel): @@ -68,6 +71,9 @@ class ChatResponse(BaseModel): trace_id: str | None = None execution_path: list[str] | None = None is_unanswerable: bool = False + engine: str | None = None # "deep" when answered by the Deep Agent + verification: str | None = None # Deep Agent answer verification status (matched, re_executed, ...) + confidence: str | None = None # Deep Agent self-reported confidence (high / medium / low) # --------------------------------------------------------------------------- @@ -210,6 +216,8 @@ async def chat(request: ChatRequest) -> ChatResponse: "hitl_enabled": request.hitl_enabled, "connection_id": request.connection_id, } + if request.engine is not None: # only sent when chosen, so older agent services keep working + tool_arguments["engine"] = request.engine result = await _call_agent_mcp(tool_arguments) return ChatResponse(**result) diff --git a/backend/tests/test_agent_engine.py b/backend/tests/test_agent_engine.py new file mode 100644 index 0000000..9e89b57 --- /dev/null +++ b/backend/tests/test_agent_engine.py @@ -0,0 +1,30 @@ +from unittest.mock import AsyncMock, patch + +DEEP_RESPONSE = { + "thread_id": "deep-123", "status": "completed", "summary": "466 טיסות", "sql_query": "SELECT 1", + "sql_explanation": "x", "trace_id": "t", "execution_path": ["deep:execute_trino_query"], + "is_unanswerable": False, "engine": "deep", "verification": "matched", "confidence": "high", +} + + +@patch("app.routers.agent._call_agent_mcp", new_callable=AsyncMock, return_value=DEEP_RESPONSE) +def test_chat_forwards_engine_and_returns_verification(mock_call, client): + response = client.post("/api/agent/chat", json={"query": "כמה טיסות", "connection_id": 89, "engine": "deep"}) + assert response.status_code == 200, response.text + assert mock_call.await_args.args[0]["engine"] == "deep" + body = response.json() + assert body["engine"] == "deep" and body["verification"] == "matched" and body["confidence"] == "high" + + +@patch("app.routers.agent._call_agent_mcp", new_callable=AsyncMock, + return_value={"thread_id": "t", "status": "completed", "summary": "ok"}) +def test_chat_omits_engine_when_not_chosen(mock_call, client): + response = client.post("/api/agent/chat", json={"query": "q", "connection_id": 89}) + assert response.status_code == 200, response.text + assert "engine" not in mock_call.await_args.args[0] + assert response.json()["engine"] is None + + +def test_chat_rejects_unknown_engine(client): + response = client.post("/api/agent/chat", json={"query": "q", "connection_id": 89, "engine": "gpt"}) + assert response.status_code == 422