LangGraph: prebuilt nodes that are Restate steps, and langgraph.agent(graph=...)

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

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])))

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)

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"]