diff --git a/examples-agents/README.md b/examples-agents/README.md index 57a3a9a..dfb1d31 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 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 new file mode 100644 index 0000000..d0a7732 --- /dev/null +++ b/examples-agents/intermediate/refund_graph.py @@ -0,0 +1,110 @@ +""" +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/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. ``python -m restate.agents chat refunds/alice`` +shows the wait for approval as it comes, and the answer. +""" + +import asyncio +from typing import Any, NotRequired + +from langchain_core.messages import AIMessage +from langgraph.graph import END, START, MessagesState, StateGraph +from pydantic import BaseModel + +import restate +from restate.agents import langgraph +from restate.agents.langgraph import nodes + + +class Refund(BaseModel): + order: str + amount: float + + +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 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}." + + +# 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 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", ""))]} + + +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( + "approve", + nodes.wait_for_signal( + "approval", input=lambda s: {"refund": s["refund"], "policy": s["policy"]}, output="approval" + ), +) +graph.add_node( + "pay", + 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") +graph.add_edge("read", "check") +graph.add_edge("check", "approve") +graph.add_conditional_edges("approve", approved, ["pay", "answer"]) +graph.add_edge("pay", "answer") +graph.add_edge("answer", END) + +# 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, clerk, policy, refunds]))) diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index 5d1707b..2d84b6f 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 @@ -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" @@ -408,19 +408,19 @@ 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) + 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) -> 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 + 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 ) 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) + 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 - _untyped(c.ctx).send( - target, *_inputs(handler, value), key=target_key, delay=delay, idempotency_key=c.idempotency_key + 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 ) 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..d678d60 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -28,7 +28,9 @@ 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, 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 @@ -56,5 +58,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..b1b8b99 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) @@ -257,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] = (), @@ -269,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] = (), @@ -280,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] = (), @@ -298,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) @@ -328,6 +345,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 +365,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, a call, 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``, @@ -357,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``. @@ -365,8 +397,13 @@ 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: + _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)) obj = AgentObject(name, **options) @@ -383,3 +420,23 @@ async def answer(ctx: SessionContext, text: str) -> Any: obj.chat(answer, items=ITEMS, output=output) 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: + """Create an agent whose ``chat`` runs ``graph`` (see ``_graph.py``).""" + obj = AgentObject(name, **options) + + async def answer(ctx: SessionContext, text: str) -> Any: + 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 new file mode 100644 index 0000000..e63b24c --- /dev/null +++ b/python/restate/agents/langgraph/_graph.py @@ -0,0 +1,201 @@ +# +# 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. +- 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 + 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 + +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, Entry, 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}) + + +_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 _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} + + +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. + """ + 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) + + +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: 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 between: + item = typing.cast(BaseMessage, entry.item) + if entry.prompt is not None: + 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}"})) + 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.""" + + async def run(conversation: Conversation) -> str: + running = conversation.running() + assert running is not None and running.prompt is not None + 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", [])) + 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, 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, 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, known, messages) + return _answer(conversation, running) + + return run diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py new file mode 100644 index 0000000..25be0a9 --- /dev/null +++ b/python/restate/agents/langgraph/nodes.py @@ -0,0 +1,390 @@ +# +# 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, 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 +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")) + +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") + + + @desk.handler() + async def refund(ctx: sessions.SessionContext, refund: Refund) -> State: + 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") + +Each node does its step once. A crash re-runs a node from the checkpoint it +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. +""" + +from __future__ import annotations + +import typing +from collections.abc import Awaitable, Callable +from datetime import timedelta +from typing import Any + +from langgraph.config import get_config +from langgraph.runtime import Runtime +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 +from restate.exceptions import TerminalError +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" +_DURABILITY = "__pregel_durability" + + +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))" + ) + # 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]) + node = namespace.rsplit("|", 1)[-1].split(":", 1)[0] + 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 + 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: + return adapter.validate_python(value) + except ValidationError as error: + raise TerminalError(f"the input of {name} is not valid: {error}") from error + + 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 {} + + +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: + # 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 + ) + 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. 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. + + 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). + + When it runs, the node fails with a terminal error if the input does not + validate. + """ + 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 = 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(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, + *, + 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", "call", "send", "subagent", "wait_for_signal", "workflow"] 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"] diff --git a/tests/agents_langgraph_graph.py b/tests/agents_langgraph_graph.py new file mode 100644 index 0000000..6751b27 --- /dev/null +++ b/tests/agents_langgraph_graph.py @@ -0,0 +1,320 @@ +# +# 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 threading +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, 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 +from pydantic import BaseModel # noqa: E402 + +import restate # 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 + + +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) + routes = {"refund": "approve", "stop": "approve", "crash": "first", "plan": "plan", "tool": "call_tool"} + routes |= {"history": "history", "trim": "trim", "humans": "humans"} + return routes.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}")]} + + +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 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", "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", "trim", "humans"): + builder.add_edge(last, END) + + +class Plan(BaseModel): + title: str + + +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) + + +@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_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_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_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} + + +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/agents_langgraph_nodes.py b/tests/agents_langgraph_nodes.py new file mode 100644 index 0000000..575168a --- /dev/null +++ b/tests/agents_langgraph_nodes.py @@ -0,0 +1,347 @@ +# +# 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 +from pydantic import BaseModel # 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}" + + +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 + + +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 + + +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() +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() + await until(crashed) + return f"found: {question}" + + +def approval_then_shipping() -> Any: + graph = StateGraph(State) + # 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") + 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: + await until(ready) + 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"), + "invoice": billing_graph(12.5), + "bad_invoice": billing_graph("a lot"), + "call_send": call_then_send(), +} + +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, billing, ledger, notices, 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", {"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() == { + "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" + + +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) + + +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 + + +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_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) + 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) diff --git a/tests/examples_agents.py b/tests/examples_agents.py index 5d23079..fe1a430 100644 --- a/tests/examples_agents.py +++ b/tests/examples_agents.py @@ -233,6 +233,44 @@ async def test_intermediate_triage(): assert (await ingress.post(f"{agent}/messages", {})).json()["messages"] == [] +async def test_intermediate_refund_graph(): + """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: + # 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") + + @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: + 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) + 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": 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): """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