From 55e3a902b42cd1b7ba9a0c98d3321076e09eae28 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 19:22:01 +0100 Subject: [PATCH 1/9] restate.agents.langgraph.nodes: prebuilt nodes that are Restate steps nodes.wait_for_signal, nodes.workflow and nodes.subagent make nodes for a graph of your own, as the tools do for an agent: each reads its input from the state and writes what it gets back to the state field `output`. The graph runs in a handler of an AgentObject, with ainvoke(context=ctx) and SqlCheckpointSaver(ctx.db). Each node's idempotency key is LangGraph's task namespace, which LangGraph derives from the checkpoint the step started from: a node re-run after a crash gets the same one, so it joins the wait, the workflow or the other agent's answer it started before. A wait is listed by the agent's pending view, and its signal view resolves it. The helpers the tools and the nodes share in _tools.py lose their leading underscore (handler_of, input_type_of, inputs_of, untyped). --- python/restate/agents/_tools.py | 32 +-- python/restate/agents/langgraph/__init__.py | 3 +- python/restate/agents/langgraph/nodes.py | 223 ++++++++++++++++++++ tests/agents_langgraph_nodes.py | 220 +++++++++++++++++++ 4 files changed, 461 insertions(+), 17 deletions(-) create mode 100644 python/restate/agents/langgraph/nodes.py create mode 100644 tests/agents_langgraph_nodes.py diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 5d1707b..185bb12 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -300,24 +300,24 @@ async def _once(name: str, input: object, call: ToolCall, body: Callable[[ToolCa return result -def _untyped(ctx: SessionContext) -> Any: +def untyped(ctx: SessionContext) -> Any: """Return ``ctx`` untyped, to call a handler whose input type is known only at run time.""" return ctx -def _inputs(handler: Handler[Any, Any], value: Any) -> tuple[Any, ...]: +def inputs_of(handler: Handler[Any, Any], value: Any) -> tuple[Any, ...]: """Return the positional arguments for calling ``handler``: ``(value,)`` if it takes an input.""" return (value,) if handler.arity >= 2 else () -def _handler(target: Callable[..., Any]) -> Handler[Any, Any]: +def handler_of(target: Callable[..., Any]) -> Handler[Any, Any]: try: return handler_from_callable(target) except ValueError: raise TypeError(f"{target!r} is not a handler of a service, an object or a workflow") from None -def _input_type(handler: Handler[Any, Any]) -> Any: +def input_type_of(handler: Handler[Any, Any]) -> Any: """Return the type of a handler's input, or ``None`` if it takes none.""" if handler.arity < 2: return None @@ -408,10 +408,10 @@ def call( Raises: TypeError: If ``target`` is not a handler. """ - handler = _handler(target) + handler = handler_of(target) service = handler.service_tag.name model_keys = service if _keyed(handler) and key is None else None - shape = _input_of(_input_type(handler), model_keys) + shape = _input_of(input_type_of(handler), model_keys) async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: chosen, value = arguments["key"], arguments["value"] @@ -419,8 +419,8 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: async def made(c: ToolCall) -> Any: target_key = chosen if key is None else await configured(key, c.ctx) target_key = target_key if _keyed(handler) else None - return await _untyped(c.ctx).call( - target, *_inputs(handler, value), key=target_key, idempotency_key=c.idempotency_key + return await untyped(c.ctx).call( + target, *inputs_of(handler, value), key=target_key, idempotency_key=c.idempotency_key ) return await _once(f"{service}/{handler.name}", arguments["raw"], tool_call, made) @@ -461,18 +461,18 @@ def send( Raises: TypeError: If ``target`` is not a handler. """ - handler = _handler(target) + handler = handler_of(target) service = handler.service_tag.name model_keys = service if _keyed(handler) and key is None else None - shape = _input_of(_input_type(handler), model_keys) + shape = _input_of(input_type_of(handler), model_keys) async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: chosen, value = arguments["key"], arguments["value"] async def made(c: ToolCall) -> str: target_key = (chosen if key is None else await configured(key, c.ctx)) if _keyed(handler) else None - _untyped(c.ctx).send( - target, *_inputs(handler, value), key=target_key, delay=delay, idempotency_key=c.idempotency_key + untyped(c.ctx).send( + target, *inputs_of(handler, value), key=target_key, delay=delay, idempotency_key=c.idempotency_key ) return f"sent to {service}{'' if target_key is None else f'/{target_key}'}/{handler.name}" @@ -511,17 +511,17 @@ def workflow( Raises: TypeError: If ``target`` is not a handler. """ - handler = _handler(target) + handler = handler_of(target) service = handler.service_tag.name - shape = _input_of(_input_type(handler), None) + shape = _input_of(input_type_of(handler), None) async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: value = arguments["value"] async def made(c: ToolCall) -> Any: workflow_id = f"{c.ctx.key}-{c.call_id}" if key is None else await configured(key, c.ctx) - return await _untyped(c.ctx).call( - target, *_inputs(handler, value), key=workflow_id, idempotency_key=c.idempotency_key + return await untyped(c.ctx).call( + target, *inputs_of(handler, value), key=workflow_id, idempotency_key=c.idempotency_key ) return await _once(service, arguments["raw"], tool_call, made) diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py index 194dc33..e9a161b 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -28,7 +28,8 @@ async def get_weather(city: str) -> str: ``agent`` creates an agent, ``prompt`` answers a prompt from a handler of your own, and ``SqlCheckpointSaver`` stores the checkpoints of a graph of your own -in ``ctx.db``. +in ``ctx.db``. ``nodes`` has prebuilt nodes for such a graph that are Restate +steps: wait for a signal, run a workflow, ask another agent. This differs from the older ``restate.ext.langchain`` integration, whose middleware makes a LangChain ``create_agent`` agent durable inside any regular diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py new file mode 100644 index 0000000..7767626 --- /dev/null +++ b/python/restate/agents/langgraph/nodes.py @@ -0,0 +1,223 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""Prebuilt LangGraph nodes that are Restate steps: wait for a signal, run a workflow, ask another agent. + +Add them to a graph of your own, as the tools of ``restate.agents`` are added +to an agent. Each one reads its input from the graph's state, and writes what +it gets back to the state field ``output``:: + + from restate.agents.langgraph import SqlCheckpointSaver, nodes + + graph = StateGraph(State) + graph.add_node("approve", nodes.wait_for_signal("approval", input=lambda s: s["refund"], output="approved")) + graph.add_node("ship", nodes.workflow(shipping.run, input=lambda s: s["order"], output="receipt")) + +Run the graph in a handler of an agent (an ``AgentObject``), with its context +as LangGraph's run context and a checkpointer in its database:: + + desk = AgentObject("desk") + + + @desk.handler() + async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: + app = graph.compile(checkpointer=SqlCheckpointSaver(ctx.db)) + config: RunnableConfig = {"configurable": {"thread_id": "refunds"}} + resuming = (await app.aget_state(config)).next + return await app.ainvoke(None if resuming else {"refund": refund}, config, context=ctx, durability="sync") + +Each node does its step once. A crash re-runs a node from the checkpoint it +started from, and LangGraph gives it the same task id there, so the node +joins the workflow, the agent's answer or the wait it started before rather +than start another. A wait is listed by the agent's ``pending`` view, and its +``signal`` view resolves it. +""" + +from __future__ import annotations + +import typing +from collections.abc import Awaitable, Callable +from typing import Any + +from langgraph.config import get_config +from langgraph.runtime import Runtime + +from restate.agents import _signals +from restate.agents._tools import handler_of, input_type_of, inputs_of, untyped +from restate.sessions import SessionContext + + +class Node(typing.Protocol): + """A LangGraph node: it takes the state and the run's runtime, and returns the state's update.""" + + def __call__(self, state: Any, *, runtime: Runtime[Any]) -> Awaitable[dict[str, Any]]: ... + + +_TASK_ID = "__pregel_task_id" +_CHECKPOINTER = "__pregel_checkpointer" + + +class _Step(typing.NamedTuple): + """A node's run: the agent's context, the node's name, and the key that is the same each time it re-runs.""" + + ctx: SessionContext + node: str + idempotency_key: str + + +def _step(runtime: Runtime[Any], kind: str) -> _Step: + """Return the run of a node of ``kind``, from LangGraph's runtime and config.""" + ctx = typing.cast("SessionContext | None", runtime.context) + if ctx is None: + raise TypeError(f"nodes.{kind} runs in a graph invoked with the handler's context: ainvoke(..., context=ctx)") + configurable = get_config().get("configurable", {}) + if configurable.get(_CHECKPOINTER) is None: + raise TypeError( + f"nodes.{kind} runs in a graph with a checkpointer, so that a retried attempt resumes it: " + "compile(checkpointer=SqlCheckpointSaver(ctx.db))" + ) + # The task's namespace holds its id, which LangGraph derives from the checkpoint the step started from: + # a node re-run from that checkpoint gets the same one. It holds the subgraph's path too. + namespace = typing.cast(str, configurable.get("checkpoint_ns") or configurable[_TASK_ID]) + node = namespace.rsplit("|", 1)[-1].split(":", 1)[0] + return _Step(ctx, node, f"langgraph:{ctx.key}:{namespace}") + + +def _update(output: str | None, value: object) -> dict[str, Any]: + return {output: value} if output is not None else {} + + +def wait_for_signal( + signal: str, + *, + output: str | None = None, + input: Callable[[Any], Any] | None = None, +) -> Node: + """Create a node that waits for a signal to the agent, and writes its value to ``output``. + + A person's approval, say, or another system's result. The graph waits + without holding a process, for as long as it takes. The agent's + ``pending`` view lists the wait, with ``input``'s value, and its + ``signal`` view resolves it:: + + graph.add_node("approve", nodes.wait_for_signal("approval", input=lambda s: s["refund"], output="approved")) + + Args: + signal: The signal's name. + output: The state field that gets the signal's value. None to write + nothing. + input: What the graph waits for (such as the refund to approve), from + the state, for ``pending`` to show. None to show nothing. + + Returns: + The node, for ``StateGraph.add_node``. + """ + + async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: + step = _step(runtime, "wait_for_signal") + request: _signals.WaitRequest = {"id": step.idempotency_key, "name": signal} + if input is not None: + request["input"] = input(state) + value = await step.ctx.call( + _signals.wait, request, key=_signals.signals_of(step.ctx), idempotency_key=step.idempotency_key + ) + return _update(output, value) + + return node + + +def workflow( + target: Callable[..., Any], + *, + output: str | None = None, + input: Callable[[Any], Any] | None = None, + key: Callable[[Any], str] | None = None, +) -> Node: + """Create a node that runs a workflow to completion, and writes its result to ``output``:: + + graph.add_node("ship", nodes.workflow(shipping.run, input=lambda s: s["order"], output="receipt")) + + Args: + target: The workflow's main handler. + output: The state field that gets the workflow's result. None to + write nothing. + input: The workflow's input, from the state. Required if the handler + takes one. + key: The workflow id, from the state. Defaults to one of the node's + own, made of the agent's key and the node's task. + + Returns: + The node, for ``StateGraph.add_node``. + + Raises: + TypeError: If ``target`` is not a handler, or it takes an input and + ``input`` is not given (or the other way around). + """ + handler = handler_of(target) + takes_input = input_type_of(handler) is not None + if takes_input != (input is not None): + name = f"{handler.service_tag.name}.{handler.name}" + raise TypeError( + f"{name} {'takes' if takes_input else 'takes no'} input: {'pass' if takes_input else 'drop'} input=" + ) + + async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: + step = _step(runtime, "workflow") + workflow_id = key(state) if key is not None else step.idempotency_key + value = input(state) if input is not None else None + result = await untyped(step.ctx).call( + target, *inputs_of(handler, value), key=workflow_id, idempotency_key=step.idempotency_key + ) + return _update(output, result) + + return node + + +def subagent( + agent: object, + *, + message: Callable[[Any], str], + output: str | None = None, + key: Callable[[Any], str] | None = None, + handler: str = "chat", +) -> Node: + """Create a node that sends a message to another agent, and writes its answer to ``output``:: + + graph.add_node("research", nodes.subagent(researcher, message=lambda s: s["question"], output="findings")) + + By default the message goes to the agent's key ``.``, + so each run of the graph continues the same conversation with it. + + Args: + agent: An agent, or any virtual object with a handler that takes text. + message: The message, from the state. + output: The state field that gets the agent's answer. None to write + nothing. + key: The agent's key, from the state. Defaults to + ``.``. + handler: The handler that takes the message. + + Returns: + The node, for ``StateGraph.add_node``. + """ + agent_name = typing.cast(str, getattr(agent, "name")) # noqa: B009 (an SDK's definition) + + async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: + step = _step(runtime, "subagent") + agent_key = key(state) if key is not None else f"{step.ctx.key}.{step.node}" + answer = await step.ctx.generic_call( + agent_name, handler, message(state), key=agent_key, idempotency_key=step.idempotency_key + ) + return _update(output, answer) + + return node + + +__all__ = ["Node", "subagent", "wait_for_signal", "workflow"] diff --git a/tests/agents_langgraph_nodes.py b/tests/agents_langgraph_nodes.py new file mode 100644 index 0000000..1492b1c --- /dev/null +++ b/tests/agents_langgraph_nodes.py @@ -0,0 +1,220 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""LangGraph's prebuilt Restate nodes (``restate.agents.langgraph.nodes``), against a real Restate server. + +A graph of your own, run by a handler of an agent, waits for a signal, runs a +workflow and asks another agent. A crash in a parallel branch, while a node's +step is under way, re-runs the node: it joins the same step. +""" + +# pylint: disable=missing-function-docstring,missing-class-docstring + +import asyncio +import collections +import threading +import typing +import uuid +from typing import Any, Required, TypedDict + +import httpx2 +import pytest + +pytest.importorskip("langgraph") + +from langchain_core.runnables import RunnableConfig # noqa: E402 +from langgraph.graph import END, START, StateGraph # noqa: E402 + +import restate # noqa: E402 +from restate import HarnessEnvironment, sessions # noqa: E402 +from restate.agents import AgentObject # noqa: E402 +from restate.agents.langgraph import SqlCheckpointSaver, nodes # noqa: E402 + + +class State(TypedDict, total=False): + refund: Required[str] + approved: Any + receipt: str + findings: str + + +shipping = restate.Workflow("nodes_shipping") +shipped: collections.Counter[str] = collections.Counter() + + +@shipping.main() +async def ship(ctx: restate.WorkflowContext, order: str) -> str: + shipped[order] += 1 + return f"shipped {order}" + + +researcher = restate.VirtualObject("nodes_researcher") +asked: collections.Counter[str] = collections.Counter() +researching = threading.Event() +crashed = threading.Event() + + +@researcher.handler() +async def chat(ctx: restate.ObjectContext, question: str) -> str: + asked[question] += 1 + if question.startswith("slow"): + # Answer only once the graph's attempt has crashed, so the crash cuts the node's call short. + researching.set() + while not crashed.is_set(): + await asyncio.sleep(0.01) + return f"found: {question}" + + +def approval_then_shipping() -> Any: + graph = StateGraph(State) + graph.add_node("approve", nodes.wait_for_signal("approval", input=lambda s: s["refund"], output="approved")) + graph.add_node("ship", nodes.workflow(ship, input=lambda s: s["refund"], output="receipt")) + graph.add_node("research", nodes.subagent(researcher, message=lambda s: f"why {s['refund']}", output="findings")) + graph.add_edge(START, "approve") + graph.add_edge("approve", "ship") + graph.add_edge("ship", "research") + graph.add_edge("research", END) + return graph + + +crashes: collections.Counter[str] = collections.Counter() +go_crash = threading.Event() + + +def crash_after(ready: threading.Event) -> Any: + """A node that fails its first run once ``ready`` is set, while the node beside it is under way.""" + + async def crash(state: State) -> dict[str, Any]: + if crashes[state["refund"]] == 0: + while not ready.is_set(): + await asyncio.sleep(0.01) + crashes[state["refund"]] += 1 + crashed.set() + raise RuntimeError("the process crashed") + return {} + + return crash + + +def beside_a_crash(step: str) -> Any: + """``step`` and a node that crashes while it is under way, at once.""" + graph = StateGraph(State) + if step == "research": + graph.add_node(step, nodes.subagent(researcher, message=lambda s: f"slow {s['refund']}", output="findings")) + graph.add_node("crash", typing.cast(Any, crash_after(researching))) + else: + graph.add_node(step, nodes.wait_for_signal("approval", input=lambda s: s["refund"], output="approved")) + graph.add_node("crash", typing.cast(Any, crash_after(go_crash))) + graph.add_edge(START, step) + graph.add_edge(START, "crash") + graph.add_edge([step, "crash"], END) + return graph + + +GRAPHS = { + "approval": approval_then_shipping(), + "research_crash": beside_a_crash("research"), + "wait_crash": beside_a_crash("approve"), +} + +desk = AgentObject("nodes_desk") + + +class Run(TypedDict): + graph: str + refund: str + + +@desk.handler() +async def run(ctx: sessions.SessionContext, request: Run) -> State: + app = GRAPHS[request["graph"]].compile(checkpointer=SqlCheckpointSaver(ctx.db)) + config: RunnableConfig = {"configurable": {"thread_id": request["graph"]}} + resuming = (await app.aget_state(config)).next + start: State | None = None if resuming else {"refund": request["refund"]} + return typing.cast(State, await app.ainvoke(start, config, context=ctx, durability="sync")) + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness( + restate.app([desk, shipping, researcher]), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield harness + + +pytestmark = [pytest.mark.anyio] + + +async def post(harness: HarnessEnvironment, path: str, body: Any = None) -> httpx2.Response: + async with httpx2.AsyncClient(timeout=60) as ingress: + return await ingress.post(f"{harness.ingress_url}/nodes_desk/{path}", json=body) + + +async def pending(harness: HarnessEnvironment, key: str) -> list[dict[str, Any]]: + return typing.cast(list[dict[str, Any]], (await post(harness, f"{key}/pending")).json()) + + +async def until_waiting(harness: HarnessEnvironment, key: str) -> list[dict[str, Any]]: + for _ in range(400): + if waits := await pending(harness, key): + return waits + await asyncio.sleep(0.05) + raise AssertionError("nothing waits") + + +async def test_a_graph_waits_for_a_signal_runs_a_workflow_and_asks_another_agent(harness: HarnessEnvironment): + key = uuid.uuid4().hex + running = asyncio.ensure_future(post(harness, f"{key}/run", {"graph": "approval", "refund": key})) + waits = await until_waiting(harness, key) + assert [(w["name"], w.get("input")) for w in waits] == [("approval", key)] + resolved = await post(harness, f"{key}/signal", {"name": "approval", "value": {"by": "manager"}}) + assert resolved.json() == {"resolved": waits[0]["id"]} + assert (await running).json() == { + "refund": key, + "approved": {"by": "manager"}, + "receipt": f"shipped {key}", + "findings": f"found: why {key}", + } + assert shipped[key] == 1 + assert asked[f"why {key}"] == 1 + + +async def test_a_node_cut_off_by_a_crash_joins_its_call_to_another_agent(harness: HarnessEnvironment): + key = uuid.uuid4().hex + state = (await post(harness, f"{key}/run", {"graph": "research_crash", "refund": key})).json() + assert crashes[key] == 1 + assert state["findings"] == f"found: slow {key}" + assert asked[f"slow {key}"] == 1 + + +async def test_a_node_cut_off_by_a_crash_waits_for_the_same_signal(harness: HarnessEnvironment): + key = uuid.uuid4().hex + running = asyncio.ensure_future(post(harness, f"{key}/run", {"graph": "wait_crash", "refund": key})) + first = await until_waiting(harness, key) + go_crash.set() + for _ in range(400): + if crashes[key] == 1: + break + await asyncio.sleep(0.05) + assert crashes[key] == 1 + await asyncio.sleep(0.5) # the retried attempt waits again + assert await pending(harness, key) == first + await post(harness, f"{key}/signal", {"name": "approval", "value": "yes"}) + assert (await running).json()["approved"] == "yes" + + +def test_a_workflow_node_takes_input_exactly_when_its_handler_does(): + with pytest.raises(TypeError, match="takes input: pass input="): + nodes.workflow(ship) From 3938ef8c65eecac1b86e6f9dfe6e81283c744922 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 19:30:51 +0100 Subject: [PATCH 2/9] nodes: refuse to run unless the graph runs with durability="sync" LangGraph's default durability, "async", only submits a step's checkpoint save, and the next step's tasks start right away ("exit" saves nothing between steps). A crash before the save makes the resumed run create the checkpoint again with a new id, so the node gets a new task id and a new idempotency key, and makes its step twice. Only "sync" stores the checkpoint before the step runs. --- python/restate/agents/langgraph/nodes.py | 14 ++++++++++++-- tests/agents_langgraph_nodes.py | 10 ++++++++++ 2 files changed, 22 insertions(+), 2 deletions(-) diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index 7767626..cca59f9 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -21,7 +21,8 @@ graph.add_node("ship", nodes.workflow(shipping.run, input=lambda s: s["order"], output="receipt")) Run the graph in a handler of an agent (an ``AgentObject``), with its context -as LangGraph's run context and a checkpointer in its database:: +as LangGraph's run context, a checkpointer in its database, and +``durability="sync"`` (each step's checkpoint stored before the step runs):: desk = AgentObject("desk") @@ -34,7 +35,8 @@ async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: return await app.ainvoke(None if resuming else {"refund": refund}, config, context=ctx, durability="sync") Each node does its step once. A crash re-runs a node from the checkpoint it -started from, and LangGraph gives it the same task id there, so the node +started from, which ``durability="sync"`` stored before the node ran, and +LangGraph gives it the same task id there, so the node joins the workflow, the agent's answer or the wait it started before rather than start another. A wait is listed by the agent's ``pending`` view, and its ``signal`` view resolves it. @@ -62,6 +64,7 @@ def __call__(self, state: Any, *, runtime: Runtime[Any]) -> Awaitable[dict[str, _TASK_ID = "__pregel_task_id" _CHECKPOINTER = "__pregel_checkpointer" +_DURABILITY = "__pregel_durability" class _Step(typing.NamedTuple): @@ -83,6 +86,13 @@ def _step(runtime: Runtime[Any], kind: str) -> _Step: f"nodes.{kind} runs in a graph with a checkpointer, so that a retried attempt resumes it: " "compile(checkpointer=SqlCheckpointSaver(ctx.db))" ) + # Only "sync" stores a step's checkpoint before its tasks run. Otherwise a crash before the save makes the + # resumed run create the checkpoint again, with a new id, and so new task ids: the step would be made twice. + if configurable.get(_DURABILITY) != "sync": + raise TypeError( + f"nodes.{kind} runs in a graph that stores each step's checkpoint before the step runs, so that a " + 'retried attempt gets the same task ids: ainvoke(..., durability="sync")' + ) # The task's namespace holds its id, which LangGraph derives from the checkpoint the step started from: # a node re-run from that checkpoint gets the same one. It holds the subgraph's path too. namespace = typing.cast(str, configurable.get("checkpoint_ns") or configurable[_TASK_ID]) diff --git a/tests/agents_langgraph_nodes.py b/tests/agents_langgraph_nodes.py index 1492b1c..573752f 100644 --- a/tests/agents_langgraph_nodes.py +++ b/tests/agents_langgraph_nodes.py @@ -215,6 +215,16 @@ async def test_a_node_cut_off_by_a_crash_waits_for_the_same_signal(harness: Harn assert (await running).json()["approved"] == "yes" +async def test_a_node_refuses_to_run_unless_each_step_is_stored_before_it_runs(): + from langgraph.checkpoint.memory import InMemorySaver # pylint: disable=import-outside-toplevel + + app = approval_then_shipping().compile(checkpointer=InMemorySaver()) + config: RunnableConfig = {"configurable": {"thread_id": "t"}} + for durability in ({}, {"durability": "async"}, {"durability": "exit"}): + with pytest.raises(TypeError, match='durability="sync"'): + await app.ainvoke({"refund": "r"}, config, context=object(), **durability) + + def test_a_workflow_node_takes_input_exactly_when_its_handler_does(): with pytest.raises(TypeError, match="takes input: pass input="): nodes.workflow(ship) From 73c649f82b4e211781804bade443a09dfdec6767 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 19:46:39 +0100 Subject: [PATCH 3/9] Example: a refund desk as a LangGraph graph whose steps are Restate's examples-agents/intermediate/refund_graph.py runs a graph of your own in a handler of an AgentObject: it asks a policy agent about the refund (nodes.subagent), waits for a manager's approval (nodes.wait_for_signal), routes on the answer with an ordinary conditional edge, and runs the payout workflow (nodes.workflow). tests/examples_agents.py serves it with a scripted policy agent, and checks both an approved refund, which is paid, and a rejected one, which ends unpaid. --- examples-agents/README.md | 1 + examples-agents/intermediate/refund_graph.py | 100 +++++++++++++++++++ tests/examples_agents.py | 31 ++++++ 3 files changed, 132 insertions(+) create mode 100644 examples-agents/intermediate/refund_graph.py diff --git a/examples-agents/README.md b/examples-agents/README.md index 57a3a9a..4316433 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -112,6 +112,7 @@ frameworks. A tool such as `call`, `send`, `workflow`, `subagent` or | [`team.py`](intermediate/team.py) | A team on three frameworks: an editor (OpenAI) hands research to a researcher (LangGraph) and drafting to a writer (Pydantic AI), with `subagent` | | [`router.py`](intermediate/router.py) | A front desk of your own: a handler asks a model which team handles a message (`prompt`), then calls that team's agent | | [`triage.py`](intermediate/triage.py) | Agents as steps of a workflow: classify a ticket, draft a reply, and wait for a person's review | +| [`refund_graph.py`](intermediate/refund_graph.py) | A LangGraph graph of your own whose steps are Restate's: ask a policy agent, wait for a manager's approval, run the payout workflow (`langgraph.nodes`) | | [`loop.py`](intermediate/loop.py) | An agent loop of your own, with no framework: ask the model, run its tool calls, repeat; durable on the agents' conversation | ## Tips diff --git a/examples-agents/intermediate/refund_graph.py b/examples-agents/intermediate/refund_graph.py new file mode 100644 index 0000000..170fd2d --- /dev/null +++ b/examples-agents/intermediate/refund_graph.py @@ -0,0 +1,100 @@ +""" +A LangGraph graph of your own, whose steps are Restate's: a refund desk asks a +policy agent about the refund, waits for a manager's approval, and runs the +payout workflow. Each step happens once, even if the process crashes in the +middle; the wait for approval holds no process, for as long as it takes:: + + OPENAI_API_KEY=... uv run --extra langgraph --extra serve --with langchain-openai examples-agents/intermediate/refund_graph.py + npx @restatedev/restate-server # Restate, in another terminal + npx @restatedev/restate -y deployments register --force http://localhost:9080 + curl localhost:8080/refunds/alice/request --json '{"order": "103", "amount": 30}' # waits for the approval + +Then, in another terminal, see what the desk waits for, and approve it:: + + curl localhost:8080/refunds/alice/pending + curl localhost:8080/refunds/alice/signal --json '{"name": "approval", "value": "approved"}' + +Any other value rejects the refund, and the graph ends without the payout. +""" + +import asyncio +from typing import Any, Required, TypedDict, cast + +from langchain_core.runnables import RunnableConfig +from langgraph.graph import END, START, StateGraph + +import restate +from restate import sessions +from restate.agents import AgentObject, langgraph +from restate.agents.langgraph import SqlCheckpointSaver, nodes + + +class Refund(TypedDict): + order: str + amount: float + + +class State(TypedDict, total=False): + refund: Required[Refund] + policy: str + approval: str + receipt: str + + +# The payout, as a workflow: it runs once per refund. +payouts = restate.Workflow("payouts") + + +@payouts.main() +async def pay(_ctx: restate.WorkflowContext, refund: Refund) -> str: + return f"Refunded ${refund['amount']:g} for order {refund['order']}." + + +# An agent that knows the refund policy; the graph asks it about each refund. +policy = langgraph.agent( + "policy", + system="You know the refund policy: refunds up to $50 within 30 days. Say in one sentence whether a refund fits.", +) + + +def approved(state: State) -> str: + return "pay" if state.get("approval") == "approved" else END + + +graph = StateGraph(State) +graph.add_node( + "check", + nodes.subagent( + policy, + message=lambda s: f"A refund of ${s['refund']['amount']:g} for order {s['refund']['order']}.", + output="policy", + ), +) +graph.add_node("approve", nodes.wait_for_signal("approval", input=lambda s: s, output="approval")) +graph.add_node( + "pay", nodes.workflow(pay, input=lambda s: s["refund"], key=lambda s: s["refund"]["order"], output="receipt") +) +graph.add_edge(START, "check") +graph.add_edge("check", "approve") +graph.add_conditional_edges("approve", approved, ["pay", END]) +graph.add_edge("pay", END) + +# The graph runs in a handler of an agent: its key's database keeps the graph's checkpoints, and its `pending` and +# `signal` views show and resolve the waits. +refunds = AgentObject("refunds", description="A refund desk that runs a LangGraph graph.") + + +@refunds.handler() +async def request(ctx: sessions.SessionContext, refund: Refund) -> State: + # Untyped: LangGraph types a run's context as a state-like value, and the context here is the handler's. + app: Any = graph.compile(checkpointer=SqlCheckpointSaver(ctx.db)) + config: RunnableConfig = {"configurable": {"thread_id": refund["order"]}} + # A retried attempt (after a crash) resumes the run it cut off, rather than start another. + resuming = (await app.aget_state(config)).next + state: State | None = None if resuming else {"refund": refund} + # The nodes make their Restate calls with the handler's context, and need each step stored before it runs. + return cast(State, await app.ainvoke(state, config, context=ctx, durability="sync")) + + +if __name__ == "__main__": + asyncio.run(restate.serve(restate.app([payouts, policy, refunds]))) diff --git a/tests/examples_agents.py b/tests/examples_agents.py index 5d23079..111cd52 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -233,6 +233,37 @@ async def test_intermediate_triage(): assert (await ingress.post(f"{agent}/messages", {})).json()["messages"] == [] +async def test_intermediate_refund_graph(): + """The graph asks the policy agent (scripted here), waits for the approval, and pays; a rejection ends it unpaid.""" + pytest.importorskip("langgraph") + + def setup(module: ModuleType) -> None: + policy = restate.VirtualObject("policy") # in place of the example's agent, which needs a model + + @policy.handler() + async def chat(_ctx: restate.ObjectContext, message: str) -> str: + return f"Fits the policy: {message}" + + setattr(module, "policy", policy) + + async with serving("intermediate/refund_graph.py", setup) as ingress: + for order, answer, receipt in (("103", "approved", "Refunded $30 for order 103."), ("104", "no", None)): + requested = asyncio.create_task(ingress.post("refunds/alice/request", {"order": order, "amount": 30})) + waiting: list[dict[str, Any]] = [] + for _ in range(100): + if waiting := (await ingress.post("refunds/alice/pending")).json(): + break + await asyncio.sleep(0.1) + policy = f"Fits the policy: A refund of $30 for order {order}." + refund = {"order": order, "amount": 30} + assert [(w["name"], w["input"]) for w in waiting] == [("approval", {"refund": refund, "policy": policy})] + await ingress.post("refunds/alice/signal", {"name": "approval", "value": answer}) + state = (await requested).json() + assert state == {"refund": refund, "policy": policy, "approval": answer} | ( + {"receipt": receipt} if receipt else {} + ) + + async def test_intermediate_loop(monkeypatch: pytest.MonkeyPatch): """The loop, on a scripted model: it looks two orders up at once, waits for approval, and refunds; a crash after the refund ran retries the turn, which asks the model again without refunding again; and a crash after From ccae5933697876272ebe9b75899c17f17c87040a Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 20:09:36 +0100 Subject: [PATCH 4/9] langgraph.agent(name, graph=...): chat runs a graph of your own The graph's state has LangGraph's messages (MessagesState); a StateGraph or a compiled graph (such as LangChain's create_agent makes) is accepted, with no model, system or tools of ours. Each message is a run of the graph on the key's thread, with the message as a new HumanMessage whose id is the prompt's: - The checkpoints are in the key's database (SqlCheckpointSaver, history "latest"), stored before each step (durability "sync"), and the run's context is the handler's, so the prebuilt nodes work in it. A retried attempt resumes the run it cut off, or takes the finished run's answer. - After each step, the AI and tool messages it added are copied into the agent's conversation by id, so the messages, sync and watch views show the run as it goes. LangGraph streams a step's state only once its checkpoint is stored, so the conversation never gets ahead of the checkpoints; a copy a crash cut off is made at the next step or at the end. A test checks that ordering. - The answer is the run's last AI message; output= validates it. - After a stopped run, the next message starts a run of its own. The refund example becomes a desk you chat with: a clerk agent (output=Refund) reads the request, a policy agent checks it, a manager approves, and the payout workflow runs. --- examples-agents/README.md | 2 +- examples-agents/intermediate/refund_graph.py | 112 +++++----- python/restate/agents/langgraph/__init__.py | 3 +- python/restate/agents/langgraph/_agent.py | 40 ++++ python/restate/agents/langgraph/_graph.py | 127 +++++++++++ tests/agents_langgraph_graph.py | 218 +++++++++++++++++++ tests/examples_agents.py | 33 +-- 7 files changed, 470 insertions(+), 65 deletions(-) create mode 100644 python/restate/agents/langgraph/_graph.py create mode 100644 tests/agents_langgraph_graph.py diff --git a/examples-agents/README.md b/examples-agents/README.md index 4316433..dfb1d31 100644 --- a/examples-agents/README.md +++ b/examples-agents/README.md @@ -112,7 +112,7 @@ frameworks. A tool such as `call`, `send`, `workflow`, `subagent` or | [`team.py`](intermediate/team.py) | A team on three frameworks: an editor (OpenAI) hands research to a researcher (LangGraph) and drafting to a writer (Pydantic AI), with `subagent` | | [`router.py`](intermediate/router.py) | A front desk of your own: a handler asks a model which team handles a message (`prompt`), then calls that team's agent | | [`triage.py`](intermediate/triage.py) | Agents as steps of a workflow: classify a ticket, draft a reply, and wait for a person's review | -| [`refund_graph.py`](intermediate/refund_graph.py) | A LangGraph graph of your own whose steps are Restate's: ask a policy agent, wait for a manager's approval, run the payout workflow (`langgraph.nodes`) | +| [`refund_graph.py`](intermediate/refund_graph.py) | A refund desk you chat with, as a LangGraph graph of your own (`langgraph.agent(graph=...)`) whose steps are Restate's: ask a clerk and a policy agent, wait for a manager's approval, run the payout workflow (`langgraph.nodes`) | | [`loop.py`](intermediate/loop.py) | An agent loop of your own, with no framework: ask the model, run its tool calls, repeat; durable on the agents' conversation | ## Tips diff --git a/examples-agents/intermediate/refund_graph.py b/examples-agents/intermediate/refund_graph.py index 170fd2d..99b78b9 100644 --- a/examples-agents/intermediate/refund_graph.py +++ b/examples-agents/intermediate/refund_graph.py @@ -1,100 +1,112 @@ """ -A LangGraph graph of your own, whose steps are Restate's: a refund desk asks a -policy agent about the refund, waits for a manager's approval, and runs the -payout workflow. Each step happens once, even if the process crashes in the -middle; the wait for approval holds no process, for as long as it takes:: +A refund desk you chat with, as a LangGraph graph of your own whose steps are +Restate's: it asks a clerk agent to read the request, asks a policy agent +about it, waits for a manager's approval, runs the payout workflow, and +answers. Each step happens once, even if the process crashes in the middle; +the wait for approval holds no process, for as long as it takes:: OPENAI_API_KEY=... uv run --extra langgraph --extra serve --with langchain-openai examples-agents/intermediate/refund_graph.py npx @restatedev/restate-server # Restate, in another terminal npx @restatedev/restate -y deployments register --force http://localhost:9080 - curl localhost:8080/refunds/alice/request --json '{"order": "103", "amount": 30}' # waits for the approval + curl localhost:8080/refunds/alice/chat -d 'Order 103 arrived broken, I want my $30 back' # waits for the approval Then, in another terminal, see what the desk waits for, and approve it:: curl localhost:8080/refunds/alice/pending curl localhost:8080/refunds/alice/signal --json '{"name": "approval", "value": "approved"}' -Any other value rejects the refund, and the graph ends without the payout. +Any other value rejects the refund. ``python -m restate.agents chat refunds/alice`` +shows each step of the graph as it comes. """ import asyncio -from typing import Any, Required, TypedDict, cast +from typing import Any, NotRequired -from langchain_core.runnables import RunnableConfig -from langgraph.graph import END, START, StateGraph +from langchain_core.messages import AIMessage +from langgraph.graph import END, START, MessagesState, StateGraph +from pydantic import BaseModel import restate -from restate import sessions -from restate.agents import AgentObject, langgraph -from restate.agents.langgraph import SqlCheckpointSaver, nodes +from restate.agents import langgraph +from restate.agents.langgraph import nodes -class Refund(TypedDict): +class Refund(BaseModel): order: str amount: float -class State(TypedDict, total=False): - refund: Required[Refund] - policy: str - approval: str - receipt: str +class Desk(MessagesState): + """The graph's state: the conversation (LangGraph's messages), and the refund it works on.""" + refund: NotRequired[dict[str, Any]] + policy: NotRequired[str] + approval: NotRequired[str] + receipt: NotRequired[str] -# The payout, as a workflow: it runs once per refund. + +# The payout, as a workflow: it runs once per order. payouts = restate.Workflow("payouts") @payouts.main() async def pay(_ctx: restate.WorkflowContext, refund: Refund) -> str: - return f"Refunded ${refund['amount']:g} for order {refund['order']}." + return f"Refunded ${refund.amount:g} for order {refund.order}." -# An agent that knows the refund policy; the graph asks it about each refund. +# Two agents the graph asks: one reads the request into a Refund, one knows the policy. +clerk = langgraph.agent("clerk", system="Read the refund request: which order, and how much.", output=Refund) policy = langgraph.agent( "policy", system="You know the refund policy: refunds up to $50 within 30 days. Say in one sentence whether a refund fits.", ) -def approved(state: State) -> str: - return "pay" if state.get("approval") == "approved" else END +def asked(state: Desk) -> str: + return str(state["messages"][-1].text) + + +def about(state: Desk) -> str: + refund = state.get("refund", {}) + return f"A refund of ${float(refund['amount']):g} for order {refund['order']}." + + +def answer(state: Desk) -> dict[str, list[AIMessage]]: + if state.get("approval") != "approved": + return {"messages": [AIMessage(f"The refund was not approved: {about(state)}")]} + return {"messages": [AIMessage(state.get("receipt", ""))]} -graph = StateGraph(State) +def approved(state: Desk) -> str: + return "pay" if state.get("approval") == "approved" else "answer" + + +graph = StateGraph(Desk) +graph.add_node("read", nodes.subagent(clerk, message=asked, output="refund")) +graph.add_node("check", nodes.subagent(policy, message=about, output="policy")) graph.add_node( - "check", - nodes.subagent( - policy, - message=lambda s: f"A refund of ${s['refund']['amount']:g} for order {s['refund']['order']}.", - output="policy", + "approve", + nodes.wait_for_signal( + "approval", input=lambda s: {"refund": s["refund"], "policy": s["policy"]}, output="approval" ), ) -graph.add_node("approve", nodes.wait_for_signal("approval", input=lambda s: s, output="approval")) graph.add_node( - "pay", nodes.workflow(pay, input=lambda s: s["refund"], key=lambda s: s["refund"]["order"], output="receipt") + "pay", + nodes.workflow( + pay, input=lambda s: Refund.model_validate(s["refund"]), key=lambda s: s["refund"]["order"], output="receipt" + ), ) -graph.add_edge(START, "check") +graph.add_node("answer", answer) +graph.add_edge(START, "read") +graph.add_edge("read", "check") graph.add_edge("check", "approve") -graph.add_conditional_edges("approve", approved, ["pay", END]) -graph.add_edge("pay", END) - -# The graph runs in a handler of an agent: its key's database keeps the graph's checkpoints, and its `pending` and -# `signal` views show and resolve the waits. -refunds = AgentObject("refunds", description="A refund desk that runs a LangGraph graph.") - +graph.add_conditional_edges("approve", approved, ["pay", "answer"]) +graph.add_edge("pay", "answer") +graph.add_edge("answer", END) -@refunds.handler() -async def request(ctx: sessions.SessionContext, refund: Refund) -> State: - # Untyped: LangGraph types a run's context as a state-like value, and the context here is the handler's. - app: Any = graph.compile(checkpointer=SqlCheckpointSaver(ctx.db)) - config: RunnableConfig = {"configurable": {"thread_id": refund["order"]}} - # A retried attempt (after a crash) resumes the run it cut off, rather than start another. - resuming = (await app.aget_state(config)).next - state: State | None = None if resuming else {"refund": refund} - # The nodes make their Restate calls with the handler's context, and need each step stored before it runs. - return cast(State, await app.ainvoke(state, config, context=ctx, durability="sync")) +# An agent whose chat runs the graph: each key's database keeps its checkpoints and its conversation. +refunds = langgraph.agent("refunds", graph=graph, description="A refund desk that runs a LangGraph graph.") if __name__ == "__main__": - asyncio.run(restate.serve(restate.app([payouts, policy, refunds]))) + asyncio.run(restate.serve(restate.app([payouts, clerk, policy, refunds]))) diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py index e9a161b..2338c81 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -57,5 +57,6 @@ async def get_weather(city: str) -> str: from restate.agents.langgraph._agent import AnyTool, ChatModel, agent, prompt from restate.agents.langgraph._checkpointer import SqlCheckpointSaver +from restate.agents.langgraph._graph import AnyGraph -__all__ = ["AnyTool", "ChatModel", "SqlCheckpointSaver", "agent", "prompt"] +__all__ = ["AnyGraph", "AnyTool", "ChatModel", "SqlCheckpointSaver", "agent", "prompt"] diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 13cace8..1e0c1cc 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -64,6 +64,7 @@ from restate.agents._output import answer_with, asking from restate.agents._tools import Tool, ToolCall from restate.agents._turn import tool_call_ended, tool_call_started +from restate.agents.langgraph._graph import AnyGraph, check, run_of from restate.sessions import SessionContext M = typing.TypeVar("M", bound=BaseModel) @@ -328,6 +329,7 @@ def ask(text: str) -> Awaitable[str]: def agent( name: str, *, + graph: AnyGraph | None = None, model: PerSession[ChatModel | None] = None, system: PerSession[str | None] = None, tools: PerSession[Sequence[AnyTool]] = (), @@ -347,8 +349,21 @@ def agent( assistant = langgraph.agent("assistant", system=lambda ctx: f"You help {ctx.key}.") + Or bring a graph of your own, built on LangGraph's ``MessagesState``: + ``chat`` then runs it, with the message as a new ``HumanMessage``, and + answers with its last AI message. Its checkpoints are in the key's + database, and its run's context is the handler's, so the prebuilt + ``nodes`` (a signal, a workflow, another agent) work in it:: + + desk = langgraph.agent("desk", graph=builder) + Args: name: The object's name, as Restate serves it. + graph: A graph of your own, instead of the built-in tool-calling + loop: a ``StateGraph`` whose state has ``messages``, or a + compiled one (such as LangChain's ``create_agent`` makes). It has + its own model and tools, so it takes no ``model``, ``system`` or + ``tools``. model: ``"openai/"`` (needs ``langchain-openai``) or ``"anthropic/"`` (needs ``langchain-anthropic``), or a LangChain chat model that calls tools. Defaults to ``$RESTATE_AGENT_MODEL``, @@ -365,8 +380,15 @@ def agent( The agent, to serve with ``restate.serve(restate.app([agent]))``, or any ASGI server. Raises: + TypeError: If ``graph`` is given with ``model``, ``system`` or + ``tools``, or its state has no ``messages``. TerminalError: If ``model`` names a provider other than OpenAI or Anthropic. """ + if graph is not None: + if model is not None or system is not None or tools != (): + raise TypeError("a graph of your own has its own model and tools: pass no model, system or tools") + check(graph) + return _graph_agent(name, graph, output, **options) if model is None or isinstance(model, str): _provider_of_model(model_of(model)) obj = AgentObject(name, **options) @@ -383,3 +405,21 @@ async def answer(ctx: SessionContext, text: str) -> Any: obj.chat(answer, items=ITEMS, output=output) return obj + + +def _graph_agent( + name: str, graph: AnyGraph, output: type[BaseModel] | None, **options: Unpack[AgentOptions] +) -> AgentObject: + """Create an agent whose ``chat`` runs ``graph`` (see ``_graph.py``).""" + obj = AgentObject(name, **options) + + async def answer(ctx: SessionContext, text: str) -> Any: + def ask(text: str) -> Awaitable[str]: + return answer_once(ctx, text, ITEMS, run_of(ctx, graph)) + + if output is None: + return await ask(text) + return await answer_with(output, text, ask) + + obj.chat(answer, items=ITEMS, output=output) + return obj diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py new file mode 100644 index 0000000..32a788d --- /dev/null +++ b/python/restate/agents/langgraph/_graph.py @@ -0,0 +1,127 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +# pyright: reportMissingTypeStubs=false, reportUnknownMemberType=false +"""An agent whose ``chat`` runs a LangGraph graph of your own: ``langgraph.agent(name, graph=...)``. + +The graph's state has LangGraph's ``messages`` (``MessagesState``), and each +message to the agent is a run of the graph on the key's thread, with the +message as a new ``HumanMessage``: + +- The graph's checkpoints are in the key's database (``SqlCheckpointSaver``), + stored before each step runs (``durability="sync"``), and the run's + context is the handler's: the prebuilt ``nodes`` work in it. A retried + attempt resumes the run it cut off. +- After each step, the AI and tool messages the step added are copied into + the agent's conversation, by their ids, so that the ``messages``, ``sync`` + and ``watch`` views show the run as it goes, as they do for the built-in + agent. + + The conversation never gets ahead of the checkpoints: with + ``durability="sync"``, LangGraph streams a step's state only once its + checkpoint is stored, so a message is copied only once the graph has it for + good (no crash can make the step run again and say something else). The + checkpoints can get ahead of the conversation, if a crash comes between a + step's checkpoint and its copy: the next step's state, or the end of the + run, copies what was missed, as messages are copied by id. The + conversation is a log: a message the graph removes or edits later stays as + it was copied. +- The answer is the text of the run's last AI message. +- A stopped run is not resumed: the next message starts a run of its own, + from the state the graph had stored. +""" + +from __future__ import annotations + +import typing +from collections.abc import Sequence +from typing import Any + +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage +from langchain_core.runnables import RunnableConfig +from langgraph.graph import StateGraph +from langgraph.pregel import Pregel + +from restate.agents._conversation import Conversation, Run +from restate.agents._turn import stepped +from restate.agents.langgraph._checkpointer import SqlCheckpointSaver +from restate.sessions import SessionContext + +AnyGraph: typing.TypeAlias = StateGraph[Any, Any, Any, Any] | Pregel[Any, Any, Any, Any] +"""A graph of your own: a ``StateGraph`` (uncompiled), or a compiled one such as LangChain's ``create_agent`` makes.""" + +THREAD: RunnableConfig = {"configurable": {"thread_id": "chat"}} +"""The graph's thread: there is one per key, as the key's database is its own.""" + + +def check(graph: AnyGraph) -> None: + """Raise ``TypeError`` if ``graph`` has no ``messages`` in its state.""" + if "messages" not in graph.channels: + raise TypeError("the graph's state has no messages: build it on LangGraph's MessagesState") + + +def _compiled(graph: AnyGraph, ctx: SessionContext) -> Any: + """Return ``graph`` compiled with its checkpoints in the key's database (untyped: its context is ``ctx``).""" + saver = SqlCheckpointSaver(ctx.db, history="latest") + if isinstance(graph, StateGraph): + return graph.compile(checkpointer=saver) + return graph.copy(update={"checkpointer": saver}) + + +def _after(messages: Sequence[BaseMessage], prompt_id: str) -> list[BaseMessage]: + """Return the messages after the prompt ``prompt_id``: those of its run.""" + at = next((i for i, m in enumerate(messages) if m.id == prompt_id), None) + return list(messages[at + 1 :]) if at is not None else [] + + +def _answer(messages: Sequence[BaseMessage]) -> str: + """Return the text of the last AI message of a run.""" + replies = [m for m in messages if isinstance(m, AIMessage)] + return str(replies[-1].text) if replies else "" + + +async def _copied(ctx: SessionContext, conversation: Conversation, run: Sequence[BaseMessage]) -> None: + """Copy the AI and tool messages of ``run`` that the conversation does not have yet, and publish them.""" + known = {typing.cast(BaseMessage, entry.item).id for entry in conversation.entries} + new = [m for m in run if isinstance(m, AIMessage | ToolMessage) and (m.id is None or m.id not in known)] + for message in new: + if isinstance(message, AIMessage): + await conversation.append_reply(message) + else: + await conversation.append(message) + if new: + await stepped(ctx) + + +def run_of(ctx: SessionContext, graph: AnyGraph) -> Run: + """Return the run of a prompt: the graph's, on the key's thread, from the prompt to its answer.""" + + async def run(conversation: Conversation) -> str: + running = conversation.running() + assert running is not None and running.prompt is not None + # The prompt's id is the conversation's, so that a retried attempt finds it in the graph's state. + prompt_id = f"prompt:{running.prompt}" + app = _compiled(graph, ctx) + snapshot = await app.aget_state(THREAD) + messages = typing.cast(list[BaseMessage], (snapshot.values or {}).get("messages", [])) + asked = any(m.id == prompt_id for m in messages) + if asked and not snapshot.next: + # The run had finished, but its answer was not stored: that is the answer, rather than run again. + await _copied(ctx, conversation, _after(messages, prompt_id)) + return _answer(_after(messages, prompt_id)) + # A run cut off by a crash is resumed. A new message (or one after a stopped run) starts a run. + prompt = HumanMessage(str(typing.cast(BaseMessage, running.item).text), id=prompt_id) + start = None if asked else {"messages": [prompt]} + async for state in app.astream(start, THREAD, context=ctx, durability="sync", stream_mode="values"): + messages = typing.cast(list[BaseMessage], typing.cast(dict[str, Any], state).get("messages", [])) + await _copied(ctx, conversation, _after(messages, prompt_id)) + return _answer(_after(messages, prompt_id)) + + return run diff --git a/tests/agents_langgraph_graph.py b/tests/agents_langgraph_graph.py new file mode 100644 index 0000000..c6b19a2 --- /dev/null +++ b/tests/agents_langgraph_graph.py @@ -0,0 +1,218 @@ +# +# Copyright (c) 2023-2026 - Restate Software, Inc., Restate GmbH +# +# This file is part of the Restate SDK for Python, +# which is released under the MIT license. +# +# You can find a copy of the license in file LICENSE in the root +# directory of this repository or package, or at +# https://github.com/restatedev/sdk-typescript/blob/main/LICENSE +# +"""``langgraph.agent(name, graph=...)``: chat runs a graph of your own, against a real Restate server. + +The graph is scripted: it routes on the first word of the message. +""" + +# pylint: disable=missing-function-docstring,missing-class-docstring + +import asyncio +import collections +import typing +import uuid +from typing import Any, NotRequired, TypedDict + +import httpx2 +import pytest + +pytest.importorskip("langgraph") + +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage # noqa: E402 +from langgraph.graph import END, START, MessagesState, StateGraph # noqa: E402 +from pydantic import BaseModel # noqa: E402 + +import restate # noqa: E402 +from restate import HarnessEnvironment # noqa: E402 +from restate.agents import langgraph # noqa: E402 +from restate.agents.langgraph import nodes # noqa: E402 + + +class Desk(MessagesState): + approval: NotRequired[str] + + +def _asked(state: Desk) -> tuple[str, str]: + """The first word of the message the run answers, and the rest.""" + last = next(m for m in reversed(state["messages"]) if isinstance(m, HumanMessage)) + word, _, rest = str(typing.cast(BaseMessage, last).text).partition(" ") + return word, rest + + +def route(state: Desk) -> str: + word, _ = _asked(state) + return {"refund": "approve", "stop": "approve", "crash": "first", "plan": "plan"}.get(word, "greet") + + +def greet(state: Desk) -> dict[str, Any]: + return {"messages": [AIMessage(f"hello {_asked(state)[1]}")]} + + +def confirm(state: Desk) -> dict[str, Any]: + return {"messages": [AIMessage(f"approved: {state.get('approval')}")]} + + +firsts: collections.Counter[str] = collections.Counter() +crashes: collections.Counter[str] = collections.Counter() + + +def first(state: Desk) -> dict[str, Any]: + rest = _asked(state)[1] + firsts[rest] += 1 + return {"messages": [AIMessage(f"working on {rest}")]} + + +def second(state: Desk) -> dict[str, Any]: + rest = _asked(state)[1] + crashes[rest] += 1 + if crashes[rest] == 1: + raise RuntimeError("the process crashed") + return {"messages": [AIMessage(f"done {rest}")]} + + +def plan(state: Desk) -> dict[str, Any]: + return {"messages": [AIMessage(f'{{"title": "{_asked(state)[1]}"}}')]} + + +builder = StateGraph(Desk) +for name, node in {"greet": greet, "confirm": confirm, "first": first, "second": second, "plan": plan}.items(): + builder.add_node(name, typing.cast(Any, node)) +builder.add_node("approve", nodes.wait_for_signal("approval", output="approval")) +builder.add_conditional_edges(START, route, ["approve", "first", "plan", "greet"]) +builder.add_edge("approve", "confirm") +builder.add_edge("first", "second") +for last in ("greet", "confirm", "second", "plan"): + builder.add_edge(last, END) + + +class Plan(BaseModel): + title: str + + +desk = langgraph.agent("graph_desk", graph=builder) +compiled = langgraph.agent("graph_compiled", graph=builder.compile()) +planner = langgraph.agent("graph_planner", graph=builder, output=Plan) + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness( + restate.app([desk, compiled, planner]), restate_image="ghcr.io/restatedev/restate:latest" + ) as harness: + yield harness + + +pytestmark = [pytest.mark.anyio] + + +async def post(harness: HarnessEnvironment, path: str, **kwargs: Any) -> httpx2.Response: + async with httpx2.AsyncClient(timeout=60) as ingress: + return await ingress.post(f"{harness.ingress_url}/{path}", **kwargs) + + +async def chat(harness: HarnessEnvironment, agent: str, key: str, text: str) -> httpx2.Response: + return await post(harness, f"{agent}/{key}/chat", content=text, headers={"content-type": "text/plain"}) + + +async def messages(harness: HarnessEnvironment, agent: str, key: str) -> list[tuple[str, str]]: + page = (await post(harness, f"{agent}/{key}/messages", json={})).json() + return [(m["role"], " ".join(str(p.get("text", "")) for p in m["parts"])) for m in page["messages"]] + + +async def until_waiting(harness: HarnessEnvironment, key: str) -> list[dict[str, Any]]: + for _ in range(200): + if waits := (await post(harness, f"graph_desk/{key}/pending")).json(): + return typing.cast(list[dict[str, Any]], waits) + await asyncio.sleep(0.05) + raise AssertionError("nothing waits") + + +@pytest.mark.parametrize("agent", ["graph_desk", "graph_compiled"]) +async def test_chat_runs_the_graph_and_the_conversation_goes_on(harness: HarnessEnvironment, agent: str): + key = uuid.uuid4().hex + assert (await chat(harness, agent, key, f"hi {key}")).json() == f"hello {key}" + assert (await chat(harness, agent, key, "hi again")).json() == "hello again" + assert await messages(harness, agent, key) == [ + ("user", f"hi {key}"), + ("assistant", f"hello {key}"), + ("user", "hi again"), + ("assistant", "hello again"), + ] + + +async def test_a_prebuilt_node_waits_for_a_signal_in_chat(harness: HarnessEnvironment): + key = uuid.uuid4().hex + answering = asyncio.ensure_future(chat(harness, "graph_desk", key, f"refund {key}")) + [wait] = await until_waiting(harness, key) + assert wait["name"] == "approval" + await post(harness, f"graph_desk/{key}/signal", json={"name": "approval", "value": "yes"}) + assert (await answering).json() == "approved: yes" + + +async def test_a_crash_resumes_the_run_and_its_messages_are_copied_once(harness: HarnessEnvironment): + key = uuid.uuid4().hex + assert (await chat(harness, "graph_desk", key, f"crash {key}")).json() == f"done {key}" + assert (firsts[key], crashes[key]) == (1, 2) + assert await messages(harness, "graph_desk", key) == [ + ("user", f"crash {key}"), + ("assistant", f"working on {key}"), + ("assistant", f"done {key}"), + ] + + +async def test_after_a_stopped_run_the_next_message_starts_a_run_of_its_own(harness: HarnessEnvironment): + key = uuid.uuid4().hex + answering = asyncio.ensure_future(chat(harness, "graph_desk", key, f"stop {key}")) + await until_waiting(harness, key) + await post(harness, f"graph_desk/{key}/stop") + assert (await answering).status_code == 409 + assert (await post(harness, f"graph_desk/{key}/pending")).json() == [] + assert (await chat(harness, "graph_desk", key, f"hi {key}")).json() == f"hello {key}" + + +async def test_output_validates_the_last_ai_message(harness: HarnessEnvironment): + key = uuid.uuid4().hex + assert (await chat(harness, "graph_planner", key, f"plan {key}")).json() == {"title": key} + + +async def test_langgraph_streams_a_step_only_once_its_checkpoint_is_stored(): + """What copying messages relies on: with durability="sync", a streamed state is stored, even by a slow saver.""" + from langgraph.checkpoint.memory import InMemorySaver # pylint: disable=import-outside-toplevel + + class Slow(InMemorySaver): + async def aput(self, *args: Any, **kwargs: Any) -> Any: + await asyncio.sleep(0.1) + return await super().aput(*args, **kwargs) + + saver = Slow() + app: Any = builder.compile(checkpointer=saver) + config: Any = {"configurable": {"thread_id": "t"}} + states = app.astream({"messages": [HumanMessage("hi slow")]}, config, durability="sync", stream_mode="values") + async for state in states: + stored = await saver.aget_tuple(config) + assert stored is not None + assert len(stored.checkpoint["channel_values"]["messages"]) == len(state["messages"]) + + +def test_a_graph_has_its_own_model_and_its_state_has_messages(): + with pytest.raises(TypeError, match="pass no model, system or tools"): + langgraph.agent("both", graph=builder, system="Hi.") + + class Plain(TypedDict): + count: int + + with pytest.raises(TypeError, match="MessagesState"): + langgraph.agent("plain", graph=StateGraph(Plain)) diff --git a/tests/examples_agents.py b/tests/examples_agents.py index 111cd52..fe1a430 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -234,34 +234,41 @@ async def test_intermediate_triage(): async def test_intermediate_refund_graph(): - """The graph asks the policy agent (scripted here), waits for the approval, and pays; a rejection ends it unpaid.""" + """Chat runs the graph: it reads the request and asks about the policy (both agents scripted here), waits for the + approval, and pays; a rejection is answered without paying.""" pytest.importorskip("langgraph") def setup(module: ModuleType) -> None: - policy = restate.VirtualObject("policy") # in place of the example's agent, which needs a model + # In place of the example's agents, which need a model: virtual objects of the same names. + clerk, policy = restate.VirtualObject("clerk"), restate.VirtualObject("policy") - @policy.handler() - async def chat(_ctx: restate.ObjectContext, message: str) -> str: + @clerk.handler(name="chat") + async def read(_ctx: restate.ObjectContext, message: str) -> dict[str, Any]: + return {"order": message.split()[1], "amount": 30} + + @policy.handler(name="chat") + async def check(_ctx: restate.ObjectContext, message: str) -> str: return f"Fits the policy: {message}" + setattr(module, "clerk", clerk) setattr(module, "policy", policy) async with serving("intermediate/refund_graph.py", setup) as ingress: - for order, answer, receipt in (("103", "approved", "Refunded $30 for order 103."), ("104", "no", None)): - requested = asyncio.create_task(ingress.post("refunds/alice/request", {"order": order, "amount": 30})) + cases = (("103", "approved", "Refunded $30 for order 103."), ("104", "no", None)) + for order, approval, receipt in cases: + chat = asyncio.create_task(ingress.post("refunds/alice/chat", f"Order {order} arrived broken")) waiting: list[dict[str, Any]] = [] for _ in range(100): if waiting := (await ingress.post("refunds/alice/pending")).json(): break await asyncio.sleep(0.1) - policy = f"Fits the policy: A refund of $30 for order {order}." + about = f"A refund of $30 for order {order}." refund = {"order": order, "amount": 30} - assert [(w["name"], w["input"]) for w in waiting] == [("approval", {"refund": refund, "policy": policy})] - await ingress.post("refunds/alice/signal", {"name": "approval", "value": answer}) - state = (await requested).json() - assert state == {"refund": refund, "policy": policy, "approval": answer} | ( - {"receipt": receipt} if receipt else {} - ) + assert [(w["name"], w["input"]) for w in waiting] == [ + ("approval", {"refund": refund, "policy": f"Fits the policy: {about}"}) + ] + await ingress.post("refunds/alice/signal", {"name": "approval", "value": approval}) + assert (await chat).json() == (receipt or f"The refund was not approved: {about}") async def test_intermediate_loop(monkeypatch: pytest.MonkeyPatch): From f7f4bc5f4a460b169115e676e481f9bc2dd17cd8 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 20:14:13 +0100 Subject: [PATCH 5/9] nodes.workflow: validate plain data from the state into the handler's input Graph state is plain data, while a workflow's handler often takes a Pydantic model (or a dataclass, or a TypedDict). A dict from the state was passed through as it was, so the handler's serde failed on every attempt and Restate retried forever. The node now validates the value with Pydantic's TypeAdapter into the handler's input type; a value of that type already is passed through. A value that does not validate fails the node with a terminal error naming the handler and the field, since a retry would read the same state. The refund example passes the state's dict again. --- examples-agents/intermediate/refund_graph.py | 4 +- python/restate/agents/langgraph/nodes.py | 37 ++++++++++++++-- tests/agents_langgraph_nodes.py | 44 +++++++++++++++++++- 3 files changed, 77 insertions(+), 8 deletions(-) diff --git a/examples-agents/intermediate/refund_graph.py b/examples-agents/intermediate/refund_graph.py index 99b78b9..980eac4 100644 --- a/examples-agents/intermediate/refund_graph.py +++ b/examples-agents/intermediate/refund_graph.py @@ -92,9 +92,7 @@ def approved(state: Desk) -> str: ) graph.add_node( "pay", - nodes.workflow( - pay, input=lambda s: Refund.model_validate(s["refund"]), key=lambda s: s["refund"]["order"], output="receipt" - ), + nodes.workflow(pay, input=lambda s: s["refund"], key=lambda s: s["refund"]["order"], output="receipt"), ) graph.add_node("answer", answer) graph.add_edge(START, "read") diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index cca59f9..4c65aa1 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -50,9 +50,11 @@ async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: from langgraph.config import get_config from langgraph.runtime import Runtime +from pydantic import TypeAdapter, ValidationError from restate.agents import _signals from restate.agents._tools import handler_of, input_type_of, inputs_of, untyped +from restate.exceptions import TerminalError from restate.sessions import SessionContext @@ -100,6 +102,27 @@ def _step(runtime: Runtime[Any], kind: str) -> _Step: return _Step(ctx, node, f"langgraph:{ctx.key}:{namespace}") +def _validator(input_type: Any, name: str) -> Callable[[Any], Any]: + """Return what makes a value from the state (a dict, say) a value of a handler's ``input_type``. + + Graph state is plain data, while a handler may take a Pydantic model, a + dataclass or a ``TypedDict``: the value is validated into it, and one of + that type already is passed through. A value that does not validate fails + the node with a terminal error, as a retry would read the same state. + """ + if input_type is None or input_type is Any: + return lambda value: value + adapter: TypeAdapter[Any] = TypeAdapter(input_type) + + def validated(value: Any) -> Any: + try: + return adapter.validate_python(value) + except ValidationError as error: + raise TerminalError(f"the input of {name} is not valid: {error}") from error + + return validated + + def _update(output: str | None, value: object) -> dict[str, Any]: return {output: value} if output is not None else {} @@ -159,7 +182,8 @@ def workflow( output: The state field that gets the workflow's result. None to write nothing. input: The workflow's input, from the state. Required if the handler - takes one. + takes one. Plain data (a dict, say) is validated into the + handler's input type, such as a Pydantic model. key: The workflow id, from the state. Defaults to one of the node's own, made of the agent's key and the node's task. @@ -169,11 +193,16 @@ def workflow( Raises: TypeError: If ``target`` is not a handler, or it takes an input and ``input`` is not given (or the other way around). + + When it runs, the node fails with a terminal error if the input does not + validate. """ handler = handler_of(target) - takes_input = input_type_of(handler) is not None + input_type = input_type_of(handler) + takes_input = input_type is not None + name = f"{handler.service_tag.name}.{handler.name}" + validated = _validator(input_type, name) if takes_input != (input is not None): - name = f"{handler.service_tag.name}.{handler.name}" raise TypeError( f"{name} {'takes' if takes_input else 'takes no'} input: {'pass' if takes_input else 'drop'} input=" ) @@ -181,7 +210,7 @@ def workflow( async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "workflow") workflow_id = key(state) if key is not None else step.idempotency_key - value = input(state) if input is not None else None + value = validated(input(state)) if input is not None else None result = await untyped(step.ctx).call( target, *inputs_of(handler, value), key=workflow_id, idempotency_key=step.idempotency_key ) diff --git a/tests/agents_langgraph_nodes.py b/tests/agents_langgraph_nodes.py index 573752f..494b2ea 100644 --- a/tests/agents_langgraph_nodes.py +++ b/tests/agents_langgraph_nodes.py @@ -31,6 +31,7 @@ from langchain_core.runnables import RunnableConfig # noqa: E402 from langgraph.graph import END, START, StateGraph # noqa: E402 +from pydantic import BaseModel # noqa: E402 import restate # noqa: E402 from restate import HarnessEnvironment, sessions # noqa: E402 @@ -55,6 +56,31 @@ async def ship(ctx: restate.WorkflowContext, order: str) -> str: return f"shipped {order}" +class Invoice(BaseModel): + order: str + amount: float + + +billing = restate.Workflow("nodes_billing") +billed: collections.Counter[str] = collections.Counter() + + +@billing.main() +async def bill(ctx: restate.WorkflowContext, invoice: Invoice) -> str: + billed[invoice.order] += 1 + return f"billed ${invoice.amount:g} for {invoice.order}" + + +def billing_graph(amount: object) -> Any: + """A workflow node whose input, from the state, is a dict: the workflow takes an ``Invoice``.""" + graph = StateGraph(State) + node = nodes.workflow(bill, input=lambda s: {"order": s["refund"], "amount": amount}, output="receipt") + graph.add_node("bill", node) + graph.add_edge(START, "bill") + graph.add_edge("bill", END) + return graph + + researcher = restate.VirtualObject("nodes_researcher") asked: collections.Counter[str] = collections.Counter() researching = threading.Event() @@ -122,6 +148,8 @@ def beside_a_crash(step: str) -> Any: "approval": approval_then_shipping(), "research_crash": beside_a_crash("research"), "wait_crash": beside_a_crash("approve"), + "invoice": billing_graph(12.5), + "bad_invoice": billing_graph("a lot"), } desk = AgentObject("nodes_desk") @@ -149,7 +177,7 @@ def anyio_backend(): @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([desk, shipping, researcher]), restate_image="ghcr.io/restatedev/restate:latest" + restate.app([desk, shipping, billing, researcher]), restate_image="ghcr.io/restatedev/restate:latest" ) as harness: yield harness @@ -225,6 +253,20 @@ async def test_a_node_refuses_to_run_unless_each_step_is_stored_before_it_runs() await app.ainvoke({"refund": "r"}, config, context=object(), **durability) +async def test_a_workflow_node_validates_plain_data_from_the_state_into_the_handlers_input(harness: HarnessEnvironment): + key = uuid.uuid4().hex + state = (await post(harness, f"{key}/run", {"graph": "invoice", "refund": key})).json() + assert state["receipt"] == f"billed $12.5 for {key}" + + +async def test_input_that_does_not_validate_fails_the_node_for_good(harness: HarnessEnvironment): + key = uuid.uuid4().hex + failed = await post(harness, f"{key}/run", {"graph": "bad_invoice", "refund": key}) + assert failed.status_code == 500 + assert "the input of nodes_billing.bill is not valid" in failed.text and "amount" in failed.text + assert billed[key] == 0 + + def test_a_workflow_node_takes_input_exactly_when_its_handler_does(): with pytest.raises(TypeError, match="takes input: pass input="): nodes.workflow(ship) From 203c090268fb229cd5a15d60a8d2c20ab068cd43 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 20:19:34 +0100 Subject: [PATCH 6/9] langgraph: a new run carries what closed the last one; call and send nodes The graph agent: when a message starts a run after one that was stopped, refused or failed, the run's input starts with what closed that run in the conversation: the results of its cut-off tool calls, so that the graph's history has no tool call without a result (which a model provider refuses), and the note that ends it, so that the graph's model knows how the turn ended, as the built-in agent's does. The previous message comes too, if the graph never stored it. Each carried item has an id of its own, so it is carried once. nodes.call and nodes.send call or send to a handler, as the tools do. A key is required exactly for a keyed handler (a virtual object's or a workflow's), and plain data from the state is validated into the handler's input, as for nodes.workflow. _tools.keyed is shared with them. --- python/restate/agents/_tools.py | 10 +- python/restate/agents/langgraph/__init__.py | 3 +- python/restate/agents/langgraph/_agent.py | 2 +- python/restate/agents/langgraph/_graph.py | 48 ++++++- python/restate/agents/langgraph/nodes.py | 145 ++++++++++++++++++-- tests/agents_langgraph_graph.py | 70 +++++++++- tests/agents_langgraph_nodes.py | 53 ++++++- 7 files changed, 299 insertions(+), 32 deletions(-) diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 185bb12..2d84b6f 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -374,7 +374,7 @@ def validate(args: dict[str, Any]) -> tuple[str | None, Any]: return _Input(schema=schema, validate=validate) -def _keyed(handler: Handler[Any, Any]) -> bool: +def keyed(handler: Handler[Any, Any]) -> bool: return handler.service_tag.kind != "service" @@ -410,7 +410,7 @@ def call( """ handler = handler_of(target) service = handler.service_tag.name - model_keys = service if _keyed(handler) and key is None else None + model_keys = service if keyed(handler) and key is None else None shape = _input_of(input_type_of(handler), model_keys) async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: @@ -418,7 +418,7 @@ async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: async def made(c: ToolCall) -> Any: target_key = chosen if key is None else await configured(key, c.ctx) - target_key = target_key if _keyed(handler) else None + target_key = target_key if keyed(handler) else None return await untyped(c.ctx).call( target, *inputs_of(handler, value), key=target_key, idempotency_key=c.idempotency_key ) @@ -463,14 +463,14 @@ def send( """ handler = handler_of(target) service = handler.service_tag.name - model_keys = service if _keyed(handler) and key is None else None + model_keys = service if keyed(handler) and key is None else None shape = _input_of(input_type_of(handler), model_keys) async def run(arguments: dict[str, Any], tool_call: ToolCall) -> Any: chosen, value = arguments["key"], arguments["value"] async def made(c: ToolCall) -> str: - target_key = (chosen if key is None else await configured(key, c.ctx)) if _keyed(handler) else None + target_key = (chosen if key is None else await configured(key, c.ctx)) if keyed(handler) else None untyped(c.ctx).send( target, *inputs_of(handler, value), key=target_key, delay=delay, idempotency_key=c.idempotency_key ) diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py index 2338c81..d678d60 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -29,7 +29,8 @@ async def get_weather(city: str) -> str: ``agent`` creates an agent, ``prompt`` answers a prompt from a handler of your own, and ``SqlCheckpointSaver`` stores the checkpoints of a graph of your own in ``ctx.db``. ``nodes`` has prebuilt nodes for such a graph that are Restate -steps: wait for a signal, run a workflow, ask another agent. +steps: wait for a signal, run a workflow, call or send to a handler, ask another +agent. This differs from the older ``restate.ext.langchain`` integration, whose middleware makes a LangChain ``create_agent`` agent durable inside any regular diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 1e0c1cc..08afc3e 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -353,7 +353,7 @@ def agent( ``chat`` then runs it, with the message as a new ``HumanMessage``, and answers with its last AI message. Its checkpoints are in the key's database, and its run's context is the handler's, so the prebuilt - ``nodes`` (a signal, a workflow, another agent) work in it:: + ``nodes`` (a signal, a workflow, a call, another agent) work in it:: desk = langgraph.agent("desk", graph=builder) diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py index 32a788d..18e731e 100644 --- a/python/restate/agents/langgraph/_graph.py +++ b/python/restate/agents/langgraph/_graph.py @@ -35,7 +35,12 @@ it was copied. - The answer is the text of the run's last AI message. - A stopped run is not resumed: the next message starts a run of its own, - from the state the graph had stored. + from the state the graph had stored. What closed the stopped (or refused, + or failed) run in the conversation comes first in the new run's input: the + results of its cut-off tool calls, so that no tool call is left without one + (a model provider refuses such a history), and the note that ends it (such + as "Cancelled: ..."), so that the graph's model knows, as the built-in + agent's does. So does that run's message, if the graph never stored it. """ from __future__ import annotations @@ -49,7 +54,7 @@ from langgraph.graph import StateGraph from langgraph.pregel import Pregel -from restate.agents._conversation import Conversation, Run +from restate.agents._conversation import Conversation, Entry, Run from restate.agents._turn import stepped from restate.agents.langgraph._checkpointer import SqlCheckpointSaver from restate.sessions import SessionContext @@ -100,6 +105,38 @@ async def _copied(ctx: SessionContext, conversation: Conversation, run: Sequence await stepped(ctx) +def _prompt_id(entry: Entry) -> str: + """Return the id of a prompt's ``HumanMessage`` in the graph: the conversation's id of the prompt.""" + return f"prompt:{entry.prompt}" + + +def _closing(conversation: Conversation, running: Entry, messages: Sequence[BaseMessage]) -> list[BaseMessage]: + """Return what the conversation has of the previous prompt's run that the graph does not. + + That is what closed the run, if it was stopped, refused or failed: the + results of its cut-off tool calls and the note that ends it, which only + the conversation has (copied messages have the graph's ids; these have + none). And the previous prompt itself, if the graph never stored it. Each + gets an id of its own, so that it is carried once. + """ + earlier = [e for e in conversation.entries if e.prompt is not None and e.n < running.n] + if not earlier: + return [] + previous = earlier[-1] + known = {m.id for m in messages} + carried: list[BaseMessage] = [] + for entry in conversation.entries: + if not previous.n <= entry.n < running.n: + continue + item = typing.cast(BaseMessage, entry.item) + if entry.prompt is not None: + if _prompt_id(entry) not in known: + carried.append(HumanMessage(str(item.text), id=_prompt_id(entry))) + elif item.id is None and f"closed:{entry.n}" not in known: + carried.append(item.model_copy(update={"id": f"closed:{entry.n}"})) + return carried + + def run_of(ctx: SessionContext, graph: AnyGraph) -> Run: """Return the run of a prompt: the graph's, on the key's thread, from the prompt to its answer.""" @@ -107,7 +144,7 @@ async def run(conversation: Conversation) -> str: running = conversation.running() assert running is not None and running.prompt is not None # The prompt's id is the conversation's, so that a retried attempt finds it in the graph's state. - prompt_id = f"prompt:{running.prompt}" + prompt_id = _prompt_id(running) app = _compiled(graph, ctx) snapshot = await app.aget_state(THREAD) messages = typing.cast(list[BaseMessage], (snapshot.values or {}).get("messages", [])) @@ -116,9 +153,10 @@ async def run(conversation: Conversation) -> str: # The run had finished, but its answer was not stored: that is the answer, rather than run again. await _copied(ctx, conversation, _after(messages, prompt_id)) return _answer(_after(messages, prompt_id)) - # A run cut off by a crash is resumed. A new message (or one after a stopped run) starts a run. + # A run cut off by a crash is resumed. A new message (or one after a stopped run) starts a run, after what + # closed the previous run. prompt = HumanMessage(str(typing.cast(BaseMessage, running.item).text), id=prompt_id) - start = None if asked else {"messages": [prompt]} + start = None if asked else {"messages": [*_closing(conversation, running, messages), prompt]} async for state in app.astream(start, THREAD, context=ctx, durability="sync", stream_mode="values"): messages = typing.cast(list[BaseMessage], typing.cast(dict[str, Any], state).get("messages", [])) await _copied(ctx, conversation, _after(messages, prompt_id)) diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index 4c65aa1..b648816 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -8,7 +8,7 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -"""Prebuilt LangGraph nodes that are Restate steps: wait for a signal, run a workflow, ask another agent. +"""Prebuilt LangGraph nodes that are Restate steps: wait for a signal, run a workflow, call, send, ask another agent. Add them to a graph of your own, as the tools of ``restate.agents`` are added to an agent. Each one reads its input from the graph's state, and writes what @@ -46,6 +46,7 @@ async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: import typing from collections.abc import Awaitable, Callable +from datetime import timedelta from typing import Any from langgraph.config import get_config @@ -53,7 +54,7 @@ async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: from pydantic import TypeAdapter, ValidationError from restate.agents import _signals -from restate.agents._tools import handler_of, input_type_of, inputs_of, untyped +from restate.agents._tools import handler_of, input_type_of, inputs_of, keyed, untyped from restate.exceptions import TerminalError from restate.sessions import SessionContext @@ -123,6 +124,36 @@ def validated(value: Any) -> Any: return validated +class _Target(typing.NamedTuple): + """A handler a node calls: its definition, its name for errors, and what validates its input.""" + + handler: Any + name: str + validated: Callable[[Any], Any] + + +def _target(target: Callable[..., Any], input: object, key: object, *, keys: bool) -> _Target: + """Check a node's ``target`` against its ``input`` and ``key``, and return it. + + ``input`` must be given exactly when the handler takes one; with ``keys``, + ``key`` must be given exactly when the handler is keyed (a virtual + object's or a workflow's). + """ + handler = handler_of(target) + input_type = input_type_of(handler) + name = f"{handler.service_tag.name}.{handler.name}" + takes_input = input_type is not None + if takes_input != (input is not None): + raise TypeError( + f"{name} {'takes' if takes_input else 'takes no'} input: {'pass' if takes_input else 'drop'} input=" + ) + if keys and keyed(handler) != (key is not None): + raise TypeError( + f"{name} {'is' if keyed(handler) else 'is not'} keyed: {'pass' if keyed(handler) else 'drop'} key=" + ) + return _Target(handler, name, _validator(input_type, name)) + + def _update(output: str | None, value: object) -> dict[str, Any]: return {output: value} if output is not None else {} @@ -197,28 +228,114 @@ def workflow( When it runs, the node fails with a terminal error if the input does not validate. """ - handler = handler_of(target) - input_type = input_type_of(handler) - takes_input = input_type is not None - name = f"{handler.service_tag.name}.{handler.name}" - validated = _validator(input_type, name) - if takes_input != (input is not None): - raise TypeError( - f"{name} {'takes' if takes_input else 'takes no'} input: {'pass' if takes_input else 'drop'} input=" - ) + called = _target(target, input, key, keys=False) async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "workflow") workflow_id = key(state) if key is not None else step.idempotency_key - value = validated(input(state)) if input is not None else None + value = called.validated(input(state)) if input is not None else None + result = await untyped(step.ctx).call( + target, *inputs_of(called.handler, value), key=workflow_id, idempotency_key=step.idempotency_key + ) + return _update(output, result) + + return node + + +def call( + target: Callable[..., Any], + *, + output: str | None = None, + input: Callable[[Any], Any] | None = None, + key: Callable[[Any], str] | None = None, +) -> Node: + """Create a node that calls a handler, and writes its result to ``output``:: + + graph.add_node("status", nodes.call(orders.status, input=lambda s: s["order"], output="status")) + + Args: + target: A handler of a service, a virtual object or a workflow. + output: The state field that gets the handler's result. None to + write nothing. + input: The handler's input, from the state. Required if the handler + takes one. Plain data (a dict, say) is validated into the + handler's input type, such as a Pydantic model. + key: The key to call, from the state. Required for a virtual + object's or a workflow's handler, and only for those. + + Returns: + The node, for ``StateGraph.add_node``. + + Raises: + TypeError: If ``target`` is not a handler, or ``input`` or ``key`` + is given when the handler takes none (or the other way around). + + When it runs, the node fails with a terminal error if the input does not + validate. + """ + called = _target(target, input, key, keys=True) + + async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: + step = _step(runtime, "call") + value = called.validated(input(state)) if input is not None else None result = await untyped(step.ctx).call( - target, *inputs_of(handler, value), key=workflow_id, idempotency_key=step.idempotency_key + target, + *inputs_of(called.handler, value), + key=key(state) if key is not None else None, + idempotency_key=step.idempotency_key, ) return _update(output, result) return node +def send( + target: Callable[..., Any], + *, + input: Callable[[Any], Any] | None = None, + key: Callable[[Any], str] | None = None, + delay: timedelta | None = None, +) -> Node: + """Create a node that sends to a handler, without waiting for it to run:: + + graph.add_node("remind", nodes.send(checkins.remind, input=lambda s: s["checkin"], delay=timedelta(days=1))) + + Args: + target: A handler of a service, a virtual object or a workflow. + input: The handler's input, from the state. Required if the handler + takes one. Plain data (a dict, say) is validated into the + handler's input type, such as a Pydantic model. + key: The key to send to, from the state. Required for a virtual + object's or a workflow's handler, and only for those. + delay: How long to wait before the handler runs. + + Returns: + The node, for ``StateGraph.add_node``. It writes nothing to the state. + + Raises: + TypeError: If ``target`` is not a handler, or ``input`` or ``key`` + is given when the handler takes none (or the other way around). + + When it runs, the node fails with a terminal error if the input does not + validate. + """ + called = _target(target, input, key, keys=True) + + async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: + step = _step(runtime, "send") + value = called.validated(input(state)) if input is not None else None + untyped(step.ctx).send( + target, + *inputs_of(called.handler, value), + key=key(state) if key is not None else None, + delay=delay, + idempotency_key=step.idempotency_key, + ) + return {} + + return node + + def subagent( agent: object, *, @@ -259,4 +376,4 @@ async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: return node -__all__ = ["Node", "subagent", "wait_for_signal", "workflow"] +__all__ = ["Node", "call", "send", "subagent", "wait_for_signal", "workflow"] diff --git a/tests/agents_langgraph_graph.py b/tests/agents_langgraph_graph.py index c6b19a2..5919dc5 100644 --- a/tests/agents_langgraph_graph.py +++ b/tests/agents_langgraph_graph.py @@ -17,6 +17,7 @@ import asyncio import collections +import threading import typing import uuid from typing import Any, NotRequired, TypedDict @@ -26,13 +27,16 @@ pytest.importorskip("langgraph") -from langchain_core.messages import AIMessage, BaseMessage, HumanMessage # noqa: E402 +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage # noqa: E402 +from langchain_core.tools import tool # noqa: E402 from langgraph.graph import END, START, MessagesState, StateGraph # noqa: E402 +from langgraph.prebuilt import ToolNode # noqa: E402 from pydantic import BaseModel # noqa: E402 import restate # noqa: E402 from restate import HarnessEnvironment # noqa: E402 from restate.agents import langgraph # noqa: E402 +from restate.agents.internal import CANCELLED # noqa: E402 from restate.agents.langgraph import nodes # noqa: E402 @@ -49,7 +53,8 @@ def _asked(state: Desk) -> tuple[str, str]: def route(state: Desk) -> str: word, _ = _asked(state) - return {"refund": "approve", "stop": "approve", "crash": "first", "plan": "plan"}.get(word, "greet") + routes = {"refund": "approve", "stop": "approve", "crash": "first", "plan": "plan", "tool": "call_tool"} + return routes.get(word, "history" if word == "history" else "greet") def greet(state: Desk) -> dict[str, Any]: @@ -78,18 +83,47 @@ def second(state: Desk) -> dict[str, Any]: return {"messages": [AIMessage(f"done {rest}")]} +tool_started = threading.Event() + + +@tool +async def blocking() -> str: + """Runs until the turn is stopped.""" + tool_started.set() + await asyncio.sleep(300) + return "never" + + +def call_tool(state: Desk) -> dict[str, Any]: + call = {"id": f"call-{_asked(state)[1]}", "name": "blocking", "args": {}} + return {"messages": [AIMessage("", tool_calls=[call])]} + + +def history(state: Desk) -> dict[str, Any]: + """What the graph's state holds besides the messages it was sent: AI messages and tool results.""" + seen = [ + f"tool:{m.tool_call_id}:{m.text}" if isinstance(m, ToolMessage) else f"ai:{m.text}" + for m in state["messages"] + if isinstance(m, AIMessage | ToolMessage) and (m.text or isinstance(m, ToolMessage)) + ] + return {"messages": [AIMessage(" | ".join(seen))]} + + def plan(state: Desk) -> dict[str, Any]: return {"messages": [AIMessage(f'{{"title": "{_asked(state)[1]}"}}')]} builder = StateGraph(Desk) -for name, node in {"greet": greet, "confirm": confirm, "first": first, "second": second, "plan": plan}.items(): +steps = {"greet": greet, "confirm": confirm, "first": first, "second": second, "plan": plan, "history": history} +for name, node in {**steps, "call_tool": call_tool}.items(): builder.add_node(name, typing.cast(Any, node)) +builder.add_node("tools", ToolNode([blocking])) builder.add_node("approve", nodes.wait_for_signal("approval", output="approval")) -builder.add_conditional_edges(START, route, ["approve", "first", "plan", "greet"]) +builder.add_conditional_edges(START, route, ["approve", "first", "plan", "greet", "history", "call_tool"]) +builder.add_edge("call_tool", "tools") builder.add_edge("approve", "confirm") builder.add_edge("first", "second") -for last in ("greet", "confirm", "second", "plan"): +for last in ("greet", "confirm", "second", "plan", "history", "tools"): builder.add_edge(last, END) @@ -183,6 +217,32 @@ async def test_after_a_stopped_run_the_next_message_starts_a_run_of_its_own(harn assert (await chat(harness, "graph_desk", key, f"hi {key}")).json() == f"hello {key}" +async def test_the_graph_sees_how_a_stopped_run_ended(harness: HarnessEnvironment): + key = uuid.uuid4().hex + answering = asyncio.ensure_future(chat(harness, "graph_desk", key, f"stop {key}")) + await until_waiting(harness, key) + await post(harness, f"graph_desk/{key}/stop") + assert (await answering).status_code == 409 + assert (await chat(harness, "graph_desk", key, "history")).json() == f"ai:{CANCELLED}" + + +async def test_a_tool_call_cut_off_by_a_stop_gets_a_result_in_the_graph(harness: HarnessEnvironment): + key = uuid.uuid4().hex + tool_started.clear() + answering = asyncio.ensure_future(chat(harness, "graph_desk", key, f"tool {key}")) + for _ in range(200): + if tool_started.is_set(): + break + await asyncio.sleep(0.05) + await post(harness, f"graph_desk/{key}/stop") + assert (await answering).status_code == 409 + # The graph's history has no tool call without a result, which a model provider would refuse. + expected = f"tool:call-{key}:{CANCELLED} | ai:{CANCELLED}" + assert (await chat(harness, "graph_desk", key, "history")).json() == expected + # Carried once: the next run does not add them again. + assert (await chat(harness, "graph_desk", key, "history")).json() == f"{expected} | ai:{expected}" + + async def test_output_validates_the_last_ai_message(harness: HarnessEnvironment): key = uuid.uuid4().hex assert (await chat(harness, "graph_planner", key, f"plan {key}")).json() == {"title": key} diff --git a/tests/agents_langgraph_nodes.py b/tests/agents_langgraph_nodes.py index 494b2ea..0c6d92b 100644 --- a/tests/agents_langgraph_nodes.py +++ b/tests/agents_langgraph_nodes.py @@ -81,6 +81,37 @@ def billing_graph(amount: object) -> Any: return graph +ledger = restate.VirtualObject("nodes_ledger") + + +@ledger.handler() +async def add(ctx: restate.ObjectContext, invoice: Invoice) -> str: + total = (await ctx.get("total", type_hint=float) or 0) + invoice.amount + ctx.set("total", total) + return f"{ctx.key()} owes ${total:g}" + + +notices = restate.Service("nodes_notices") +noticed: collections.Counter[str] = collections.Counter() + + +@notices.handler() +async def notice(_ctx: restate.Context, invoice: Invoice) -> None: + noticed[invoice.order] += 1 + + +def call_then_send() -> Any: + """A call to a virtual object, then a send to a service: both with a dict from the state as the input.""" + graph = StateGraph(State) + invoice = typing.cast(Any, lambda s: {"order": s["refund"], "amount": 4}) + graph.add_node("add", nodes.call(add, input=invoice, key=lambda s: s["refund"], output="receipt")) + graph.add_node("notice", nodes.send(notice, input=invoice)) + graph.add_edge(START, "add") + graph.add_edge("add", "notice") + graph.add_edge("notice", END) + return graph + + researcher = restate.VirtualObject("nodes_researcher") asked: collections.Counter[str] = collections.Counter() researching = threading.Event() @@ -150,6 +181,7 @@ def beside_a_crash(step: str) -> Any: "wait_crash": beside_a_crash("approve"), "invoice": billing_graph(12.5), "bad_invoice": billing_graph("a lot"), + "call_send": call_then_send(), } desk = AgentObject("nodes_desk") @@ -177,7 +209,8 @@ def anyio_backend(): @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([desk, shipping, billing, researcher]), restate_image="ghcr.io/restatedev/restate:latest" + restate.app([desk, shipping, billing, ledger, notices, researcher]), + restate_image="ghcr.io/restatedev/restate:latest", ) as harness: yield harness @@ -267,6 +300,24 @@ async def test_input_that_does_not_validate_fails_the_node_for_good(harness: Har assert billed[key] == 0 +async def test_a_graph_calls_an_object_and_sends_to_a_service(harness: HarnessEnvironment): + key = uuid.uuid4().hex + state = (await post(harness, f"{key}/run", {"graph": "call_send", "refund": key})).json() + assert state["receipt"] == f"{key} owes $4" + for _ in range(100): + if noticed[key]: + break + await asyncio.sleep(0.05) + assert noticed[key] == 1 + + +def test_a_call_takes_a_key_exactly_when_its_handler_is_keyed(): + with pytest.raises(TypeError, match="is keyed: pass key="): + nodes.call(add, input=lambda s: s) + with pytest.raises(TypeError, match="is not keyed: drop key="): + nodes.send(notice, input=lambda s: s, key=lambda s: "k") + + def test_a_workflow_node_takes_input_exactly_when_its_handler_does(): with pytest.raises(TypeError, match="takes input: pass input="): nodes.workflow(ship) From f7b518c93dee27f64c973b41d11de456b50988d7 Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 20:24:12 +0100 Subject: [PATCH 7/9] langgraph: fixes from a review of the whole PR - nodes: a handler input type Pydantic does not know (a msgspec Struct) no longer makes building the node fail: the value goes to the handler's serde as it is. - nodes.wait_for_signal: what the graph waits for is turned into JSON, so a model in the state (a workflow's result, say) no longer fails each attempt. - prompt(ctx, text, graph=...): a handler of your own prompts the graph as chat does; agent(graph=...) answers through it. - The graph agent copies a run's messages looking at the run's own entries only, not the whole conversation at each step. - Docs: nodes lead with langgraph.agent(graph=...); output= with a graph is the graph's own prompts' to ask for; the example does not claim the chat shows steps that add no messages. - Tests wait at most 20 seconds for their events, so a broken one fails rather than hangs; new tests for the above. --- examples-agents/intermediate/refund_graph.py | 2 +- python/restate/agents/langgraph/_agent.py | 47 +++++++++++++------- python/restate/agents/langgraph/_graph.py | 13 +++--- python/restate/agents/langgraph/nodes.py | 23 +++++++--- tests/agents_langgraph_graph.py | 22 ++++++++- tests/agents_langgraph_nodes.py | 36 ++++++++++++--- 6 files changed, 109 insertions(+), 34 deletions(-) diff --git a/examples-agents/intermediate/refund_graph.py b/examples-agents/intermediate/refund_graph.py index 980eac4..d0a7732 100644 --- a/examples-agents/intermediate/refund_graph.py +++ b/examples-agents/intermediate/refund_graph.py @@ -16,7 +16,7 @@ curl localhost:8080/refunds/alice/signal --json '{"name": "approval", "value": "approved"}' Any other value rejects the refund. ``python -m restate.agents chat refunds/alice`` -shows each step of the graph as it comes. +shows the wait for approval as it comes, and the answer. """ import asyncio diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 08afc3e..b1b8b99 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -258,6 +258,7 @@ async def prompt( ctx: SessionContext, text: str, *, + graph: AnyGraph | None = None, model: ChatModel | None = None, system: str | None = None, tools: Sequence[AnyTool] = (), @@ -270,6 +271,7 @@ async def prompt( ctx: SessionContext, text: str, *, + graph: AnyGraph | None = None, model: ChatModel | None = None, system: str | None = None, tools: Sequence[AnyTool] = (), @@ -281,6 +283,7 @@ async def prompt( ctx: SessionContext, text: str, *, + graph: AnyGraph | None = None, model: ChatModel | None = None, system: str | None = None, tools: Sequence[AnyTool] = (), @@ -299,27 +302,40 @@ async def summarize(ctx: sessions.SessionContext, topic: str) -> str: Args: ctx: The handler's context. text: The prompt. + graph: A graph of your own to run, as ``agent(graph=...)`` runs it + for ``chat``, instead of the built-in tool-calling loop. It takes + no ``model``, ``system`` or ``tools``. model: ``"openai/"`` or ``"anthropic/"``, or a LangChain chat model that calls tools. Defaults to ``$RESTATE_AGENT_MODEL``, else ``"openai/gpt-6-sol"``. system: The system prompt. tools: ``restate.agents`` tools, or LangChain tools. output: A Pydantic model, for an answer that is one rather than - text. The system prompt asks for its JSON. + text. The system prompt asks for its JSON; with ``graph``, the + graph's own prompts must, and its last AI message is validated. Returns: The agent's answer: its text, or a value of ``output``. Raises: + TypeError: If ``graph`` is given with ``model``, ``system`` or + ``tools``, or its state has no ``messages``. TerminalError: If the agent's reply is not a value of ``output``, twice. """ + if graph is not None: + _check_graph(graph, model, system, tools) + run = run_of(ctx, graph) - # Without a way to ask for ``output`` natively with any provider's model, the system prompt asks for it. - asked = system if output is None else asking(system, output) + def ask(text: str) -> Awaitable[str]: + return answer_once(ctx, text, ITEMS, run) - def ask(text: str) -> Awaitable[str]: - return answer_once(ctx, text, ITEMS, _run(ctx, model, asked, tools)) + else: + # Without a way to ask for ``output`` natively with any provider's model, the system prompt asks for it. + asked = system if output is None else asking(system, output) + + def ask(text: str) -> Awaitable[str]: + return answer_once(ctx, text, ITEMS, _run(ctx, model, asked, tools)) if output is None: return await ask(text) @@ -372,7 +388,8 @@ def agent( tools: ``restate.agents`` tools, or LangChain tools. output: A Pydantic model, for answers that are one rather than text. The system prompt asks for its JSON, which ``chat`` - returns. + returns; with ``graph``, the graph's own prompts must, and its + last AI message is validated. **options: The object's options, such as ``description``, ``background`` or ``invocation_retry_policy``: see ``AgentOptions``. @@ -385,9 +402,7 @@ def agent( TerminalError: If ``model`` names a provider other than OpenAI or Anthropic. """ if graph is not None: - if model is not None or system is not None or tools != (): - raise TypeError("a graph of your own has its own model and tools: pass no model, system or tools") - check(graph) + _check_graph(graph, model, system, tools) return _graph_agent(name, graph, output, **options) if model is None or isinstance(model, str): _provider_of_model(model_of(model)) @@ -407,6 +422,13 @@ async def answer(ctx: SessionContext, text: str) -> Any: return obj +def _check_graph(graph: AnyGraph, model: object, system: object, tools: object) -> None: + """Raise ``TypeError`` if ``graph`` comes with a model, a system prompt or tools, or its state has no messages.""" + if model is not None or system is not None or tools != (): + raise TypeError("a graph of your own has its own model and tools: pass no model, system or tools") + check(graph) + + def _graph_agent( name: str, graph: AnyGraph, output: type[BaseModel] | None, **options: Unpack[AgentOptions] ) -> AgentObject: @@ -414,12 +436,7 @@ def _graph_agent( obj = AgentObject(name, **options) async def answer(ctx: SessionContext, text: str) -> Any: - def ask(text: str) -> Awaitable[str]: - return answer_once(ctx, text, ITEMS, run_of(ctx, graph)) - - if output is None: - return await ask(text) - return await answer_with(output, text, ask) + return await prompt(ctx, text, graph=graph, output=output) obj.chat(answer, items=ITEMS, output=output) return obj diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py index 18e731e..70edab5 100644 --- a/python/restate/agents/langgraph/_graph.py +++ b/python/restate/agents/langgraph/_graph.py @@ -92,9 +92,12 @@ def _answer(messages: Sequence[BaseMessage]) -> str: return str(replies[-1].text) if replies else "" -async def _copied(ctx: SessionContext, conversation: Conversation, run: Sequence[BaseMessage]) -> None: - """Copy the AI and tool messages of ``run`` that the conversation does not have yet, and publish them.""" - known = {typing.cast(BaseMessage, entry.item).id for entry in conversation.entries} +async def _copied(ctx: SessionContext, conversation: Conversation, running: Entry, run: Sequence[BaseMessage]) -> None: + """Copy the AI and tool messages of ``run`` (the prompt ``running``'s) that the conversation does not have yet. + + Only the run's own entries (after its prompt) can hold them, so only those are looked at. Then publish. + """ + known = {typing.cast(BaseMessage, e.item).id for e in conversation.entries if e.n > running.n} new = [m for m in run if isinstance(m, AIMessage | ToolMessage) and (m.id is None or m.id not in known)] for message in new: if isinstance(message, AIMessage): @@ -151,7 +154,7 @@ async def run(conversation: Conversation) -> str: asked = any(m.id == prompt_id for m in messages) if asked and not snapshot.next: # The run had finished, but its answer was not stored: that is the answer, rather than run again. - await _copied(ctx, conversation, _after(messages, prompt_id)) + await _copied(ctx, conversation, running, _after(messages, prompt_id)) return _answer(_after(messages, prompt_id)) # A run cut off by a crash is resumed. A new message (or one after a stopped run) starts a run, after what # closed the previous run. @@ -159,7 +162,7 @@ async def run(conversation: Conversation) -> str: start = None if asked else {"messages": [*_closing(conversation, running, messages), prompt]} async for state in app.astream(start, THREAD, context=ctx, durability="sync", stream_mode="values"): messages = typing.cast(list[BaseMessage], typing.cast(dict[str, Any], state).get("messages", [])) - await _copied(ctx, conversation, _after(messages, prompt_id)) + await _copied(ctx, conversation, running, _after(messages, prompt_id)) return _answer(_after(messages, prompt_id)) return run diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index b648816..25be0a9 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -20,8 +20,13 @@ graph.add_node("approve", nodes.wait_for_signal("approval", input=lambda s: s["refund"], output="approved")) graph.add_node("ship", nodes.workflow(shipping.run, input=lambda s: s["order"], output="receipt")) -Run the graph in a handler of an agent (an ``AgentObject``), with its context -as LangGraph's run context, a checkpointer in its database, and +The simplest is an agent whose ``chat`` runs the graph, built on LangGraph's +``MessagesState``; it sets everything below up:: + + desk = langgraph.agent("desk", graph=graph) + +Or run the graph in a handler of an agent (an ``AgentObject``) yourself, with +its context as LangGraph's run context, a checkpointer in its database, and ``durability="sync"`` (each step's checkpoint stored before the step runs):: desk = AgentObject("desk") @@ -29,7 +34,7 @@ @desk.handler() async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: - app = graph.compile(checkpointer=SqlCheckpointSaver(ctx.db)) + app: Any = graph.compile(checkpointer=SqlCheckpointSaver(ctx.db)) # Any: LangGraph types no Restate context config: RunnableConfig = {"configurable": {"thread_id": "refunds"}} resuming = (await app.aget_state(config)).next return await app.ainvoke(None if resuming else {"refund": refund}, config, context=ctx, durability="sync") @@ -51,7 +56,8 @@ async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: from langgraph.config import get_config from langgraph.runtime import Runtime -from pydantic import TypeAdapter, ValidationError +from pydantic import PydanticSchemaGenerationError, TypeAdapter, ValidationError +from pydantic_core import to_jsonable_python from restate.agents import _signals from restate.agents._tools import handler_of, input_type_of, inputs_of, keyed, untyped @@ -113,7 +119,11 @@ def _validator(input_type: Any, name: str) -> Callable[[Any], Any]: """ if input_type is None or input_type is Any: return lambda value: value - adapter: TypeAdapter[Any] = TypeAdapter(input_type) + try: + adapter: TypeAdapter[Any] = TypeAdapter(input_type) + except PydanticSchemaGenerationError: + # A type Pydantic does not know, such as a msgspec Struct: the handler's serde takes the value as it is. + return lambda value: value def validated(value: Any) -> Any: try: @@ -188,7 +198,8 @@ async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "wait_for_signal") request: _signals.WaitRequest = {"id": step.idempotency_key, "name": signal} if input is not None: - request["input"] = input(state) + # As JSON, which the companion stores: the state may hold models (a workflow's result, say). + request["input"] = to_jsonable_python(input(state)) value = await step.ctx.call( _signals.wait, request, key=_signals.signals_of(step.ctx), idempotency_key=step.idempotency_key ) diff --git a/tests/agents_langgraph_graph.py b/tests/agents_langgraph_graph.py index 5919dc5..eed7eba 100644 --- a/tests/agents_langgraph_graph.py +++ b/tests/agents_langgraph_graph.py @@ -34,7 +34,7 @@ from pydantic import BaseModel # noqa: E402 import restate # noqa: E402 -from restate import HarnessEnvironment # noqa: E402 +from restate import HarnessEnvironment, sessions # noqa: E402 from restate.agents import langgraph # noqa: E402 from restate.agents.internal import CANCELLED # noqa: E402 from restate.agents.langgraph import nodes # noqa: E402 @@ -132,6 +132,14 @@ class Plan(BaseModel): desk = langgraph.agent("graph_desk", graph=builder) + + +@desk.handler() +async def ask(ctx: sessions.SessionContext, text: str) -> str: + """A handler of your own that prompts the graph, as chat does.""" + return await langgraph.prompt(ctx, text, graph=builder) + + compiled = langgraph.agent("graph_compiled", graph=builder.compile()) planner = langgraph.agent("graph_planner", graph=builder, output=Plan) @@ -243,6 +251,18 @@ async def test_a_tool_call_cut_off_by_a_stop_gets_a_result_in_the_graph(harness: assert (await chat(harness, "graph_desk", key, "history")).json() == f"{expected} | ai:{expected}" +async def test_a_handler_of_your_own_prompts_the_graph(harness: HarnessEnvironment): + key = uuid.uuid4().hex + assert (await post(harness, f"graph_desk/{key}/ask", json=f"hi {key}")).json() == f"hello {key}" + assert (await chat(harness, "graph_desk", key, "hi again")).json() == "hello again" + assert await messages(harness, "graph_desk", key) == [ + ("user", f"hi {key}"), + ("assistant", f"hello {key}"), + ("user", "hi again"), + ("assistant", "hello again"), + ] + + async def test_output_validates_the_last_ai_message(harness: HarnessEnvironment): key = uuid.uuid4().hex assert (await chat(harness, "graph_planner", key, f"plan {key}")).json() == {"title": key} diff --git a/tests/agents_langgraph_nodes.py b/tests/agents_langgraph_nodes.py index 0c6d92b..575168a 100644 --- a/tests/agents_langgraph_nodes.py +++ b/tests/agents_langgraph_nodes.py @@ -112,6 +112,15 @@ def call_then_send() -> Any: return graph +async def until(event: threading.Event) -> None: + """Wait until ``event`` is set, for at most 20 seconds, so that a broken test fails rather than hangs.""" + for _ in range(2000): + if event.is_set(): + return + await asyncio.sleep(0.01) + raise AssertionError("the event never came") + + researcher = restate.VirtualObject("nodes_researcher") asked: collections.Counter[str] = collections.Counter() researching = threading.Event() @@ -124,14 +133,15 @@ async def chat(ctx: restate.ObjectContext, question: str) -> str: if question.startswith("slow"): # Answer only once the graph's attempt has crashed, so the crash cuts the node's call short. researching.set() - while not crashed.is_set(): - await asyncio.sleep(0.01) + await until(crashed) return f"found: {question}" def approval_then_shipping() -> Any: graph = StateGraph(State) - graph.add_node("approve", nodes.wait_for_signal("approval", input=lambda s: s["refund"], output="approved")) + # A model as what the graph waits for: pending shows it as JSON. + waiting_for = typing.cast(Any, lambda s: Invoice(order=s["refund"], amount=1)) + graph.add_node("approve", nodes.wait_for_signal("approval", input=waiting_for, output="approved")) graph.add_node("ship", nodes.workflow(ship, input=lambda s: s["refund"], output="receipt")) graph.add_node("research", nodes.subagent(researcher, message=lambda s: f"why {s['refund']}", output="findings")) graph.add_edge(START, "approve") @@ -150,8 +160,7 @@ def crash_after(ready: threading.Event) -> Any: async def crash(state: State) -> dict[str, Any]: if crashes[state["refund"]] == 0: - while not ready.is_set(): - await asyncio.sleep(0.01) + await until(ready) crashes[state["refund"]] += 1 crashed.set() raise RuntimeError("the process crashed") @@ -239,7 +248,7 @@ async def test_a_graph_waits_for_a_signal_runs_a_workflow_and_asks_another_agent key = uuid.uuid4().hex running = asyncio.ensure_future(post(harness, f"{key}/run", {"graph": "approval", "refund": key})) waits = await until_waiting(harness, key) - assert [(w["name"], w.get("input")) for w in waits] == [("approval", key)] + assert [(w["name"], w.get("input")) for w in waits] == [("approval", {"order": key, "amount": 1.0})] resolved = await post(harness, f"{key}/signal", {"name": "approval", "value": {"by": "manager"}}) assert resolved.json() == {"resolved": waits[0]["id"]} assert (await running).json() == { @@ -311,6 +320,21 @@ async def test_a_graph_calls_an_object_and_sends_to_a_service(harness: HarnessEn assert noticed[key] == 1 +def test_a_node_takes_a_handler_whose_input_pydantic_does_not_know(): + import msgspec # pylint: disable=import-outside-toplevel + + class Parcel(msgspec.Struct): + weight: float + + parcels = restate.Workflow("nodes_parcels") + + @parcels.main() + async def weigh(_ctx: restate.WorkflowContext, parcel: Parcel) -> float: + return parcel.weight + + nodes.workflow(weigh, input=lambda s: {"weight": 1}) # the handler's serde takes the value as it is + + def test_a_call_takes_a_key_exactly_when_its_handler_is_keyed(): with pytest.raises(TypeError, match="is keyed: pass key="): nodes.call(add, input=lambda s: s) From 73c9121477a290d36e8a8c2ab168988867d3ad0a Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 20:30:41 +0100 Subject: [PATCH 8/9] Graph agent: a graph that trims its history keeps its answer A run was found by its prompt's place in the graph's state, and its messages and answer were those after it. A graph that trims its history (a summarizing one) removes the prompt as well, and the run then copied nothing and answered with nothing; and the next message brought the trimmed prompt back. Each checkpoint of a run now names its prompt in its metadata, which is how a retried attempt finds the run it cut off. The messages copied are the graph's AI and tool messages that the conversation does not have (not those it gave the graph), and the answer is the last AI message the run copied, read from the conversation, which is a log. The previous prompt is carried into a new run only if its run copied nothing, so a prompt the graph removed is not brought back. --- python/restate/agents/langgraph/_graph.py | 85 ++++++++++++++++------- tests/agents_langgraph_graph.py | 30 ++++++-- 2 files changed, 85 insertions(+), 30 deletions(-) diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py index 70edab5..e63b24c 100644 --- a/python/restate/agents/langgraph/_graph.py +++ b/python/restate/agents/langgraph/_graph.py @@ -34,6 +34,9 @@ conversation is a log: a message the graph removes or edits later stays as it was copied. - The answer is the text of the run's last AI message. +- The graph may trim its history, its input and answer included (as a + summarizing graph does): a run is found by its checkpoints' metadata, + which name its prompt, and its answer is read from the conversation. - A stopped run is not resumed: the next message starts a run of its own, from the state the graph had stored. What closed the stopped (or refused, or failed) run in the conversation comes first in the new run's input: the @@ -80,30 +83,55 @@ def _compiled(graph: AnyGraph, ctx: SessionContext) -> Any: return graph.copy(update={"checkpointer": saver}) -def _after(messages: Sequence[BaseMessage], prompt_id: str) -> list[BaseMessage]: - """Return the messages after the prompt ``prompt_id``: those of its run.""" - at = next((i for i, m in enumerate(messages) if m.id == prompt_id), None) - return list(messages[at + 1 :]) if at is not None else [] +_PROMPT = "restate_prompt" +"""The key, in each checkpoint's metadata, of the prompt whose run made it: the run's input may be trimmed away.""" +_CARRIED = ("prompt:", "closed:") +"""The prefixes of the ids of the messages that the conversation gave the graph, rather than the graph made.""" -def _answer(messages: Sequence[BaseMessage]) -> str: - """Return the text of the last AI message of a run.""" - replies = [m for m in messages if isinstance(m, AIMessage)] - return str(replies[-1].text) if replies else "" +def _known(conversation: Conversation) -> set[str]: + """Return the ids of the graph's messages that the conversation has: those copied from it.""" + return {i for e in conversation.entries if (i := typing.cast(BaseMessage, e.item).id) is not None} -async def _copied(ctx: SessionContext, conversation: Conversation, running: Entry, run: Sequence[BaseMessage]) -> None: - """Copy the AI and tool messages of ``run`` (the prompt ``running``'s) that the conversation does not have yet. - Only the run's own entries (after its prompt) can hold them, so only those are looked at. Then publish. +def _answer(conversation: Conversation, running: Entry) -> str: + """Return the run's answer: the text of the last AI message it copied. + + Read from the conversation, which is a log, as the graph may have + trimmed its own history (the run's input and answer included). + """ + copied = [ + item + for e in conversation.entries + if e.n > running.n and isinstance(item := typing.cast(BaseMessage, e.item), AIMessage) and item.id is not None + ] + return str(copied[-1].text) if copied else "" + + +async def _copied( + ctx: SessionContext, conversation: Conversation, known: set[str], messages: Sequence[BaseMessage] +) -> None: + """Copy the AI and tool messages of ``messages`` (the graph's state) that the conversation does not have yet. + + ``known`` is the ids the conversation has, which this adds to; the + messages the conversation gave the graph are not copied back. Then + publish. """ - known = {typing.cast(BaseMessage, e.item).id for e in conversation.entries if e.n > running.n} - new = [m for m in run if isinstance(m, AIMessage | ToolMessage) and (m.id is None or m.id not in known)] + new = [ + m + for m in messages + if isinstance(m, AIMessage | ToolMessage) + and m.id is not None + and m.id not in known + and not m.id.startswith(_CARRIED) + ] for message in new: if isinstance(message, AIMessage): await conversation.append_reply(message) else: await conversation.append(message) + known.add(typing.cast(str, message.id)) if new: await stepped(ctx) @@ -119,21 +147,24 @@ def _closing(conversation: Conversation, running: Entry, messages: Sequence[Base That is what closed the run, if it was stopped, refused or failed: the results of its cut-off tool calls and the note that ends it, which only the conversation has (copied messages have the graph's ids; these have - none). And the previous prompt itself, if the graph never stored it. Each - gets an id of its own, so that it is carried once. + none). And the previous prompt itself, if the graph never stored it: if + its run copied nothing (a run that copied a message had stored its + prompt). A prompt the graph stored and then removed (trimming its + history, say) is not brought back. Each gets an id of its own, so that + it is carried once. """ earlier = [e for e in conversation.entries if e.prompt is not None and e.n < running.n] if not earlier: return [] previous = earlier[-1] known = {m.id for m in messages} + between = [e for e in conversation.entries if previous.n <= e.n < running.n] + copied = any(e.prompt is None and typing.cast(BaseMessage, e.item).id is not None for e in between) carried: list[BaseMessage] = [] - for entry in conversation.entries: - if not previous.n <= entry.n < running.n: - continue + for entry in between: item = typing.cast(BaseMessage, entry.item) if entry.prompt is not None: - if _prompt_id(entry) not in known: + if not copied and _prompt_id(entry) not in known: carried.append(HumanMessage(str(item.text), id=_prompt_id(entry))) elif item.id is None and f"closed:{entry.n}" not in known: carried.append(item.model_copy(update={"id": f"closed:{entry.n}"})) @@ -146,23 +177,25 @@ def run_of(ctx: SessionContext, graph: AnyGraph) -> Run: async def run(conversation: Conversation) -> str: running = conversation.running() assert running is not None and running.prompt is not None - # The prompt's id is the conversation's, so that a retried attempt finds it in the graph's state. prompt_id = _prompt_id(running) app = _compiled(graph, ctx) + # Each checkpoint of the run names its prompt, so that a retried attempt finds the run it cut off. + config: RunnableConfig = {**THREAD, "metadata": {_PROMPT: prompt_id}} snapshot = await app.aget_state(THREAD) messages = typing.cast(list[BaseMessage], (snapshot.values or {}).get("messages", [])) - asked = any(m.id == prompt_id for m in messages) + known = _known(conversation) + asked: bool = typing.cast(dict[str, Any], snapshot.metadata or {}).get(_PROMPT) == prompt_id if asked and not snapshot.next: # The run had finished, but its answer was not stored: that is the answer, rather than run again. - await _copied(ctx, conversation, running, _after(messages, prompt_id)) - return _answer(_after(messages, prompt_id)) + await _copied(ctx, conversation, known, messages) + return _answer(conversation, running) # A run cut off by a crash is resumed. A new message (or one after a stopped run) starts a run, after what # closed the previous run. prompt = HumanMessage(str(typing.cast(BaseMessage, running.item).text), id=prompt_id) start = None if asked else {"messages": [*_closing(conversation, running, messages), prompt]} - async for state in app.astream(start, THREAD, context=ctx, durability="sync", stream_mode="values"): + async for state in app.astream(start, config, context=ctx, durability="sync", stream_mode="values"): messages = typing.cast(list[BaseMessage], typing.cast(dict[str, Any], state).get("messages", [])) - await _copied(ctx, conversation, running, _after(messages, prompt_id)) - return _answer(_after(messages, prompt_id)) + await _copied(ctx, conversation, known, messages) + return _answer(conversation, running) return run diff --git a/tests/agents_langgraph_graph.py b/tests/agents_langgraph_graph.py index eed7eba..6751b27 100644 --- a/tests/agents_langgraph_graph.py +++ b/tests/agents_langgraph_graph.py @@ -27,7 +27,7 @@ pytest.importorskip("langgraph") -from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, ToolMessage # noqa: E402 +from langchain_core.messages import AIMessage, BaseMessage, HumanMessage, RemoveMessage, ToolMessage # noqa: E402 from langchain_core.tools import tool # noqa: E402 from langgraph.graph import END, START, MessagesState, StateGraph # noqa: E402 from langgraph.prebuilt import ToolNode # noqa: E402 @@ -54,7 +54,8 @@ def _asked(state: Desk) -> tuple[str, str]: def route(state: Desk) -> str: word, _ = _asked(state) routes = {"refund": "approve", "stop": "approve", "crash": "first", "plan": "plan", "tool": "call_tool"} - return routes.get(word, "history" if word == "history" else "greet") + routes |= {"history": "history", "trim": "trim", "humans": "humans"} + return routes.get(word, "greet") def greet(state: Desk) -> dict[str, Any]: @@ -109,21 +110,36 @@ def history(state: Desk) -> dict[str, Any]: return {"messages": [AIMessage(" | ".join(seen))]} +def trim(state: Desk) -> dict[str, Any]: + """Trim the history to the answer alone, as a summarizing graph does.""" + removed = [RemoveMessage(id=m.id) for m in state["messages"] if m.id is not None] + return {"messages": [*removed, AIMessage("trimmed")]} + + +def humans(state: Desk) -> dict[str, Any]: + """The messages the graph has from the user.""" + said = [str(m.text) for m in state["messages"] if isinstance(m, HumanMessage)] + return {"messages": [AIMessage(" | ".join(said))]} + + def plan(state: Desk) -> dict[str, Any]: return {"messages": [AIMessage(f'{{"title": "{_asked(state)[1]}"}}')]} builder = StateGraph(Desk) steps = {"greet": greet, "confirm": confirm, "first": first, "second": second, "plan": plan, "history": history} +steps |= {"trim": trim, "humans": humans} for name, node in {**steps, "call_tool": call_tool}.items(): builder.add_node(name, typing.cast(Any, node)) builder.add_node("tools", ToolNode([blocking])) builder.add_node("approve", nodes.wait_for_signal("approval", output="approval")) -builder.add_conditional_edges(START, route, ["approve", "first", "plan", "greet", "history", "call_tool"]) +builder.add_conditional_edges( + START, route, ["approve", "first", "plan", "greet", "history", "call_tool", "trim", "humans"] +) builder.add_edge("call_tool", "tools") builder.add_edge("approve", "confirm") builder.add_edge("first", "second") -for last in ("greet", "confirm", "second", "plan", "history", "tools"): +for last in ("greet", "confirm", "second", "plan", "history", "tools", "trim", "humans"): builder.add_edge(last, END) @@ -263,6 +279,12 @@ async def test_a_handler_of_your_own_prompts_the_graph(harness: HarnessEnvironme ] +async def test_a_message_the_graph_trimmed_is_not_brought_back(harness: HarnessEnvironment): + key = uuid.uuid4().hex + assert (await chat(harness, "graph_desk", key, f"trim {key}")).json() == "trimmed" + assert (await chat(harness, "graph_desk", key, "humans")).json() == "humans" + + async def test_output_validates_the_last_ai_message(harness: HarnessEnvironment): key = uuid.uuid4().hex assert (await chat(harness, "graph_planner", key, f"plan {key}")).json() == {"title": key} From 669b763b1ff1c5eb620d8671a9b78110cc2b07bd Mon Sep 17 00:00:00 2001 From: igalshilman Date: Fri, 6 Nov 2026 20:42:53 +0100 Subject: [PATCH 9/9] Checkpointer test: c may fail before b starts The test assumed that b ran in the first attempt, beside c, which fails. When c fails before b starts, b runs in the retry instead (seen on CI, Python 3.13). Either way a and b run once, c twice, and d last, which is what the test now checks. --- tests/agents_langgraph_checkpointer.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/agents_langgraph_checkpointer.py b/tests/agents_langgraph_checkpointer.py index e6b7015..c06db48 100644 --- a/tests/agents_langgraph_checkpointer.py +++ b/tests/agents_langgraph_checkpointer.py @@ -102,8 +102,10 @@ async def go(ctx: sessions.SessionContext) -> None: with pytest.raises(RuntimeError, match="c crashed"): await attempt(backend, go) await attempt(backend, go) - # a ran once; b's write, stored while c failed, is not made again; c runs again. - assert ran == ["a", "b", "c", "c", "d"] or ran == ["a", "c", "b", "c", "d"] + # a ran once; b ran once (its write, stored if it finished before c failed, is not made again; if c failed + # before b started, b runs in the retry); c runs again; d runs last. + assert (ran[0], ran[-1]) == ("a", "d") + assert {step: ran.count(step) for step in "abcd"} == {"a": 1, "b": 1, "c": 2, "d": 1} assert sorted(results[-1]["log"]) == ["a", "b", "c", "d"]