diff --git a/README.md b/README.md index c434eb9..5dddf69 100644 --- a/README.md +++ b/README.md @@ -43,8 +43,9 @@ app = restate.app(services=[greeter]) ``` `restate.app` takes services, virtual objects and workflows, and anything else that is a -`restate.Servable`: session objects, and agents, which bring along the `signals` object their tools -wait on. Each is served once, however many things bring it along. +`restate.Servable`: session objects, and agents, each of which brings along its own companion +(`.prebuilt.restate`), the object its tools wait on. Each is served once, however many +things bring it along. Serve it with `restate.serve`, which needs the `serve` extra (`pip install "restate-sdk[serve]"`). It listens on `$PORT`, or 9080, with the HTTP/2 that Restate connects with, and diff --git a/docs/agents/README.md b/docs/agents/README.md index d98ba8b..60442b8 100644 --- a/docs/agents/README.md +++ b/docs/agents/README.md @@ -35,9 +35,14 @@ building it, and how to port it. - **A turn whose invocation dies is closed, never finished:** `settle`, or the next message, gives each open call a result and ends the turn as `ended`. -- **The companion** (the `signals` object) is the control plane: the turn - being answered, the tools' waits, and a revision clients follow. Nothing in - it is needed to complete a turn. +- **The companion** is the control plane: the turn being answered, the + tools' waits, and a revision clients follow. Each agent has its own, the + virtual object `.prebuilt.restate`, keyed as the agent (key + `alice` of `assistant` has key `alice` of `assistant.prebuilt.restate`). + It is ingress-private: clients reach it through the agent's views. Nothing + in it is needed to complete a turn. Upgrading from the shared `signals` + object of earlier versions leaves its state behind: see + [The companion is the agent's own](turns.md#the-companion-is-the-agents-own). ```mermaid flowchart LR @@ -50,10 +55,11 @@ flowchart LR T[("Key's database: turns, items, snapshots, call records")] V["views: messages, sync"] -- read --> T X["private: settle, compact, steered"] --> T + W["views: status, stop, signal, steer, watch"] end - H -- "start, close, publish" --> S[("Companion signals, key assistant:alice: turn, waits, steered messages, revision")] + H -- "start, close, publish" --> S[("Companion assistant.prebuilt.restate, key alice: turn, waits, steered messages, revision")] F -- "wait, for a tool or an approval" --> S - C -- "signal resolves a wait" --> S + W -- "ask, signal resolves a wait, steer" --> S S -- "settle a dead turn, and steered for each steered message" --> X ``` @@ -78,10 +84,10 @@ then the same. | Module | What it is | | --- | --- | -| [`_turn.py`](../../python/restate/agents/_turn.py) | the turn's wrapper: start and close in the companion, cancellation, attempts | -| [`_conversation.py`](../../python/restate/agents/_conversation.py) | the conversation: its tables, `Conversation`, `Turn`, `answer_once`, `compact`, `settle` | +| [`_turn.py`](../../python/restate/agents/_turn.py) | the agent's turn: start, a publish after each commit of its steps (and after its answer, for a turn no `close` ends), and close in the companion; cancellation, attempts | +| [`_conversation.py`](../../python/restate/agents/_conversation.py) | the conversation, which only stores (no companion, no publish): its tables, `Conversation`, `Turn`, `answer_once`, `compact`, `settle` | | [`_messages.py`](../../python/restate/agents/_messages.py) | the conversation as clients read it (`messages`, `sync`) | -| [`_signals.py`](../../python/restate/agents/_signals.py) | the companion | +| [`_signals.py`](../../python/restate/agents/_signals.py) | the companion: `make_companion`, one per agent | | [`_steering.py`](../../python/restate/agents/_steering.py) | the messages steered to a turn: `take_steers`, at its steps | | [`_object.py`](../../python/restate/agents/_object.py) | the agent object: `chat`, the views, the private handlers | | [`_calls.py`](../../python/restate/agents/_calls.py) | the call records: a tool's Restate calls, made once | diff --git a/docs/agents/access-patterns.md b/docs/agents/access-patterns.md index 8436626..a147705 100644 --- a/docs/agents/access-patterns.md +++ b/docs/agents/access-patterns.md @@ -61,7 +61,7 @@ at a time, so each get is charged to the right one). - **Probes:** single statements, each alone in a new process. - The agent's defaults: `window=1`, `budget=32_000`, no summaries. -Not covered: the companion (`signals`), the call records' writes (no tool +Not covered: the companion (`.prebuilt.restate`), the call records' writes (no tool here makes Restate calls), snapshots and `compact`, and a real server (see [Open questions](#open-questions)). diff --git a/docs/agents/contexts.md b/docs/agents/contexts.md new file mode 100644 index 0000000..f7ed8e8 --- /dev/null +++ b/docs/agents/contexts.md @@ -0,0 +1,182 @@ +# The agent's contexts: `AgentContext` and `AgentViewContext` + +An agent's handler gets an `AgentContext`, its view an `AgentViewContext`. +They are the session contexts with a name that says which layer they belong +to, and they let a type checker catch a handler reaching up a layer — a +`wait_for_signal` from a plain session object, say — where today only the +running turn finds out. + +They are a **signpost, not a fence**. Nothing new exists at runtime, nothing +is checked at runtime by them, and `NeedsAgent` / `NoCompanion` stay exactly +as they are: the guard that actually stops a wait from a session object. + +## Why + +`restate.agents` sits on `restate.sessions`. The lower tier is a session +object with handlers of its own: `ctx.db`, the Restate tools that need +nothing more (`tool`, `call`, `send`, `workflow`, `subagent`, and the call +records of `_calls.py`), and a framework driven by hand — pi-durable over +`openSqlKvStorage(ctx.db)` in TypeScript +([`packages/template/src/16_pi_durable.ts`](https://github.com/restatedev/sdk-sessions)), +`withPi` for more than a prompt. The upper tier is an agent: a conversation +of turns, the views, steering, signals, and the companion +(`.prebuilt.restate`) that holds them. + +Both tiers hand a handler the same `SessionContext`, so the agent's own APIs +accept a context from either. Reaching up a layer is only found out when it +runs, and for a wait that means after a model round-trip, in the middle of a +turn, looking like a tool failure. + +Naming the contexts puts the boundary in the signature, where the person +writing the handler reads it. + +## What they are + +| | Python | TypeScript | +| --- | --- | --- | +| The handler's | `AgentContext(SessionContext, Protocol)` | `interface AgentContext extends SessionContext` | +| The view's | `AgentViewContext(ViewContext, Protocol)` | `interface AgentViewContext extends ViewContext` | + +**Python.** The session contexts are `typing.Protocol`s +(`restate/sessions/_api.py`), so the agent's are protocols that extend them +with one private marker member, declared and never present at runtime: + +```python +class AgentContext(SessionContext, Protocol): + """The context of an agent's handler: a session context, from an object that has a companion.""" + + @property + def _agent(self) -> None: + """A marker, never present at runtime and never read: it is what makes a plain session context + fail to be an ``AgentContext`` for a type checker.""" +``` + +The marker is a fiction, and the rules that keep it honest are: + +- **Nothing reads it.** No `ctx._agent`, no `getattr`, no `hasattr`, no + `isinstance` (a protocol with a non-method member cannot be + `runtime_checkable` anyway). +- **Nothing constructs it.** `restate/sessions/` does not change at all: + `_context_of` keeps building `AttemptContext` and `AttemptViewContext`, + and the agent's handlers receive those, as now. +- **No cast is needed at the boundary.** `SessionObject.handler` and `.view` + are bound to `Callable[..., Awaitable[Any]]`, so a function annotated + `AgentContext` registers with no complaint, and the runtime passes what it + always passed. + +**TypeScript.** Plain `extends` with no extra member and no brand. Two +interfaces that are structurally identical are mutually assignable, so this +does not stop anything — it names the layer in the signature, which is what +it is for. + +## What changes + +### Python + +1. **Declare** `AgentContext` and `AgentViewContext` in the agents package + (`python/restate/agents/_api.py`, beside the other public shapes), export + them from `restate.agents`, and name them in + `python/restate/agents/internal.py` for integrations. +2. **Narrow the agent's own surface** — what an agent registers, and what + cannot run without the companion: + + | Where | What | + | --- | --- | + | `_object.py` | `chat()`'s `answer`, the built-in views, `_stopped`, and the private handlers (`erase`, `settle`, `compact`, `steered`, `clear`, `sync`, `messages`, `steer`) | + | `_turn.py` | `chat_turn`, `publish`, `publish_steps`, `publish_end`, `cleared` | + +3. **Move the one strict bound**: `A = TypeVar("A", bound=Callable[[SessionContext, str], Awaitable[Any]])` + in `agents/_object.py` becomes `AgentContext`. Parameter types are + contravariant, so this is the one place where a user's `answer` annotated + `AgentContext` is rejected today. +4. **Leave the shared tier on `SessionContext`**: `tool()`, `call`, `send`, + `workflow`, `subagent`, `run_tool`, everything in `_calls.py`, and + **`ToolCall.ctx`**. A tool is a value, passed to a framework, which may be + run from either tier — the type cannot reach through it, and + `wait_for_signal`'s tool keeps `NeedsAgent` as its guard. +5. **Keep `waits_in(ctx: SessionContext, ...)`** as it is: it is the + backstop, and it must go on taking the wider type to raise `NeedsAgent` + rather than fail to compile. + +### What stays wide, and why + +A first draft of this note also listed the conversation, the turn's helpers, +the signal waits and the langgraph nodes. None of them can narrow, and the +reason is the same each time: **they are the lower tier's entry points, or +the probes that find out which tier they are on.** + +- **`_conversation.py` — `answer_once`, `settle`, `compact`.** A handler of a + plain session object drives them: `tests/agents_conversation.py` registers + `answer_once` on a `sessions.SessionObject`, and each framework's public + `prompt(ctx: SessionContext, ...)` — the documented "for a handler of your + own" — funnels into it. `settings_of` returns "the defaults, for a session + object that is no agent" by design. +- **`_turn.py` — `settings_of`, `stop_signal_of`, `inbox_of`, `settled`, + `answering`, `tool_call_started`.** `answer_once` calls them with a + session context, so they cascade from the point above. +- **`_tools.py` — `wait_for`**, and **`_steering.py` — `take_steers`**. Both + are reached from a framework's `prompt()`. `take_steers` is written to be + safe on a plain context: it returns `[]` before touching the companion + unless it proves it is in a chat turn. +- **`_signals.py` — `companion_at`, `companion_in`, `deliver`.** + `companion_at` is the probe (`"or None for another session object's"`), + `waits_in` calls it, and `deliver` takes any Restate handler's context on + purpose. +- **`langgraph/nodes.py`, `langgraph/_graph.py`.** The signal node takes its + context from the graph runtime as `SessionContext | None`. Narrowing it + would need a cast asserting "this is an agent" at the very spot whose job + is to find out that it is not. + +So the narrowed surface is the agent's registration surface, and the guard +for everything reached through a framework stays `NeedsAgent`. That is the +honest reach of a type that is only a signpost. + +### TypeScript + +The same split, and the same line: declare the two interfaces in +`packages/agents/src/common.ts`, export them from the package, and type the +agents' own surface with them — `sessionObject()`'s and `agent()`'s handlers +and views, `chatTurn`, `publish`, `tookSteers`, and the object's private +handlers. The tools (`tool`, `call`, `send`, `workflow`, `subagent`) and +`ToolCall`/`ToolContext` stay on `SessionContext`, as in Python. + +TypeScript could narrow more than Python can, since nothing there is +structurally barred: `waitFor`, `companionHere` and `takeSteers` would +compile either way. They stay wide so that the two SDKs say the same thing +about the same function, and so that `companionHere` stays what `waits_in` +is — the raiser of `NoCompanion`, which needs the wider type to be reached +at all. + +One asymmetry worth knowing: in TypeScript every object from +`@restatedev/agents`' `sessionObject()` carries a companion, including a +hand-driven one such as `packages/template/src/16_pi_durable.ts`. So +`AgentContext` on its handlers is accurate there in a way that Python's +`SessionObject` / `AgentObject` split is not. + +## What must hold + +- `restate/sessions/` and `packages/sessions/` are untouched: `git diff` + shows no file under either. +- A handler annotated `SessionContext` still registers on an agent and still + type-checks: a function taking the supertype satisfies a callback given + the subtype. No user code breaks — the tutorial in modal-template annotates + `sessions.SessionContext` throughout. +- A handler annotated `AgentContext` registers on a plain `SessionObject` + too: nothing prevents it, and nothing should. The protocol says what the + handler expects, not what the object is. +- `pyright` and `mypy` pass (both are dev dependencies, with `[tool.pyright]` + in `pyproject.toml`), and `tsc -p .` passes in sdk-sessions. +- `NeedsAgent` and `NoCompanion` keep their tests: the wait from a plain + session object still fails for good, with the same message. +- No new runtime behaviour: no class, no attribute, no check. A diff that + adds one has gone wrong. + +## Non-goals + +- Signals for session objects. A session has no awakeable to park on, no + `pending` to list a wait, and no `signal` view to resolve it, and this does + not change that. +- Making the marker real. It stays a fiction; the runtime guard is + `NeedsAgent`. +- A brand in TypeScript. Structural assignability is accepted: the names + guide the reader. diff --git a/docs/agents/conversation.md b/docs/agents/conversation.md index c620565..4d959e8 100644 --- a/docs/agents/conversation.md +++ b/docs/agents/conversation.md @@ -10,6 +10,20 @@ The code is [`_conversation.py`](../../python/restate/agents/_conversation.py) view) and [`_calls.py`](../../python/restate/agents/_calls.py) (the call records). +The conversation only stores. It knows nothing of the companion: an agent's +handlers (its `chat` turn, and handlers you add) publish after each write of +steps that commits, by setting what follows a commit (`after_steps`, through +`publish_steps`), which `answer_once` hands to +`Conversation.hold` as it is ([Who does what](turns.md#who-does-what)). A +handler you add publishes once more after the commit that stored the turn's +answer, by setting what follows it (`after_answer`, through `publish_end`), +which `answer_once` runs itself; a `chat` turn sets nothing there, and ends +with the companion's `close`. A +handler of your own on a plain `SessionObject` that prompts stores its turn +the same way, and publishes nothing. Writing the held rows before a call +record or a Restate call (`holding`, `flush_held`) is the conversation's own +durability, and applies to both. + ## Conversation and turns - **A conversation is a key's database.** Each key of an agent (`alice` of @@ -175,7 +189,7 @@ CREATE TABLE restate_agent_steers ( -- the messages steered to a turn t | Where | What | Written by | | --- | --- | --- | | `restate_agent_calls` (schema `restate-agents.calls`), in the key's database | one row: the invocation's Restate calls made by tools, as JSON ([Tool calls, once](#tool-calls-once-the-call-records)) | the tools built on Restate (`call`, `send`, `workflow`, `subagent`, `wait_for_signal`) and `wait_for` | -| the companion's state: the `signals` object, keyed `:` | `turn`: the turn being answered (`Started`: its invocation, since when, whether the agent `settles`); `waiting`: the tools' waits, each with its awakeable; `kept`: signals that came before their wait; `revision`; `watching`: the `watch` calls waiting for a newer revision | the companion's handlers (`start`, `close`, `wait`, `signal`, `publish`, `watch`) | +| the companion's state: the agent's own object `.prebuilt.restate`, keyed as the agent | `turn`: the turn being answered (`Started`: its invocation, since when, whether it takes steered messages and whether `start` publishes it); `waiting`: the tools' waits, each with its awakeable; `kept`: signals that came before their wait; `revision`; `watching`: the `watch` calls waiting for a newer revision | the companion's handlers (`start`, `close`, `wait`, `signal`, `publish`, `watch`) | | `restate_langgraph_checkpoints`, `restate_langgraph_channel_values`, `restate_langgraph_writes` (schema `restate_agents.langgraph`) | a LangGraph graph's checkpoints, channel values by version, and pending writes, per thread and namespace | `SqlCheckpointSaver`, for `langgraph.agent(graph=...)`; pruned to what the latest checkpoint of each namespace needs after each run | | `restate_pydantic_ai_plan`, `restate_pydantic_ai_spills`, `restate_pydantic_ai_journal` (schema `restate_agents.pydantic_ai`) | the harness's plan (`plan_store()`), spilled tool results (`overflow_store()`), and each tool call's journal of plan reads and writes, for the current invocation | the Pydantic AI stores | | `restate_pydantic_ai_memory`, `..._receipts`, `..._generation` (schema `restate_agents.pydantic_ai.memory`), in a memory object's key | the harness's memory files, the receipts of the changes, and the version counter: one key per namespace (a user), shared by that user's conversations | `memory()`'s session object | @@ -524,8 +538,10 @@ invocation (the next invocation's first call replaces it): again under the same id; a new call with the same input (the model asking again after the first's answer) waits anew, and never gets an earlier call's approval or denial. -- **Forgetting.** At each publish (after each write of steps that commits) - and before each tool call, a record whose Restate call returned is +- **Forgetting.** After each write of steps that commits (once, before the + publish an agent's turn makes there, and in a plain session object's turn + too, which publishes nothing) and before each tool call, a record whose + Restate call returned is forgotten once the transcript has its result under the latest id it was made with: a success, or the very error the call returned. Another error in the transcript (one the framework made up for a diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index a59d026..f10b1c2 100644 --- a/docs/agents/frameworks.md +++ b/docs/agents/frameworks.md @@ -232,7 +232,7 @@ sequenceDiagram autonumber participant M as Model participant G as _Restate - participant S as Companion (signals) + participant S as Companion participant C as Client M->>G: call delete_account(id=7), which requires approval G->>S: wait for the call's id diff --git a/docs/agents/porting.md b/docs/agents/porting.md index 372b966..c566e2c 100644 --- a/docs/agents/porting.md +++ b/docs/agents/porting.md @@ -9,14 +9,16 @@ against sdk-sessions at `6d5e9a2` (after v0.3.0). ## Where the two SDKs stand Both SDKs share the turn's outer life: a key answers one message at a time -(`chat` is exclusive), a turn starts and closes in the companion (`signals`, -keyed `:`, with the same handler names), `stop` cancels the +(`chat` is exclusive), a turn starts and closes in the companion (with the +same handler names), `stop` cancels the invocation, `status`, `pending`, `messages`, `sync` and `watch` mean the same, and a tool's Restate calls are made once (`calls.ts` and `_calls.py` are the -same design). They differ in who owns the conversation: +same design). They differ in whose companion it is, and in who owns the +conversation: | | sdk-sessions (TypeScript) today | This SDK (Python) | | --- | --- | --- | +| The companion | `signals`, shared by every agent, keyed `:` | the agent's own, `.prebuilt.restate`, keyed as the agent, ingress-private ([The companion is the agent's own](turns.md#the-companion-is-the-agents-own)) | | The conversation | each framework's own tables: Claude's `restate_claude_messages` and `restate_claude_answers`, DeepSeek's sessions and events, opencode's and pi's stores, a LangGraph checkpointer | one store for every framework: turn rows and items ([The conversation](conversation.md)) | | A turn in storage | an answer per prompt id, in the framework's tables | a row per turn, with a status, an answer, and what it carries | | What clients read | `restate_agent_log` (`log.ts`), kept in step with the framework's transcript at each publish (`transcribing`) | the turn rows and items themselves: the full and summary views | @@ -30,7 +32,7 @@ The map of the code: | Concern | Python | TypeScript today | | --- | --- | --- | | The turn's wrapper: start, close, cancellation, attempt end | `_turn.py` (`chat_turn`, `answering`) | `common.ts` (`chatHandler`) | -| The companion | `_signals.py` | `signals.ts` | +| The companion | `_signals.py` (`make_companion`, one per agent) | `signals.ts` (one, shared) | | The views and private handlers | `_object.py` (`settle`, `compact`, `erase`) | `object.ts` | | The conversation | `_conversation.py` | — (each framework's store) | | What clients read | `_messages.py` (from the turns) | `log.ts`, `Transcript` in `common.ts` | @@ -61,20 +63,20 @@ Each step stands on its own, with its tests, and the earlier ones are small. (`Started.publishes`, a followed turn) is published as `start` starts it, and `start` returns what `publish` returns (`{revision, steering?}`): the turn makes no publish of its own to start. A turn that - doesn't say so gets nothing back, and is not published: so a TypeScript - turn that this SDK's companion holds publishes its start as before, and - this SDK's turn publishes its own start when the TypeScript companion - returns nothing. `publish` (and `start`) clear `watching` only when a - watch waits. -2. **The companion's `settle`.** A turn's start says whether its agent - settles (`Started.settles`); when `awaitTurn` finds the invocation gone, - the companion closes the turn and sends such an agent `settle(turn)`, a - private handler that closes the turn in the conversation - ([Dead turns](turns.md#dead-turns)). The handler names stay shared, and an - agent that doesn't say it settles is sent nothing, so the two SDKs' agents - can share a deployment's companion while one has it and the other - doesn't. Port it with the conversation (step 3), since `settle` closes the - turn there. + doesn't say so (a background agent's) gets nothing back, and is not + published. `publish` (and `start`) clear `watching` only when a watch + waits. +2. **The agent's own companion, and its `settle`.** Each agent makes its + own companion, `.prebuilt.restate` (Restate refuses names + that start with `restate`), keyed as the agent, ingress-private, which + knows its agent rather than parsing it from its key. Since the two are + always the same code, the companion sends every turn whose invocation is + gone to its agent's `settle(turn)`, a private handler that closes the turn + in the conversation ([Dead turns](turns.md#dead-turns)); this SDK has no + `Started.settles` flag any more. Port the companion with the + conversation (step 3), since `settle` closes the turn there. Until then, + the two SDKs' companions are different objects: they share handler names + and payloads, not state. 3. **The conversation store.** The tables, `Conversation`, `Turn` and `answerOnce`, as [The conversation](conversation.md) and [The turn model](turns.md#the-turn-model) describe them, with an `Items` @@ -172,7 +174,7 @@ tests. Test names are this SDK's. | The turn model: scope, ownership, the window rule by status, asking again, a turn's end, bounded loading, summaries, dead turns and `settle`, the summary view | `tests/agents_conversation.py`, `tests/agents_messages.py`, `tests/agents_object.py` | | Snapshots, `asked_at`, `wait_for` (a wait per tool call, approvals not outcomes, an approval and its tool's call apart) | `tests/agents_snapshots.py` | | Call records: made once, rejoined under a new id, kept until stored, forgotten after | `tests/agents_tools.py` (`test_a_returned_call_is_kept_until_its_result_is_stored_and_forgotten_once_it_is`) | -| The companion: start, close, `awaitTurn`, the `settle` gate, waits and signals | `tests/agents_signals.py`, `tests/agents_object.py` | +| The companion: start, close, `awaitTurn`, `settle`, waits and signals; one per agent, keyed as it, ingress-private | `tests/agents_signals.py`, `tests/agents_companion.py`, `tests/agents_object.py` | | Per framework: answer and the full view after it, a crash in a tool, a stored answer not asked again, refused, stopped, `output=`, several prompts | each framework's `tests/agents_*.py` | | The LangGraph checkpointer: resume, history, prune, `DeltaChannel`, many channels | `tests/agents_langgraph_checkpointer.py` | diff --git a/docs/agents/prebuilt.md b/docs/agents/prebuilt.md new file mode 100644 index 0000000..e73d1eb --- /dev/null +++ b/docs/agents/prebuilt.md @@ -0,0 +1,180 @@ +# The agent's prebuilt object: `.prebuilt.restate` + +Every agent brings one object of its own, `.prebuilt.restate`, keyed +as the agent. It replaces `signals`, the single object both SDKs ship today, +and it is the agent's alone: a plain session object has none, and needs none. + +This note is the design for both SDKs. It was written against sdk-python +`c56a55b` and sdk-sessions `3b451c4` (the branches), with sdk-python +`4730b10` and sdk-sessions `5871ed4` (`main`) as the shipped state. + +## Where the two SDKs stand + +| | The object | Its key | Where | +| --- | --- | --- | --- | +| Shipped (both `main`s) | `signals`, one for the whole endpoint | `:` (`assistant:alice`) | sdk-python `_signals.py` at `4730b10`; sdk-sessions `packages/agents/src/signals.ts` | +| The branches, before this | `.agentprebuilt.restate`, one per agent | the agent's own key (`alice`) | sdk-python `_signals.py`; sdk-sessions `packages/agents/src/companion.ts` | +| **Decided here** | **`.prebuilt.restate`, one per agent** | the agent's own key | the same files, renamed | + +The move from one shared object to one per agent is already made on both +branches; what is left is the name, and taking the word `signals` out of the +code with it. + +## Why the shared `signals` has to go + +1. **Two endpoints, one name.** A Python endpoint and a TypeScript endpoint + registered with the same Restate cluster each register a service called + `signals`. There is one service of that name in a cluster, so the second + registration takes it over: whichever SDK registered last serves the + companion for both SDKs' agents, with its own handler set and its own + payload details. Nothing warns about it. This is the collision that + started this note, and no per-endpoint fix removes it — only a name that + cannot be shared does. +2. **Serving a subset shares the name too.** One agent per image + (sdk-sessions' `RESTATE_SERVE_ONLY`, or a `restate.app([one])` of its + own in Python) means every image also serves `signals`, and every one of + them claims the same service: the same clash inside a single language. +3. **It sits in the user's namespace.** `signals` is an ordinary name that a + user may want for a service of their own. `.prebuilt.restate` + cannot be mistaken for one, and cannot be taken by accident. +4. **The composite key hides the agent.** `:` means one state + table for every agent in the endpoint, and every query, cleanup or + inspection has to parse keys apart. With one object per agent, the key is + the agent's own key, and `service_name` tells the agent. +5. **Deployment follows the agent.** A companion served with its agent is + versioned and rolled out with it; a shared one is a single service that + every agent's deployment fights over. + +## The name + +`.prebuilt.restate`: the agent's name, then the marker. Restate +refuses service names that *start* with `restate`, so the marker is a +suffix and ends with it. + +- **Not `.agentprebuilt.restate`** (what the branches carried): `agent` says + nothing the agent's own name does not already say, and one marker is + enough. Shorter reads better in a service list, in a state query and in an + error message. +- **One prebuilt object per name.** The suffix is the agent's single reserved + name. If anything else ever needs to live beside an agent, it goes *inside* + this object rather than beside it under a second suffix. +- **The name is derived, never configurable**: `companion_name(agent)` / + `companionName(agent)`, from the agent's name alone, in both SDKs. That is + what keeps the two SDKs' names identical for the same agent. +- **Collisions.** An agent named `a` and an agent named `a.prebuilt` do not + collide (`a.prebuilt.restate` against `a.prebuilt.prebuilt.restate`). An + agent named `support.prebuilt.restate` would own agent `support`'s + companion name, and **nothing refuses it**: naming an agent that is a + deliberate act, and a guard was weighed and left out. Within one endpoint + `Endpoint.bind` still catches it (`Service … already exists`); across + endpoints, or across the two SDKs, it would take the companion over, which + is the one case this name does not close. + +## Only agents have one + +A plain `SessionObject` gets no prebuilt object. Sessions have no signals, no +turn and no revision — the database is a session's only state — and nothing +in `restate.sessions` or `@restatedev/sessions` changes here. A +`wait_for_signal` from a plain session object's handler keeps failing for +good, with `NeedsAgent` (Python) and `NoCompanion` (TypeScript), which both +branches already do (sdk-python `4268641`). + +**The word `signals` goes; the word `signal` stays.** The public vocabulary — +the `signal` view, `pending`, `wait_for_signal`, `waitFor` — is unchanged. +What disappears is the *object* called `signals`, and the module and test +names that follow from it. + +## What the object holds + +Unchanged by this rename, and listed here so the rename can be checked +against it. + +| State | What it is | Written by | +| --- | --- | --- | +| `waiting` | the tools' waits, oldest first, each with its awakeable | `listen`, `resolve`, `forget`, `close` | +| `kept` | signals that arrived with nobody waiting, by name (`keep`) | `resolve`, `listen` | +| `turn` | the message the agent is answering (`Started`) | `start`, `close` | +| `inbox` | messages steered to the turn being answered, taken or not | `steer`, `taken`, `close` | +| `revision`, `watching` | the conversation's revision, and the `watch`es waiting for a newer one | every publish; `subscribe`, `unsubscribe` | + +| Handlers | | +| --- | --- | +| Signals | `wait`, `resolve`, `pending`, `listen`, `forget` | +| The turn | `start`, `close`, `ended`, `awaitTurn` | +| Steering | `steer`, `taken` | +| Notifications | `publish`, `notifications`, `watch`, `subscribe`, `unsubscribe` | +| The rest | `status`, `reset` | + +All of them stay ingress-private: clients reach them through the agent's own +views (`status`, `pending`, `signal`, `steer`, `watch`, `sync`). The handler +names, `awaitTurn` among them, stay as they are — they are the wire the two +SDKs share. + +## No migration + +Nothing carries the companion's state over, and nothing needs to. The +shared `signals` and `.agentprebuilt.restate` are both preview, and +neither per-agent companion has been released: the new object starts empty +for every key, and whatever a running `signals` holds — waits, the turn +being answered, the steering inbox, the revision — is dropped with it. + +The conversation itself is untouched. It lives in the key's SQL database: +messages, turns, statuses and cursors all survive the rename, since the +companion holds none of them. + +## The change, file by file + +The wire name is the only thing that changes in behaviour. In the code, keep +**companion** as the word for the thing — `companion_of` / `companionOf`, +`COMPANION_SUFFIX`, `companion.ts` — and let `prebuilt` be the wire name +alone. + +### sdk-python + +| File | Change | +| --- | --- | +| `python/restate/agents/_signals.py` | `COMPANION_SUFFIX = ".prebuilt.restate"`, and the module docstring; **rename the file to `_companion.py`** (it matches `companion.ts`, and it is the last place called `signals`) | +| `python/restate/agents/internal.py` | the import, and the paragraph naming the companion | +| `python/restate/agents/_object.py`, `servable.py`, `endpoint.py` | the name in the docstrings | +| `docs/agents/README.md`, `turns.md`, `conversation.md`, `porting.md`, `access-patterns.md`, `README.md` (root) | the name, and `porting.md`'s table row, which still describes `signals` as the TypeScript state | +| `tests/agents_companion.py` (11), `tests/agents_signals.py` (4), `agents_object.py`, `agents_steering.py`, `agents_snapshots.py`, `agents_tools.py`, `servable.py` | the expected names; **rename `tests/agents_signals.py` to `tests/agents_companion_server.py`**, the against-a-server twin of `agents_companion.py` | +| `scripts/steering_traffic.py`, `scripts/check_no_extras.py` | the name | + +### sdk-sessions + +| File | Change | +| --- | --- | +| `packages/agents/src/companion.ts` | `COMPANION_SUFFIX = '.prebuilt.restate'`, and the header comment | +| `packages/agents/src/object.ts` | the name in the comment | +| `packages/agents/README.md` (8), `docs/README.md`, `turns.md`, `conversation.md`, `access-patterns.md`, `plan.md` | the name; `plan.md` row D2 records the decision | +| `packages/agents/test/companion.test.ts` (10), `webhook.test.ts`, `tools.test.ts`, `layers.test.ts`, `calls.test.ts`, `steering.test.ts`, `support.ts` | the expected names | +| `packages/e2e/restate-tools.ts` | the `/services` assertion, the `state` query by `service_name`, and the per-handler ingress checks | + +`packages/agents/src/signals.ts` is already gone on the branch, replaced by +`companion.ts`; on `main` it is what this whole note deletes. + +## What must hold afterwards + +- `companion_name("one") == "one.prebuilt.restate"` in both SDKs, from the + agent's name alone, with no normalisation beyond it (`Odd-Name_1` keeps its + spelling). +- An endpoint serving agents `one` and `two` serves exactly `one`, + `one.prebuilt.restate`, `two`, `two.prebuilt.restate`, and no `signals`. +- Serving one agent of several serves it with its own companion and nothing + else: `RESTATE_SERVE_ONLY=two` in sdk-sessions, `restate.app([two])` in + sdk-python (`__restate_services__`). +- Every companion handler is refused by the ingress, by name, in the server + tests of both SDKs. +- The companion's state keys are the agent's keys (`alice`), not + `:` — asserted against `state` in the e2e. +- Both repos' e2e assert that no service named `signals` is registered. + +## Open questions + +1. **The Python file names** (`_companion.py`, `agents_companion_server.py`): + confirm, or keep `_signals.py` as the file name and change only the + strings. +2. **Which repo lands first.** The name must be identical in both before + either is released; sdk-python's branch is the one the TypeScript port + follows (`docs/agents/porting.md`), so it goes first, and sdk-sessions + follows in the same week. diff --git a/docs/agents/testing.md b/docs/agents/testing.md index 0d42873..aee1c1e 100644 --- a/docs/agents/testing.md +++ b/docs/agents/testing.md @@ -40,7 +40,7 @@ run in two tiers. | Tier | How | Used by | | --- | --- | --- | -| In-process | `sessions_object.Host`: an attempt's context over an in-memory `Backend`, with no Restate | `agents_conversation`, `agents_messages`, `agents_tools`, `agents_steering`, and the crash tests that monkeypatch a hook | +| In-process | `sessions_object.Host`: an attempt's context over an in-memory `Backend`, with no Restate | `agents_conversation`, `agents_messages`, `agents_tools`, `agents_steering`, `agents_companion`, and the crash tests that monkeypatch a hook | | Real server | `restate.create_test_harness`: Restate in Docker, with the app on a thread of the test process | each framework's file, `agents_object`, `agents_signals`, `agents_steering_server`, `agents_webhook`, `examples_agents` | Each framework's file runs its agent on a scripted model, which answers from diff --git a/docs/agents/turns.md b/docs/agents/turns.md index 5499abc..f8ca288 100644 --- a/docs/agents/turns.md +++ b/docs/agents/turns.md @@ -11,7 +11,8 @@ The code lives in: - [`_turn.py`](../../python/restate/agents/_turn.py): `chat_turn`, the turn; - [`_conversation.py`](../../python/restate/agents/_conversation.py): the conversation, and `answer_once`, which answers a prompt; -- [`_signals.py`](../../python/restate/agents/_signals.py): the companion; +- [`_signals.py`](../../python/restate/agents/_signals.py): the companion, one per + agent (`make_companion`); - [`_steering.py`](../../python/restate/agents/_steering.py): the messages steered to a turn, which it takes at its steps; - [`_object.py`](../../python/restate/agents/_object.py): the views (`status`, @@ -24,10 +25,59 @@ The code lives in: The TypeScript SDK's agents follow the same turn and companion (`sdk-sessions/packages/agents/TURN-SEMANTICS.md`). How the conversation is -stored and loaded, and the `settle` and `compact` handlers, are this SDK's -(see [Agents of the TypeScript SDK](#agents-of-the-typescript-sdk), and +stored and loaded, the `settle` and `compact` handlers, and each agent +having a companion of its own, are this SDK's (see +[The companion is the agent's own](#the-companion-is-the-agents-own), and [Porting to sdk-sessions](porting.md)). +## Who does what + +The conversation stores: `answer_once` runs a prompt's turn and writes its +steps in the key's database, and nothing else. The agent's turn publishes: +`chat_turn` starts the turn with the companion, has each write of the turn's +steps that commits publish (`publish_steps` sets what follows a commit, +`after_steps`, which `answer_once` runs as an opaque callback), and closes +the turn. A handler of your own that prompts on a plain `SessionObject` +(`pydantic_ai.prompt(ctx, …)`, say) stores its turn and publishes nothing: +it has no companion, and nothing follows it. + +A handler you add to an agent (`@assistant.handler`) that prompts publishes +each step its turns store, as `chat` does, and publishes once more when a +turn's answer is stored: every handler of an `AgentObject` calls +`publish_steps` and `publish_end` (what follows the answer's commit, +`after_answer`, which `answer_once` runs as the other opaque callback). So a +follower of such a turn sees its steps arrive and then sees it end, instead +of hearing nothing more. The end publish comes after the commit that stored +the answer, never before it: a follower told earlier would read a +conversation that has no answer yet. A turn closed for good (failed, +refused, stopped) publishes its end too — its close stores the turn's +answer; a turn whose attempt failed publishes none, as it has not ended and +the next attempt carries it on. A background agent's turn publishes neither +its steps nor its end. + +A `chat` turn publishes no end of its own: `chat_turn` sets `after_answer` +to nothing, because the turn ends with the companion's `close`, which +publishes the turn's last change with it. So the publishes of a `chat` turn +are what they always were: `start`, one per commit of its steps, `close`. + +A turn of another handler is not a turn in the companion, though: no `start` +nor `close`, so `status` does not report it and `stop` does not cancel it. +There is no public way to make it one: `chat_turn` is private (`_turn.py`), +and needs the handler declared with `explicit_cancellation=True` and the +agent's `settle` handler, which only `chat()` adds. For a followed turn, +send the message to the agent's `chat` (`deliver`, or `ctx.generic_send`). + +Waiting for a signal (a `wait_for_signal` tool, `wait_for`, an approval, +LangGraph's `interrupt()` and `nodes.wait_for_signal`) and steering need the +companion, so an agent. From a plain session object every wait fails the +turn once, for good (`NeedsAgent`, a `TerminalError`), with a message that +says to use `.agent(...)` or `AgentObject`, and the turn is +closed `failed`. A signal tool's wait included: unlike another +`TerminalError` a tool raises, it is never handed to the model as the +tool's result, since the model cannot fix how the object is set up (as in +the TypeScript SDK). A `publish` from a plain session object is a +programming error, which raises. + ## Words - **Agent, key.** A session object of `restate.agents`, such as `assistant`, @@ -50,8 +100,12 @@ stored and loaded, and the `settle` and `compact` handlers, are this SDK's - **Attempt.** One run of a turn's invocation. A crash, a lost connection or a suspension ends an attempt; Restate starts another, which runs the handler again from its first line and resumes from what the database has. -- **Companion.** The agent's `signals` object, keyed `:` (such - as `assistant:alice`). It is a journaled virtual object, so it can wait +- **Companion.** The agent's own virtual object, named + `.prebuilt.restate` and keyed as the agent: key `alice` of + `assistant` has the companion `assistant.prebuilt.restate`, key + `alice`. Whatever makes the agent makes it, and the agent serves it + alongside itself. It is ingress-private: clients reach it through the + agent's views. It is a journaled virtual object, so it can wait (awakeables, timers) and be written while the turn runs. It keeps the running turn, the waits of its tools, and the conversation's revision. Nothing in it is needed to complete a turn: it is the control plane. @@ -110,9 +164,9 @@ sequenceDiagram participant C as Client participant A as Agent (chat) participant DB as Key's database - participant S as Companion (signals) + participant S as Companion C->>A: chat("Plan my week") - A->>S: start(turn, settles, publishes) + A->>S: start(turn, steers, publishes) S-)S: awaitTurn(turn) Note over S: status.answering = turn, publish S-->>A: the revision, and the messages steered to the turn, if any @@ -140,14 +194,14 @@ sequenceDiagram 1. **Start.** The turn calls the companion's `start(turn)`: from here, `status` says the agent answers it (`answering: {invocationId, since}`). - It also says the agent has `settle` (see [Dead turns](#dead-turns)). A - retried attempt calls `start` again with the same invocation, which + It also says whether the turn takes steered messages (`steers`), and + whether `start` publishes it (`publishes`: a followed turn, not a + background agent's). A retried attempt calls `start` again with the same invocation, which changes nothing (it keeps when the turn started). `start` sends `awaitTurn(turn)`, the turn's watch. It also publishes, as the turn's first revision, and returns what `publish` returns: the turn makes no - publish of its own to start. A companion that returns nothing from - `start` (an older deployment, or the TypeScript SDK's) gets a `publish` - from the turn instead. + publish of its own to start. A background agent's turn is not published + at its start, and `start` returns nothing. 2. **The prompt.** `answer_once` reads the prompt's answer by its id, one indexed read: a retried attempt of an answered turn returns it, and reads nothing more; one of a turn that was closed (cancelled, failed, refused, or @@ -159,7 +213,8 @@ sequenceDiagram is stored as it comes: a reply with tool calls in one transaction, before its tools run; the tools' results before the model is asked again. A final reply, which no tool waits for, is held and commits with the - answer. After each write that commits, the turn publishes once, and + answer. After each write that commits, the turn (`chat_turn`, not + `answer_once`, see [Who does what](#who-does-what)) publishes once, and waits for it: so a follower sees the calls as their tools start, and their results before the model's next request returns. The companion answers each `publish` with the messages steered to the turn that wait: @@ -230,9 +285,9 @@ sequenceDiagram participant O as Operator participant A as Agent (chat) participant DB as Key's database - participant S as Companion (signals) + participant S as Companion participant T as Agent (settle) - A->>S: start(turn, settles) + A->>S: start(turn) S-)S: awaitTurn(turn) A->>DB: the prompt, a reply with a tool call O->>A: kill (admin API) @@ -275,7 +330,7 @@ sequenceDiagram autonumber participant N as Next turn (chat) participant DB as Key's database - participant S as Companion (signals) + participant S as Companion participant T as Agent (settle) N->>S: start(next turn) Note over S: a turn before it never closed: close it, send settle @@ -652,7 +707,11 @@ shared with Restate's other SDKs). A handler of your own may prompt more than once (`openai.prompt(ctx, …)` twice, say): each prompt is a turn of the conversation, ``, then `:2`. A retried attempt makes them in the same order, so each finds its -own answer, and the companion's turn is the invocation, all of them. +own answer. In an agent's `chat` turn, the companion's turn is the +invocation, all of them. A handler you add to an agent publishes each +step its turns store, as `chat` does, but only `chat` starts and closes a +turn in the companion. A handler of a plain `SessionObject` stores its turns +and publishes none of them. ## What a client sees @@ -681,27 +740,36 @@ calls as their tools start, their results before the model's next request returns, its end, and a killed turn end. A final reply, held until the answer, shows with the turn's end. -## Agents of the TypeScript SDK - -The companion's handler names are shared with the TypeScript SDK, so one -deployment's `signals` may hold the turns of the other's agents. An agent -says it has `settle` as it starts a turn (`Started.settles`, which `status` -does not show), and only such a turn is sent it: - -- a TypeScript agent's turn that this SDK's companion holds is closed in the - companion, and its agent is sent nothing; -- this SDK's agent whose turn the TypeScript companion holds has a dead turn - closed by its next message. - -A turn says it takes steered messages the same way (`Started.steers`): the -companion holds messages only for such a turn, and every other gets them as -turns of their own. And it says whether `start` publishes -(`Started.publishes`): this SDK's companion then publishes and returns what -`publish` returns; a TypeScript companion returns nothing from `start`, and -this SDK's turn then publishes its start itself. A TypeScript agent's turn -does not say so, and this SDK's companion publishes nothing at its `start`. A TypeScript companion, which has no `steer`, answers -`publish` with no `steering`, so this SDK's turns take none from it; and the -`steer` view fails there, as the companion has no such handler. +## The companion is the agent's own + +Each agent has a companion of its own: the virtual object +`.prebuilt.restate`, with the agent's name as it is (Restate +refuses names that start with `restate`, so the suffix ends the name +instead). It is keyed as the agent, and it knows its agent: it sends +`settle` and `steered` to the agent it was made for, at its own key. An +agent and its companion are always the same code in the same deployment, so +the companion sends `settle` for every dead turn, and `start` always returns +what it published for a turn that `publishes`. + +A turn still says whether it takes steered messages (`Started.steers`: a +framework that cannot take them, or a background agent, does not), and +whether `start` publishes it (`Started.publishes`: a background agent's turn +is not published). The handlers' names (`awaitTurn` among them) and payloads +are those of the TypeScript SDK's companion, but nothing is shared with it +any more: the TypeScript SDK's agents keep their `signals` object, which this +SDK's agents no longer serve or call. + +**Upgrading** from a version whose agents shared the `signals` object +(keyed `:`): the new companions start empty, and nothing in +`signals` carries over. Turns in flight, pending waits, steered messages +held for a turn, and the revision are left behind. A turn in flight as you +deploy runs on with the deployment it started on, and with `signals`: the +new `status` and `stop` do not see it, and the new `signal` cannot resolve +its tools' waits (while the old deployment is registered, its `signals` +still takes `resolve` at `/signals/:/resolve`). The revision +starts again, so a follower should start over with a `sync` without one. +Let the agents' turns end before you deploy, and remove the old deployment +once nothing runs on it. ## Guarantees and limits diff --git a/python/restate/agents/__init__.py b/python/restate/agents/__init__.py index ac20ef8..25247cd 100644 --- a/python/restate/agents/__init__.py +++ b/python/restate/agents/__init__.py @@ -45,7 +45,9 @@ async def get_weather(city: str) -> str: This package holds what all frameworks share: -- ``AgentObject``, a session object set up for agents. +- ``AgentObject``, a session object set up for agents, and ``AgentContext`` + and ``AgentViewContext``, which its handlers and views take: the session + contexts under a name that says which layer they belong to. - Tools that work with any framework: ``tool`` for your own functions, and ``call``, ``send``, ``workflow``, ``subagent`` and ``wait_for_signal``, which are built on Restate. @@ -68,6 +70,8 @@ async def get_weather(city: str) -> str: """ from restate.agents._api import ( + AgentContext, + AgentViewContext, Answering, ChatMessage, Messages, @@ -90,8 +94,10 @@ async def get_weather(city: str) -> str: from restate.agents._webhook import Event, webhook, webhook_message __all__ = [ + "AgentContext", "AgentObject", "AgentOptions", + "AgentViewContext", "Answering", "ChatMessage", "Event", diff --git a/python/restate/agents/_api.py b/python/restate/agents/_api.py index 3abfef6..e856142 100644 --- a/python/restate/agents/_api.py +++ b/python/restate/agents/_api.py @@ -8,18 +8,63 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -"""The JSON types that an agent's clients send and receive. - -These cover the conversation's messages and pages of them, the agent's -status, and signals. Some keys are camelCase (``invocationId``, -``timeoutSeconds``): they are wire-format names, shared with Restate's other -SDKs so that the same client can follow an agent written in any of them. +"""The shapes an agent's handlers and clients are written to. + +The contexts come first: ``AgentContext`` and ``AgentViewContext``, which +name the layer a handler or a view belongs to. The rest are the JSON types +an agent's clients send and receive, which cover the conversation's messages +and pages of them, the agent's status, and signals. Some keys are camelCase +(``invocationId``, ``timeoutSeconds``): they are wire-format names, shared +with Restate's other SDKs so that the same client can follow an agent +written in any of them. """ from __future__ import annotations import typing -from typing import Any, Literal, TypedDict +from typing import Any, Literal, Protocol, TypedDict + +from restate.sessions import SessionContext, ViewContext + + +class AgentContext(SessionContext, Protocol): + """The context of an agent's handler: a session context, from an object that has a companion. + + The agent's own surface takes it: the ``answer`` of ``chat()``, the + agent's handlers, and what only an agent does with a context + (``chat_turn``, ``publish``). The layer a handler belongs to is then in + its signature, where the person writing it reads it, rather than found + out in the middle of a turn, after a model round-trip, looking like a + tool failure. It is a signpost, not a fence: nothing of it exists at + runtime, the handler is passed the context a session object always + passes, and ``NeedsAgent`` stays the guard that stops a wait from a + session object. + + The tier below keeps ``SessionContext``: the tools, which are values a + framework may run from either, and the conversation a handler of one's + own prompts with (``answer_once``, and a framework's ``prompt``), which + a plain session object drives too. + """ + + @property + def _agent(self) -> None: + """A marker, never present at runtime and never read: it is what makes a plain session context + fail to be an ``AgentContext`` for a type checker.""" + + +class AgentViewContext(ViewContext, Protocol): + """The context of an agent's view: a view context, from an object that has a companion. + + What ``AgentContext`` is to a handler, for the views that read what the + agent is doing, or tell its companion to (``status``, ``pending``, + ``signal``, ``steer``, ``watch``). Its marker is the same fiction. + """ + + @property + def _agent(self) -> None: + """A marker, never present at runtime and never read: it is what makes a plain view context + fail to be an ``AgentViewContext`` for a type checker.""" + Role = Literal["user", "assistant"] """Who wrote a message.""" diff --git a/python/restate/agents/_calls.py b/python/restate/agents/_calls.py index 0d508aa..80d7f01 100644 --- a/python/restate/agents/_calls.py +++ b/python/restate/agents/_calls.py @@ -26,8 +26,8 @@ ``run_workflow``.) A call is kept until the framework's transcript has its result under the -latest id it was made with (``settle`` checks this at each publish and before -each tool call). An error in the transcript does not count as the result, +latest id it was made with (``settle`` checks this after each write of the +turn's steps that commits, and before each tool call). An error in the transcript does not count as the result, unless the call itself returned that error. A framework's own wait (``wait_for``) is its tool call's alone: the framework diff --git a/python/restate/agents/_common.py b/python/restate/agents/_common.py index e8138b9..1a301fd 100644 --- a/python/restate/agents/_common.py +++ b/python/restate/agents/_common.py @@ -86,6 +86,47 @@ async def flush_held(ctx: SessionContext) -> None: await flush() +_AFTER_STEPS: ScopeKey[Callable[[], Awaitable[None]]] = ScopeKey("what follows each commit of the turn's steps") + + +def after_steps(ctx: SessionContext, then: Callable[[], Awaitable[None]]) -> None: + """Have ``then`` run after each write of the attempt's turn's steps that commits (see ``Conversation.hold``). + + The caller of ``answer_once`` sets it, before the turn runs: an agent's + ``chat`` turn publishes there, for the clients that follow it. The + conversation itself only stores: without it, a commit is followed by + nothing. + """ + _AFTER_STEPS.set(ctx, then) + + +def after_steps_of(ctx: SessionContext) -> Callable[[], Awaitable[None]] | None: + """Return what follows each commit of the attempt's turn's steps (``after_steps``), or ``None``.""" + return _AFTER_STEPS.get(ctx) + + +_AFTER_ANSWER: ScopeKey[Callable[[], Awaitable[None]]] = ScopeKey("what follows the commit of the turn's answer") + + +def after_answer(ctx: SessionContext, then: Callable[[], Awaitable[None]] | None) -> None: + """Have ``then`` run after the commit that stores the attempt's turn's answer, or ``None`` for nothing. + + The caller of ``answer_once`` sets it, before the turn runs: a handler of + an agent publishes there, so that a follower of the turn hears it end, as + it heard each of its steps. A ``chat`` turn sets ``None``: its end is the + companion's ``close``, which publishes the turn's last change with it. + """ + if then is None: + _AFTER_ANSWER.delete(ctx) + else: + _AFTER_ANSWER.set(ctx, then) + + +def after_answer_of(ctx: SessionContext) -> Callable[[], Awaitable[None]] | None: + """Return what follows the commit of the attempt's turn's answer (``after_answer``), or ``None``.""" + return _AFTER_ANSWER.get(ctx) + + MODEL_ENV = "RESTATE_AGENT_MODEL" """The environment variable that names the model agents use when none is given.""" @@ -294,6 +335,16 @@ def __init__(self, tool: str, cause: BaseException) -> None: self.__cause__ = cause +class NeedsAgent(TerminalError): + """A step that only an agent takes (waiting for a signal: it needs the agent's companion), from a handler of a + session object that is no agent. + + It fails the turn for good, and closes it ``failed``. Unlike another + ``TerminalError`` that a tool raises, it is never the tool's result for + the model: the model cannot fix how the object is set up. + """ + + class ErrorSet: """A set of errors, by identity, that holds them for as long as it lives (an attempt, or a run). diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index 62b9ca6..0b4010d 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -106,6 +106,8 @@ UNANSWERED, Check, Window, + after_answer_of, + after_steps_of, check_window, cut_off_result, holding, @@ -114,7 +116,7 @@ refusal_of, stored_answer, ) -from restate.agents._turn import AttemptEnded, answering, settings_of, stepped, stop_signal_of +from restate.agents._turn import AttemptEnded, answering, settings_of, settled, stop_signal_of from restate.exceptions import TerminalError from restate.sessions import Database, DatabaseReader, DatabaseWriter, SessionContext, SqlParam, Transaction from restate.sessions.internal import Schema, definition_of, ensure_schemas @@ -487,8 +489,9 @@ class _Ledger: sees or that outlives the attempt: a publish, a Restate call of the library, the attempt failing. - The run publishes after each write of steps that commits (``published``), - not after the held rows: they are no step that anyone waits for. + What the run's caller asks for after each write of steps that commits + (``committed``: an agent's ``chat`` turn publishes there) follows that + write, not the held rows: they are no step that anyone waits for. """ def __init__(self, db: Database, items: Items, start: Start, transcript: list[Entry], next: int, seq: int) -> None: @@ -519,8 +522,8 @@ def __init__(self, db: Database, items: Items, start: Start, transcript: list[En self.held: list[list[SqlParam]] = [] """The rows of the trajectories given positions but not written yet, in order of position (``_INSERT_STEP``'s parameters): the next write writes them first.""" - self.published: Callable[[], Awaitable[None]] | None = None - """What the run makes after each write of steps that commits (a publish, ``Conversation.hold``), if anything.""" + self.committed: Callable[[], Awaitable[None]] | None = None + """What follows each write of steps that commits (``Conversation.hold``), if anything.""" def holds_back(self, entries: Sequence[Entry]) -> bool: """Return whether ``entries``, which are read, are held rather than written: no tool call nor result.""" @@ -1018,12 +1021,12 @@ def _check_writable(self) -> None: async def _stored(self, made: Callable[[int], list[Entry]]) -> int: """Store the entries that ``made`` makes (``_insert``), and return their position; then, if they committed - (they are not held), make what the run makes after a commit (``_Ledger.published``). + (they are not held), run what follows a commit (``_Ledger.committed``). - That is a publish, outside the lock, awaited before the caller goes on: - the companion returns the messages steered to the turn from it, which - the caller's next step takes. Writes that run at once each make their - own. + It runs outside the lock, awaited before the caller goes on: for an + agent's ``chat`` turn it is a publish, which also learns of the + messages steered to the turn, for the caller's next step to take. + Writes that run at once each run it. """ added: list[Entry] = [] @@ -1032,9 +1035,9 @@ def making(n: int) -> list[Entry]: return added n = await self._insert(making) - published = self._ledger.published - if published is not None and not self._ledger.holds_back(added): - await published() + committed = self._ledger.committed + if committed is not None and not self._ledger.holds_back(added): + await committed() return n @_shielded @@ -1451,7 +1454,7 @@ def compaction_asked(self) -> bool: """Whether the run asked for the history to be compacted after the turn (``compact_after``).""" return self._ledger.compaction - def hold(self, published: Callable[[], Awaitable[None]] | None = None) -> None: + def hold(self, committed: Callable[[], Awaitable[None]] | None = None) -> None: """Hold the steps that no tool waits for, rather than write each at once: the next write writes them. For a run that writes them (``flush``) before each publish and Restate @@ -1459,13 +1462,14 @@ def hold(self, published: Callable[[], Awaitable[None]] | None = None) -> None: A conversation that is written otherwise (step by step, from handlers of its own) writes each step as it comes. - ``published``, if given, is awaited after each write of steps that + ``committed``, if given, is awaited after each write of steps that commits (a reply with tool calls, the tools' results, with the rows - held before them): ``answer_once`` publishes there. Held rows make - none, nor does ``flush``. + held before them): ``answer_once`` passes what its caller set + (``after_steps``), where an agent's ``chat`` turn publishes. Held rows + run none, nor does ``flush``. """ self._ledger.holds = True - self._ledger.published = published + self._ledger.committed = committed @_shielded async def flush(self) -> None: @@ -1665,14 +1669,36 @@ async def answer_once( # With the tool calls' (``_calls.py``), which the turn reads as it ends, whether or not a tool ran: a key's # first turn applies both in one commit. await ensure_schemas(ctx.db, [SCHEMA, CALLS_SCHEMA]) + # What follows the commit that stored the turn's answer (`after_answer`): an agent's handlers publish there, + # a `chat` turn sets nothing (its `close` publishes its end), and so does a handler of one's own. It is read + # before the answer below, so that a retry which finds the turn answered publishes its end too: an attempt + # that died between the commit and the publish left it unpublished, and its followers waiting. + at_end = after_answer_of(ctx) if (found := await answer_in(ctx.db, prompt_id)) is not None: - # Answered (or closed) already: a retry answers, or fails, as the turn ended. + # Answered (or closed) already: a retry answers, or fails, as the turn ended. The end publishes again, + # which costs a revision and tells a follower what it may already know. answer, status = found + if at_end is not None: + await at_end() return stored_answer(answer, status) conversation = await Conversation.load(ctx.db, items, window, budget, snapshots=settings.snapshots) # The turn's final reply commits with its answer; the library's publishes and Restate calls write what the run - # holds back first. Each write of steps that commits publishes, unless the agent runs in the background. - conversation.hold(lambda: stepped(ctx)) + # holds back first. Each write of steps that commits forgets the calls whose results it stored (`_calls.py`), + # then runs what the caller set (`after_steps`): an agent's `chat` turn publishes there; a handler of one's own + # sets nothing, and its turn is only stored. + after = after_steps_of(ctx) + + async def committed() -> None: + await settled(ctx) + if after is not None: + await after() + + async def ended() -> None: + # The turn's calls are forgotten already: `answering` does it as its work ends, failed for good or not. + if at_end is not None: + await at_end() + + conversation.hold(committed) holding(ctx, conversation.flush) async def work() -> str: @@ -1692,7 +1718,15 @@ def stored(call_id: str, failed: str | None) -> bool: running = conversation.running() return running is not None and running.id == prompt_id and running.stored(call_id, failed) - answered = await answering(ctx, stored, work) + try: + answered = await answering(ctx, stored, work) + except TerminalError: + # The turn ended for good, closed with the error (`_answer`): its answer is stored, so what follows the + # answer runs, and a follower hears the turn end. An attempt that failed otherwise ends no turn, and runs + # nothing: the next attempt carries it on. + await ended() + raise + await ended() definition = definition_of(ctx) if settings.snapshots: due = settings.compact_history is not None and conversation.compaction_asked diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index 39577e9..f19b774 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -25,6 +25,8 @@ from restate.agents import _messages, _signals from restate.agents._api import ( + AgentContext, + AgentViewContext, Messages, MessagesRequest, Notifications, @@ -51,14 +53,13 @@ from restate.agents._conversation import SCHEMA, Items, next_start, restart, steer_stored from restate.agents._conversation import compact as compact_turns from restate.agents._conversation import settle as settle_turn -from restate.agents._turn import AGENT, AgentInfo, chat_turn, cleared, publish +from restate.agents._turn import AGENT, AgentInfo, chat_turn, cleared, publish, publish_end, publish_steps from restate.object import VirtualObject from restate.retry_policy import InvocationRetryPolicy -from restate.sessions import Migration, SessionContext, SessionObject, ViewContext +from restate.sessions import Migration, SessionContext, SessionObject from restate.sessions.internal import hide_entries, reset -T = typing.TypeVar("T") -A = typing.TypeVar("A", bound=Callable[[SessionContext, str], Awaitable[Any]]) +A = typing.TypeVar("A", bound=Callable[[AgentContext, str], Awaitable[Any]]) """The function that answers a chat message: it takes the message and returns the answer.""" RETRIES = InvocationRetryPolicy( @@ -278,7 +279,7 @@ def __init__( raise ValueError(f"budget must be a number of tokens (0 or more), not {budget!r}") self._settings = AgentInfo(background=background, window=check_window(window), budget=budget) AGENT.on(self, self._settings) - _signals.with_companion(self) + self._companion = _signals.with_companion(self) self._building = True try: self._built_in() @@ -286,9 +287,9 @@ def __init__( self._building = False def __restate_services__(self) -> tuple[VirtualObject, ...]: - """What ``restate.app`` serves for an agent (see ``restate.Servable``): itself, and ``signals``, - the object its tools wait on, shared by every agent.""" - return (self, _signals.signals) + """What ``restate.app`` serves for an agent (see ``restate.Servable``): itself, and its own companion, + ``.prebuilt.restate``, which its tools wait on.""" + return (self, self._companion.service) def _reserved(self, name: str) -> str | None: if name not in BUILT_IN or getattr(self, "_building", False): @@ -299,13 +300,20 @@ def _reserved(self, name: str) -> str | None: ) def _attempt_fn(self, fn: Callable[..., Awaitable[Any]], view: bool) -> Callable[..., Awaitable[Any]]: - """Wrap a handler so that a model provider's refusal fails its invocation for good, as in ``chat``.""" + """Wrap a handler so that a model provider's refusal fails its invocation for good, as in ``chat``, the + steps its turns store publish (``publish_steps``), as in ``chat``, and each turn's answer publishes once + more (``publish_end``): for a handler of your own that prompts. + + A ``chat`` turn is wrapped the same way, and then keeps its end for + its ``close`` (see ``chat_turn``): it publishes no end of its own.""" if view: return fn # Keep fn's signature: the SDK reads from it whether the handler takes an input. @functools.wraps(fn) - async def answering(ctx: SessionContext, *args: Any) -> Any: + async def answering(ctx: AgentContext, *args: Any) -> Any: + publish_steps(ctx) + publish_end(ctx) try: return await fn(ctx, *args) except Exception as error: @@ -317,35 +325,48 @@ async def answering(ctx: SessionContext, *args: Any) -> Any: return answering def _built_in(self) -> None: - """Register the views every agent has; each one asks the agent's companion ``signals`` object.""" + """Register the views every agent has; each one asks the agent's companion.""" + companion = self._companion - async def signal(ctx: ViewContext, signal: Signal) -> Resolved: + async def signal(ctx: AgentViewContext, signal: Signal) -> Resolved: """Resolve a wait of the agent's tools with ``value``: the wait with ``id``, or the oldest one for ``name``.""" - return await _companion(ctx, _signals.resolve, "signal")(signal) - - async def pending(ctx: ViewContext, _: Nothing = None) -> list[Waiting]: + return await ctx.call( + companion.resolve, + signal, + key=ctx.key, + idempotency_key=f"{ctx.request.id}:signal", + name=_signals.ENTRY_NAME, + ) + + async def pending(ctx: AgentViewContext, _: Nothing = None) -> list[Waiting]: """List the signals the agent's tools are waiting for, oldest first.""" - return await _companion(ctx, _signals.pending)() + return await ctx.call(companion.pending, key=ctx.key, name=_signals.ENTRY_NAME) - async def status(ctx: ViewContext, _: Nothing = None) -> Status: + async def status(ctx: AgentViewContext, _: Nothing = None) -> Status: """Return what the agent is doing: the message it is answering, if any, and what its tools wait for.""" - return await _companion(ctx, _signals.status)() + return await ctx.call(companion.status, key=ctx.key, name=_signals.ENTRY_NAME) - async def stop(ctx: ViewContext, _: Nothing = None) -> Stopped: + async def stop(ctx: AgentViewContext, _: Nothing = None) -> Stopped: """Cancel the invocation that is answering the current message, if there is one.""" - return await _stopped(ctx, "stop") + return await _stopped(ctx, companion, "stop") - async def notifications(ctx: ViewContext, _: Nothing = None) -> Notifications: + async def notifications(ctx: AgentViewContext, _: Nothing = None) -> Notifications: """Return the conversation's revision, which increases whenever ``messages`` or ``status`` change.""" - return await _companion(ctx, _signals.notifications)() + return await ctx.call(companion.notifications, key=ctx.key, name=_signals.ENTRY_NAME) - async def watch(ctx: ViewContext, request: Watch) -> Notifications: + async def watch(ctx: AgentViewContext, request: Watch) -> Notifications: """Wait for a revision newer than ``revision``, then return the current one. It returns anyway after ``timeoutSeconds`` (50 by default, at most 60). The wait happens in Restate and holds no process. """ - return await _companion(ctx, _signals.watch, "watch")(request) + return await ctx.call( + companion.watch, + request, + key=ctx.key, + idempotency_key=f"{ctx.request.id}:watch", + name=_signals.ENTRY_NAME, + ) for fn in (signal, pending, status, stop, notifications): self.view(fn) @@ -383,7 +404,7 @@ def chat( def chat( self, - answer: Callable[[SessionContext, str], Awaitable[Any]] | None = None, + answer: Callable[[AgentContext, str], Awaitable[Any]] | None = None, /, *, items: Items, @@ -448,8 +469,8 @@ def chat( if answer is None: def decorate( - fn: Callable[[SessionContext, str], Awaitable[Any]], - ) -> Callable[[SessionContext, str], Awaitable[Any]]: + fn: Callable[[AgentContext, str], Awaitable[Any]], + ) -> Callable[[AgentContext, str], Awaitable[Any]]: return self.chat( fn, items=items, @@ -476,12 +497,12 @@ def decorate( AGENT.on(self, self._settings) answering = answer - async def chat(ctx: SessionContext, text: str) -> Any: + async def chat(ctx: AgentContext, text: str) -> Any: return await chat_turn(ctx, text, answering) chat.__annotations__["return"] = output if output is not None else str - async def steered(ctx: SessionContext, message: Steer) -> None: + async def steered(ctx: AgentContext, message: Steer) -> None: """Answer a message sent to ``steer`` as a turn of its own, unless the turn it was sent to took it. The companion sends it as the message is sent, so that it keeps its @@ -501,32 +522,19 @@ async def steered(ctx: SessionContext, message: Steer) -> None: return answer -def _companion( - ctx: ViewContext, handler: Callable[..., Awaitable[T]], step: str | None = None -) -> Callable[..., Awaitable[T]]: - """Return a function that calls ``handler`` on the companion of the agent of ``ctx``. - - With ``step``, the call gets an idempotency key made of the request id and - ``step``, so a restarted view joins the same call. Reads have no key. - """ - key = None if step is None else f"{ctx.request.id}:{step}" - companion = _signals.signals_of(ctx) - untyped = typing.cast(Any, ctx) - - def call(*args: Any) -> Awaitable[T]: - made = untyped.call(handler, *args, key=companion, idempotency_key=key, name=_signals.ENTRY_NAME) - return typing.cast(Awaitable[T], made) - - return call - - -async def _stopped(ctx: ViewContext, step: str) -> Stopped: +async def _stopped(ctx: AgentViewContext, companion: _signals.Companion, step: str) -> Stopped: """Cancel the invocation answering the current message, if any, and return its id. - The status is read with an idempotency key, so a restarted view cancels the - same invocation again (a no-op) rather than the next message's. + The status is read with an idempotency key, made of the request id and + ``step``, so a restarted view cancels the same invocation again (a no-op) + rather than the next message's. The companion's other calls from a view + that change something (``signal``, ``steer``) or wait (``watch``) have + such a key too, so a restarted view joins the same call; reads have none. """ - turn = (await _companion(ctx, _signals.status, step)())["answering"] + status = await ctx.call( + companion.status, key=ctx.key, idempotency_key=f"{ctx.request.id}:{step}", name=_signals.ENTRY_NAME + ) + turn = status["answering"] if turn is not None: ctx.cancel(turn["invocationId"]) return {"stopped": turn["invocationId"] if turn is not None else None} @@ -535,8 +543,9 @@ async def _stopped(ctx: ViewContext, step: str) -> Stopped: def _followed(agent: AgentObject, items: Items) -> None: """Register the views that ``chat()`` adds (``messages``, ``steer``, ``clear``, ``sync``), and the private handlers.""" + companion = _signals.companion_of(agent) - async def messages(ctx: ViewContext, request: MessagesRequest | None = None) -> Messages: + async def messages(ctx: AgentViewContext, request: MessagesRequest | None = None) -> Messages: """Read the conversation: a page of it, or what changed since a cursor. Send ``{"limit": n, "before": id}`` for a page, or ``{"cursor": c}`` @@ -545,7 +554,7 @@ async def messages(ctx: ViewContext, request: MessagesRequest | None = None) -> """ return await _messages.page(ctx.db, items, request) - async def steer(ctx: ViewContext, text: str) -> Steered: + async def steer(ctx: AgentViewContext, text: str) -> Steered: """Send a message to the agent while it answers another: the turn takes it at its next step. The message is the request's body, as plain text or as a JSON string, @@ -557,9 +566,12 @@ async def steer(ctx: ViewContext, text: str) -> Steered: Either way it is answered once: send it with an idempotency key, so that a retried request is the same message. """ - return await _companion(ctx, _signals.steer, "steer")({"id": ctx.request.id, "text": text}) + message: Steer = {"id": ctx.request.id, "text": text} + return await ctx.call( + companion.steer, message, key=ctx.key, idempotency_key=f"{ctx.request.id}:steer", name=_signals.ENTRY_NAME + ) - async def erase(ctx: SessionContext) -> None: + async def erase(ctx: AgentContext) -> None: """Empty the key's database and its signals, once the handlers queued before it are done.""" # Read first, which also makes the reset apply the conversation's schema again. await SCHEMA.ensure(ctx.db) @@ -569,13 +581,13 @@ async def erase(ctx: SessionContext) -> None: await restart(tx, start) await cleared(ctx) await ctx.call( - _signals.reset, - key=_signals.signals_of(ctx), + companion.reset, + key=ctx.key, idempotency_key=f"{ctx.request.id}:reset", name=_signals.ENTRY_NAME, ) - async def settle(ctx: SessionContext, invocation_id: str) -> bool: + async def settle(ctx: AgentContext, invocation_id: str) -> bool: """Close the latest turn, unanswered, if invocation ``invocation_id`` ran it and is gone. The companion sends this once it finds the turn's invocation ended @@ -588,11 +600,11 @@ async def settle(ctx: SessionContext, invocation_id: str) -> bool: await publish(ctx) return True - async def compact(ctx: SessionContext, _: Nothing = None) -> bool: + async def compact(ctx: AgentContext, _: Nothing = None) -> bool: """Compact the history: summarize the turns past the budget, or compact a snapshot. Return whether it did.""" return await compact_turns(ctx, items) - async def clear(ctx: ViewContext, _: Nothing = None) -> Stopped: + async def clear(ctx: AgentViewContext, _: Nothing = None) -> Stopped: """Forget the conversation, and empty the key's database. It stops the message being answered, as ``stop`` does. Messages sent @@ -601,11 +613,11 @@ async def clear(ctx: ViewContext, _: Nothing = None) -> Stopped: that follows the conversation keeps reading, and drops the messages before ``first``. """ - stopped = await _stopped(ctx, "clear") + stopped = await _stopped(ctx, companion, "clear") await ctx.call(erase, key=ctx.key, idempotency_key=f"{ctx.request.id}:erase") return stopped - async def sync(ctx: ViewContext, request: SyncRequest | None = None) -> Synced: + async def sync(ctx: AgentViewContext, request: SyncRequest | None = None) -> Synced: """Return the conversation since ``cursor``, what the agent is doing, and the revision. With ``revision``, it first waits for a newer revision, or for the @@ -617,12 +629,14 @@ async def sync(ctx: ViewContext, request: SyncRequest | None = None) -> Synced: revision = request.get("revision") if revision is not None: watch: Watch = {"revision": revision, "timeoutSeconds": request.get("timeoutSeconds", 50)} - await _companion(ctx, _signals.watch, "watch")(watch) + await ctx.call( + companion.watch, watch, key=ctx.key, idempotency_key=f"{ctx.request.id}:watch", name=_signals.ENTRY_NAME + ) # Read the revision first, then the page in an invocation of its own, which starts after it: a view # reads the database as committed when it started, and this one started before the revision (by # the whole wait, when it waited). So the page has every change the revision counts. - now = await _companion(ctx, _signals.notifications)() - status = await _companion(ctx, _signals.status)() + now = await ctx.call(companion.notifications, key=ctx.key, name=_signals.ENTRY_NAME) + status = await ctx.call(companion.status, key=ctx.key, name=_signals.ENTRY_NAME) wanted: MessagesRequest = {"cursor": request.get("cursor", 0), "limit": request.get("limit", 100)} if "view" in request: wanted["view"] = request["view"] diff --git a/python/restate/agents/_signals.py b/python/restate/agents/_signals.py index 2aac338..d2df3cb 100644 --- a/python/restate/agents/_signals.py +++ b/python/restate/agents/_signals.py @@ -8,26 +8,33 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -"""The ``signals`` virtual object: each agent key's companion. - -The companion of an agent's key is keyed ``:`` (such as -``assistant:alice``). It holds the signals the agent's tools wait for (``wait`` -returns once ``resolve`` gives the wait its value), the message the agent is +"""Each agent's companion: a virtual object of its own, ``.prebuilt.restate``. + +Whatever makes an agent (a framework's ``agent()``, ``AgentObject``) makes its +companion too (``with_companion``), and the agent serves it alongside itself. +The companion of agent ``assistant`` is ``assistant.prebuilt.restate``, +keyed as the agent: key ``alice`` of the agent has key ``alice`` of its +companion. It holds the signals the agent's tools wait for (``wait`` returns +once ``resolve`` gives the wait its value), the message the agent is answering, the messages steered to it that it has not taken yet (``steer``), -and the conversation's revision, which ``watch`` waits on. Its handler names, -including the camelCase ``awaitTurn``, are wire names shared with Restate's -other SDKs. +and the conversation's revision, which ``watch`` waits on. + +The companion is ingress-private: only its agent's handlers and views call it, +and clients reach it through the agent's views (``status``, ``pending``, +``signal``, ``steer``, ``watch``, ...). Its handlers' names, the camelCase +``awaitTurn`` among them, and payloads are those of the TypeScript SDK's +companion. An agent's own handler cannot wait on an awakeable, since it has no journal to replay one from. Instead it calls ``wait``, a regular journaled handler of -this object, which waits without holding a process. +its companion, which waits without holding a process. """ from __future__ import annotations import json import typing -from collections.abc import Callable, Sequence +from collections.abc import Awaitable, Callable, Sequence from dataclasses import dataclass from datetime import UTC, datetime, timedelta from typing import Any, TypedDict @@ -47,19 +54,24 @@ Waiting, Watch, ) +from restate.agents._common import NeedsAgent from restate.sessions import SessionContext, ViewContext from restate.sessions._link import Link from restate.sessions.internal import definition_of T = typing.TypeVar("T") -D = typing.TypeVar("D") -signals = restate.VirtualObject("signals") -"""The companion object of every agent, keyed ``:`` (see ``signals_of``). +COMPANION_SUFFIX = ".prebuilt.restate" +"""What an agent's name is followed by to name its companion: ``.prebuilt.restate``. + +Restate refuses service names that start with ``restate``, so the companion's +name ends with it instead.""" + + +def companion_name(agent: str) -> str: + """Return the name of the companion of agent ``agent``: ``.prebuilt.restate``.""" + return f"{agent}{COMPANION_SUFFIX}" -``restate.app`` serves it alongside the agents: each agent lists it (``AgentObject.__restate_services__``). Its ingress-private handlers -are called only by the agents' own handlers. -""" ENTRY_NAME = "signals" """The journal entry name of every call and send of an agent to its companion. @@ -115,19 +127,15 @@ class _Watching(TypedDict): class Started(Answering, total=False): - """A turn as its agent starts it: and whether the agent closes it itself once its invocation is gone.""" + """A turn as its agent starts it: and whether it takes steered messages, and whether ``start`` publishes it.""" - settles: bool - """Whether the agent has the private ``settle`` handler, which the companion sends a turn whose invocation is - gone. An agent of another SDK, whose turns this companion may hold too (its handler names are shared), does - not say so, and is sent nothing: its next message closes such a turn.""" steers: bool - """Whether the turn takes steered messages at its steps (``steer``). A turn that does not say so (of another - SDK, or of an agent whose framework cannot take them) gets none: each is answered as a turn of its own.""" + """Whether the turn takes steered messages at its steps (``steer``). A turn that does not (of an agent whose + framework cannot take them, or that runs in the background) gets none: each is answered as a turn of its own.""" publishes: bool - """Whether ``start`` publishes the turn's start, and returns what ``publish`` returns: a followed turn of this - SDK's agents, which then makes no publish of its own. A turn that does not say so (of another SDK, or of an - agent that runs in the background) is not published by ``start``, which returns nothing.""" + """Whether ``start`` publishes the turn's start, and returns what ``publish`` returns: a followed turn, which + then makes no publish of its own. A turn that does not (of an agent that runs in the background) is not + published by ``start``, which returns nothing.""" class _Held(TypedDict): @@ -197,7 +205,7 @@ async def _publish(ctx: ObjectContext) -> Notifications: return notified -async def _closed(ctx: ObjectContext, turn: str, gone: bool = False) -> None: +async def _closed(ctx: ObjectContext, agent: str, turn: str, gone: bool = False) -> None: """Close ``turn`` (an invocation id): forget it and its tools' waits, and publish the change. The waits are rejected, which ends their ``wait`` calls. A turn that ends @@ -205,8 +213,8 @@ async def _closed(ctx: ObjectContext, turn: str, gone: bool = False) -> None: Closing a turn twice changes nothing and publishes nothing. If the turn's invocation is ``gone`` and the turn did not close itself, the - agent is sent ``settle``, which closes the turn in its conversation too: if - it said it has that handler as it started the turn. + agent (``agent``, whose companion this is) is sent ``settle``, which closes + the turn in its conversation too. The messages steered to the turn are forgotten here: each that it did not take was queued as a turn of its own as it was sent (``steer``), and is @@ -216,9 +224,8 @@ async def _closed(ctx: ObjectContext, turn: str, gone: bool = False) -> None: is_answered = answering is not None and answering["invocationId"] == turn if is_answered: ctx.clear("turn") - if gone and typing.cast(Started, answering).get("settles", False): - agent, key = ctx.key().split(":", 1) - ctx.generic_send(agent, "settle", json.dumps(turn).encode(), key=key) + if gone: + ctx.generic_send(agent, "settle", json.dumps(turn).encode(), key=ctx.key()) inbox = await _inbox(ctx) dropped = any(one["turn"] == turn for one in inbox) if dropped: @@ -241,252 +248,6 @@ def _set_inbox(ctx: ObjectContext, inbox: list[_Inboxed]) -> None: ctx.clear("inbox") -@signals.handler(ingress_private=True) -async def steer(ctx: ObjectContext, message: Steer) -> Steered: - """Take a message steered to the agent (its ``steer`` view): for the turn being answered, or as a turn of its own. - - The message is always sent to the agent's private ``steered`` handler at - once, which keeps its place in the key's queue, in order with the - messages sent after it: once the turn being answered ends, ``steered`` - answers the message as a turn of its own, unless that turn took it. A - turn that takes steered messages (``Started.steers``) also finds it here, - at its next step, through ``publish``, until it takes it (``taken``) or - ends. - """ - message_id, text = message["id"], message["text"] - agent, key = ctx.key().split(":", 1) - placeholder = json.dumps({"id": message_id, "text": text}).encode() - ctx.generic_send(agent, "steered", placeholder, key=key, idempotency_key=f"steer:{message_id}") - answering = await _get(ctx, "turn", typing.cast(Started | None, None)) - if answering is None or not answering.get("steers", False): - return {"id": message_id, "turn": None} - turn, inbox = answering["invocationId"], await _inbox(ctx) - if any(one["id"] == message_id for one in inbox): - return {"id": message_id, "turn": turn} - # Taken or not: the turn has room for so many in all. - if sum(one["turn"] == turn for one in inbox) >= STEERS: - return {"id": message_id, "turn": None} - since = iso(await ctx.time()) - _set_inbox(ctx, [*inbox, {"id": message_id, "text": text, "since": since, "turn": turn}]) - await _publish(ctx) - return {"id": message_id, "turn": turn} - - -@signals.handler(ingress_private=True) -async def taken(ctx: ObjectContext, ids: list[str]) -> None: - """Mark the steered messages ``ids`` taken, which their turn has stored, and publish the change, if any. - - The turn sends it with the transaction that stores them, so it comes once - they are stored, and once only. Each is kept without its text until the - turn ends, as one of its ``STEERS``; one already forgotten (its turn - closed first) is left so. - """ - inbox = await _inbox(ctx) - marked: list[_Inboxed] = [ - {"id": one["id"], "since": one["since"], "turn": one["turn"]} if one["id"] in ids else one for one in inbox - ] - if marked != inbox: - _set_inbox(ctx, marked) - await _publish(ctx) - - -@signals.handler(kind="shared", inactivity_timeout=PARKED) -async def wait(ctx: ObjectSharedContext, request: WaitRequest) -> Any: - """Wait for the signal ``name`` as wait ``id``, and return its value.""" - awakeable, value = ctx.awakeable(type_hint=object) - since = iso(await ctx.time()) - waiting: _WaitingOn = {"id": request["id"], "name": request["name"], "since": since, "awakeable": awakeable} - if "input" in request: - waiting["input"] = request["input"] - ctx.object_send(listen, ctx.key(), waiting) - try: - return await value - except TerminalError as error: - # The agent's invocation was cancelled, so nothing waits any more. - if error.status_code == 409: - ctx.object_send(forget, ctx.key(), request["id"]) - raise - - -@signals.handler() -async def resolve(ctx: ObjectContext, signal: Signal) -> Resolved: - """Resolve the wait ``id``, or else the oldest wait for signal ``name``, with ``value``.""" - wait_id, name = signal.get("id"), signal.get("name") - if wait_id is None and name is None: - raise TerminalError("a signal needs an id or a name") - value = signal.get("value") - waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) - found = next( - (at for at, wait in enumerate(waiting) if (wait["name"] == name if wait_id is None else wait["id"] == wait_id)), - None, - ) - if found is None: - if signal.get("keep", False) and name is not None and wait_id is None: - kept = await _get(ctx, "kept", typing.cast(dict[str, list[Any]], {})) - ctx.set("kept", {**kept, name: [*kept.get(name, []), value]}) - return {"resolved": None} - resolved = waiting.pop(found) - ctx.resolve_awakeable(resolved["awakeable"], value) - ctx.set("waiting", waiting) - await _publish(ctx) - return {"resolved": resolved["id"]} - - -@signals.handler(kind="shared") -async def pending(ctx: ObjectSharedContext) -> list[Waiting]: - """List the waits, oldest first.""" - return [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))] - - -@signals.handler(ingress_private=True) -async def listen(ctx: ObjectContext, waiting: _WaitingOn) -> None: - """Register a wait, resolving it right away if a signal was kept for it.""" - kept = await _get(ctx, "kept", typing.cast(dict[str, list[Any]], {})) - values = kept.get(waiting["name"], []) - if values: - ctx.resolve_awakeable(waiting["awakeable"], values[0]) - ctx.set("kept", {**kept, waiting["name"]: values[1:]}) - return - ctx.set("waiting", [*await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])), waiting]) - await _publish(ctx) - - -@signals.handler(ingress_private=True) -async def start(ctx: ObjectContext, turn: Started) -> Published | None: - """Record that the agent's ``chat`` started answering ``turn``, and publish it if the turn ``publishes``. - - A retry of the same invocation keeps the original start time. A previous - turn that never closed itself (because its invocation was killed) is - closed now. ``awaitTurn`` then watches the new turn, so that it is closed - even if no other turn follows it. - - A turn that ``publishes`` gets what ``publish`` returns (the turn's first - publish, made here rather than by a call of its own); any other gets - ``None``, and publishes itself if it is followed. - """ - before = await _get(ctx, "turn", typing.cast(Started | None, None)) - if before is None or before["invocationId"] != turn["invocationId"]: - if before is not None: - await _closed(ctx, before["invocationId"], gone=True) - ctx.set("turn", turn) - ctx.object_send(await_turn, ctx.key(), turn["invocationId"]) - if not turn.get("publishes", False): - return None - return await _published(ctx) - - -@signals.handler( - name="awaitTurn", - kind="shared", - ingress_private=True, - inactivity_timeout=PARKED, - # Keep nothing once it is done: there is one at a time per turn. - journal_retention=timedelta(0), - idempotency_retention=timedelta(0), -) -async def await_turn(ctx: ObjectSharedContext, turn: str) -> None: - """Wait for ``turn`` to end, without holding a process, and close it if it failed. - - Closing ends the turn's waits, and ``status`` stops reporting it. A turn - that ended normally or was cancelled has closed itself already; one that - did not (killed, out of retries, lost) is closed here. - - It waits at most ``WATCH``: if the turn is still running then, a new - watch takes over, which runs on the current deployment. Cancelling the - watch itself (an operator might) closes nothing. - """ - attached = ctx.attach_invocation(turn, type_hint=object) - # Cancelling this watch raises here, and closes nothing; the turn failing does not raise here. - completed, _ = await restate.wait_completed(attached, ctx.sleep(WATCH)) - if attached in completed: - try: - await attached - except TerminalError: - ctx.object_send(ended, ctx.key(), turn) - return - answering = await _get(ctx, "turn", typing.cast(Answering | None, None)) - if answering is not None and answering["invocationId"] == turn: - ctx.object_send(await_turn, ctx.key(), turn) - - -@signals.handler(kind="shared") -async def status(ctx: ObjectSharedContext) -> Status: - """Return what the agent is doing: the message it is answering, if any, its tools' waits, and the messages - steered to it that it has not taken yet.""" - answering = await _get(ctx, "turn", typing.cast(Started | None, None)) - shown: Status = { - "answering": _shown(answering), - "waiting": [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))], - } - waiting = _waiting(await _inbox(ctx), answering) - steering: list[Steering] = [{"id": one["id"], "since": one["since"]} for one in waiting] - if steering: - shown["steering"] = steering - return shown - - -# Its journal is not kept once it is done: a client follows with one watch after another. -@signals.handler(kind="shared", inactivity_timeout=PARKED, journal_retention=timedelta(0)) -async def watch(ctx: ObjectSharedContext, request: Watch) -> Notifications: - """Wait for a revision newer than ``revision``, or the timeout, and return the current revision. - - The timeout is ``timeoutSeconds`` (50 by default, at most 60). The wait - happens in Restate and holds no process; a caller that retries with the - same idempotency key joins the same wait. - """ - awakeable, changed = ctx.awakeable(type_hint=object) - watching: _Watching = {"revision": request["revision"], "awakeable": awakeable} - ctx.object_send(subscribe, ctx.key(), watching) - seconds = min(max(float(request.get("timeoutSeconds", 50)), 0.0), 60.0) - which, value = await restate.select(changed=changed, timeout=ctx.sleep(timedelta(seconds=seconds))) - if which == "changed": - return typing.cast(Notifications, value) - ctx.object_send(unsubscribe, ctx.key(), awakeable) - return {"revision": await _get(ctx, "revision", 0)} - - -@signals.handler(ingress_private=True) -async def close(ctx: ObjectContext, turn: str) -> None: - """Close ``turn`` (an invocation id), which has been answered or ended otherwise.""" - await _closed(ctx, turn) - - -@signals.handler(ingress_private=True) -async def ended(ctx: ObjectContext, turn: str) -> None: - """Close ``turn`` (an invocation id), which failed: if it did not close itself, the agent closes it too.""" - await _closed(ctx, turn, gone=True) - - -@signals.handler(ingress_private=True) -async def reset(ctx: ObjectContext) -> None: - """Forget all waits, kept signals and steered messages after the agent's ``clear``, and publish the change. - Each steered message was queued as a turn of its own as it was sent, before the ``clear``.""" - ctx.clear("waiting") - ctx.clear("kept") - ctx.clear("inbox") - await _publish(ctx) - - -@signals.handler(ingress_private=True) -async def forget(ctx: ObjectContext, wait_id: str) -> None: - """Forget wait ``wait_id``, which the agent no longer waits for.""" - waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) - ctx.set("waiting", [wait for wait in waiting if wait["id"] != wait_id]) - await _publish(ctx) - - -@signals.handler(ingress_private=True) -async def publish(ctx: ObjectContext) -> Published: - """Increase the conversation's revision after a change, and wake up every waiting ``watch``. - - It returns the new revision, and the messages steered to the turn being - answered that it has not taken yet, if any: a turn publishes as it starts - and after each step it stores that a tool waits for (a reply with tool - calls, the tools' results), so it learns of them at no cost of its own. - """ - return await _published(ctx) - - async def _published(ctx: ObjectContext) -> Published: """Publish (``_publish``), and return the new revision and the messages the turn being answered has not taken.""" published: Published = {**await _publish(ctx)} @@ -500,82 +261,399 @@ async def _published(ctx: ObjectContext) -> Published: return published -@signals.handler(kind="shared") -async def notifications(ctx: ObjectSharedContext) -> Notifications: - """Return the conversation's revision.""" - return {"revision": await _get(ctx, "revision", 0)} +@dataclass(frozen=True) +class Companion: + """An agent's companion object (``service``), and its handlers, to call with the agent's key. + From a handler or view of the agent:: -@signals.handler(ingress_private=True) -async def subscribe(ctx: ObjectContext, watching: _Watching) -> None: - """Resolve the watch's awakeable once there is a revision newer than its ``revision``, right away if there is one.""" - revision = await _get(ctx, "revision", 0) - if revision > watching["revision"]: - ctx.resolve_awakeable(watching["awakeable"], {"revision": revision}) - return - ctx.set("watching", [*await _get(ctx, "watching", typing.cast(list[_Watching], [])), watching]) + status = await ctx.call(companion_of(assistant).status, key=ctx.key) + """ + agent: str + """The name of the agent whose companion this is.""" + service: restate.VirtualObject + """The companion's virtual object, ``.prebuilt.restate``, which the agent serves alongside itself.""" + steer: Callable[[ObjectContext, Steer], Awaitable[Steered]] + taken: Callable[[ObjectContext, list[str]], Awaitable[None]] + wait: Callable[[ObjectSharedContext, WaitRequest], Awaitable[Any]] + resolve: Callable[[ObjectContext, Signal], Awaitable[Resolved]] + pending: Callable[[ObjectSharedContext], Awaitable[list[Waiting]]] + listen: Callable[[ObjectContext, _WaitingOn], Awaitable[None]] + start: Callable[[ObjectContext, Started], Awaitable[Published | None]] + await_turn: Callable[[ObjectSharedContext, str], Awaitable[None]] + status: Callable[[ObjectSharedContext], Awaitable[Status]] + watch: Callable[[ObjectSharedContext, Watch], Awaitable[Notifications]] + close: Callable[[ObjectContext, str], Awaitable[None]] + ended: Callable[[ObjectContext, str], Awaitable[None]] + reset: Callable[[ObjectContext], Awaitable[None]] + forget: Callable[[ObjectContext, str], Awaitable[None]] + publish: Callable[[ObjectContext], Awaitable[Published]] + notifications: Callable[[ObjectSharedContext], Awaitable[Notifications]] + subscribe: Callable[[ObjectContext, _Watching], Awaitable[None]] + unsubscribe: Callable[[ObjectContext, str], Awaitable[None]] + + +def make_companion(agent: str) -> Companion: + """Make the companion of agent ``agent``: the virtual object ``.prebuilt.restate``, keyed as the agent. + + It knows its agent: it sends it ``settle`` and ``steered``. It is + ingress-private: only the agent's handlers and views call it. + """ + service = restate.VirtualObject( + companion_name(agent), + description=f"The companion of agent {agent}: its turn, its tools' waits, steered messages and revision", + ingress_private=True, + ) -@signals.handler(ingress_private=True) -async def unsubscribe(ctx: ObjectContext, awakeable: str) -> None: - """Forget the ``watch`` with ``awakeable``, which no longer waits.""" - watching = await _get(ctx, "watching", typing.cast(list[_Watching], [])) - ctx.set("watching", [one for one in watching if one["awakeable"] != awakeable]) + @service.handler() + async def steer(ctx: ObjectContext, message: Steer) -> Steered: + """Take a message steered to the agent (its ``steer`` view): for the turn being answered, or as a turn of its own. + + The message is always sent to the agent's private ``steered`` handler at + once, which keeps its place in the key's queue, in order with the + messages sent after it: once the turn being answered ends, ``steered`` + answers the message as a turn of its own, unless that turn took it. A + turn that takes steered messages (``Started.steers``) also finds it here, + at its next step, through ``publish``, until it takes it (``taken``) or + ends. + """ + message_id, text = message["id"], message["text"] + placeholder = json.dumps({"id": message_id, "text": text}).encode() + ctx.generic_send(agent, "steered", placeholder, key=ctx.key(), idempotency_key=f"steer:{message_id}") + answering = await _get(ctx, "turn", typing.cast(Started | None, None)) + if answering is None or not answering.get("steers", False): + return {"id": message_id, "turn": None} + turn, inbox = answering["invocationId"], await _inbox(ctx) + if any(one["id"] == message_id for one in inbox): + return {"id": message_id, "turn": turn} + # Taken or not: the turn has room for so many in all. + if sum(one["turn"] == turn for one in inbox) >= STEERS: + return {"id": message_id, "turn": None} + since = iso(await ctx.time()) + _set_inbox(ctx, [*inbox, {"id": message_id, "text": text, "since": since, "turn": turn}]) + await _publish(ctx) + return {"id": message_id, "turn": turn} + + @service.handler() + async def taken(ctx: ObjectContext, ids: list[str]) -> None: + """Mark the steered messages ``ids`` taken, which their turn has stored, and publish the change, if any. + + The turn sends it with the transaction that stores them, so it comes once + they are stored, and once only. Each is kept without its text until the + turn ends, as one of its ``STEERS``; one already forgotten (its turn + closed first) is left so. + """ + inbox = await _inbox(ctx) + marked: list[_Inboxed] = [ + {"id": one["id"], "since": one["since"], "turn": one["turn"]} if one["id"] in ids else one for one in inbox + ] + if marked != inbox: + _set_inbox(ctx, marked) + await _publish(ctx) + + @service.handler(kind="shared", inactivity_timeout=PARKED) + async def wait(ctx: ObjectSharedContext, request: WaitRequest) -> Any: + """Wait for the signal ``name`` as wait ``id``, and return its value.""" + awakeable, value = ctx.awakeable(type_hint=object) + since = iso(await ctx.time()) + waiting: _WaitingOn = {"id": request["id"], "name": request["name"], "since": since, "awakeable": awakeable} + if "input" in request: + waiting["input"] = request["input"] + ctx.object_send(listen, ctx.key(), waiting) + try: + return await value + except TerminalError as error: + # The agent's invocation was cancelled, so nothing waits any more. + if error.status_code == 409: + ctx.object_send(forget, ctx.key(), request["id"]) + raise + + @service.handler() + async def resolve(ctx: ObjectContext, signal: Signal) -> Resolved: + """Resolve the wait ``id``, or else the oldest wait for signal ``name``, with ``value``.""" + wait_id, name = signal.get("id"), signal.get("name") + if wait_id is None and name is None: + raise TerminalError("a signal needs an id or a name") + value = signal.get("value") + waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) + found = next( + ( + at + for at, wait in enumerate(waiting) + if (wait["name"] == name if wait_id is None else wait["id"] == wait_id) + ), + None, + ) + if found is None: + if signal.get("keep", False) and name is not None and wait_id is None: + kept = await _get(ctx, "kept", typing.cast(dict[str, list[Any]], {})) + ctx.set("kept", {**kept, name: [*kept.get(name, []), value]}) + return {"resolved": None} + resolved = waiting.pop(found) + ctx.resolve_awakeable(resolved["awakeable"], value) + ctx.set("waiting", waiting) + await _publish(ctx) + return {"resolved": resolved["id"]} + + @service.handler(kind="shared") + async def pending(ctx: ObjectSharedContext) -> list[Waiting]: + """List the waits, oldest first.""" + return [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))] + + @service.handler() + async def listen(ctx: ObjectContext, waiting: _WaitingOn) -> None: + """Register a wait, resolving it right away if a signal was kept for it.""" + kept = await _get(ctx, "kept", typing.cast(dict[str, list[Any]], {})) + values = kept.get(waiting["name"], []) + if values: + ctx.resolve_awakeable(waiting["awakeable"], values[0]) + ctx.set("kept", {**kept, waiting["name"]: values[1:]}) + return + ctx.set("waiting", [*await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])), waiting]) + await _publish(ctx) + + @service.handler() + async def start(ctx: ObjectContext, turn: Started) -> Published | None: + """Record that the agent's ``chat`` started answering ``turn``, and publish it if the turn ``publishes``. + + A retry of the same invocation keeps the original start time. A previous + turn that never closed itself (because its invocation was killed) is + closed now. ``awaitTurn`` then watches the new turn, so that it is closed + even if no other turn follows it. + + A turn that ``publishes`` gets what ``publish`` returns (the turn's first + publish, made here rather than by a call of its own); any other gets + ``None``, and publishes itself if it is followed. + """ + before = await _get(ctx, "turn", typing.cast(Started | None, None)) + if before is None or before["invocationId"] != turn["invocationId"]: + if before is not None: + await _closed(ctx, agent, before["invocationId"], gone=True) + ctx.set("turn", turn) + ctx.object_send(await_turn, ctx.key(), turn["invocationId"]) + if not turn.get("publishes", False): + return None + return await _published(ctx) + + @service.handler( + name="awaitTurn", + kind="shared", + inactivity_timeout=PARKED, + # Keep nothing once it is done: there is one at a time per turn. + journal_retention=timedelta(0), + idempotency_retention=timedelta(0), + ) + async def await_turn(ctx: ObjectSharedContext, turn: str) -> None: + """Wait for ``turn`` to end, without holding a process, and close it if it failed. + + Closing ends the turn's waits, and ``status`` stops reporting it. A turn + that ended normally or was cancelled has closed itself already; one that + did not (killed, out of retries, lost) is closed here. + + It waits at most ``WATCH``: if the turn is still running then, a new + watch takes over, which runs on the current deployment. Cancelling the + watch itself (an operator might) closes nothing. + """ + attached = ctx.attach_invocation(turn, type_hint=object) + # Cancelling this watch raises here, and closes nothing; the turn failing does not raise here. + completed, _ = await restate.wait_completed(attached, ctx.sleep(WATCH)) + if attached in completed: + try: + await attached + except TerminalError: + ctx.object_send(ended, ctx.key(), turn) + return + answering = await _get(ctx, "turn", typing.cast(Answering | None, None)) + if answering is not None and answering["invocationId"] == turn: + ctx.object_send(await_turn, ctx.key(), turn) + + @service.handler(kind="shared") + async def status(ctx: ObjectSharedContext) -> Status: + """Return what the agent is doing: the message it is answering, if any, its tools' waits, and the messages + steered to it that it has not taken yet.""" + answering = await _get(ctx, "turn", typing.cast(Started | None, None)) + shown: Status = { + "answering": _shown(answering), + "waiting": [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))], + } + waiting = _waiting(await _inbox(ctx), answering) + steering: list[Steering] = [{"id": one["id"], "since": one["since"]} for one in waiting] + if steering: + shown["steering"] = steering + return shown + + # Its journal is not kept once it is done: a client follows with one watch after another. + @service.handler(kind="shared", inactivity_timeout=PARKED, journal_retention=timedelta(0)) + async def watch(ctx: ObjectSharedContext, request: Watch) -> Notifications: + """Wait for a revision newer than ``revision``, or the timeout, and return the current revision. + + The timeout is ``timeoutSeconds`` (50 by default, at most 60). The wait + happens in Restate and holds no process; a caller that retries with the + same idempotency key joins the same wait. + """ + awakeable, changed = ctx.awakeable(type_hint=object) + watching: _Watching = {"revision": request["revision"], "awakeable": awakeable} + ctx.object_send(subscribe, ctx.key(), watching) + seconds = min(max(float(request.get("timeoutSeconds", 50)), 0.0), 60.0) + which, value = await restate.select(changed=changed, timeout=ctx.sleep(timedelta(seconds=seconds))) + if which == "changed": + return typing.cast(Notifications, value) + ctx.object_send(unsubscribe, ctx.key(), awakeable) + return {"revision": await _get(ctx, "revision", 0)} + + @service.handler() + async def close(ctx: ObjectContext, turn: str) -> None: + """Close ``turn`` (an invocation id), which has been answered or ended otherwise.""" + await _closed(ctx, agent, turn) + + @service.handler() + async def ended(ctx: ObjectContext, turn: str) -> None: + """Close ``turn`` (an invocation id), which failed: if it did not close itself, the agent closes it too.""" + await _closed(ctx, agent, turn, gone=True) + + @service.handler() + async def reset(ctx: ObjectContext) -> None: + """Forget all waits, kept signals and steered messages after the agent's ``clear``, and publish the change. + Each steered message was queued as a turn of its own as it was sent, before the ``clear``.""" + ctx.clear("waiting") + ctx.clear("kept") + ctx.clear("inbox") + await _publish(ctx) + + @service.handler() + async def forget(ctx: ObjectContext, wait_id: str) -> None: + """Forget wait ``wait_id``, which the agent no longer waits for.""" + waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) + ctx.set("waiting", [wait for wait in waiting if wait["id"] != wait_id]) + await _publish(ctx) + + @service.handler() + async def publish(ctx: ObjectContext) -> Published: + """Increase the conversation's revision after a change, and wake up every waiting ``watch``. + + It returns the new revision, and the messages steered to the turn being + answered that it has not taken yet, if any: a turn publishes as it starts + and after each step it stores that a tool waits for (a reply with tool + calls, the tools' results), so it learns of them at no cost of its own. + """ + return await _published(ctx) + + @service.handler(kind="shared") + async def notifications(ctx: ObjectSharedContext) -> Notifications: + """Return the conversation's revision.""" + return {"revision": await _get(ctx, "revision", 0)} + + @service.handler() + async def subscribe(ctx: ObjectContext, watching: _Watching) -> None: + """Resolve the watch's awakeable once there is a revision newer than its ``revision``, right away if there is one.""" + revision = await _get(ctx, "revision", 0) + if revision > watching["revision"]: + ctx.resolve_awakeable(watching["awakeable"], {"revision": revision}) + return + ctx.set("watching", [*await _get(ctx, "watching", typing.cast(list[_Watching], [])), watching]) + + @service.handler() + async def unsubscribe(ctx: ObjectContext, awakeable: str) -> None: + """Forget the ``watch`` with ``awakeable``, which no longer waits.""" + watching = await _get(ctx, "watching", typing.cast(list[_Watching], [])) + ctx.set("watching", [one for one in watching if one["awakeable"] != awakeable]) + + return Companion( + agent=agent, + service=service, + steer=steer, + taken=taken, + wait=wait, + resolve=resolve, + pending=pending, + listen=listen, + start=start, + await_turn=await_turn, + status=status, + watch=watch, + close=close, + ended=ended, + reset=reset, + forget=forget, + publish=publish, + notifications=notifications, + subscribe=subscribe, + unsubscribe=unsubscribe, + ) # --- An agent's companion ------------------------------------------------------------------------- -def key_of(service: str, key: str) -> str: - """Return the companion's key for key ``key`` of agent ``service``: ``:``.""" - return f"{service}:{key}" +_COMPANION: Link[Companion] = Link("restate-agents.companion") +"""The companion of each agent.""" -def signals_of(ctx: SessionContext | ViewContext) -> str: - """Return the companion's key for the agent key that ``ctx`` runs on.""" - definition = definition_of(ctx) - if definition is None: - raise TypeError("signals_of takes the context of an agent's handler or view") - return key_of(definition.name, ctx.key) +def with_companion(definition: object) -> Companion: + """Make the companion of ``definition`` (an agent), attach it, and return it. + The agent serves it alongside itself (``AgentObject.__restate_services__``). + """ + name = typing.cast(str, getattr(definition, "name")) # noqa: B009 (a definition of the SDK's) + companion = make_companion(name) + _COMPANION.on(definition, companion) + return companion -@dataclass(frozen=True) -class Companion: - """The address of an agent key's companion in the ``signals`` object. - The companion knows what the agent is doing: the message it is answering - and what its tools wait for. +def companion_of(agent: object) -> Companion: + """Return the companion of ``agent``, whose handlers take the agent's keys, to find out what that agent is doing. + + Raises: + TypeError: If ``agent`` is not a ``restate.agents`` agent. """ + companion = _COMPANION.of(agent) + if companion is None: + raise TypeError(f"{agent!r} is not an agent of restate.agents") + return companion - key: str - """The companion's key in ``signals``: ``:``.""" - agent: str - """The agent's name.""" - agent_key: str - """The agent's key.""" +def companion_at(ctx: SessionContext | ViewContext) -> Companion | None: + """Return the companion of the agent whose handler or view ``ctx`` runs, or ``None`` for another session + object's: call it with ``key=ctx.key``.""" + definition = definition_of(ctx) + return _COMPANION.of(definition) if definition is not None else None -_COMPANION: Link[Callable[[str], Companion]] = Link("restate-agents.companion") -"""How an agent finds the companion of each of its keys.""" +def companion_in(ctx: SessionContext | ViewContext) -> Companion: + """Return the companion of the agent whose handler or view ``ctx`` runs: call it with ``key=ctx.key``. -def with_companion(definition: D) -> D: - """Attach the companion to ``definition`` (an agent): how it finds what each of its keys waits for.""" - name = typing.cast(str, getattr(definition, "name")) # noqa: B009 (a definition of the SDK's) - _COMPANION.on(definition, lambda key: Companion(key=key_of(name, key), agent=name, agent_key=key)) - return definition + Raises: + TypeError: If ``ctx`` is not the context of an agent's handler or view. + """ + companion = companion_at(ctx) + if companion is None: + raise TypeError("companion_in takes the context of an agent's handler or view") + return companion + + +def needs_agent(ctx: SessionContext, what: str) -> NeedsAgent: + """Return the error for ``what`` (a step that only an agent takes, as it needs the companion), from a handler + of a session object that is no agent: it fails the turn for good, as a retry would meet it again.""" + definition = definition_of(ctx) + name = typing.cast(str, getattr(definition, "name", "?")) + return NeedsAgent( + f"{what} needs an agent, whose companion holds its waits: {name!r} is a plain session object. " + "Use `.agent(...)` (such as `restate.agents.pydantic_ai.agent`) or `AgentObject` instead of " + "`SessionObject`" + ) -def companion_of(agent: object, key: str) -> Companion: - """Return the companion of ``agent``'s key ``key``, to find out what that agent is doing. +def waits_in(ctx: SessionContext, what: str) -> Companion: + """Return the companion of the agent whose handler ``ctx`` runs, to wait for a signal (``what``) on it. Raises: - TypeError: If ``agent`` is not a ``restate.agents`` agent. + NeedsAgent: If ``ctx`` is a handler's of a session object that is no agent (``needs_agent``). """ - companion = _COMPANION.of(agent) + companion = companion_at(ctx) if companion is None: - raise TypeError(f"{agent!r} is not an agent of restate.agents") - return companion(key) + raise needs_agent(ctx, what) + return companion def has_companion(agent: object) -> bool: @@ -598,13 +676,21 @@ def deliver( Raises: TypeError: If ``agent`` is not a ``restate.agents`` agent. """ - companion = companion_of(agent, key) # an agent of restate.agents, with a `chat` handler + agent_name = companion_of(agent).agent # an agent of restate.agents, with a `chat` handler if isinstance(ctx, restate.Context): # A Restate handler's context sends bytes: the message as a JSON string, which `chat` takes. message = json.dumps(text).encode() - ctx.generic_send(companion.agent, "chat", message, key=key, idempotency_key=idempotency_key) + ctx.generic_send(agent_name, "chat", message, key=key, idempotency_key=idempotency_key) else: - ctx.generic_send(companion.agent, "chat", text, key=key, idempotency_key=idempotency_key) - - -__all__: Sequence[str] = ["Companion", "companion_of", "deliver", "signals"] + ctx.generic_send(agent_name, "chat", text, key=key, idempotency_key=idempotency_key) + + +__all__: Sequence[str] = [ + "COMPANION_SUFFIX", + "Companion", + "companion_in", + "companion_name", + "companion_of", + "deliver", + "make_companion", +] diff --git a/python/restate/agents/_steering.py b/python/restate/agents/_steering.py index 20a7433..2bc4751 100644 --- a/python/restate/agents/_steering.py +++ b/python/restate/agents/_steering.py @@ -13,7 +13,7 @@ A client sends the message to the agent's ``steer`` view. The companion queues it at once as a turn of its own (the agent's private ``steered``), in order with the messages sent after it, and, if the turn being answered takes -steered messages, holds it for that turn too (see ``_signals.steer``). +steered messages, holds it for that turn too (see the companion's ``steer``, in ``_signals.py``). The turn learns of it from its next ``publish``, which it makes anyway as it starts and after each write of steps that commits (a reply with tool calls, @@ -66,10 +66,10 @@ def going() -> bool: turn = conversation.running() if not going() or turn is None or turn.open_calls(): return [] - companion = _signals.signals_of(ctx) + companion = _signals.companion_in(ctx) def taken(tx: Transaction, ids: list[str]) -> None: # With the transaction that stores them: so they are marked taken once stored, and once. - tx.send(_signals.taken, ids, key=companion, idempotency_key=f"taken:{ids[0]}", name=_signals.ENTRY_NAME) + tx.send(companion.taken, ids, key=ctx.key, idempotency_key=f"taken:{ids[0]}", name=_signals.ENTRY_NAME) return await turn.steer(inbox.messages, taken, going) diff --git a/python/restate/agents/_tools.py b/python/restate/agents/_tools.py index c7a5c08..de02906 100644 --- a/python/restate/agents/_tools.py +++ b/python/restate/agents/_tools.py @@ -40,7 +40,7 @@ from restate.agents import _signals from restate.agents._calls import call_of -from restate.agents._common import PerSession, ToolError, configured +from restate.agents._common import NeedsAgent, PerSession, ToolError, configured from restate.agents._turn import stop_signal_of, tool_call_started from restate.exceptions import TerminalError from restate.handler import Handler, handler_from_callable @@ -172,6 +172,9 @@ async def run(self, args: object, call: ToolCall) -> str: # nothing else. The turn still stops with its reason, from the stop signal. raise raise stop.reason from error + if isinstance(error, NeedsAgent): + # Not the model's to fix (a wait from a session object that is no agent): the turn fails with it. + raise if isinstance(error, TerminalError): # A terminal error is the tool's result, which the model can act on. return f"Error: {error.message}" @@ -688,12 +691,16 @@ async def _wait( """Wait for signal ``signal`` as tool call ``call``, once however often it runs; ``pending`` shows ``input`` if ``shown``. ``rejoin`` and ``approval`` are ``call_of``'s.""" + # Before the call is recorded: from a session object that is no agent, the turn fails for good (``NeedsAgent``). + waiting = "waiting for an approval" if approval else f"waiting for signal {signal!r}" + wait = _signals.waits_in(call.ctx, waiting).wait + async def made(c: ToolCall) -> Any: wait_id = c.idempotency_key request: _signals.WaitRequest = {"id": wait_id, "name": signal} if shown: request["input"] = input - return await c.ctx.call(_signals.wait, request, key=_signals.signals_of(c.ctx), idempotency_key=wait_id) + return await c.ctx.call(wait, request, key=c.ctx.key, idempotency_key=wait_id) return await _once(f"signal/{signal}", input, call, made, rejoin=rejoin, approval=approval) diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py index 38ec530..8b50b7b 100644 --- a/python/restate/agents/_turn.py +++ b/python/restate/agents/_turn.py @@ -14,6 +14,14 @@ to the clients that follow the conversation. The conversation itself is already in the key's database, where the ``messages`` view reads it; a publish only tells the clients to read again. + +The layers: the conversation (``answer_once``) stores a turn, and knows +nothing of the companion; the agent's ``chat`` turn (``chat_turn``) starts, +publishes and closes it. A turn of another handler of the agent publishes +each of its steps and its end (``publish_steps``, ``publish_end``), but is no +turn in the companion: no ``start``, no ``close``. A handler of one's own +that prompts, on a session object that is no agent, stores its turn and +publishes nothing. """ from __future__ import annotations @@ -26,9 +34,19 @@ from dataclasses import dataclass, field from restate.agents import _signals -from restate.agents._api import Answering, Published, Steer +from restate.agents._api import AgentContext, Answering, Published, Steer from restate.agents._calls import settle -from restate.agents._common import BUDGET, WINDOW, CompactHistory, Summarize, Window, flush_held, refusal_of +from restate.agents._common import ( + BUDGET, + WINDOW, + CompactHistory, + Summarize, + Window, + after_answer, + after_steps, + flush_held, + refusal_of, +) from restate.exceptions import TerminalError from restate.sessions import SessionContext from restate.sessions._link import Link @@ -154,7 +172,7 @@ def settings_of(ctx: SessionContext) -> AgentInfo: return _agent_of(ctx) or AgentInfo() -async def cleared(ctx: SessionContext) -> None: +async def cleared(ctx: AgentContext) -> None: """Clear whatever the agent of ``ctx`` keeps outside its database.""" agent = _agent_of(ctx) if agent is not None and agent.clear is not None: @@ -181,11 +199,9 @@ class Inbox: revision: int = -1 messages: list[Steer] = field(default_factory=list[Steer]) - def saw(self, published: Published | None) -> None: + def saw(self, published: Published) -> None: """Keep what ``publish`` returned, if it is newer than what is kept.""" - if not isinstance(published, dict): - return - revision = published.get("revision", -1) + revision = published["revision"] if revision > self.revision: self.revision = revision self.messages = list(published.get("steering", [])) @@ -221,32 +237,77 @@ async def run() -> None: return done -async def publish(ctx: SessionContext) -> None: +async def publish(ctx: AgentContext) -> None: """Publish a change to the conversation of ``ctx``: the revision increases, and the clients read it again. - While a framework is answering (inside ``answering``), the calls whose - results are stored are first forgotten. The rows the conversation holds - back are written first: the clients read them once told to. + The rows the conversation holds back are written first: the clients + read them once told to. It forgets no calls: ``answer_once`` does that + after each write of steps that commits, before it publishes there + (``settled``), and the ``settle`` handler publishes outside a turn. + + Only an agent publishes, to its companion: after each step its handlers' + turns store (``publish_steps``), after the answer of a turn that no + ``close`` ends (``publish_end``), and its ``settle``. A session object that + is no agent has no companion, and nothing follows it: its turns are only + stored. + + Raises: + TypeError: If ``ctx`` is not the context of an agent's handler. """ + companion = _signals.companion_in(ctx) await flush_held(ctx) - stored = _ANSWERING.get(ctx) - if stored is not None: - await _settled(ctx, stored) - published = await ctx.call(_signals.publish, key=_signals.signals_of(ctx), name=_signals.ENTRY_NAME) + published = await ctx.call(companion.publish, key=ctx.key, name=_signals.ENTRY_NAME) inbox = _INBOX.get(ctx) if inbox is not None: - inbox.saw(typing.cast(Published | None, published)) + inbox.saw(published) -async def stepped(ctx: SessionContext) -> None: - """Publish a step of the turn, unless the agent runs in the background. +def publish_steps(ctx: AgentContext) -> None: + """Have each write of steps that commits, in the turns that ``ctx``'s handler answers, publish: unless the + agent runs in the background, which publishes only as a ``chat`` turn ends. - ``answer_once`` has it made after each write of the turn's steps that - commits (see ``Conversation.hold``): a reply with tool calls, the tools' - results. + ``chat_turn`` calls it, and so does every handler of an agent + (``AgentObject``), a handler of your own that prompts included: its + followers see its steps as they are stored. ``answer_once`` runs it + (``after_steps``), and knows nothing of it. + """ + if not _unfollowed(ctx): + after_steps(ctx, lambda: publish(ctx)) + + +def publish_end(ctx: AgentContext) -> None: + """Have the commit that stores the answer of a turn that ``ctx``'s handler answers publish: unless the agent + runs in the background, which publishes only as a ``chat`` turn ends. + + Every handler of an agent (``AgentObject``) calls it, as it calls + ``publish_steps``: a follower of such a turn sees its steps as they are + stored, and then the turn end, instead of hearing nothing more. A turn + closed for good (failed, refused, stopped) publishes its end too: its + close stores the turn's answer. One whose attempt failed publishes none: + the turn has not ended, and the next attempt carries it on. An attempt + that died between the commit that stored the answer and this publish + leaves it to the retry, which finds the turn answered and publishes its + end before it answers. + + ``chat_turn`` undoes it (``after_answer(ctx, None)``): a ``chat`` turn + ends with the companion's ``close``, which publishes the turn's last + change with it, as one revision. ``answer_once`` runs what is set + (``after_answer``), and knows nothing of it. """ if not _unfollowed(ctx): - await publish(ctx) + after_answer(ctx, lambda: publish(ctx)) + + +async def settled(ctx: SessionContext) -> None: + """Forget the calls whose results the turn being answered (inside ``answering``) has stored (see ``_calls.py``). + + ``answer_once`` has it done after each write of the turn's steps that + commits, before what its caller set (``after_steps``), and each tool + call's start does it. Outside ``answering``, it does nothing. + """ + stored = _ANSWERING.get(ctx) + if stored is not None: + await _settled(ctx, stored) async def tool_call_started(ctx: SessionContext) -> None: @@ -255,9 +316,7 @@ async def tool_call_started(ctx: SessionContext) -> None: It publishes nothing: the turn publishes once the tool call is stored (with its reply), and once its result is. """ - stored = _ANSWERING.get(ctx) - if stored is not None: - await _settled(ctx, stored) + await settled(ctx) async def tool_call_ended(ctx: SessionContext) -> None: @@ -306,7 +365,7 @@ def _now() -> str: return _signals.iso(time.time()) -async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionContext, str], Awaitable[T]]) -> T: +async def chat_turn(ctx: AgentContext, text: str, answer: Callable[[AgentContext, str], Awaitable[T]]) -> T: """Run ``answer`` as the turn of a ``chat`` handler. While it answers, the object's ``status`` reports the turn. Cancelling the @@ -314,8 +373,14 @@ async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionCon and cancels the calls its tools made; the handler then fails with the cancellation. If the attempt ends first, the turn stops with ``AttemptEnded`` and the next attempt carries on. The turn publishes when - it starts (the companion's ``start`` does) and when it ends (its - ``close``). + it starts (the companion's ``start`` does), after each write of its steps + that commits (a reply with tool calls, the tools' results: ``answer_once`` + runs what ``after_steps`` sets, see ``Conversation.hold``), unless the + agent runs in the background, and when it ends (its ``close``, which + publishes the turn's last change with it: so the commit of its answer + publishes nothing of its own, unlike the turns of the agent's other + handlers, see ``publish_end``). The conversation itself only stores: a + handler of one's own that prompts stores its turn and publishes nothing. An error ends the turn for good if it is a ``TerminalError``, or a model provider's refusal (``refusal_of``), as a 424. Any other error fails the @@ -327,26 +392,28 @@ async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionCon AI's usage limits, retries that ran out and setup errors. """ invocation_id = ctx.request.id - companion = _signals.signals_of(ctx) + companion = _signals.companion_in(ctx) turn: Answering = {"invocationId": invocation_id, "since": _now()} settings = settings_of(ctx) steers = settings.steers and not settings.background if steers: # Before the first publish, which may return messages steered to the turn already. _INBOX.set(ctx, Inbox()) - # From here until the turn closes, the companion knows about it: `status` and `stop` ask it. - # The agent has ``settle`` (``chat()`` adds it), so the companion may send it this turn once it is gone. + # From here until the turn closes, the companion knows about it: `status` and `stop` ask it. Once the + # turn's invocation is gone, the companion sends it to the agent's `settle` (which `chat()` adds). followed = not _unfollowed(ctx) - started: _signals.Started = {**turn, "settles": True, "steers": steers, "publishes": followed} - published = await ctx.call(_signals.start, started, key=companion, name=_signals.ENTRY_NAME) - if isinstance(published, dict): - # The companion published the turn's start as it started it. - inbox = _INBOX.get(ctx) - if inbox is not None: - inbox.saw(published) - else: - # A companion that does not publish as it starts a turn (of an earlier deployment): publish here. - await stepped(ctx) + started: _signals.Started = {**turn, "steers": steers, "publishes": followed} + # A followed turn's start is its first publish; a background agent's publishes nothing (and takes no steers). + published = await ctx.call(companion.start, started, key=ctx.key, name=_signals.ENTRY_NAME) + inbox = _INBOX.get(ctx) + if inbox is not None and published is not None: + inbox.saw(published) + # Each step the turn stores that a tool waits for (a reply with tool calls, the tools' results). + publish_steps(ctx) + # Nothing follows the commit of the answer: this turn ends with the companion's `close`, which publishes + # its last change with it. So it keeps the end the handler's wrapper gave it (`publish_end`) from making a + # second publish, and a `chat` turn's calls are as they were. + after_answer(ctx, None) async def end() -> None: # Have the companion close the turn, which publishes its last change @@ -355,7 +422,7 @@ async def end() -> None: stored = _ANSWERING.get(ctx) if stored is not None: await _settled(ctx, stored) - await ctx.call(_signals.close, invocation_id, key=companion, name=_signals.ENTRY_NAME) + await ctx.call(companion.close, invocation_id, key=ctx.key, name=_signals.ENTRY_NAME) stop = Stop() _STOP.set(ctx, stop) diff --git a/python/restate/agents/internal.py b/python/restate/agents/internal.py index 43493e9..06479bd 100644 --- a/python/restate/agents/internal.py +++ b/python/restate/agents/internal.py @@ -10,7 +10,9 @@ # """Building blocks for writing a framework integration of ``restate.agents``. -An integration needs, beyond the public API: +An integration needs, beyond the public API (``AgentContext`` and +``AgentViewContext``, the contexts of an agent's handlers and views, which +an integration's ``answer`` and any handler of its own take): - ``Conversation``, ``Turn`` and ``answer_once``: the conversation every integration stores the same way, as turns, and the run of a prompt on it @@ -32,13 +34,17 @@ - ``refusal_of`` and ``from_tool``, to tell a model provider's refusal from a tool's error. - ``companion_of`` and ``deliver``, to see what an agent is doing and to send - it a message from another session object's handler. + it a message from another session object's handler. Each agent has its own + companion, the ingress-private virtual object ``.prebuilt.restate`` + (``COMPANION_SUFFIX``) keyed as the agent: ``companion_of(agent)`` returns it, + with its handlers, to call with the agent's key. - ``take_steers``, at each step of a turn: the items of the messages steered to it, which the next model request gets (an integration says it calls it with ``chat(..., steers=True)``, and its ``Items`` may make a steered message's item, with ``steered``). """ +from restate.agents._api import AgentContext, AgentViewContext from restate.agents._common import ( CANCELLED, CHECK_MODEL, @@ -73,7 +79,7 @@ answer_once, ) from restate.agents._output import answer_with, asking -from restate.agents._signals import Companion, companion_of, deliver, signals +from restate.agents._signals import COMPANION_SUFFIX, Companion, companion_of, deliver from restate.agents._steering import take_steers from restate.agents._tools import run_tool, wait_for from restate.agents._turn import ( @@ -89,11 +95,14 @@ __all__ = [ "CANCELLED", "CHECK_MODEL", + "COMPANION_SUFFIX", "DEFAULT_MODEL", "ENDED", "MODEL_ENV", "REFUSED", "UNANSWERED", + "AgentContext", + "AgentViewContext", "AttemptEnded", "Call", "Companion", @@ -125,7 +134,6 @@ "refusal", "refusal_of", "run_tool", - "signals", "stop_signal_of", "stored_answer", "take_steers", diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py index f2b2112..232b8e1 100644 --- a/python/restate/agents/langgraph/_graph.py +++ b/python/restate/agents/langgraph/_graph.py @@ -217,7 +217,7 @@ async def _copied( messages whose ids are not ``known``. ``known`` is the ids the conversation has, which this adds to; the messages the conversation gave the graph are not copied back. Tool messages that follow one another (a step's results) are copied in one - write. Each write that commits publishes (``Conversation.hold``). + write. Each write that commits publishes, in an agent's ``chat`` turn (``Conversation.hold``). """ at = next((i for i, m in enumerate(messages) if m.id == prompt), None) new = [ @@ -322,6 +322,7 @@ async def _resumed(ctx: SessionContext, snapshot: StateSnapshot) -> dict[str, An pending = _interrupts(snapshot) if not pending: return None + companion = _signals.waits_in(ctx, "resuming an interrupt") # What the turn holds back is written before it waits: the waits outlive the attempt. await flush_held(ctx) waits: dict[str, Any] = {} @@ -332,7 +333,7 @@ async def _resumed(ctx: SessionContext, snapshot: StateSnapshot) -> dict[str, An "name": INTERRUPT, "input": to_jsonable_python(interrupted.value), } - waits[interrupted.id] = ctx.call(_signals.wait, request, key=_signals.signals_of(ctx), idempotency_key=wait_id) + waits[interrupted.id] = ctx.call(companion.wait, request, key=ctx.key, idempotency_key=wait_id) values = await asyncio.gather(*waits.values()) return dict(zip(waits, values, strict=True)) diff --git a/python/restate/agents/langgraph/nodes.py b/python/restate/agents/langgraph/nodes.py index 184015b..bd09961 100644 --- a/python/restate/agents/langgraph/nodes.py +++ b/python/restate/agents/langgraph/nodes.py @@ -205,6 +205,7 @@ def wait_for_signal( async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: step = _step(runtime, "wait_for_signal") + wait = _signals.waits_in(step.ctx, f"waiting for signal {signal!r}").wait await flush_held(step.ctx) # The wait's id starts with the invocation's (the same in each attempt), as the waits of tools do, so that a # turn that ends without it (killed, say) has it closed: the companion closes the waits whose ids start @@ -214,7 +215,7 @@ async def node(state: Any, *, runtime: Runtime[Any]) -> dict[str, Any]: 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=wait_id) + value = await step.ctx.call(wait, request, key=step.ctx.key, idempotency_key=wait_id) return _update(output, value) return node diff --git a/python/restate/agents/openai/_agent.py b/python/restate/agents/openai/_agent.py index 970bac7..4a4a10d 100644 --- a/python/restate/agents/openai/_agent.py +++ b/python/restate/agents/openai/_agent.py @@ -299,8 +299,8 @@ def _current(conversation: Conversation, root: Agent[Any], everyone: dict[str, A class _Hooks(RunHooks[Any]): - """Store each step of a run as it comes (each write that commits publishes), and settle the earlier calls as - each tool call starts.""" + """Store each step of a run as it comes (in an agent's ``chat`` turn, each write that commits publishes), and + settle the earlier calls as each tool call starts.""" def __init__(self, ctx: SessionContext, conversation: Conversation, config: RunConfig) -> None: self.ctx = ctx diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index be5bc85..2ad0f9f 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -158,6 +158,7 @@ CompactHistory, ErrorSet, ModelError, + NeedsAgent, PerSession, check_built_in, check_own, @@ -350,8 +351,8 @@ class _Restate(AbstractCapability[Any]): """Make a Pydantic AI run durable in a Restate handler. It stores the requests before each model request and each response as it - arrives (each write that commits publishes, for clients that follow the - conversation), and settles the earlier calls as each tool call starts. A + arrives (in an agent's ``chat`` turn, each write that commits publishes, + for clients that follow the conversation), and settles the earlier calls as each tool call starts. A tool's ``restate.TerminalError`` becomes its answer to the model; any other error fails the attempt. It records the errors of model requests, the only ones that can be a refusal. @@ -455,7 +456,7 @@ async def on_tool_execute_error( stop = stop_signal_of(self.ctx) if stop is not None and stop.reason is not None: raise stop.reason from error - if isinstance(error, TerminalError) and not is_cancellation(error): + if isinstance(error, TerminalError) and not is_cancellation(error) and not isinstance(error, NeedsAgent): return f"Error: {error.message}" # Mark the error as the tool's, so that it is never mistaken for the model provider refusing. from_tool(self.ctx, error) diff --git a/python/restate/endpoint.py b/python/restate/endpoint.py index 47892b4..4fd8ab7 100644 --- a/python/restate/endpoint.py +++ b/python/restate/endpoint.py @@ -132,7 +132,7 @@ def app( A restate ASGI application that hosts the given services. ``services`` are services, virtual objects and workflows, or anything else ``restate.Servable``: - a session object, an agent (which brings the ``signals`` object its tools wait on), or a group of + a session object, an agent (which brings its own companion, ``.prebuilt.restate``), or a group of them. Each is served with what it brings along, once:: app = restate.app([greeter, notes, assistant]) diff --git a/python/restate/servable.py b/python/restate/servable.py index c317c98..ed2a58d 100644 --- a/python/restate/servable.py +++ b/python/restate/servable.py @@ -12,9 +12,9 @@ A service, a virtual object and a workflow are servable: each serves itself. So is anything else that says what to serve for it, such as an agent, which -serves itself and the ``signals`` object its tools wait on:: +serves itself and its own companion, the object its tools wait on:: - app = restate.app([greeter, assistant]) # greeter, assistant, and signals + app = restate.app([greeter, assistant]) # greeter, assistant, and assistant.prebuilt.restate """ from __future__ import annotations @@ -37,7 +37,7 @@ class Servable(typing.Protocol): ``__restate_services__`` returns the things to serve. A service, a virtual object or a workflow lists itself, and is served as it is; anything else it lists is served the same way, in turn. So a library's object can bring what - it needs along (an agent lists itself and ``signals``), and a group of + it needs along (an agent lists itself and its companion), and a group of services can be one servable that lists them:: class Team: diff --git a/scripts/check_no_extras.py b/scripts/check_no_extras.py index 92fc262..1f7906b 100644 --- a/scripts/check_no_extras.py +++ b/scripts/check_no_extras.py @@ -111,7 +111,7 @@ async def add(ctx: sessions.SessionContext, text: str) -> None: manifest = asyncio.run(discover(restate.app([notes, restate.agents.AgentObject("assistant")]))) served = sorted(s["name"] for s in manifest["services"]) - check("sessions and agents serve", served == ["assistant", "notes", "signals"], served) + check("sessions and agents serve", served == ["assistant", "assistant.prebuilt.restate", "notes"], served) refused("restate.agents.tool", lambda: restate.agents.tool(lambda city: city), "pydantic") refused("restate.agents.langgraph", lambda: __import__("restate.agents.langgraph"), "restate-sdk[langgraph]") diff --git a/scripts/resurrect.py b/scripts/resurrect.py index 7851dd2..2271d9b 100644 --- a/scripts/resurrect.py +++ b/scripts/resurrect.py @@ -659,8 +659,8 @@ def kind(self) -> str: "pydantic-denial-approves", "a deferred call's approval signal with any text runs the tool, a denial's reason included", f"{AGENTS}/pydantic_ai/_agent.py", - " return ToolApproved() if value.strip().lower() in _APPROVED else ToolDenied(value)", - " return ToolApproved()", + " if said in _APPROVED:\n return ToolApproved()", + " if True:\n return ToolApproved()", PYDANTIC_AI, ), Mutant( @@ -717,7 +717,7 @@ def kind(self) -> str: "steer-taken-never-sent", "the companion keeps a message the turn took, and status shows it waiting", f"{AGENTS}/_steering.py", - ' tx.send(_signals.taken, ids, key=companion, idempotency_key=f"taken:{ids[0]}")\n', + ' tx.send(companion.taken, ids, key=ctx.key, idempotency_key=f"taken:{ids[0]}")\n', " pass\n", STEERING, ), diff --git a/scripts/steering_traffic.py b/scripts/steering_traffic.py index 5a51dea..f0011dc 100644 --- a/scripts/steering_traffic.py +++ b/scripts/steering_traffic.py @@ -42,6 +42,11 @@ from typing import Any +def _companion(service: str) -> bool: + """Whether ``service`` is an agent's companion, ``.prebuilt.restate``.""" + return service.endswith(".prebuilt.restate") + + async def query(client: Any, admin: str, sql: str) -> list[dict[str, Any]]: response = await client.post(f"{admin}/query", json={"query": sql}, headers={"accept": "application/json"}) response.raise_for_status() @@ -58,7 +63,7 @@ async def cost(client: Any, admin: str, invocation: str) -> dict[str, Any]: f"SELECT target_service_name, target_handler_name FROM sys_invocation WHERE invoked_by_id = '{invocation}'", ) companion = collections.Counter( - row["target_handler_name"] for row in called if row["target_service_name"] == "signals" + row["target_handler_name"] for row in called if _companion(row["target_service_name"]) ) return { "journal": len(journal), diff --git a/tests/agents_adk.py b/tests/agents_adk.py index b50701f..a318d56 100644 --- a/tests/agents_adk.py +++ b/tests/agents_adk.py @@ -36,6 +36,7 @@ from restate import HarnessEnvironment, sessions # noqa: E402 from restate.agents import adk, tool # noqa: E402 from restate.agents._api import Messages # noqa: E402 +from restate.agents._signals import with_companion # noqa: E402 from restate.agents.internal import CANCELLED # noqa: E402 @@ -46,6 +47,9 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] +REVISION = b'{"revision": 1}' +"""What the agent's companion answers a turn's ``start`` and each ``publish`` with: the revision.""" + class AuthenticationError(Exception): """As a model SDK names the error of a missing or wrong API key.""" @@ -606,11 +610,11 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: return await adk.prompt(ctx, text, model=Answering()) backend, invocation = Backend(), request() - host = Host(backend, service="asker", answers={"publish": b"null"}) + host = Host(backend, service="asker", answers={"publish": REVISION}) with pytest.raises(RuntimeError, match="crashed after storing"): await invoke(asker, "ask", host, invocation, "hi") # The retry (the same invocation) returns the answer that was stored, without asking the model again. - host = Host(backend, service="asker", answers={"publish": b"null"}) + host = Host(backend, service="asker", answers={"publish": REVISION}) assert await invoke(asker, "ask", host, invocation, "hi") == "answer 1" assert replies == ["asked"] @@ -654,7 +658,7 @@ async def stored(ctx: sessions.SessionContext) -> list[bool]: backend = Backend() for text in ("first", "second"): - host[:] = [Syncing(backend, service="asker", answers={"publish": b"null"})] + host[:] = [Syncing(backend, service="asker", answers={"publish": REVISION})] assert await invoke(asker, "ask", host[0], request(), text) == "hello" # The second turn (the conversation's tables exist): from the model's response on, one commit, which has the # final response, the end of the agent and the answer. @@ -691,7 +695,9 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: backend = Backend() async def asked(text: str) -> Any: - return await invoke(asker, "ask", Host(backend, service="asker", answers={"publish": b"null"}), request(), text) + return await invoke( + asker, "ask", Host(backend, service="asker", answers={"publish": REVISION}), request(), text + ) assert await asked("system x") == "system: Mine." assert await asked("weather Rome") == "tool said: sunny in Rome" @@ -718,7 +724,7 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: backend = backend or Backend() async def asked(text: str, invocation: Any = None) -> Any: - host = Host(backend, service="asker", answers={"publish": b"null"}) + host = Host(backend, service="asker", answers={"publish": REVISION}) return await invoke(asker, "ask", host, invocation or request(), text) asked.backend = backend # type: ignore[attr-defined] @@ -1055,7 +1061,7 @@ async def taken(ctx: sessions.SessionContext) -> bool: return await steer_stored(ctx.db, "s1") backend = Backend() - host = Host(backend, service="asker", answers={"publish": b"null"}) + host = Host(backend, service="asker", answers={"publish": REVISION}) answer = await invoke(asker, "ask", host, request(), "weather") assert answer == ('{"summary": "sunny"}' if output else "") # No request after the tool's result: the turn left the message to a turn of its own. @@ -1074,7 +1080,7 @@ async def test_a_steered_message_starts_an_invocation_of_its_own_which_a_retry_r from restate.agents._turn import _INBOX, Inbox # pylint: disable=import-outside-toplevel from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel - inbox = Inbox(revision=0) + inbox = Inbox(revision=1) # as the companion's start returned it, which each publish answers again crashed: list[bool] = [] def get_weather(city: str) -> str: # pylint: disable=redefined-outer-name @@ -1096,6 +1102,7 @@ async def generate_content_async( steps = SequentialAgent(name="steps", sub_agents=[LlmAgent(name="step", model=Crashing(), tools=[get_weather])]) asker = sessions.SessionObject("asker") + with_companion(asker) # the companion it takes steered messages from, as an agent's @asker.handler async def ask(ctx: sessions.SessionContext, text: str) -> str: @@ -1106,7 +1113,7 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: backend, invocation = Backend(), request() async def attempt() -> Any: - host = Host(backend, service="asker", answers={"publish": b"null"}) + host = Host(backend, service="asker", answers={"publish": REVISION}) return await invoke(asker, "ask", host, invocation, "weather Rome") with pytest.raises(RuntimeError, match="crashed after the steered message was stored"): diff --git a/tests/agents_companion.py b/tests/agents_companion.py new file mode 100644 index 0000000..4fcbb61 --- /dev/null +++ b/tests/agents_companion.py @@ -0,0 +1,110 @@ +# +# 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 +# +"""Each agent's own companion, in process: named for its agent, keyed as its agent, and served with it only. + +The companion's handlers are tested against a server in ``agents_signals.py``, and through an agent's views in +``agents_object.py``.""" + +# pylint: disable=missing-function-docstring,import-error + +import pytest +from agents_conversation import ITEMS +from sessions_database import Backend +from sessions_object import Host, request + +import restate +from restate import sessions +from restate.agents import AgentObject +from restate.agents._conversation import Conversation, answer_once +from restate.agents._signals import COMPANION_SUFFIX, companion_name, companion_of +from restate.endpoint import Endpoint +from restate.sessions.internal import invoke + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + +REVISION = b'{"revision": 1}' + + +def agent(name: str) -> AgentObject: + """An agent that echoes each message.""" + obj = AgentObject(name) + + @obj.chat(items=ITEMS) + async def answer(ctx: sessions.SessionContext, text: str) -> str: + async def run(conversation: Conversation) -> str: + await conversation.turn.append_reply({"role": "assistant", "text": f"echo: {text}"}) + return f"echo: {text}" + + return await answer_once(ctx, text, ITEMS, run) + + return obj + + +def test_each_agent_has_a_companion_of_its_own_named_for_it(): + assert COMPANION_SUFFIX == ".prebuilt.restate" + first, second = agent("first"), agent("second") + assert companion_of(first).service.name == companion_name("first") == "first.prebuilt.restate" + assert companion_of(second).service.name == "second.prebuilt.restate" + assert companion_of(first).service is not companion_of(second).service + assert (companion_of(first).agent, companion_of(second).agent) == ("first", "second") + # The agent's name is taken as it is. + assert companion_of(AgentObject("Odd-Name_1")).service.name == "Odd-Name_1.prebuilt.restate" + + +def test_an_app_serves_each_agent_with_its_companion_and_no_shared_signals(): + first, second = agent("first"), agent("second") + endpoint = Endpoint().bind(first, second) + assert list(endpoint.services) == [ + "first", + "first.prebuilt.restate", + "second", + "second.prebuilt.restate", + ] + assert "signals" not in endpoint.services + # The companion is ingress-private, all of it: clients reach it through its agent's views. + for name in ("first.prebuilt.restate", "second.prebuilt.restate"): + companion = endpoint.services[name] + assert companion.ingress_private is True + assert all(not handler.ingress_private for handler in companion.handlers.values()) + # Your own services may be named as you like, "signals" too. + assert "signals" in Endpoint().bind(first, restate.VirtualObject("signals")).services + + +async def test_a_turn_calls_its_own_agents_companion_with_the_agents_key_and_not_the_others(): + first, second = agent("first"), agent("second") + host = Host(Backend(), service="first", key="alice", answers={"start": REVISION, "close": b"null"}) + assert await invoke(first, "chat", host, request(), "hi") == "echo: hi" + called = [(service, handler, key) for service, handler, _, key, _ in host.calls] + assert called == [ + ("first.prebuilt.restate", "start", "alice"), + ("first.prebuilt.restate", "close", "alice"), + ] + # Nothing reaches the other agent's companion. + other = companion_of(second).service.name + assert other not in [service for service, *_ in host.calls] + [send.service for send in host.sends] + + +async def test_the_views_ask_the_agents_own_companion_at_the_agents_key(): + first = agent("first") + status = b'{"answering": null, "waiting": []}' + view = Host(Backend(), service="first", key="bob", shared=True, answers={"status": status, "watch": REVISION}) + assert await invoke(first, "status", view, request(), None) == {"answering": None, "waiting": []} + assert await invoke(first, "watch", view, request(), {"revision": 0}) == {"revision": 1} + assert [(service, handler, key) for service, handler, _, key, _ in view.calls] == [ + ("first.prebuilt.restate", "status", "bob"), + ("first.prebuilt.restate", "watch", "bob"), + ] diff --git a/tests/agents_conversation.py b/tests/agents_conversation.py index 1cbbc81..2df72d2 100644 --- a/tests/agents_conversation.py +++ b/tests/agents_conversation.py @@ -40,6 +40,7 @@ last_of, ) from restate.agents._common import Check +from restate.agents._signals import with_companion from restate.agents.internal import CANCELLED, from_tool, refusal_of from restate.sessions.internal import invoke @@ -51,8 +52,11 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] -PUBLISHED = {"publish": b"null"} -"""What the signals object answers when the turn publishes its progress.""" +REVISION = b'{"revision": 1}' +"""What the agent's companion answers a turn's ``start`` and each ``publish`` with: the revision.""" + +PUBLISHED = {"publish": REVISION} +"""What the agent's companion answers when the turn publishes its progress.""" class Items: @@ -980,8 +984,12 @@ def generic_call(self, service, handler, arg, key=None, idempotency_key=None, he def _turn_rows() -> sessions.SessionObject: - """An object whose ``ask`` answers with the run that ``runs`` holds, and whose ``rows`` reads each turn back.""" + """An object whose ``ask`` answers with the run that ``runs`` holds, and whose ``rows`` reads each turn back. + + It has a companion, as an agent has, which ``publish`` publishes to: its turns, outside an agent's ``chat``, + publish nothing themselves.""" obj = sessions.SessionObject("asker") + with_companion(obj) @obj.handler async def rows(ctx: sessions.SessionContext) -> list[dict[str, Any]]: @@ -1064,7 +1072,7 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: async def test_a_reply_held_back_is_stored_before_a_publish_that_follows_it(): - from restate.agents._turn import stepped # pylint: disable=import-outside-toplevel + from restate.agents._turn import publish # pylint: disable=import-outside-toplevel obj = _turn_rows() @@ -1073,7 +1081,8 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: async def run(conversation: Conversation) -> str: await conversation.turn.append_reply({"role": "assistant", "text": "on it"}) # The clients are told to read the reply: it is stored by then, even if the process crashes after. - await stepped(ctx) + # ``publish`` takes an agent's context; this object stands in for one, so the test says so. + await publish(typing.cast(Any, ctx)) await conversation.turn.append_reply({"role": "assistant", "text": "done"}) return "done" @@ -1223,3 +1232,60 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: assert await invoke(obj, "ask", host, request(), text) == "hello" syncs.append(host.syncs) assert syncs[0] - syncs[1] == 1 + + +# --- Layers: the conversation stores, an agent's turn publishes ------------------------------------------------- + + +def _no_companion(*_: Any) -> Any: + raise AssertionError("a session object that is no agent has no companion to call, nor anything to publish") + + +async def test_answer_once_in_a_plain_session_object_stores_a_turn_of_tool_calls_and_publishes_nothing( + monkeypatch: pytest.MonkeyPatch, +): + """The conversation only stores: its commits publish nothing unless an agent's ``chat`` turn asks for it.""" + from restate.agents import _signals, _turn # pylint: disable=import-outside-toplevel + + monkeypatch.setattr(_turn, "publish", _no_companion) + monkeypatch.setattr(_signals, "companion_at", _no_companion) + obj = sessions.SessionObject("desk") + + @obj.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + async def run(conversation: Conversation) -> str: + calls = [{"id": "t1", "name": "lookup", "input": {}}] + await conversation.turn.append_reply({"role": "assistant", "text": "", "calls": calls}) + await conversation.turn.append({"result": "t1", "output": "found"}) + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + return await answer_once(ctx, text, ITEMS, run) + + @obj.handler + async def rows(ctx: sessions.SessionContext) -> list[Any]: + conversation = await Conversation.load(ctx.db, ITEMS, "all") + return [[turn.answer, len(turn.items)] for turn in conversation.turns] + + backend = Backend() + host = Host(backend, service="desk") + assert await invoke(obj, "ask", host, request(), "hi") == "done" + assert host.calls == [] + assert await invoke(obj, "rows", Host(backend, service="desk"), request()) == [["done", 3]] + + +async def test_a_publish_from_a_session_object_that_is_no_agent_raises(): + from restate.agents.internal import publish # pylint: disable=import-outside-toplevel + + obj = sessions.SessionObject("desk") + + @obj.handler + async def tell(ctx: sessions.SessionContext) -> None: + # The cast is the test: ``publish`` takes an agent's context, and a plain one passed off as one still + # raises when it runs. The names are a signpost; the companion is the fence. + await publish(typing.cast(Any, ctx)) + + host = Host(Backend(), service="desk") + with pytest.raises(TypeError, match="companion_in takes the context of an agent's handler"): + await invoke(obj, "tell", host, request()) + assert host.calls == [] diff --git a/tests/agents_langgraph.py b/tests/agents_langgraph.py index 69cabeb..4cf3654 100644 --- a/tests/agents_langgraph.py +++ b/tests/agents_langgraph.py @@ -48,6 +48,9 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] +REVISION = b'{"revision": 1}' +"""What the agent's companion answers a turn's ``start`` and each ``publish`` with: the revision.""" + class AuthenticationError(Exception): """As a model SDK names the error of a missing or wrong API key.""" @@ -551,7 +554,9 @@ async def ask(ctx: sessions.SessionContext, text: str) -> str: backend = Backend() async def asked(text: str) -> Any: - return await invoke(asker, "ask", Host(backend, service="asker", answers={"publish": b"null"}), request(), text) + return await invoke( + asker, "ask", Host(backend, service="asker", answers={"publish": REVISION}), request(), text + ) return asked diff --git a/tests/agents_object.py b/tests/agents_object.py index 2be2f71..5141361 100644 --- a/tests/agents_object.py +++ b/tests/agents_object.py @@ -95,10 +95,14 @@ async def forget(ctx: sessions.SessionContext) -> None: assistant = AgentObject("assistant") assistant.chat(answer, items=ITEMS, clear=forget) +other = AgentObject("other") +"""A second agent in the same app, with a companion of its own.""" +other.chat(answer, items=ITEMS) + @pytest.fixture(scope="session") async def harness(): - async with restate.create_test_harness(restate.app([assistant]), restate_image=RESTATE_IMAGE) as harness: + async with restate.create_test_harness(restate.app([assistant, other]), restate_image=RESTATE_IMAGE) as harness: yield harness @@ -212,6 +216,32 @@ async def answering() -> bool: assert (await call(harness, "chat", key, "next")).json() == "echo: next" +async def test_each_agent_has_its_own_companion_which_a_turn_of_the_other_does_not_touch(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + async with httpx2.AsyncClient(timeout=60) as admin: + served = { + service["name"] for service in (await admin.get(f"{harness.admin_api_url}/services")).json()["services"] + } + assert {"assistant.prebuilt.restate", "other.prebuilt.restate"} <= served and "signals" not in served + # A turn of the assistant, which waits for an approval at the same key. + invocation = await send(harness, key, "refund o7") + + async def waits() -> bool: + return len((await call(harness, "pending", key)).json()) == 1 + + await until(waits) + assert (await call(harness, "notifications", key)).json()["revision"] > 0 + async with httpx2.AsyncClient(timeout=60) as ingress: + other_status = (await ingress.post(f"{harness.ingress_url}/other/{key}/status")).json() + other_revision = (await ingress.post(f"{harness.ingress_url}/other/{key}/notifications")).json() + unresolved = (await ingress.post(f"{harness.ingress_url}/other/{key}/signal", json={"name": "approval"})).json() + # The other agent's companion, at the same key, knows nothing of it, and has no wait to resolve. + assert other_status == {"answering": None, "waiting": []} and other_revision == {"revision": 0} + assert unresolved == {"resolved": None} + await call(harness, "signal", key, {"name": "approval", "value": "approved"}) + assert (await attach(harness, invocation)).json() == "refund: approved" + + async def test_clear_forgets_the_conversation_and_its_ids_go_on(harness: HarnessEnvironment): key = str(uuid.uuid4()) await call(harness, "chat", key, "one") diff --git a/tests/agents_openai.py b/tests/agents_openai.py index 83a0289..837377d 100644 --- a/tests/agents_openai.py +++ b/tests/agents_openai.py @@ -69,6 +69,9 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] +REVISION = b'{"revision": 1}' +"""What the agent's companion answers a turn's ``start`` and each ``publish`` with: the revision.""" + class AuthenticationError(Exception): """As a model SDK names the error of a missing or wrong API key.""" @@ -440,8 +443,8 @@ async def test_a_handoff_holds_when_the_model_sees_no_earlier_turn_whole(harness # --- In process ----------------------------------------------------------------------------------- -PUBLISHED = {"publish": b"null"} -"""What the signals object answers when the turn publishes its progress (at each tool call).""" +PUBLISHED = {"publish": REVISION} +"""What the agent's companion answers when the turn publishes its progress.""" def asker(**options: Any) -> sessions.SessionObject: @@ -1313,11 +1316,11 @@ async def test_a_message_steered_while_the_model_writes_its_answer_is_the_next_t class Companion(dict[str, bytes]): - """An agent's companion, in this process: it answers a turn's start and close, and each publish with the + """An agent's companion, in this process: it answers a turn's close, and its start and each publish with the messages steered to the turn (``steering``), as the companion does until the turn closes.""" def __init__(self) -> None: - super().__init__(start=b"null") + super().__init__() self.revision = 0 self.sent = 0 self.steering: list[dict[str, str]] = [] @@ -1329,7 +1332,7 @@ def __missing__(self, handler: str) -> bytes: # The messages the turn did not take are each answered as a turn of their own. self.steering.clear() return b"null" - assert handler == "publish", f"the companion has no answer for {handler}" + assert handler in ("start", "publish"), f"the companion has no answer for {handler}" self.revision += 1 return json.dumps({"revision": self.revision, "steering": self.steering}).encode() diff --git a/tests/agents_publishes.py b/tests/agents_publishes.py index 7544981..7456668 100644 --- a/tests/agents_publishes.py +++ b/tests/agents_publishes.py @@ -10,10 +10,12 @@ # """When a followed turn publishes, and what a follower then sees: in this process, for each framework. -A turn publishes as it starts (the companion's ``start`` does), after each -write of steps that commits (a reply with tool calls, the tools' results), and -as it closes. A follower reads the ``messages`` view once told: at each -publish, the companion here reads it as a follower would, at once. +A ``chat`` turn publishes as it starts (the companion's ``start`` does), +after each write of steps that commits (a reply with tool calls, the tools' +results), and as it closes. A turn of another handler of the agent has no +``start`` nor ``close``: it publishes each of its steps, and once more when +its answer is stored. A follower reads the ``messages`` view once told: at +each publish, the companion here reads it as a follower would, at once. """ # pylint: disable=import-error,import-outside-toplevel @@ -28,7 +30,8 @@ import pytest -from restate.agents import tool +from restate import sessions +from restate.agents import _signals, tool from restate.agents._object import AgentObject pytestmark = pytest.mark.anyio @@ -51,6 +54,8 @@ class Seen: """Each companion call that answers with a revision, or a model request, and the results the view showed then.""" steering: list[dict[str, str]] = field(default_factory=list[dict[str, str]]) """The messages steered to the turn, which each publish returns.""" + answers: list[tuple[str, str]] = field(default_factory=list[tuple[str, str]]) + """Each companion call, and the text the assistant's last message showed then (``""`` for none).""" names: list[tuple[str, str | None]] = field(default_factory=list[tuple[str, str | None]]) """Each call the agent made, by handler, with the name of its journal entry.""" @@ -62,45 +67,59 @@ class Companion: """The agent's companion in this process: each publish (and a ``start`` that publishes) reads the ``messages`` view, as a follower told of it would, and returns a new revision with the messages steered to the turn.""" - def __init__(self, obj: AgentObject, backend: Any, seen: Seen, *, start_publishes: bool = True) -> None: + def __init__(self, obj: AgentObject, backend: Any, seen: Seen) -> None: self.obj = obj self.backend = backend self.seen = seen - self.start_publishes = start_publishes self.revision = 0 + self.answered = "" + """The text of the assistant's last message, as the latest read of the view showed it.""" async def results(self) -> int: - """Return how many tool results the ``messages`` view shows.""" + """Return how many tool results the ``messages`` view shows, and record the assistant's last message.""" from sessions_object import Host, request from restate.sessions.internal import invoke view = Host(self.backend, service=self.obj.name).view() page = await invoke(self.obj, "messages", view, request(), {"limit": 1000}) + said = [ + part.get("text", "") + for message in page["messages"] + if message["role"] == "assistant" + for part in message["parts"] + if part["type"] == "text" + ] + self.answered = said[-1] if said else "" return sum("output" in part for message in page["messages"] for part in message["parts"]) async def answer(self, handler: str, arg: Any) -> bytes: - if handler == "start" and not (self.start_publishes and json.loads(arg).get("publishes")): - self.seen.events.append(("start", await self.results())) + def seen(event: str, results: int) -> None: + self.seen.events.append((event, results)) + self.seen.answers.append((event, self.answered)) + + if handler == "start" and not json.loads(arg).get("publishes"): + seen("start", await self.results()) return b"null" if handler in ("start", "publish"): self.revision += 1 - self.seen.events.append((handler, await self.results())) + seen(handler, await self.results()) return json.dumps({"revision": self.revision, "steering": self.seen.steering}).encode() assert handler == "close", f"the companion has no answer for {handler}" - self.seen.events.append(("close", await self.results())) + seen("close", await self.results()) return b"null" -def follow(obj: AgentObject, seen: Seen, *, start_publishes: bool = True) -> Callable[[str], Any]: - """Return a ``chat`` of ``obj`` in this process, each call a turn, whose follower records what it sees in ``seen``.""" +def follow(obj: AgentObject, seen: Seen, handler: str = "chat") -> Callable[[str], Any]: + """Return a ``chat`` of ``obj`` in this process, each call a turn, whose follower records what it sees in ``seen``; + or a call of another of its handlers, ``handler``.""" from sessions_database import Backend from sessions_object import Answer, Host, request from restate.sessions.internal import invoke backend = Backend() - companion = Companion(obj, backend, seen, start_publishes=start_publishes) + companion = Companion(obj, backend, seen) class Deferred(Answer): def __init__(self, handler: str, arg: Any) -> None: @@ -118,7 +137,7 @@ def generic_call(self, service, handler, arg, key=None, idempotency_key=None, he return Deferred(handler, arg) async def chat(text: str) -> Any: - return await invoke(obj, "chat", Followed(backend, service=obj.name), request(), text) + return await invoke(obj, handler, Followed(backend, service=obj.name), request(), text) return chat @@ -139,12 +158,10 @@ async def lookup(about: str) -> str: # --- The frameworks' agents, each with a model that calls ``lookup`` CALLS times at once, then answers ---------- -def pydantic_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: +def pydantic_model(seen: Seen, calls: int = CALLS) -> Any: from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart, UserPromptPart from pydantic_ai.models.function import AgentInfo, FunctionModel - from restate.agents import pydantic_ai - async def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: seen.events.append(("model", 0)) if any(isinstance(m, ModelResponse) for m in messages) or not calls: @@ -154,15 +171,19 @@ async def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: parts=[ToolCallPart("lookup", {"about": f"t{i}"}, tool_call_id=f"c{i}") for i in range(calls)] ) - return pydantic_ai.agent("followed", model=FunctionModel(model), tools=[lookup_of(seen, steer)], **options) + return FunctionModel(model) -def openai_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: +def pydantic_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: + from restate.agents import pydantic_ai + + return pydantic_ai.agent("followed", model=pydantic_model(seen, calls), tools=[lookup_of(seen, steer)], **options) + + +def openai_model(seen: Seen, calls: int = CALLS) -> Any: from agents import Model, ModelResponse, Usage from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText - from restate.agents import openai - class Scripted(Model): async def get_response(self, system_instructions, input, *args, **kwargs): # noqa: A002 seen.events.append(("model", 0)) @@ -191,17 +212,21 @@ async def get_response(self, system_instructions, input, *args, **kwargs): # no def stream_response(self, *args, **kwargs): raise NotImplementedError - return openai.agent("followed", model=Scripted(), tools=[lookup_of(seen, steer)], **options) + return Scripted() -def adk_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: +def openai_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: + from restate.agents import openai + + return openai.agent("followed", model=openai_model(seen, calls), tools=[lookup_of(seen, steer)], **options) + + +def adk_model(seen: Seen, calls: int = CALLS) -> Any: from google.adk.models import BaseLlm from google.adk.models.llm_request import LlmRequest from google.adk.models.llm_response import LlmResponse from google.genai import types - from restate.agents import adk - class Scripted(BaseLlm): model: str = "scripted" @@ -227,10 +252,18 @@ async def generate_content_async( ] yield LlmResponse(content=types.Content(role="model", parts=parts)) - return adk.agent("followed", model=Scripted(), tools=[lookup_of(seen, steer)], **options) + return Scripted() + + +def adk_agent(seen: Seen, steer: str | None = None, *, calls: int = CALLS, **options: Any) -> AgentObject: + from restate.agents import adk + + return adk.agent("followed", model=adk_model(seen, calls), tools=[lookup_of(seen, steer)], **options) AGENTS = {"pydantic_ai": pydantic_agent, "openai": openai_agent, "adk": adk_agent} +MODELS = {"pydantic_ai": pydantic_model, "openai": openai_model, "adk": adk_model} +"""Each framework's scripted model, for its ``prompt``.""" def make(framework: str, seen: Seen, steer: str | None = None, **options: Any) -> AgentObject: @@ -254,11 +287,12 @@ async def test_the_calls_to_the_companion_are_named_for_the_agents_metadata_to_h agent = make(framework, seen) chat = follow(agent, seen) await chat("hello") + companion = _signals.companion_of(agent).service.name assert [handler for handler, _ in seen.names] == [ - "signals/start", - "signals/publish", - "signals/publish", - "signals/close", + f"{companion}/start", + f"{companion}/publish", + f"{companion}/publish", + f"{companion}/close", ] assert {name for _, name in seen.names} == {"signals"} # Restate's UI hides them, with the database's entries; an agent's own metadata still wins. @@ -299,15 +333,6 @@ async def test_a_message_steered_while_the_tools_run_is_taken_at_the_next_step(f assert seen.companion_calls() == ["start", "publish", "publish", "close"] -@pytest.mark.parametrize("framework", list(AGENTS)) -async def test_a_companion_whose_start_returns_nothing_still_gets_the_turns_first_publish(framework: str): - seen = Seen() - chat = follow(make(framework, seen, calls=0), seen, start_publishes=False) - await chat("hello") - # An older companion: the turn publishes its start itself. - assert seen.companion_calls() == ["start", "publish", "close"] - - @pytest.mark.parametrize("framework", list(AGENTS)) async def test_a_background_agent_publishes_only_its_end(framework: str): seen = Seen() @@ -317,6 +342,206 @@ async def test_a_background_agent_publishes_only_its_end(framework: str): assert seen.companion_calls() == ["start", "close"] +@pytest.mark.parametrize("framework", list(AGENTS)) +@pytest.mark.parametrize("background", [False, True]) +async def test_a_handler_of_your_own_on_an_agent_publishes_each_committed_step_and_its_answer( + framework: str, background: bool +): + """No turn in the companion (no ``start`` nor ``close``, which only ``chat`` makes), but its steps publish as + ``chat``'s do, and its answer publishes once more, so that the agent's followers see the turn through to its + end; a background agent's publish nothing.""" + import importlib + + seen = Seen() + agent = make(framework, seen, background=background) + framework_module = importlib.import_module(f"restate.agents.{framework}") + + @agent.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await framework_module.prompt(ctx, text, model=MODELS[framework](seen), tools=[lookup_of(seen)]) + + assert await follow(agent, seen, "ask")("hello") == "answered after hello" + # The reply with the calls (no results yet); the results; the answer, which the final reply commits with. + assert [e for e in seen.events if e[0] != "model"] == ( + [] if background else [("publish", 0), ("publish", CALLS), ("publish", CALLS)] + ) + # The answer is stored before the publish that tells the followers of it: a follower that reads then sees it. + assert seen.answers == ( + [] if background else [("publish", ""), ("publish", ""), ("publish", "answered after hello")] + ) + + +# --- The end of a turn, with a run of our own: the exact publishes, whichever handler answers ------------------ + + +class Items: + """The items of the runs below: a message (``role`` and ``text``), a tool call (``call``), or its result.""" + + def dump(self, item: dict[str, Any]) -> str: + return json.dumps(item) + + def load(self, stored: str) -> dict[str, Any]: + return json.loads(stored) + + def prompt(self, text: str) -> dict[str, Any]: + return {"role": "user", "text": text} + + def reply(self, text: str, agent: str | None = None) -> dict[str, Any]: + return {"role": "assistant", "text": text} + + def results(self, results: Any) -> list[dict[str, Any]]: + return [{"result": call.id, "text": text} for call, text in results] + + def read(self, item: dict[str, Any]) -> Any: + from restate.agents.internal import Call, Parsed, Result + + if "result" in item: + return Parsed(results=[Result(item["result"], item["text"])]) + if "call" in item: + return Parsed(calls=[Call(item["call"], "lookup", {})]) + return Parsed((item["role"], item["text"])) + + +ITEMS = Items() + + +def ending(fails: BaseException | None = None, *, background: bool = False) -> AgentObject: + """An agent whose ``chat`` and whose own ``ask`` handler answer alike, with a run of our own. + + The run stores a reply with a tool call, then the call's result — two + commits, each a step a tool waits for — and then answers ``"done"``, whose + reply is held and commits with the answer. With ``fails``, it raises that + instead of answering, once the result is stored. + """ + from restate.agents.internal import Conversation, answer_once + + async def run(conversation: Conversation) -> str: + turn = conversation.turn + await turn.append_reply({"call": "c1"}) + await turn.append({"result": "c1", "text": "found"}) + if fails is not None: + raise fails + return "done" + + async def answer(ctx: sessions.SessionContext, text: str) -> str: + return await answer_once(ctx, text, ITEMS, run) + + agent = AgentObject("ending", background=background) + agent.chat(answer, items=ITEMS) + + @agent.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + """A handler of your own that prompts: a turn with no start and no close.""" + return await answer(ctx, text) + + return agent + + +async def test_a_turn_of_a_handler_of_your_own_publishes_each_step_and_then_its_answer(): + """Each commit of the turn's steps publishes, and the commit that stores the answer publishes once more: a + follower sees the steps arrive, and then the turn end.""" + seen = Seen() + assert await follow(ending(), seen, "ask")("hello") == "done" + assert seen.companion_calls() == ["publish", "publish", "publish"] + # The reply with the call (no result yet); the result; the answer, which is stored before its publish. + assert seen.events == [("publish", 0), ("publish", 1), ("publish", 1)] + assert seen.answers == [("publish", ""), ("publish", ""), ("publish", "done")] + + +async def test_a_turn_of_a_handler_of_your_own_that_nobody_follows_publishes_nothing(): + """A background agent is unwatched: neither its steps nor the end of its turn publish.""" + seen = Seen() + assert await follow(ending(background=True), seen, "ask")("hello") == "done" + assert seen.companion_calls() == [] + + +async def test_a_turn_of_a_handler_of_your_own_that_fails_for_good_publishes_its_end(): + """A turn closed for good has its answer stored (its error, as the turn's last message), so its end publishes + as an answered turn's does: a follower hears it end, instead of hearing nothing more.""" + from restate.exceptions import TerminalError + + seen = Seen() + with pytest.raises(TerminalError) as raised: + await follow(ending(TerminalError("the tool is gone")), seen, "ask")("hello") + assert raised.value.message == "the tool is gone" + assert seen.companion_calls() == ["publish", "publish", "publish"] + assert seen.answers[-1] == ("publish", "Error: the tool is gone") + + +async def test_a_turn_of_a_handler_of_your_own_whose_attempt_fails_publishes_no_end(): + """A retryable error ends no turn: the attempt fails, the next one carries the turn on, and nothing publishes + an end that never happened. The steps stored before it published, as they commit.""" + seen = Seen() + with pytest.raises(RuntimeError, match="the model is unreachable"): + await follow(ending(RuntimeError("the model is unreachable")), seen, "ask")("hello") + assert seen.companion_calls() == ["publish", "publish"] + assert seen.answers[-1] == ("publish", "") + + +@pytest.mark.parametrize("fails", [None, "terminal"]) +async def test_a_chat_turn_publishes_as_it_starts_after_each_step_and_as_it_closes_only(fails: str | None): + """The same run from ``chat``: its end is the companion's ``close``, which publishes the turn's last change + with it, so the commit of its answer publishes nothing of its own. Its calls are as they were.""" + from restate.exceptions import TerminalError + + seen = Seen() + chat = follow(ending(TerminalError("the tool is gone") if fails else None), seen) + if fails: + with pytest.raises(TerminalError): + await chat("hello") + else: + assert await chat("hello") == "done" + assert seen.companion_calls() == ["start", "publish", "publish", "close"] + assert seen.answers[-1] == ("close", "Error: the tool is gone" if fails else "done") + + +def no_companion(*_: Any) -> Any: + raise AssertionError("a session object that is no agent has no companion to call, nor anything to publish") + + +@pytest.mark.parametrize("framework", list(AGENTS)) +async def test_a_plain_session_objects_prompt_with_tool_calls_stores_its_turn_and_publishes_nothing( + framework: str, monkeypatch: pytest.MonkeyPatch +): + """A handler of one's own, on a session object that is no agent, prompts with a framework: a turn of tool calls. + + The conversation stores the turn; nothing publishes, nor looks for a companion, and the handler makes no Restate + call at all.""" + import importlib + + from sessions_database import Backend + from sessions_object import Host, request + + from restate.agents import _signals, _turn + from restate.sessions.internal import invoke + + monkeypatch.setattr(_turn, "publish", no_companion) + monkeypatch.setattr(_signals, "companion_at", no_companion) + seen = Seen() + framework_module = importlib.import_module(f"restate.agents.{framework}") + desk = sessions.SessionObject("desk") + + @desk.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await framework_module.prompt(ctx, text, model=MODELS[framework](seen), tools=[lookup_of(seen)]) + + @desk.handler + async def stored(ctx: sessions.SessionContext) -> list[Any]: + turns = await ctx.db.values("SELECT prompt, answer FROM restate_agent_turns") + steps = await ctx.db.values("SELECT count(*) FROM restate_agent_items") + return [[list(turn) for turn in turns], steps[0][0]] + + backend = Backend() + host = Host(backend, service="desk") + assert await invoke(desk, "ask", host, request(), "hello") == "answered after hello" + assert host.calls == [] + assert [event for event, _ in seen.events] == ["model", "model"] + [turns, steps] = await invoke(desk, "stored", Host(backend, service="desk"), request()) + assert [answer for _, answer in turns] == ["answered after hello"] + # The reply with the calls, their results, and the final reply. + assert steps >= 3 + + class _State: """A companion's context in this process, for its state only: what it writes, in order.""" @@ -351,3 +576,41 @@ async def test_a_publish_clears_the_watches_only_if_there_are_any(): watched = _State(revision=4, watching=[{"revision": 4, "awakeable": "a1"}]) await _signals._publish(typing.cast(Any, watched)) # pyright: ignore[reportPrivateUsage] assert watched.resolved == ["a1"] and watched.writes == [("set", "revision"), ("clear", "watching")] + + +@pytest.mark.parametrize("framework", list(AGENTS)) +async def test_a_signal_tool_of_a_plain_session_object_fails_its_turn_for_good_and_the_model_never_sees_it( + framework: str, +): + """Waiting needs an agent: the turn fails once, terminally, and is closed ``failed``; the model is not asked + again with the error as the tool's result, as it could not fix how the object is set up.""" + import importlib + + from sessions_database import Backend + from sessions_object import Host, request + + from restate.agents import wait_for_signal + from restate.exceptions import TerminalError + from restate.sessions.internal import invoke + + seen = Seen() + waits = wait_for_signal("approval", description="Wait for an approval.", name="lookup") + framework_module = importlib.import_module(f"restate.agents.{framework}") + desk = sessions.SessionObject("desk") + + @desk.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await framework_module.prompt(ctx, text, model=MODELS[framework](seen), tools=[waits]) + + @desk.handler + async def stored(ctx: sessions.SessionContext) -> list[Any]: + return [list(row) for row in await ctx.db.values("SELECT status, answer FROM restate_agent_turns")] + + backend = Backend() + host = Host(backend, service="desk") + with pytest.raises(TerminalError) as raised: + await invoke(desk, "ask", host, request(), "refund o1") + assert raised.value.message.startswith("waiting for signal 'approval' needs an agent"), raised.value.message + assert [event for event, _ in seen.events] == ["model"] and host.calls == [] + [[status, answer]] = await invoke(desk, "stored", Host(backend, service="desk"), request()) + assert status == "failed" and "needs an agent" in answer diff --git a/tests/agents_pydantic_ai.py b/tests/agents_pydantic_ai.py index 145f907..45f8d0b 100644 --- a/tests/agents_pydantic_ai.py +++ b/tests/agents_pydantic_ai.py @@ -57,6 +57,9 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] +REVISION = b'{"revision": 1}' +"""What the agent's companion answers a turn's ``start`` and each ``publish`` with: the revision.""" + class AuthenticationError(Exception): """As a model SDK names the error of a missing or wrong API key.""" @@ -1275,7 +1278,7 @@ def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: def lookup(about: str) -> str: """Look something up; meanwhile, a client steers the turn twice.""" steering = [{"id": "s1", "text": "use Celsius"}, {"id": "s2", "text": "and be brief"}] - answers["publish"] = json.dumps({"revision": 1, "steering": steering}).encode() + answers["publish"] = json.dumps({"revision": 2, "steering": steering}).encode() return f"found {about}" mine: Any = Agent(FunctionModel(model), tools=[lookup]) @@ -1327,7 +1330,7 @@ def lookup(about: str) -> str: obj = pydantic_ai.agent("pai_reused_ids", model=FunctionModel(model), tools=[lookup]) backend = Backend() for word in ("one", "two", "three"): - host = Host(backend, service=obj.name, answers={"publish": b"null", "start": b"null", "close": b"null"}) + host = Host(backend, service=obj.name, answers={"publish": REVISION, "start": REVISION, "close": b"null"}) assert await invoke(obj, "chat", host, request(), word) == "answered" # The last request has each turn's result, its own, and none taken for interrupted. returned = [ diff --git a/tests/agents_pydantic_ai_snapshots.py b/tests/agents_pydantic_ai_snapshots.py index 766deac..0f84e23 100644 --- a/tests/agents_pydantic_ai_snapshots.py +++ b/tests/agents_pydantic_ai_snapshots.py @@ -59,8 +59,11 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] -ANSWERS = {"publish": b"null", "start": b"null", "close": b"null"} -"""What the signals object answers when a turn starts, publishes its progress and closes.""" +REVISION = b'{"revision": 1}' +"""What the agent's companion answers a turn's ``start`` and each ``publish`` with: the revision.""" + +ANSWERS = {"publish": REVISION, "start": REVISION, "close": b"null"} +"""What the agent's companion answers when a turn starts, publishes its progress and closes.""" SUMMARY = "SUMMARY of what was said" LONG = "word " * 200 @@ -517,8 +520,8 @@ async def own(ctx: sessions.SessionContext, text: str) -> str: assert ids == ["pai_ids/alice"] * 3 -STEERING = json.dumps({"revision": 1, "steering": [{"id": "s1", "text": "use Celsius"}]}).encode() -"""What the signals object answers a publish with once a client steered the turn: the message waiting for it.""" +STEERING = json.dumps({"revision": 2, "steering": [{"id": "s1", "text": "use Celsius"}]}).encode() +"""What the agent's companion answers a publish with once a client steered the turn: the message waiting for it.""" def _returned(message: ModelMessage) -> bool: @@ -633,7 +636,8 @@ def planner(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: with pytest.raises(RuntimeError, match="crashed"): await invoke(chat.obj, "chat", chat.host(), invocation, "plan one") # The retried attempt takes a message steered meanwhile, right after the stored feedback. - host = Host(chat.backend, service=chat.name, answers={**ANSWERS, "publish": STEERING}) + # The companion lists it from the retry's start, as from each publish. + host = Host(chat.backend, service=chat.name, answers={**ANSWERS, "start": STEERING, "publish": STEERING}) assert await invoke(chat.obj, "chat", host, invocation, "plan one") == Plan(title="use Celsius") assert await chat("plan two") == Plan(title="plan two") # The model got each feedback, and the message, once; the conversation stored each once. diff --git a/tests/agents_signals.py b/tests/agents_signals.py index 98a5aff..ee39515 100644 --- a/tests/agents_signals.py +++ b/tests/agents_signals.py @@ -8,14 +8,19 @@ # directory of this repository or package, or at # https://github.com/restatedev/sdk-typescript/blob/main/LICENSE # -"""The agents' companion object, signals, against a real Restate server.""" +"""An agent's companion object, against a real Restate server. + +The companion is ingress-private: the tests call it as its agent's handlers do, from a service (``driver``). +""" # pylint: disable=missing-function-docstring import asyncio +import json import typing import uuid from collections.abc import Awaitable, Callable +from typing import Any import httpx2 import pytest @@ -23,9 +28,8 @@ import restate from support import RESTATE_IMAGE # pylint: disable=import-error from restate import HarnessEnvironment -from restate.agents import _signals -from restate.agents._api import Published, Signal, Status, Steer, Steered, Watch -from restate.agents._signals import STEERS, Started, WaitRequest +from restate.agents._api import Published, Signal, Status, Steered, Watch +from restate.agents._signals import STEERS, Started, WaitRequest, make_companion @pytest.fixture(scope="session") @@ -35,46 +39,56 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] -# The agent's side: its handlers call the companion's internal handlers, which the ingress refuses. +companion = make_companion("assistant") +"""The companion of the agent ``assistant`` (below), ``assistant.prebuilt.restate``.""" + +# The agent's side: its handlers call the companion's handlers, which the ingress refuses. driver = restate.Service("driver") @driver.handler() -async def start_turn(ctx: restate.Context, request: dict[str, str]) -> Published | None: - turn: Started = {"invocationId": request["invocation"], "since": "2026-11-04T00:00:00.000Z"} - if request.get("settles"): - turn["settles"] = True - if request.get("steers"): - turn["steers"] = True - if request.get("publishes"): - turn["publishes"] = True - return await ctx.object_call(_signals.start, key=request["key"], arg=turn) +async def call(ctx: restate.Context, request: dict[str, Any]) -> Any: + """Call the companion's ``handler`` at ``key``, with ``arg`` if given, and return what it returns.""" + arg = json.dumps(request["arg"]).encode() if "arg" in request else b"" + answer = await ctx.generic_call(companion.service.name, request["handler"], arg, key=request["key"]) + return json.loads(answer) if answer else None @driver.handler() -async def close_turn(ctx: restate.Context, request: dict[str, str]) -> None: - await ctx.object_call(_signals.close, key=request["key"], arg=request["invocation"]) +async def send(ctx: restate.Context, request: dict[str, Any]) -> str: + """Send ``arg`` to the companion's ``handler`` at ``key``, and return the invocation's id.""" + arg = json.dumps(request["arg"]).encode() + sent = ctx.generic_send(companion.service.name, request["handler"], arg, key=request["key"]) + return await sent.invocation_id() -@driver.handler() -async def steer_to(ctx: restate.Context, request: dict[str, str]) -> Steered: - message = {"id": request["id"], "text": request["text"]} - return await ctx.object_call(_signals.steer, key=request["key"], arg=typing.cast(Steer, message)) +async def companion_call(harness: HarnessEnvironment, handler: str, k: str, *arg: Any) -> Any: + request: dict[str, Any] = {"handler": handler, "key": k} + if arg: + request["arg"] = arg[0] + return await harness.client.service_call(call, arg=request) -@driver.handler() -async def publish_to(ctx: restate.Context, key: str) -> Published: - return await ctx.object_call(_signals.publish, key=key, arg=None) +async def companion_send(harness: HarnessEnvironment, handler: str, k: str, arg: Any) -> str: + return await harness.client.service_call(send, arg={"handler": handler, "key": k, "arg": arg}) -@driver.handler() -async def taken_by(ctx: restate.Context, request: dict[str, typing.Any]) -> None: - await ctx.object_call(_signals.taken, key=request["key"], arg=request["ids"]) +async def start_turn(harness: HarnessEnvironment, k: str, invocation: str, **flags: bool) -> Published | None: + turn: Started = {"invocationId": invocation, "since": "2026-11-04T00:00:00.000Z"} + turn.update(typing.cast(Started, flags)) + return await companion_call(harness, "start", k, turn) -@driver.handler() -async def reset_of(ctx: restate.Context, key: str) -> None: - await ctx.object_call(_signals.reset, key=key, arg=None) +async def close_turn(harness: HarnessEnvironment, k: str, invocation: str) -> None: + await companion_call(harness, "close", k, invocation) + + +async def status_of(harness: HarnessEnvironment, k: str) -> Status: + return typing.cast(Status, await companion_call(harness, "status", k)) + + +async def revision_of(harness: HarnessEnvironment, k: str) -> int: + return (await companion_call(harness, "notifications", k))["revision"] settled: list[tuple[str, str]] = [] @@ -84,7 +98,8 @@ async def reset_of(ctx: restate.Context, key: str) -> None: """The steered messages the companion queued as turns of their own: the agent's key, and the message.""" assistant = restate.VirtualObject("assistant") -"""The agent of the companions' keys (``assistant:``), which records the turns it is sent to settle.""" +"""The companion's agent, keyed as its companion, which records the turns it is sent to settle and the messages it +is sent to answer.""" @assistant.handler() @@ -109,7 +124,7 @@ async def fail(ctx: restate.Context) -> None: @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([_signals.signals, driver, failing, assistant]), restate_image=RESTATE_IMAGE + restate.app([companion.service, driver, failing, assistant]), restate_image=RESTATE_IMAGE ) as harness: yield harness @@ -128,7 +143,7 @@ async def until(check: Callable[[], Awaitable[bool]]) -> None: def key() -> str: - return f"assistant:{uuid.uuid4()}" + return str(uuid.uuid4()) def waiting(wait_id: str, name: str) -> WaitRequest: @@ -139,151 +154,154 @@ def watching(revision: int, timeout: float) -> Watch: return {"revision": revision, "timeoutSeconds": timeout} -def signal(**fields: typing.Any) -> Signal: +def signal(**fields: Any) -> Signal: return typing.cast(Signal, fields) +async def test_the_companion_is_named_for_its_agent_and_refuses_the_ingress(harness: HarnessEnvironment): + assert companion.service.name == "assistant.prebuilt.restate" + async with httpx2.AsyncClient(timeout=60) as admin: + listed = (await admin.get(f"{harness.admin_api_url}/services")).json()["services"] + served = {service["name"]: service for service in listed} + assert served["assistant.prebuilt.restate"]["public"] is False + assert "signals" not in served + async with httpx2.AsyncClient(timeout=60) as ingress: + refused = await ingress.post(f"{harness.ingress_url}/assistant.prebuilt.restate/{key()}/status") + assert refused.is_client_error + + async def test_a_wait_returns_the_value_of_the_signal_that_resolves_it(harness: HarnessEnvironment): k = key() - sent = await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:1", "approval")) + sent = await companion_send(harness, "wait", k, waiting("inv:1", "approval")) async def listed() -> bool: - return len(await harness.client.object_call(_signals.pending, key=k, arg=None)) == 1 + return len(await companion_call(harness, "pending", k)) == 1 await until(listed) - [pending] = await harness.client.object_call(_signals.pending, key=k, arg=None) + [pending] = await companion_call(harness, "pending", k) assert pending["id"] == "inv:1" and pending["name"] == "approval" and pending["since"].endswith("Z") - resolved = await harness.client.object_call( - _signals.resolve, key=k, arg=signal(**{"id": "inv:1", "value": "approved"}) - ) + resolved = await companion_call(harness, "resolve", k, signal(id="inv:1", value="approved")) assert resolved == {"resolved": "inv:1"} - answer = await attach(harness, sent.invocation_id) + answer = await attach(harness, sent) assert answer.json() == "approved" - assert await harness.client.object_call(_signals.pending, key=k, arg=None) == [] + assert await companion_call(harness, "pending", k) == [] async def test_a_signal_resolves_the_oldest_wait_of_its_name_or_is_dropped_or_kept(harness: HarnessEnvironment): k = key() # No one waits: dropped, unless kept for the next wait. - assert await harness.client.object_call(_signals.resolve, key=k, arg=signal(**{"name": "go", "value": 1})) == { - "resolved": None - } - await harness.client.object_call(_signals.resolve, key=k, arg=signal(**{"name": "go", "value": 2, "keep": True})) - kept = await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:2", "go")) - assert (await attach(harness, kept.invocation_id)).json() == 2 + assert await companion_call(harness, "resolve", k, signal(name="go", value=1)) == {"resolved": None} + await companion_call(harness, "resolve", k, signal(name="go", value=2, keep=True)) + kept = await companion_send(harness, "wait", k, waiting("inv:2", "go")) + assert (await attach(harness, kept)).json() == 2 - first = await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:3", "go")) - second = await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:4", "go")) + first = await companion_send(harness, "wait", k, waiting("inv:3", "go")) + second = await companion_send(harness, "wait", k, waiting("inv:4", "go")) async def both() -> bool: - return len(await harness.client.object_call(_signals.pending, key=k, arg=None)) == 2 + return len(await companion_call(harness, "pending", k)) == 2 await until(both) - assert await harness.client.object_call(_signals.resolve, key=k, arg=signal(**{"name": "go", "value": "a"})) in ( + assert await companion_call(harness, "resolve", k, signal(name="go", value="a")) in ( {"resolved": "inv:3"}, {"resolved": "inv:4"}, ) - await harness.client.object_call(_signals.resolve, key=k, arg=signal(**{"name": "go", "value": "b"})) - answers = {(await attach(harness, w.invocation_id)).json() for w in (first, second)} + await companion_call(harness, "resolve", k, signal(name="go", value="b")) + answers = {(await attach(harness, w)).json() for w in (first, second)} assert answers == {"a", "b"} async def test_watch_returns_a_newer_revision_or_the_current_one_at_its_timeout(harness: HarnessEnvironment): k = key() - now = (await harness.client.object_call(_signals.notifications, key=k, arg=None))["revision"] - timed_out = await harness.client.object_call(_signals.watch, key=k, arg=watching(now, 0.2)) + now = await revision_of(harness, k) + timed_out = await companion_call(harness, "watch", k, watching(now, 0.2)) assert timed_out == {"revision": now} - watcher = asyncio.ensure_future(harness.client.object_call(_signals.watch, key=k, arg=watching(now, 30))) + watcher = asyncio.ensure_future(companion_call(harness, "watch", k, watching(now, 30))) await asyncio.sleep(0.3) # A wait that starts publishes a change. - await harness.client.object_send(_signals.wait, key=k, arg=waiting("inv:5", "x")) + await companion_send(harness, "wait", k, waiting("inv:5", "x")) assert (await watcher)["revision"] > now async def test_a_turn_is_answered_until_it_closes_and_a_failed_one_is_closed_for_it(harness: HarnessEnvironment): k = key() - await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_running"}) - status = typing.cast(Status, await harness.client.object_call(_signals.status, key=k, arg=None)) + await start_turn(harness, k, "inv_running") + status = await status_of(harness, k) assert status["answering"] is not None and status["answering"]["invocationId"] == "inv_running" - await harness.client.service_call(close_turn, arg={"key": k, "invocation": "inv_running"}) - assert (await harness.client.object_call(_signals.status, key=k, arg=None))["answering"] is None + await close_turn(harness, k, "inv_running") + assert (await status_of(harness, k))["answering"] is None # A turn whose invocation failed without closing itself: awaitTurn closes it. killed = await harness.client.service_send(fail, arg=None) - await harness.client.service_call(start_turn, arg={"key": k, "invocation": killed.invocation_id}) + await start_turn(harness, k, killed.invocation_id) async def closed() -> bool: - return (await harness.client.object_call(_signals.status, key=k, arg=None))["answering"] is None + return (await status_of(harness, k))["answering"] is None await until(closed) -async def test_a_turn_whose_invocation_is_gone_is_settled_by_its_agent_if_the_agent_settles( - harness: HarnessEnvironment, -): +async def test_a_turn_whose_invocation_is_gone_is_settled_by_its_agent_at_the_same_key(harness: HarnessEnvironment): k = key() - agent_key = k.split(":", 1)[1] + gone = await harness.client.service_send(fail, arg=None) + await start_turn(harness, k, gone.invocation_id) + # The status shows the turn as the agent's views give it, with no field of the companion's own. + answering = (await status_of(harness, k))["answering"] + assert answering is None or set(answering) == {"invocationId", "since"} async def closed() -> bool: - return (await harness.client.object_call(_signals.status, key=k, arg=None))["answering"] is None + return (await status_of(harness, k))["answering"] is None - # An agent of another SDK does not say it settles its turns: it is sent nothing. - other = await harness.client.service_send(fail, arg=None) - await harness.client.service_call(start_turn, arg={"key": k, "invocation": other.invocation_id}) - await until(closed) - # The status shows the turn as every SDK's companion does, with no field of this one's own. - mine = await harness.client.service_send(fail, arg=None) - await harness.client.service_call(start_turn, arg={"key": k, "invocation": mine.invocation_id, "settles": "yes"}) - answering = (await harness.client.object_call(_signals.status, key=k, arg=None))["answering"] - assert answering is None or set(answering) == {"invocationId", "since"} await until(closed) async def sent() -> bool: - return (agent_key, mine.invocation_id) in settled + return (k, gone.invocation_id) in settled await until(sent) - assert (agent_key, other.invocation_id) not in settled # --- Steered messages ----------------------------------------------------------------------------- async def steer(harness: HarnessEnvironment, k: str, message_id: str, text: str) -> Steered: - return await harness.client.service_call(steer_to, arg={"key": k, "id": message_id, "text": text}) + return await companion_call(harness, "steer", k, {"id": message_id, "text": text}) async def published(harness: HarnessEnvironment, k: str) -> Published: - return await harness.client.service_call(publish_to, arg=k) + return await companion_call(harness, "publish", k) + + +async def taken(harness: HarnessEnvironment, k: str, ids: list[str]) -> None: + await companion_call(harness, "taken", k, ids) async def queued(k: str, message_id: str) -> bool: - agent_key = k.split(":", 1)[1] - return any(at == agent_key and message["id"] == message_id for at, message in placeholders) + return any(at == k and message["id"] == message_id for at, message in placeholders) async def test_a_steered_message_is_queued_as_a_turn_and_held_for_a_turn_that_steers_until_taken( harness: HarnessEnvironment, ): k = key() - await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_steers", "steers": "yes"}) + await start_turn(harness, k, "inv_steers", steers=True) assert await steer(harness, k, "s1", "use Celsius") == {"id": "s1", "turn": "inv_steers"} # Queued as a turn of its own at once, in case the turn does not take it. await until(lambda: queued(k, "s1")) assert [m for _, m in placeholders if m["id"] == "s1"] == [{"id": "s1", "text": "use Celsius"}] # The turn learns of it from its next publish, and status shows it. assert (await published(harness, k)).get("steering") == [{"id": "s1", "text": "use Celsius"}] - status = typing.cast(Status, await harness.client.object_call(_signals.status, key=k, arg=None)) + status = await status_of(harness, k) assert [one["id"] for one in status.get("steering", [])] == ["s1"] # Once the turn stored it, it is gone. - await harness.client.service_call(taken_by, arg={"key": k, "ids": ["s1"]}) + await taken(harness, k, ["s1"]) assert "steering" not in await published(harness, k) - assert "steering" not in await harness.client.object_call(_signals.status, key=k, arg=None) + assert "steering" not in await status_of(harness, k) # Sent again, it is still the turn's, and not held again; marked taken again, nothing changes or publishes. assert await steer(harness, k, "s1", "use Celsius") == {"id": "s1", "turn": "inv_steers"} assert "steering" not in await published(harness, k) - now = (await harness.client.object_call(_signals.notifications, key=k, arg=None))["revision"] - await harness.client.service_call(taken_by, arg={"key": k, "ids": ["s1"]}) - assert (await harness.client.object_call(_signals.notifications, key=k, arg=None))["revision"] == now + now = await revision_of(harness, k) + await taken(harness, k, ["s1"]) + assert await revision_of(harness, k) == now async def test_a_turn_that_does_not_steer_holds_nothing_and_a_closed_turn_forgets_what_it_left( @@ -292,29 +310,29 @@ async def test_a_turn_that_does_not_steer_holds_nothing_and_a_closed_turn_forget k = key() # No turn, or one that takes no steered messages: only the turn of its own. assert await steer(harness, k, "s1", "hello") == {"id": "s1", "turn": None} - await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_plain"}) + await start_turn(harness, k, "inv_plain") assert await steer(harness, k, "s2", "hello again") == {"id": "s2", "turn": None} assert "steering" not in await published(harness, k) await until(lambda: queued(k, "s2")) - await harness.client.service_call(close_turn, arg={"key": k, "invocation": "inv_plain"}) + await close_turn(harness, k, "inv_plain") # A turn that steers ends without taking it: the message is left to its own turn. - await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_left", "steers": "yes"}) + await start_turn(harness, k, "inv_left", steers=True) assert (await steer(harness, k, "s3", "too late"))["turn"] == "inv_left" - await harness.client.service_call(close_turn, arg={"key": k, "invocation": "inv_left"}) - await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_next", "steers": "yes"}) + await close_turn(harness, k, "inv_left") + await start_turn(harness, k, "inv_next", steers=True) assert "steering" not in await published(harness, k) await until(lambda: queued(k, "s3")) # The agent's clear forgets them too. assert (await steer(harness, k, "s4", "forget me"))["turn"] == "inv_next" - await harness.client.service_call(reset_of, arg=k) + await companion_call(harness, "reset", k) assert "steering" not in await published(harness, k) async def test_a_turn_takes_at_most_its_share_of_steered_messages(harness: HarnessEnvironment): k = key() - await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_busy", "steers": "yes"}) + await start_turn(harness, k, "inv_busy", steers=True) for n in range(STEERS): assert (await steer(harness, k, f"s{n}", f"message {n}"))["turn"] == "inv_busy" # Past them, a message is answered as a turn of its own only. @@ -324,7 +342,7 @@ async def test_a_turn_takes_at_most_its_share_of_steered_messages(harness: Harne assert (await steer(harness, k, "s0", "message 0"))["turn"] == "inv_busy" assert len((await published(harness, k)).get("steering", [])) == STEERS # The messages the turn took still count: it has taken its share, and the next is a turn of its own. - await harness.client.service_call(taken_by, arg={"key": k, "ids": [f"s{n}" for n in range(STEERS)]}) + await taken(harness, k, [f"s{n}" for n in range(STEERS)]) assert "steering" not in await published(harness, k) assert await steer(harness, k, "later", "and more") == {"id": "later", "turn": None} assert "steering" not in await published(harness, k) @@ -335,28 +353,21 @@ async def test_a_turn_that_publishes_is_published_as_it_starts_and_gets_what_pub ): k = key() - async def revision() -> int: - return (await harness.client.object_call(_signals.notifications, key=k, arg=None))["revision"] - - # A turn that does not say it publishes (of another SDK, or a background agent's): nothing is published. - before = await revision() - assert await harness.client.service_call(start_turn, arg={"key": k, "invocation": "inv_quiet"}) is None - assert await revision() == before - await harness.client.service_call(close_turn, arg={"key": k, "invocation": "inv_quiet"}) + # A turn that does not publish (a background agent's): nothing is published. + before = await revision_of(harness, k) + assert await start_turn(harness, k, "inv_quiet") is None + assert await revision_of(harness, k) == before + await close_turn(harness, k, "inv_quiet") # A followed turn: its start is its first publish, with the new revision. - before = await revision() - started = await harness.client.service_call( - start_turn, arg={"key": k, "invocation": "inv_loud", "steers": "yes", "publishes": "yes"} - ) + before = await revision_of(harness, k) + started = await start_turn(harness, k, "inv_loud", steers=True, publishes=True) assert started == {"revision": before + 1} # A retried attempt starts it again: it publishes again, with the messages steered to it meanwhile. await steer(harness, k, "s1", "use Celsius") - again = await harness.client.service_call( - start_turn, arg={"key": k, "invocation": "inv_loud", "steers": "yes", "publishes": "yes"} - ) - assert again is not None and again["revision"] == await revision() + again = await start_turn(harness, k, "inv_loud", steers=True, publishes=True) + assert again is not None and again["revision"] == await revision_of(harness, k) assert again.get("steering") == [{"id": "s1", "text": "use Celsius"}] - status = typing.cast(Status, await harness.client.object_call(_signals.status, key=k, arg=None)) + status = await status_of(harness, k) assert status["answering"] is not None and status["answering"]["invocationId"] == "inv_loud" - await harness.client.service_call(close_turn, arg={"key": k, "invocation": "inv_loud"}) + await close_turn(harness, k, "inv_loud") diff --git a/tests/agents_snapshots.py b/tests/agents_snapshots.py index b47327a..1d29c0d 100644 --- a/tests/agents_snapshots.py +++ b/tests/agents_snapshots.py @@ -47,6 +47,7 @@ ) from restate.agents._tools import ToolCall, _once from restate.agents._turn import AGENT, AgentInfo +from restate.agents._signals import with_companion from restate.agents.internal import wait_for from restate.exceptions import TerminalError from restate.sessions.internal import invoke, reset @@ -59,6 +60,14 @@ def anyio_backend(): pytestmark = [pytest.mark.anyio] + +def waiter() -> sessions.SessionObject: + """A session object whose handlers wait for signals: at the companion it has, as an agent's.""" + obj = sessions.SessionObject("waiter") + with_companion(obj) + return obj + + SUMMARY = {"role": "user", "text": "q0 and q1 were asked and answered"} """What the made-up framework's compaction leaves in place of the history it compacted.""" @@ -398,7 +407,7 @@ async def run(conversation: Conversation) -> str: async def test_a_framework_waits_for_a_signal_for_its_own_tool_call_once(): - obj = sessions.SessionObject("waiter") + obj = waiter() @obj.handler async def deferred(ctx: sessions.SessionContext, call_id: str) -> Any: @@ -412,7 +421,12 @@ def host() -> Host: first = host() assert await invoke(obj, "deferred", first, invocation, "call-1") == "approved" [(service, handler, arg, key, idempotency_key)] = first.calls - assert (service, handler, key, idempotency_key) == ("signals", "wait", "waiter:alice", f"{invocation.id}:call-1") + assert (service, handler, key, idempotency_key) == ( + "waiter.prebuilt.restate", + "wait", + "alice", + f"{invocation.id}:call-1", + ) assert json.loads(arg) == {"id": f"{invocation.id}:call-1", "name": "approval", "input": {"refund": 5}} # The tool call runs again (after a crash, say): it joins the same wait. again = host() @@ -428,7 +442,7 @@ def host() -> Host: async def test_waits_opened_at_once_each_keep_their_record(): from restate.agents._calls import _load # pylint: disable=import-outside-toplevel - obj = sessions.SessionObject("waiter") + obj = waiter() @obj.handler async def deferred(ctx: sessions.SessionContext, call_ids: list[str]) -> Any: @@ -451,7 +465,7 @@ async def _lookup(ctx: sessions.SessionContext, _: Any = None) -> str: async def _closed_result(wait: Callable[[sessions.SessionContext], Awaitable[None]]) -> str: """Return what a turn stopped right after ``wait`` stores as the result of its tool call ``c1``.""" - obj = sessions.SessionObject("waiter") + obj = waiter() AGENT.on(obj, AgentInfo()) contexts: list[sessions.SessionContext] = [] @@ -511,7 +525,7 @@ async def lookup(call: ToolCall) -> Any: async def test_a_wait_with_nothing_to_show_shows_no_input(): - obj = sessions.SessionObject("waiter") + obj = waiter() @obj.handler async def deferred(ctx: sessions.SessionContext, call_id: str) -> Any: diff --git a/tests/agents_steering.py b/tests/agents_steering.py index 1851e4c..8f480d6 100644 --- a/tests/agents_steering.py +++ b/tests/agents_steering.py @@ -29,6 +29,7 @@ from restate.agents import _messages from restate.agents._api import Steer from restate.agents._common import Check +from restate.agents._signals import with_companion from restate.agents._conversation import SCHEMA, STEERED, Conversation, answer_once, steer_stored from restate.agents._steering import take_steers from restate.agents._turn import _INBOX, _STOP, Inbox, Stop @@ -50,14 +51,15 @@ class Asker: def __init__(self, run: Run, *, window: int | None = None, check: Check | None = None): self.obj = sessions.SessionObject("asker") + with_companion(self.obj) # the companion it calls, as an agent's self.backend = Backend() self.inbox: list[Steer] = [] self.hosts: list[Host] = [] @self.obj.handler async def ask(ctx: sessions.SessionContext, text: str) -> str: - # As ``chat_turn`` sets it, from what the companion's ``publish`` returned. - _INBOX.set(ctx, Inbox(revision=0, messages=list(self.inbox))) + # As ``chat_turn`` sets it, from what the companion's ``start`` returned (each publish answers the same). + _INBOX.set(ctx, Inbox(revision=1, messages=list(self.inbox))) async def ran(conversation: Conversation) -> str: return await run(ctx, conversation) @@ -120,9 +122,9 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: # The transaction that stored it sent the companion its id, once. [sent] = asker.taken() assert (sent.service, json.loads(sent.arg), sent.key, sent.idempotency_key) == ( - "signals", + "asker.prebuilt.restate", ["s1"], - "asker:alice", + "alice", "taken:s1", ) diff --git a/tests/agents_tools.py b/tests/agents_tools.py index a0d7358..1378fba 100644 --- a/tests/agents_tools.py +++ b/tests/agents_tools.py @@ -22,7 +22,7 @@ import restate from restate import sessions from restate.sessions import internal -from restate.agents._common import ToolError +from restate.agents._common import NeedsAgent, ToolError from restate.agents._object import AgentObject from restate.agents._tools import Tool, ToolCall, call, subagent, tool, wait_for_signal, workflow from restate.context import Request @@ -153,7 +153,7 @@ async def test_signal_waits_at_the_companion_of_the_agent_with_the_tool_calls_ke req = request() assert await run_tool(host, "approval", {"order": "o1", "amount": 5}, "w1", req) == "approved" [(service, handler, arg, key, idempotency_key)] = host.calls - assert (service, handler, key, idempotency_key) == ("signals", "wait", "runner:alice", f"{req.id}:w1") + assert (service, handler, key, idempotency_key) == ("runner.prebuilt.restate", "wait", "alice", f"{req.id}:w1") wait = typing.cast(dict[str, Any], json.loads(arg)) assert wait == {"id": f"{req.id}:w1", "name": "approval", "input": {"order": "o1", "amount": 5}} @@ -265,3 +265,89 @@ async def records(ctx: sessions.SessionContext) -> list[Any]: assert [c[4] for c in made] == [f"{req.id}:t1", f"{req.id}:t1"] # Its result is stored, and so its record is forgotten. assert await internal.invoke(obj, "records", Host(backend, service="crashed_caller"), req) == [] + + +# --- Waits need an agent ------------------------------------------------------------------------------------------ + +NEEDS_AGENT = "needs an agent, whose companion holds its waits: 'desk' is a plain session object" + + +async def test_a_wait_for_a_signal_from_a_session_object_that_is_no_agent_fails_for_good_with_what_to_do(): + """A ``wait_for_signal`` tool's call, and ``wait_for``, fail with a terminal error that is not the tool's result + (``NeedsAgent``): a retry would fail again, and the model cannot fix it. They make no call, nor record one.""" + from restate.agents._tools import wait_for # pylint: disable=import-outside-toplevel + from restate.exceptions import TerminalError # pylint: disable=import-outside-toplevel + + approval = wait_for_signal("approval", description="Wait for an approval.") + desk = sessions.SessionObject("desk") + waited: list[str] = [] + + @desk.handler + async def tool_call(ctx: sessions.SessionContext) -> str: + try: + return await approval.run({}, ToolCall(ctx, "c1")) + except NeedsAgent as error: + waited.append(error.message) + raise + + @desk.handler + async def deferred(ctx: sessions.SessionContext) -> None: + try: + await wait_for(ctx, "c2", "refund", {"order": "o1"}) + except TerminalError as error: + waited.append(error.message) + raise + + backend = Backend() + host = Host(backend, service="desk") + with pytest.raises(TerminalError): + await internal.invoke(desk, "tool_call", host, request()) + with pytest.raises(TerminalError): + await internal.invoke(desk, "deferred", host, request()) + assert len(waited) == 2, waited + assert waited[0].startswith(f"waiting for signal 'approval' {NEEDS_AGENT}"), waited + assert "`.agent(...)`" in waited[0] and "`AgentObject`" in waited[0] + assert waited[1].startswith(f"waiting for signal 'refund' {NEEDS_AGENT}") + assert host.calls == [] + + +def _langgraph_desk(node: Any) -> sessions.SessionObject: + """A session object that is no agent, whose ``ask`` prompts a graph that runs ``node``.""" + from langgraph.graph import END, START, MessagesState, StateGraph # pylint: disable=import-outside-toplevel + + from restate.agents import langgraph # pylint: disable=import-outside-toplevel + + builder = StateGraph(MessagesState) + builder.add_node("wait", node) + builder.add_edge(START, "wait") + builder.add_edge("wait", END) + desk = sessions.SessionObject("desk") + + @desk.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await langgraph.prompt(ctx, text, graph=builder) + + return desk + + +def _interrupting(state: Any) -> dict[str, Any]: + from langgraph.types import interrupt # pylint: disable=import-outside-toplevel + + return {"messages": [("ai", f"approved: {interrupt('approve?')}")]} + + +@pytest.mark.parametrize("waits", ["node", "interrupt"]) +async def test_a_langgraph_wait_from_a_session_object_that_is_no_agent_fails_the_turn_for_good(waits: str): + """``nodes.wait_for_signal`` and ``interrupt()``: the turn fails once, terminally, with what to do.""" + pytest.importorskip("langgraph") + from restate.agents.langgraph import nodes # pylint: disable=import-outside-toplevel + from restate.exceptions import TerminalError # pylint: disable=import-outside-toplevel + + node = nodes.wait_for_signal("approval") if waits == "node" else _interrupting + what = "waiting for signal 'approval'" if waits == "node" else "resuming an interrupt" + desk = _langgraph_desk(node) + host = Host(Backend(), service="desk") + with pytest.raises(TerminalError) as raised: + await internal.invoke(desk, "ask", host, request(), "refund o1") + assert raised.value.message.startswith(f"{what} {NEEDS_AGENT}"), raised.value.message + assert host.calls == [] diff --git a/tests/agents_tools_server.py b/tests/agents_tools_server.py index a270d71..bee8fe8 100644 --- a/tests/agents_tools_server.py +++ b/tests/agents_tools_server.py @@ -53,3 +53,79 @@ async def greet(ctx: sessions.SessionContext, call_id: str) -> str: async with restate.create_test_harness(restate.app([greeter, runner]), restate_image=RESTATE_IMAGE) as harness: assert await harness.client.object_call(greet, key="alice", arg="c1") == "hi" + + +# --- A session object that is no agent: the conversation stores, nothing publishes, waits fail for good ------------ + + +def _scripted(tool_name: str, replies: list[str]): + """A model that calls ``tool_name`` once, then answers with what the tool returned (kept in ``replies``).""" + from pydantic_ai.messages import ModelMessage, ModelResponse, TextPart, ToolCallPart, ToolReturnPart + from pydantic_ai.models.function import AgentInfo, FunctionModel + + async def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + returned = [p for m in messages for p in getattr(m, "parts", []) if isinstance(p, ToolReturnPart)] + if not returned: + return ModelResponse(parts=[ToolCallPart(tool_name, {}, tool_call_id="c1")]) + replies.append(str(returned[0].content)) + return ModelResponse(parts=[TextPart(f"the tool said: {returned[0].content}")]) + + return FunctionModel(model) + + +async def test_an_app_of_only_a_session_object_that_prompts_answers_a_turn_of_tool_calls_and_stores_it(): + """No agent, so no companion service at all: the turn is answered and stored, and publishes nothing.""" + from restate.agents import pydantic_ai, tool # pylint: disable=import-outside-toplevel + + @tool + async def lookup() -> str: + """Look it up.""" + return "found" + + replies: list[str] = [] + desk = sessions.SessionObject("plain_desk") + + @desk.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + return await pydantic_ai.prompt(ctx, text, model=_scripted("lookup", replies), tools=[lookup]) + + @desk.handler + async def turns(ctx: sessions.SessionContext) -> list[list[str]]: + rows = await ctx.db.values("SELECT answer FROM restate_agent_turns") + steps = await ctx.db.values("SELECT count(*) FROM restate_agent_items") + return [[str(row[0]) for row in rows], [str(steps[0][0])]] + + async with restate.create_test_harness(restate.app([desk]), restate_image=RESTATE_IMAGE) as harness: + assert await harness.client.object_call(ask, key="alice", arg="hi") == "the tool said: found" + [answers, [steps]] = await harness.client.object_call(turns, key="alice", arg=None) + assert answers == ["the tool said: found"] and int(steps) >= 3 + assert replies == ["found"] + + +async def test_a_wait_for_a_signal_from_a_session_object_that_is_no_agent_fails_once_with_what_to_do(): + """The turn fails with a terminal error, in the one attempt (not retried), and is closed ``failed``; the model is + not asked again with it.""" + from restate.agents import pydantic_ai # pylint: disable=import-outside-toplevel + from restate.agents._tools import wait_for_signal # pylint: disable=import-outside-toplevel + + approval = wait_for_signal("approval", description="Wait for an approval.") + replies: list[str] = [] + attempts: list[str] = [] + desk = sessions.SessionObject("plain_waiter") + + @desk.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + attempts.append(ctx.request.id) + return await pydantic_ai.prompt(ctx, text, model=_scripted(approval.name, replies), tools=[approval]) + + @desk.handler + async def turns(ctx: sessions.SessionContext) -> list[list[str]]: + return [[str(v) for v in row] for row in await ctx.db.values("SELECT status, answer FROM restate_agent_turns")] + + async with restate.create_test_harness(restate.app([desk]), restate_image=RESTATE_IMAGE) as harness: + with pytest.raises(Exception, match="HTTP 500"): # a terminal error, whose message the turn stored + await harness.client.object_call(ask, key="alice", arg="refund o1") + [[status, answer]] = await harness.client.object_call(turns, key="alice", arg=None) + assert len(attempts) == 1 and replies == [] + assert status == "failed" and answer.startswith("Error: waiting for signal 'approval' needs an agent"), answer + assert "`.agent(...)`" in answer diff --git a/tests/servable.py b/tests/servable.py index 9cc14af..7e09e2f 100644 --- a/tests/servable.py +++ b/tests/servable.py @@ -69,11 +69,13 @@ def test_a_group_serves_its_members(): assert names([Team(Needing("a", shared), Team(Needing("b", shared))), shared]) == ["a", "shared", "b"] -def test_an_agent_brings_signals_along_once(): +def test_an_agent_brings_its_own_companion_along(): first, second = AgentObject("first_agent"), AgentObject("second_agent") - assert names([first, second]) == ["first_agent", "signals", "second_agent"] - endpoint = Endpoint().bind(first).bind(second) # bound one at a time, signals still once - assert list(endpoint.services) == ["first_agent", "signals", "second_agent"] + served = ["first_agent", "first_agent.prebuilt.restate", "second_agent", "second_agent.prebuilt.restate"] + assert names([first, second]) == served + endpoint = Endpoint().bind(first).bind(second) # bound one at a time, each with its companion + assert list(endpoint.services) == served + assert names([first, first]) == served[:2] def test_what_cannot_be_served():