diff --git a/docs/agents/README.md b/docs/agents/README.md index 5e9737ca..d98ba8b4 100644 --- a/docs/agents/README.md +++ b/docs/agents/README.md @@ -14,9 +14,14 @@ building it, and how to port it. ## In brief -- **A conversation is a sequence of turns.** A turn is one prompt and what - answers it: its steps (the model's replies, tool calls and results) and its - answer. Turns never overlap, so at most one turn has no answer: the latest. +- **A conversation is a sequence of turns.** A turn is one prompt, the + messages steered to it, and what answers them: its steps (the model's + replies, tool calls and results) and its answer. Turns never overlap, so + at most one turn has no answer: the latest. +- **A message can join the turn being answered.** Sent to `steer` rather than + `chat`, the turn takes it at its next step, once the model's tool calls + have their results, and the model's next request has it; otherwise it is + answered as a turn of its own, in order. Either way, once. - **The turn is the unit of storage.** One row per turn (the transcript), and every other item tagged with its turn (the trajectories), in one sequence of positions that are the message ids clients see. @@ -36,7 +41,7 @@ building it, and how to port it. ```mermaid flowchart LR - C[Client] -- "chat, stop, signal" --> A + C[Client] -- "chat, steer, stop, signal" --> A C -- "status, messages, sync, watch" --> A subgraph A["Agent (session object), key alice"] H["chat: a turn"] --> F["Framework run: model, tools"] @@ -44,19 +49,19 @@ flowchart LR H -- "prompt, answer" --> T T[("Key's database: turns, items, snapshots, call records")] V["views: messages, sync"] -- read --> T - X["private: settle, compact"] --> T + X["private: settle, compact, steered"] --> T end - H -- "start, close, publish" --> S[("Companion signals, key assistant:alice: turn, waits, revision")] + H -- "start, close, publish" --> S[("Companion signals, key assistant:alice: turn, waits, steered messages, revision")] F -- "wait, for a tool or an approval" --> S C -- "signal resolves a wait" --> S - S -- "settle a dead turn" --> X + S -- "settle a dead turn, and steered for each steered message" --> X ``` ## The pages | Page | What it holds | | --- | --- | -| [Turn semantics](turns.md) | what a turn is; its life, its end, attempts, stops, dead turns and `settle`; the turn model; what clients see; the guarantees | +| [Turn semantics](turns.md) | what a turn is; its life, its end, attempts, stops, dead turns and `settle`; steering; the turn model; what clients see; the guarantees | | [The conversation](conversation.md) | the data layout (every table), a worked example, which turns a turn reads, what the model gets (the turn projection, summaries, snapshots), what a turn carries, the call records, the schemas | | [The frameworks](frameworks.md) | per framework: what it expects, how we map onto it, and what we learned | | [What we learned](learnings.md) | the research (coding agents, frameworks), the decisions and why, the lessons, the bugs a review found, the open questions | @@ -77,6 +82,7 @@ then the same. | [`_conversation.py`](../../python/restate/agents/_conversation.py) | the conversation: 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 | +| [`_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 | | [`_tools.py`](../../python/restate/agents/_tools.py) | the tools built on Restate (`call`, `send`, `workflow`, `subagent`, `wait_for_signal`), and `wait_for` | diff --git a/docs/agents/access-patterns.md b/docs/agents/access-patterns.md index 4c8137aa..38340e2f 100644 --- a/docs/agents/access-patterns.md +++ b/docs/agents/access-patterns.md @@ -290,6 +290,24 @@ In order of value for effort: a real server. - **The companion and the call records**, not exercised by this harness: their writes per turn, and the `restate_agent_calls` read at every turn. +- **Steering** adds no request to a turn that nobody steers: the companion + answers the `publish` a turn makes at each tool call's start and end with + the messages waiting for it (one state read more, of its inbox, which is + empty then), and a turn takes them at a step with no call of its own; + the `taken` that tells the companion is sent from the storing + transaction's outbox. A steered message costs the `steer` view, one call + to the companion (an inbox entry, a publish, and the placeholder sent to + `steered`), the transaction that stores it (an item and a row of + `restate_agent_steers`), and the placeholder's short invocation (open the + database, one read by id). A later turn reads the steered messages of the + transcript it loads as one range of that table per page; the summary + view, one range. Measured on a real server with + `scripts/steering_traffic.py` (a turn of one tool call): a turn nobody + steers is what it is before steering, 47 journal entries and 5 companion + calls (`start`, 3 `publish`, `close`); a turn that takes a steered + message, 57 entries and 6 calls (the outbox's `taken`); the message + itself, the `steer` view (5 entries), the companion's `steer` (13), and + the placeholder that finds it taken (3). - **Snapshots and `compact`:** a snapshot replaces the previous one (one `DELETE` and one `INSERT` of the whole compacted history), and `compact` reads the transcript from the summary's start. Both are occasional, but diff --git a/docs/agents/conversation.md b/docs/agents/conversation.md index c20c6d68..3249c7fc 100644 --- a/docs/agents/conversation.md +++ b/docs/agents/conversation.md @@ -16,15 +16,17 @@ records). `assistant`) has its own SQL database, and its conversation is the tables below in it. Nothing in it is shared with another key. - **A conversation is a sequence of turns,** in the order their messages - came. A turn is one prompt and what answers it: + came. A turn is one prompt, the messages steered to it, and what answers + them: - its **row**: the prompt, then, once the turn ends, its status, its answer and what it carries to later turns. The rows together are the *transcript*; - its **steps**: everything else it stored, the model's replies, the tool - calls and their results, a message that asked again. A turn's steps are - its *trajectory*. + calls and their results, a message that asked again, the messages + steered to it that it took ([Steering a turn](turns.md#steering-a-turn)). + A turn's steps are its *trajectory*. - **Turns never overlap**, so at most one turn has no answer: the latest - ([One message at a time](turns.md#one-message-at-a-time)). + ([One turn at a time](turns.md#one-turn-at-a-time)). - **Every row belongs to one turn,** and only the latest turn, while it has no answer, is written ([The turn model](turns.md#the-turn-model)). - **The model's input is composed from the turns in one place,** @@ -32,14 +34,15 @@ records). [turn projection](#what-the-model-sees), or [history from snapshots](#history-from-snapshots). - **Clients read the turns as stored**, in two projections of the same rows: - the full view (every step) and the summary view (each turn's message and - answer), under the same message ids. + the full view (every step) and the summary view (each turn's message, the + messages steered to it, and its answer), under the same message ids. ```mermaid erDiagram restate_agent_turns ||--o{ restate_agent_items : "its steps (items.turn = turns.n)" restate_agent_conversation_start ||--o{ restate_agent_turns : "where they start, and their summary" restate_agent_snapshots }o--|| restate_agent_turns : "stored in a turn, or between turns" + restate_agent_steers |o--|| restate_agent_items : "a steered message (steers.n = items.n)" restate_agent_turns { int n PK "position: the prompt's message id" int seq "the change it was written at" @@ -123,6 +126,12 @@ CREATE TABLE restate_agent_snapshots ( -- the latest only, for agents with turn INTEGER, items TEXT NOT NULL ); + +CREATE TABLE restate_agent_steers ( -- the messages steered to a turn that it took (appended) + n INTEGER PRIMARY KEY, -- the item's position: the message of the user it is + turn INTEGER NOT NULL, -- the turn's position + id TEXT NOT NULL UNIQUE -- the message's id: the `steer` request's +); ``` - **Positions.** The two tables share one sequence of positions (`n`), the @@ -153,6 +162,13 @@ CREATE TABLE restate_agent_snapshots ( -- the latest only, for agents with - **A snapshot** takes no position of its own, so message ids and cursors don't depend on it, and it is never shown ([History from snapshots](#history-from-snapshots)). +- **A steered message** is an item of its turn, a message of the user made + with `Items.steered` (by default `Items.prompt`), and a row of + `restate_agent_steers` at the same position, with its id: unique, so a + turn stores it once whatever its attempts do, and the agent's `steered` + finds that it was taken with one read. Earlier turns' steered messages + are one range of that table's primary key, as positions only grow. A drop + deletes those after the reply it drops. ### Around the conversation @@ -301,17 +317,18 @@ the previous turn whole, to give its graph what closed it. 1. The summary of the earliest turns, if they have one: a user's message with it, and the assistant's acknowledgement. -2. The turns before the window, as their prompt and answer: the latest of - them that the agent's `budget` holds (32 000 tokens by default, counted - as 4 characters each). +2. The turns before the window, as their prompt, the messages steered to + them, and their answer: the latest of them that the agent's `budget` + holds (32 000 tokens by default, counted as 4 characters each). 3. The window's turns, whole. The model sees the earlier turns' messages, not their tool calls: those of the window only. The `messages` view shows every turn whole, as stored. This is a contract for every integration on the turn projection: **earlier -turns reach the model only from their rows** (each as its prompt, then -`Items.reply(answer, agent)`), so a framework must never depend on an +turns reach the model only from their rows** (each as its prompt, the +messages steered to it, then `Items.reply(answer, agent)`), so a framework +must never depend on an earlier turn's steps outside the window. What it needs from them, it carries ([What a turn carries](#what-a-turn-carries)). @@ -348,7 +365,8 @@ sequenceDiagram - **What is summarized:** the turns before the next window, but the latest that half the budget holds, which stay as they are, with the summary - before. Turns stay stored, and `messages` shows them as before. + before: each turn's message, the messages steered to it, and its answer. + Turns stay stored, and `messages` shows them as before. - **Who summarizes:** the built-in agents of OpenAI, ADK and LangGraph, with their own model. An agent you configure yourself only trims the earliest turns. A LangGraph graph agent never summarizes: its model reads the @@ -537,7 +555,8 @@ invocation (the next invocation's first call replaces it): [What we learned](learnings.md#decisions)): its schemas are written for new databases, with no migration from the layout of the 1.4.x releases, whose tables no deployment uses. From the next release on, schema changes are - appended migrations. + appended migrations: the first is `restate_agent_steers`, a new table, + which a view that applies it to its own session creates at no cost. - **The database** is the session object's SQL engine (sdk-shared-core's, the same for the TypeScript SDK), one per key, stored in Restate's state. It refuses an expression more than 100 deep: match many pairs with a list diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index ce9c5f19..a0cfc039 100644 --- a/docs/agents/frameworks.md +++ b/docs/agents/frameworks.md @@ -35,7 +35,9 @@ What the core gives an integration (`restate.agents.internal`): them), `prompt(text)`, `reply(text, agent)`, `results(calls)` (to make them), `read(item)` (what an item says, the tool calls it makes, the results it holds), and optionally `summary(text)` (the summary in the - framework's own form). + framework's own form) and `steered(text, steer_id)` (a message steered + to a turn, as the framework tells it apart from a prompt; by default a + prompt's item). - **`answer_once(ctx, text, items, run, window, budget, check=)`**, which answers a prompt once. `run` gets the `Conversation`: `history()` is the model's input, and `conversation.turn` is the turn to write. @@ -59,6 +61,7 @@ At a glance: | Summaries | ours: one user message | the harness's compaction | ours: a user message, then "Understood." | ours: as ADK | none: the graph trims its own state | | Steps are stored from | the run's hooks and input filter | the `_Restate` capability | `_Sessions.append_event` | the agent's nodes | a copy of the messages after the run's input | | The framework's own approvals | not given: such a run fails its turn | durable: deferred tool calls wait | not integrated | not integrated | durable: `interrupt()` waits | +| Steered messages ([Steering a turn](turns.md#steering-a-turn)) | the input filter, at each step; spliced back into each later request | `_Steering`, innermost after `_Snapshots`: Pydantic AI's `enqueue('asap')` | `_Restate.before_model_callback`, as a request goes out, for an `LlmAgent` root: in the request and the session | the model node, kept in the run's state | none: each is a turn of its own | In every framework, a tool can wait for a person: `wait_for_signal`. @@ -109,6 +112,26 @@ In every framework, a tool can wait for a person: `wait_for_signal`. - **Approvals aren't durable yet.** A run that stops for tool calls that need an approval (`needs_approval`) fails its turn: nothing would give the approval. A tool that waits for a person uses `wait_for_signal`. +- **A steered message must be put back in every request.** The input filter + changes one request only: the SDK makes each request's input afresh, from + the run's input and its own items. So the run keeps the messages it took, + each with the item it came after (a tool call's id), and splices them back + into each later request; where a handoff's `input_filter` took that item + out, before the first item the run saw after the message, as if the filter + had kept it. The run goes on in place: its input guardrails run once, and + `max_turns`, tool choice and hooks carry on. A run also takes the messages + waiting as it starts (a retried attempt's, say), into its input. +- **An agent used as a tool runs with the run's `run_config`** (the SDK + passes it on to `Agent.as_tool`'s run), our input filter with it. A request + made while one of the run's tool calls runs is that nested run's: it gets + your own filter only. The turn stores none of its items (the call's output + is what the turn keeps; an output without its call would break the next + request) and puts none of the steered messages in it. +- **A handoff filter that drops the handoff's own output** (`input_filter= + remove_all_tools`) leaves the handoff's call without an output in the + turn: the run takes no steered message after it. Found with it, and not + steering's: that output is never stored, so the next turns' requests carry + a function call with no output, which OpenAI's APIs refuse (an open bug). ## Pydantic AI, on Pydantic AI Harness @@ -252,10 +275,29 @@ handlers have no `ctx.run`; it disables harness features), `StepPersistence` compaction); store `RetryFeedbackPart` requests. - **Compacting off the answer's path** (`compact_after()`) triggers at 70% of the window: `compact_now` at an 80% target does nothing below 80%. -- **Order matters.** `_Restate` runs innermost in `before_model_request`, so - that it sees the messages as they will be sent. Capabilities at the - agent's own layer run before it, so one that edits the newest tool results - could change them before they are stored (an open question). +- **Order matters.** `_Snapshots` runs innermost in `before_model_request` + but for `_Steering`, which it wraps, so that it sees the messages as they + will be sent; `_Restate`, which stores the tools' results, has no position + of its own. Capabilities at the agent's own layer run before them, so one + that edits the newest tool results could change them before they are + stored (an open question). +- **A steered message goes in with Pydantic AI's own steering**, + `RunContext.enqueue(..., priority='asap')` from `before_model_request`: + it is delivered once every before-hook ran, into the request and the + run's history, after what both hold, and the run goes on with its usage + limits, retries and capabilities' state. `_Steering` stores the messages + after any snapshot of the same request, so that the snapshot covers only + what came before them. A request that resumes a paused response takes + none. +- **Feedback is found by its parts.** A run merges the consecutive requests + it is given, and resumes from the last one as a request of its own, so a + stored `RetryFeedbackPart` request came back as a new request and was + stored again (a retry, then every later turn): `_Restate` now knows each + feedback part by its own JSON. +- **What other capabilities enqueue is not stored** (the harness's system + reminders, loop detection, ...): it reaches the model, but a retried + attempt and the next turns lack it. Planning records its plan again when + it is missing; the others are lost (an open question). - **The trade-offs of snapshots.** A turn reads everything since the latest compaction, up to the harness's target (80% of the model's window), where the turn projection reads up to its budget; and a snapshot repeats the @@ -328,6 +370,22 @@ sequenceDiagram guess), then applies the state deltas of the loaded turns after it. - **The summary keeps the default form**, a user event then "Understood.": the tests run on scripted models, which can't tell what Gemini accepts. +- **A steered message is taken as a model request goes out** + (`_Restate.before_model_callback`), not as the tools' results are stored: + after some tools ADK asks the model nothing more (a tool's + `skip_summarization`, an output schema's `set_model_response`), and a + message taken there would go unread. It is added to the request and to the + session, which ADK builds every later request from, so the run goes on + with it. The agent's own `before_model_callback` runs after the plugin's, + on the request with the message: a response it gives in place of the + model's answers the message too. +- **A steered message starts an invocation of its own**, as a prompt does: + after a crash ADK resumes the prompt's invocation from its events, and its + classic runner (a workflow root's) fails on a user message last among + them, on every retry. Only an `LlmAgent` root takes steered messages; with + a workflow root, each is a turn of its own. With + `include_contents='none'`, the requests after the one that takes a message + start at it, without the prompt and the tools' results before it. ## LangGraph and LangChain @@ -357,6 +415,7 @@ sequenceDiagram | Our items | the run's messages | a copy of the messages after the run's input, matched by id | | The window | the agent's | at least 1: the graph reads what closed the previous turn | | Summaries | ours | none: its model reads its own state | +| Steered messages | taken by the model node at each step, kept in the run's state | none: each is answered as a turn of its own | ### What we learned @@ -384,5 +443,13 @@ sequenceDiagram an answer the conversation stored. - **A structured response is the answer only of the run that wrote it:** the turn compares the channel's version before and after its run. +- **A steered message joins the run in place.** The model node takes it + (`take_steers`) before its request, sends it after the tool results, and + returns it with the reply, so the state keeps it for the rounds after: + running the graph again with it would reset the recursion limit. Its id is + `steer:`, as `_Items.steered` makes and stores it, so the history of + a retried attempt and of the next turns has it too; a trimmed request + never starts at one, which would leave out the turn's prompt and its tool + results. - **Subgraph checkpoints pile up:** the latest checkpoint of each subgraph namespace stays, one per turn that ran a subgraph (an open question). diff --git a/docs/agents/learnings.md b/docs/agents/learnings.md index 677008b1..f390edb5 100644 --- a/docs/agents/learnings.md +++ b/docs/agents/learnings.md @@ -77,6 +77,11 @@ Each framework's details are in [The frameworks](frameworks.md). | Q1 | **Pydantic AI's usage limits are per turn**, not per conversation. | As they were; a budget per conversation would be the harness's `Spend`, later. | | Q3 | **An `Agent` of your own with no compaction capability gets the default compaction.** | Otherwise its history would grow without bound. | | — | **The summary's form stays per framework.** OpenAI's is one user message; ADK and LangGraph keep a user message then "Understood.". | OpenAI's was checked against both its APIs; the others run on scripted models in the tests, which can't tell what Gemini or other chat models accept. | +| D6 | **Steering is a view of its own, `steer`; `chat` keeps one message per turn.** | A client that sends to `chat` gets the answer to its message as `chat`'s response, which joining another turn would take from it. A UI that wants one way in sends everything to `steer`, which starts a turn when none is answered. | +| D7 | **A steered message is queued as a turn of its own the moment it is sent** (the agent's `steered`), and the turn that takes it makes that turn do nothing. | Its place in the key's queue is kept with the messages sent after it, a turn that dies leaves no message behind, and the companion forwards nothing at a turn's close. The cost is one short invocation per steered message. | +| D8 | **The turn learns of steered messages from `publish`**, which it makes at each tool call's start and end, and takes them with an outbox send. | A turn that nobody steers makes no call more: the check is free. A call of its own at each step would cost one round trip per tool round of every turn. | +| D9 | **Each framework takes a steered message in place**, in the request about to go out (OpenAI's input filter, Pydantic AI's `enqueue`, ADK's session events, LangGraph's model node); none restarts its run. | Restarting a run from the stored steps (the crash path) reran input guardrails, reset `max_turns`, usage limits and recursion limits, and left the old run able to write. | +| D10 | **A message taken late is the next turn**: one that comes after the turn's last step is answered next, not folded into the turn. | The turn's answer stays the answer to what it was asked; folding it in would hide that answer from the summary view. | | — | **Steps retention is future work** ([its design](steps-retention.md)). | Nothing reads a turn's steps once it leaves the window, but deleting them changes what the full view and conversation search can show; it can come later, as an option. | ## Lessons from building it diff --git a/docs/agents/porting.md b/docs/agents/porting.md index 817a4d16..02d455dc 100644 --- a/docs/agents/porting.md +++ b/docs/agents/porting.md @@ -108,6 +108,20 @@ Each step stands on its own, with its tests, and the earlier ones are small. if the TypeScript agents have `output=` or guardrails that ask again. 7. **A framework's own approvals,** with `waitFor(..., { approval: true })` for any framework whose tool calls can wait for a person. +8. **Steering** ([Steering a turn](turns.md#steering-a-turn)), on the + conversation store. The companion first, as its handlers are shared: + `steer` (the placeholder sent to the agent's `steered`, idempotency key + `steer:`, and the inbox for a turn whose `start` says + `steers: true`, at most `STEERS` per turn, taken or not), `taken` + (marking them, without their text, and publishing only a change), + `publish` answering `{revision, steering?}`, `_closed` and `reset` + forgetting the inbox, and `status.steering`. Then the agent's `steer` + view and private `steered` (it answers the message as a turn unless + `restate_agent_steers` has its id), `takeSteers` at a step with the + `taken` send in the storing transaction's outbox (taking none once the + turn stopped, checked again as the store gets the ledger), and each + framework putting the messages in its next request in place: pi's + `whenBusy: "steer"` inbox may be that framework's way. ## What must hold afterwards diff --git a/docs/agents/testing.md b/docs/agents/testing.md index 23c3e67d..0d428734 100644 --- a/docs/agents/testing.md +++ b/docs/agents/testing.md @@ -40,8 +40,8 @@ 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`, 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_webhook`, `examples_agents` | +| 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 | +| 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 the first word of the prompt (`weather`, `crash`, `wait`, `refuse`, ...). The diff --git a/docs/agents/turns.md b/docs/agents/turns.md index 80fcc716..e1919343 100644 --- a/docs/agents/turns.md +++ b/docs/agents/turns.md @@ -12,9 +12,11 @@ The code lives in: - [`_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; +- [`_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`, - `stop`, `clear`, `sync`, `watch`) and the private handlers (`settle`, - `compact`, `erase`); + `stop`, `steer`, `clear`, `sync`, `watch`) and the private handlers + (`settle`, `compact`, `erase`, `steered`); - [`_messages.py`](../../python/restate/agents/_messages.py): the conversation as clients read it; - [`_calls.py`](../../python/restate/agents/_calls.py): the Restate calls of a @@ -31,16 +33,20 @@ stored and loaded, and the `settle` and `compact` handlers, are this SDK's - **Agent, key.** A session object of `restate.agents`, such as `assistant`, and one of its keys, such as `alice`: one conversation, in its own SQL database. -- **Message.** One call of the agent's `chat` handler: one Restate - invocation, whose id (`inv_1f…`) is also the message's id. -- **Turn.** One prompt, and what answers it: the unit of the conversation +- **Message.** What a client sends the agent: a call of its `chat` handler, + one Restate invocation whose id (`inv_1f…`) is also the message's id; or a + call of its `steer` view, whose request's id is the message's id. +- **Turn.** One prompt, the messages steered to it that it took, and what + answers them: the unit of the conversation ([The turn model](#the-turn-model)). Its id is the prompt's: the invocation's id, or `:2`, `:3` for later prompts of a handler that - prompts more than once. A `chat` message is one turn. -- **Invocation.** The `chat` invocation answering a message, from its start - to its end, across its attempts: what the companion follows (`status` - calls it the turn being answered, `answering`) and what `stop` cancels. - An invocation of a handler of your own can hold several turns. + prompts more than once. A `chat` message is one turn; a steered message + joins the turn being answered, or is a turn of its own. +- **Invocation.** The invocation answering a message (`chat`'s, or + `steered`'s for a steered message answered as a turn of its own), from its + start to its end, across its attempts: what the companion follows + (`status` calls it the turn being answered, `answering`) and what `stop` + cancels. An invocation of a handler of your own can hold several turns. - **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. @@ -59,14 +65,25 @@ stored and loaded, and the `settle` and `compact` handlers, are this SDK's framework's own state (ADK's session state). - **Window.** The turns the model sees whole, with their trajectories ([Which turns a turn loads](conversation.md#which-turns-a-turn-loads)). +- **Step.** The point in a turn after the model's reply has a result for + each of its tool calls, before the model's next request. +- **Steered message.** A message sent to the agent's `steer` view: the turn + being answered takes it at its next step, or else it is answered as a turn + of its own ([Steering a turn](#steering-a-turn)). -## One message at a time +## One turn at a time -A key answers one message at a time, in the order they came. `chat` is an -exclusive handler, and it keeps the key while it waits (suspended on a +A key answers one turn at a time, in the order its messages came. `chat` is +an exclusive handler, and it keeps the key while it waits (suspended on a signal, say): a message sent while the agent answers another waits in Restate's queue for the key, and its turn starts when the one before it -ended. There is no steering: a message never joins the turn being answered. +ended. + +A message sent to `steer` takes its place in the same queue as it is sent, +and is answered there, as a turn of its own, unless the turn being answered +takes it first, at its next step, and answers it together with its own +message. A message sent to `chat` never joins the turn being answered +([Steering a turn](#steering-a-turn)). So **turns never overlap.** At most one turn of a conversation has no answer, and it is the latest. A turn that finds an earlier one without an @@ -80,8 +97,9 @@ still hold the key. - **Once.** A message sent with an idempotency key is one invocation, answered once, however often it is sent. A webhook's delivery (`deliver`) is a message of this kind, keyed by the delivery. -- **To change course,** stop the turn (`stop`) and send the message that - replaces it: it is answered once the stopped turn ended, after any +- **To change course,** send the turn a message with `steer`, which it + takes at its next step; or stop the turn (`stop`) and send the message + that replaces it: it is answered once the stopped turn ended, after any messages already waiting. ## A turn's life @@ -104,6 +122,10 @@ sequenceDiagram loop each step of the framework A->>DB: append the model's reply, the tool calls' results A->>S: publish (a tool call starts, ends) + S-->>A: the revision, and the messages steered to the turn, if any + opt messages steered to the turn wait, at the step + A->>DB: take them: steps of the turn, and taken in the outbox + end end A->>DB: answered: the answer, and what the turn carries opt the earlier turns outgrow the budget @@ -132,7 +154,10 @@ sequenceDiagram 3. **Steps.** The framework answers: model calls and tool calls. Each step is stored as it comes (a reply in one transaction, each result as it comes back), and each tool call publishes as it starts and as it ends, - so a follower sees the call start and its result. + so a follower sees the call start and its result. The companion answers + each `publish` with the messages steered to the turn that wait: at its + next step, the turn takes them, and the model's next request has them + ([Steering a turn](#steering-a-turn)). 4. **Answer.** The answer is stored on the turn's row, with what the turn carries. If the turns before the next turn's window outgrow the agent's budget, the agent sends itself `compact` @@ -143,7 +168,8 @@ sequenceDiagram turn's error. A **background** agent (`background=True`, a subagent say) publishes its -turn's end only: no step, nor its start, wakes a follower. +turn's end only: no step, nor its start, wakes a follower. Its turns take no +steered messages. ## How a turn ends @@ -164,11 +190,18 @@ or else the note that ends the turn), then the note is the turn's last reply, and its answer. So the next turn's model never sees a tool call without a result, and knows how the turn ended. +**The messages steered to a turn go as the turn goes.** Those it took are +its own, as its prompt is: its answer answers them, and a turn that ends +otherwise (refused, failed, stopped, ended) ends them too, and they are not +answered again. Those it did not take were queued as turns of their own as +they were sent, and are answered after it, in order. + Whoever closes a turn in the companion, a close does the same: it clears the turn (if the agent still answers it), drops the waits of the turn's tools -(rejecting their awakeables, so a `wait` still parked ends), and publishes -once if either changed something. A second close of the same turn changes -nothing, and publishes nothing. +(rejecting their awakeables, so a `wait` still parked ends), forgets the +messages steered to the turn, and publishes once if any of these changed +something. A second close of the same turn changes nothing, and publishes +nothing. ## Dead turns @@ -308,6 +341,10 @@ sequenceDiagram once however often the tool call runs: a framework that gives a cut-off call a new id gets the earlier call's id back ([`_calls.py`](../../python/restate/agents/_calls.py)). +- **Steered messages, once.** A message the turn stored is a step of the + turn for every later attempt: in the model's history, and never taken + again (by its id). The `taken` that tells the companion is sent from the + storing transaction's outbox, again with the same key after a crash. - **Waits survive.** A tool waiting for a signal waits in the companion (`wait`), by that key: the next attempt rejoins the same wait, and the turn is still answered meanwhile. A framework's own tool call that waits @@ -363,10 +400,162 @@ sequenceDiagram tool ends cancelled instead: a framework that cancels a task and waits for it (ADK 2's runner) expects nothing else. The turn ends cancelled either way. +- A message steered to the turn ends with it if the turn took it, as the + turn's prompt does; one it had not taken is answered as a turn of its own + after it. A turn that stops while it waits to store messages takes none. - `clear` stops the turn the same way, then forgets the conversation (the - key's database, the companion's waits and kept signals). A message sent - before `clear` is answered first, then forgotten too. Message ids and - cursors carry on. + key's database, the companion's waits, kept signals and steered + messages). A message sent before `clear` is answered first (unless the + stopped turn took it), then forgotten too; one sent after it is answered + in the cleared conversation. Message ids and cursors carry on. + +## Steering a turn + +A message sent to the agent's `steer` view, rather than to `chat`, adds to +the turn being answered. The turn takes it at its next **step**: once the +model's reply has a result for each of its tool calls, before the model's +next request. That request has the message, after those results, and the +turn's answer answers the turn's messages so far. If the turn reaches no +step before it ends, the message is answered as a turn of its own, right +after it. Either way, it is answered once. + +| The message is sent | It is answered | +| --- | --- | +| while the turn has a step ahead (a tool call runs, say) | by the turn being answered, which takes it at that step | +| while the model writes the turn's final answer | as the next turn, before the messages sent after it: no step follows | +| with no turn being answered | as a turn of its own, which it starts, as `chat/send` does | +| to a turn that takes none (see the limits), or past its 16 | as a turn of its own, once the turn being answered ended | +| to a turn that then stops, fails, is refused or killed | not at all if the turn took it (it ends with the turn); otherwise as a turn of its own after it | + +The view answers `{id, turn}`: the message's id (the request's: send it +with an idempotency key, so that a retried request is the same message), +and the turn that holds it for its next step, or `null` if it is answered as +a turn of its own. + +### Taken at the next step + +```mermaid +sequenceDiagram + autonumber + participant C as Client + participant V as Agent (steer view) + participant S as Companion + participant A as Agent (chat) + participant P as Agent (steered) + participant DB as Key's database + C->>A: chat/send("Plan my week") + A->>S: start(turn, steers) + A->>DB: the model's reply, a tool call + C->>V: steer("Skip Friday") + V->>S: steer(id, text) + S-)P: steered(id, text): queued now, behind the turn + Note over S: inbox: the message, for this turn + A->>S: publish (the tool call ends) + S-->>A: revision, steering: [the message] + A->>DB: the tool's result + A->>DB: the message, a step of the turn, and taken(id) in the transaction's outbox + DB--)S: taken(id), once the transaction commits + A->>A: the model's next request has it + A->>DB: the answer + A->>S: close(turn) + Note over P: runs once the turn ended: the message is stored, so it does nothing +``` + +### Too late: the next turn + +The model writes the turn's final answer, so no step follows: the turn ends +without taking the message, which its own turn answers next, before a +message sent after it. + +```mermaid +sequenceDiagram + autonumber + participant C as Client + participant S as Companion + participant A as Agent (chat) + participant P as Agent (steered) + participant DB as Key's database + C->>A: chat/send("Plan my week") + Note over A: the tools are done: the model writes the final answer + C->>S: steer("Skip Friday"), through the steer view + S-)P: steered(id, text): queued behind the turn + Note over S: inbox: the message, for this turn + C->>A: chat/send("And next week?"): queued behind steered + A->>DB: the answer + A->>S: close(turn): the message it did not take is forgotten here + P->>DB: stored by a turn? no + P->>DB: a turn of its own: "Skip Friday", and its answer + Note over A: then "And next week?" is answered +``` + +### The turn ends otherwise + +A message the turn took is its own, as its prompt is: stopped with it here. +One it had not taken is answered as a turn of its own after it. + +```mermaid +sequenceDiagram + autonumber + participant C as Client + participant S as Companion + participant A as Agent (chat) + participant P as Agent (steered) + participant DB as Key's database + C->>S: steer("Use Celsius"), through the steer view + S-)P: steered(1): queued behind the turn + A->>DB: at a step: "Use Celsius", a step of the turn, and taken(1) + C->>S: steer("And the rain") + S-)P: steered(2): queued behind steered(1) + C->>A: stop + A->>DB: the turn ends cancelled, and "Use Celsius" with it + A->>S: close(turn): message 2, not taken, is forgotten here + P->>DB: steered(1): stored by a turn? yes: nothing to do + P->>DB: steered(2): stored? no: a turn of its own, "And the rain", and its answer +``` + +### How it holds + +- **Where it waits.** The companion's `steer` sends the agent's private + `steered` handler the message at once, so it takes its place in the key's + queue as it is sent; and, if the turn being answered takes steered + messages (its `start` says so, `Started.steers`), it holds the message for + that turn (`inbox`), 16 at most (`STEERS`). +- **How the turn learns of it.** From `publish`, which the turn makes at + each tool call's start and end anyway: the companion answers it with the + messages waiting for the turn. A turn that nobody steers makes no call + more. +- **Taken once.** At a step, the framework's integration calls + `take_steers`: the turn stores each waiting message by its id, once, as a + message of the user that is a step of the turn, and the integration puts it + in the request that is about to go out, in place: the framework's run goes + on, with its limits, guardrails and state. The transaction that stores the + messages sends the companion `taken` from its outbox, once, with the + commit; the companion marks them taken, and counts them among the turn's + 16 until it ends. A turn that stops while it waits to store messages + takes none. +- **Answered once.** Once the turn ended, `steered` runs, for each message: + it does nothing if a turn stored the message, and otherwise answers it as + a turn of its own. So each message is answered once, in the order it was + sent with the messages after it. +- **A taken message is its turn's.** It goes as the turn goes + ([How a turn ends](#how-a-turn-ends)): answered by the turn's answer, or + stopped, refused, failed or ended with it, and never answered again. A + reply dropped to ask the model again drops the messages taken after it, + which `steered` then answers. +- **What the model sees.** In the window, the message is where it was taken. + Later, an earlier turn is its message, the messages steered to it, and its + answer; budgets and summaries count them. `latest_ask` skips them: the + turn goes on from its prompt (an ADK invocation, a graph's run). +- **Limits.** A step comes once a reply's tool calls all have results: a + tool that waits a day for an approval holds the message until it + returns (stop the turn to change course sooner). A turn takes at most 16 + messages in all; past them, each is a turn of its own. A message is taken + only as a model request goes out: one that no request follows (ADK asks + the model nothing more after some tools) is a turn of its own. A LangGraph + graph (`graph=`), an ADK agent whose root is a workflow, and a background + agent take none; an OpenAI run takes none after a handoff whose + `input_filter` drops the handoff's own output + ([The frameworks](frameworks.md)). ## The turn model @@ -382,16 +571,22 @@ Turns are composed into one flat list in one place only: what the model gets. nothing else. - **Scope.** A turn's run reads and writes its own turn: its open calls, its replies, its final reply, its drops. Earlier turns are read-only context. +- **Messages.** A turn's messages of the user are its prompt, then the + messages steered to it that it took, each a step at the place it was + taken, and a message that asked again (`ask_again`). Its answer answers + them all. - **Lifecycle.** A turn is being answered (no status), then ends once, with a status: `answered` by its run, or closed as `cancelled`, `failed`, `refused` or `ended`. The status, the answer, the answer's message and - what the turn carries are stored in one update of its row. -- **Two projections.** The transcript (each turn's message and answer) and a - turn's detail (its steps), read separately: the `messages` view's - `summary` and `full` views. + what the turn carries are stored in one update of its row. The messages + steered to it that it took end with it. +- **Two projections.** The transcript (each turn's message, the messages + steered to it, and its answer) and a turn's detail (its steps), read + separately: the `messages` view's `summary` and `full` views. - **One composition.** `Conversation.history()` composes the turns into what - the model gets: the summary, the earlier turns' messages and answers, the - window's turns whole. Nothing else flattens turns. + the model gets: the summary, the earlier turns' messages (with the + messages steered to them) and answers, the window's turns whole. Nothing + else flattens turns. In code ([`_conversation.py`](../../python/restate/agents/_conversation.py)): @@ -404,12 +599,13 @@ Conversation the key's conversation, loaded for a turn knows_result(), new_results() what the model was given, across turns (to tell new results) ask(prompt_id, text) -> Turn start the next turn: refused while another has no answer -Turn one prompt, and what answers it +Turn one prompt, the messages steered to it, and what answers them id, n, status, answer, prompt its row entries, items its prompt and steps open_calls(), replies(), last_reply(), final_reply(), stored() its own state, never another turn's - latest_ask() its latest message of the user + latest_ask(), steered() its latest ask; the messages steered to it append(), append_reply(), drop(), carry(), ask_again(text) its run's writes + steer(messages, taken, going) take steered messages, at a step (take_steers) answered(answer), close(ctx, status, note) its end: once ``` @@ -433,7 +629,9 @@ invocation, a LangGraph graph's input) runs from the turn's latest message of the user (`latest_ask`). A retried attempt that finds the ask stored does not ask again. The check reads the reply to the ask too (`last`): a `TerminalError` it raises fails the turn, which is closed `failed` with it, -as `output=` does for a reply that still does not validate. +as `output=` does for a reply that still does not validate. A steered +message is a message of the user that is a step of the turn too, but no +ask: the turn goes on from its latest ask, which skips it. **Names.** A turn is the conversation's (`restate.agents.internal.Turn`). What `status` reports, `answering`, is the invocation being answered @@ -451,9 +649,9 @@ own answer, and the companion's turn is the invocation, all of them. | View | What it gives | | --- | --- | -| `status` | `answering`: the turn from its `start` to its close, or null; `waiting`: its tools' waits | +| `status` | `answering`: the turn from its `start` to its close, or null; `waiting`: its tools' waits; `steering`: the messages steered to the turn that it has not taken yet (absent if none) | | `pending` | the waits, oldest first, each with its id (`:`) for `signal` | -| `messages` | the conversation as stored. `view: "full"` (the default): every turn, with its tool calls and their outputs, and how a closed turn ended. `view: "summary"`: each turn's message and answer, under the same ids | +| `messages` | the conversation as stored. `view: "full"` (the default): every turn, with its tool calls and their outputs, the messages steered to it where it took them, and how a closed turn ended. `view: "summary"`: each turn's message, the messages steered to it, and its answer, under the same ids | | `notifications`, `watch` | the revision; `watch` waits (up to a minute) for a newer one | | `sync` | the revision, `status`, and the messages since a cursor, at once: what a follower reads | @@ -466,6 +664,7 @@ The revision increases at: - `settle`, once it closed a dead turn in the conversation; - a wait that starts, is resolved by a `signal`, or is forgotten (a cancelled turn's); +- a message steered to the turn, and the turn taking it (`taken`); - `clear`. So a follower that reads each new revision sees every turn begin, its tool @@ -483,10 +682,22 @@ does not show), and only such a turn is sent it: - 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. 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. + ## Guarantees and limits - **One turn at a time per key**, in the order messages came. A turn is - never joined by another message. + joined only by a message sent to `steer`, at a step. +- **A steered message is answered once**: by the turn that took it, or as a + turn of its own, in the order it was sent with the messages after it. One + that a turn took goes as the turn goes: stopped, refused, failed or ended + with it, it is not answered again. +- **A turn takes at most 16 steered messages**, and only as a model request + goes out, once each of the model's tool calls has a result. - **A turn is completed from its own rows**, and reads a bounded part of the conversation: the transcript the budget holds, and the window. - **`status` says the turn being answered** from its start to its close, diff --git a/python/restate/agents/__main__.py b/python/restate/agents/__main__.py index d85c00d0..d1b3e5e5 100644 --- a/python/restate/agents/__main__.py +++ b/python/restate/agents/__main__.py @@ -15,6 +15,7 @@ What you type goes to the agent, as a message it answers after the one it is answering, if any. A line that starts with ``/`` is a command:: + /steer send to the turn being answered, which takes it at its next step /signal [value] resolve the oldest wait for the signal (value: JSON, or text; default "approved") /pending what the agent's tools wait for /status what the agent is doing @@ -35,12 +36,13 @@ import os import sys import threading +import uuid from typing import Any from restate.agents._api import Part USAGE = """\ -commands: /signal [value], /pending, /status, /stop, /clear, /help; Ctrl-D leaves +commands: /steer , /signal [value], /pending, /status, /stop, /clear, /help; Ctrl-D leaves """ @@ -97,7 +99,11 @@ def _value(words: list[str]) -> object: async def _say(agent: Any, line: str) -> None: """Run a command, or send a message to the agent.""" command, *words = line.split() - if command == "/signal" and words: + if command == "/steer" and words: + steered = await agent.steer(line.split(None, 1)[1], idempotency_key=str(uuid.uuid4())) + turn = steered["turn"] + print(f"steered to {turn}" if turn else "queued: no turn takes it, it is answered next", flush=True) + elif command == "/signal" and words: resolved = await agent.signal(words[0], _value(words[1:])) print(f"resolved {resolved}" if resolved else f"nothing waits for {words[0]}", flush=True) elif command == "/pending": diff --git a/python/restate/agents/_api.py b/python/restate/agents/_api.py index fe65ca74..3abfef6f 100644 --- a/python/restate/agents/_api.py +++ b/python/restate/agents/_api.py @@ -182,6 +182,42 @@ class Notifications(TypedDict): revision: int +class Steer(TypedDict): + """A message steered to the turn being answered (``steer``): its id (the ``steer`` request's) and its text.""" + + id: str + text: str + + +class Published(Notifications, total=False): + """What the companion's ``publish`` returns to the turn that published: the new revision, and the messages + steered to that turn that it has not taken yet, if any, oldest first.""" + + steering: list[Steer] + + +class Steering(TypedDict): + """A message steered to the turn being answered that it has not taken yet, as ``status`` shows it.""" + + id: str + since: str + """When it was sent, as an ISO 8601 timestamp.""" + + +class Steered(TypedDict): + """The response of the ``steer`` view. + + ``id`` is the message's (the request's id). ``turn`` is the invocation of + the turn being answered, which takes the message at its next step if it + reaches one; ``None`` if no turn takes it. Either way the message is + answered once: in that turn, or else as a turn of its own, in the order + it was sent. + """ + + id: str + turn: str | None + + class _Watch(TypedDict): revision: int @@ -192,13 +228,18 @@ class Watch(_Watch, total=False): timeoutSeconds: float -class Status(TypedDict): - """What an agent is doing: the message it is answering, if any, and what its tools wait for, oldest first.""" - +class _Status(TypedDict): answering: Answering | None waiting: list[Waiting] +class Status(_Status, total=False): + """What an agent is doing: the message it is answering, if any, and what its tools wait for, oldest first.""" + + steering: list[Steering] + """The messages steered to the turn being answered that it has not taken yet, oldest first; absent if none.""" + + class Stopped(TypedDict): """The response of the ``stop`` view: the invocation it cancelled, or ``None`` if no message was being answered.""" diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index dad35f64..d0cf3d3b 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -88,9 +88,10 @@ import time import typing from collections.abc import Awaitable, Callable, Sequence -from dataclasses import dataclass, replace +from dataclasses import dataclass, field, replace from typing import Any, Literal, Protocol +from restate.agents._api import Steer from restate.agents._calls import outcomes from restate.agents._common import ( BUDGET, @@ -109,7 +110,7 @@ ) from restate.agents._turn import AttemptEnded, answering, settings_of, stop_signal_of from restate.exceptions import TerminalError -from restate.sessions import Database, DatabaseReader, DatabaseWriter, SessionContext +from restate.sessions import Database, DatabaseReader, DatabaseWriter, SessionContext, Transaction from restate.sessions.internal import Schema, definition_of _logger = logging.getLogger(__name__) @@ -129,6 +130,10 @@ SNAPSHOTS = "restate_agent_snapshots" """The history as a framework's own compaction left it, for agents that take their history from snapshots.""" +STEERED = "restate_agent_steers" +"""The messages steered to a turn that it took (``Turn.steer``): which items of the trajectories they are, by +position, and their ids, each stored once.""" + SCHEMA = Schema( "restate_agents.turns", [ @@ -169,6 +174,12 @@ turn INTEGER, items TEXT NOT NULL )""", + # Appended: a new table, which costs a view that applies it to its own session nothing. + f"""CREATE TABLE {STEERED} ( + n INTEGER PRIMARY KEY, + turn INTEGER NOT NULL, + id TEXT NOT NULL UNIQUE + )""", ], ) @@ -233,6 +244,12 @@ class Entry: """On a prompt: the ``seq`` given as its turn was answered.""" asked_at: int | None = None """On a prompt: when its turn's row was stored, in Unix milliseconds.""" + steer: str | None = None + """The id of the steered message it is, if it is one: a message of the user that its turn took at a step + (``Turn.steer``).""" + steers: list[Entry] = field(default_factory=list["Entry"]) + """On a prompt of the transcript: the messages steered to its turn that it took, in order, which the model + sees between its message and its answer once the turn is before the window.""" class Parsed(typing.NamedTuple): @@ -252,6 +269,10 @@ class Items(Protocol): A framework may also make the items that give the model the summary of the earliest turns, with ``summary(text) -> list[Item]``; by default they are a user's message with the summary and an assistant's acknowledgement. + And it may make the item of a message steered to a turn, with + ``steered(text, steer_id) -> Item``, to tell it apart from a prompt in + its runs (LangGraph's has the id ``steer:``); by default it is + a prompt's item (``prompt``). """ def dump(self, item: Item) -> str: @@ -310,6 +331,12 @@ def load_item(items: Items, stored: str) -> Item: ) from error +def steered_item(items: Items, text: str, steer_id: str) -> Item: + """Return the item of the message ``steer_id`` steered to a turn, with ``text``: ``items.steered`` if it has one.""" + own = getattr(items, "steered", None) + return own(text, steer_id) if own is not None else items.prompt(text) + + @dataclass(frozen=True) class Start: """Where a conversation starts: its first position, the last ``seq`` given, and the summary of the earliest turns.""" @@ -506,6 +533,9 @@ def carried(self) -> tuple[str | None, str | None]: _INSERT_STEP = f"INSERT INTO {STEPS} (n, turn, seq, item, reply, agent) VALUES (?, ?, ?, ?, ?, ?)" """Store a step of a turn: its position, its turn's, its ``seq``, its item, and the reply and agent it is of.""" +_INSERT_STEERED = f"INSERT INTO {STEERED} (n, turn, id) VALUES (?, ?, ?)" +"""Mark the step at a position a message steered to its turn: the position, its turn's, and the message's id.""" + _ANSWER = ( f"UPDATE {TURNS} SET status = ?, answer = ?, agent = ?, carried = ?, answered_at = ?, answered_seq = ? WHERE n = ?" ) @@ -624,16 +654,23 @@ def open_calls(self) -> list[Call]: return list(self._open) def latest_ask(self) -> Entry: - """Return the turn's latest message of the user: its prompt, or the message that asked again (``ask_again``). + """Return the turn's latest message of the user that asks for its answer: its prompt, or the message that + asked again (``ask_again``). A framework that runs from a message of its own (ADK's invocation, a - LangGraph graph's input) runs from this one. + LangGraph graph's input) runs from this one. A steered message + (``steer``) is not one: the model reads it within the run, which goes + on from the same ask. """ return next( - (e for e in reversed(self.items) if e.said is not None and e.said[0] == "user"), + (e for e in reversed(self.items) if e.steer is None and e.said is not None and e.said[0] == "user"), self.prompt, ) + def steered(self) -> list[Entry]: + """Return the messages steered to the turn that it took, in order.""" + return [entry for entry in self.items if entry.steer is not None] + def has_result(self, call_id: str) -> bool: """Return whether the turn's tool call ``call_id`` has a result.""" return call_id in self._results @@ -666,6 +703,65 @@ async def ask_again(self, text: str) -> None: """ await self.append(self._ledger.items.prompt(text)) + async def steer( + self, messages: Sequence[Steer], taken: Callable[[Transaction, list[str]], None], going: Callable[[], bool] + ) -> list[Item]: + """Store the steered ``messages`` that the turn has not taken yet as steps of the turn, and return their items. + + Each is stored once, by its id, as a message of the user (as + ``Items.steered`` makes it) after the turn's steps before: the model + reads it in its next request, and in every turn's history after, + between the turn's message and its answer. ``taken`` adds what goes + out with the transaction that stores them (the companion's + ``taken``), which commits once. + + ``going`` tells whether the turn still goes on: once it has the + ledger, a turn that stopped meanwhile takes none, which are then + answered as turns of their own. A store under way finishes even if the + turn stops: the turn then closes after it, at the positions after its + messages, and the messages are its own (stopped with it). + """ + held = {entry.steer for entry in self.steered()} + wanted = list({message["id"]: message for message in messages if message["id"] not in held}.values()) + if not wanted: + return [] + # Shielded: a stop cancels the run, but a commit that is under way still lands, and the turn's own record of + # its positions with it, so that the turn's close stores after it rather than over it. + return await asyncio.shield(self._steer(wanted, taken, going)) + + async def _steer( + self, messages: list[Steer], taken: Callable[[Transaction, list[str]], None], going: Callable[[], bool] + ) -> list[Item]: + ledger = self._ledger + items = ledger.items + async with ledger.lock: + if not going(): + return [] + self._check_writable() + # The database, not the turn: an id is stored once in the conversation, by whichever turn took it. + stored = await _steers_stored(ledger.db, [message["id"] for message in messages]) + new = [message for message in messages if message["id"] not in stored] + if not new: + return [] + n = ledger.next + added = [ + Entry(n + i, steered_item(items, message["text"], message["id"]), steer=message["id"]) + for i, message in enumerate(new) + ] + async with ledger.db.transaction() as tx: + for i, entry in enumerate(added, 1): + await tx.execute( + _INSERT_STEP, [entry.n, self.n, ledger.seq + i, items.dump(entry.item), None, None] + ) + await tx.execute(_INSERT_STEERED, [entry.n, self.n, entry.steer]) + taken(tx, [message["id"] for message in new]) + ledger.seq += len(added) + ledger.next = n + len(added) + self._add(added) + # The transcript's entry of the turn has them too, as the next turn loads it: for its size, say. + self.prompt.steers.extend(added) + return [entry.item for entry in added] + def carry(self, *, agent: str | None = None, state: str | None = None) -> None: """Keep what the turn carries to the turns after it: it is stored with its answer. @@ -702,6 +798,9 @@ async def drop(self, reply: int) -> None: was made from the reply. The positions are given again, so the new reply's messages take the dropped ones' ids; the ``seq`` is kept, so that a client's cursor still sees them. + + A steered message after the reply is dropped too: it is no longer the + turn's, so the agent's ``steered`` answers it as a turn of its own. """ self._check_writable() if reply not in self.replies(): @@ -714,9 +813,11 @@ async def drop(self, reply: int) -> None: start = replace(ledger.start, seq=ledger.seq) async with ledger.db.transaction() as tx: await tx.execute(f"DELETE FROM {STEPS} WHERE n >= ?", [reply]) + await tx.execute(f"DELETE FROM {STEERED} WHERE n >= ?", [reply]) await restart(tx, start) ledger.start = start ledger.next = min(ledger.next, reply) + self.prompt.steers = [steered for steered in self.prompt.steers if steered.n < reply] cut = bisect.bisect_left(self._entries, reply, key=_position) kept, dropped = self._entries[:cut], self._entries[cut:] # The dropped calls are not made: their ids may be made again, by the reply that replaces them. @@ -960,13 +1061,7 @@ async def load( first = start.n if window == "all" or start.summarized is None else start.summarized transcript = await _transcript(db, items, first, None if window == "all" or complete else (window, budget)) whole = transcript[window_start(transcript, window) :] - steps: list[Entry] = [] - if whole: - stored = await db.values(f"SELECT n, item, reply, agent FROM {STEPS} WHERE n > ? ORDER BY n", [whole[0].n]) - steps = [ - Entry(int(typing.cast(int, n)), load_item(items, str(item)), reply=_int(reply), agent=_str(agent)) - for n, item, reply, agent in stored - ] + steps = await _steps(db, items, whole[0].n) if whole else [] entries = sorted([*whole, *steps], key=_position) return cls(db, items, entries, start, transcript=transcript, last=last, seq=seq, budget=budget, window=window) @@ -988,13 +1083,7 @@ async def _from_snapshot(cls, db: Database, items: Items, start: Start, last: in [[carried]] = await db.values(f"SELECT MAX(n) FROM {TURNS} WHERE n < ? AND answer IS NOT NULL", [first]) transcript = await _transcript(db, items, first if carried is None else int(typing.cast(int, carried)), None) whole = [turn for turn in transcript if turn.n >= first] - steps: list[Entry] = [] - if whole: - stored = await db.values(f"SELECT n, item, reply, agent FROM {STEPS} WHERE n > ? ORDER BY n", [whole[0].n]) - steps = [ - Entry(int(typing.cast(int, n)), load_item(items, str(item)), reply=_int(reply), agent=_str(agent)) - for n, item, reply, agent in stored - ] + steps = await _steps(db, items, whole[0].n) if whole else [] entries = sorted([*whole, *steps], key=_position) return cls( db, items, entries, start, transcript=transcript, last=last, seq=seq, snapshots=True, snapshot=snapshot @@ -1093,17 +1182,11 @@ def _before(self) -> list[Entry]: start = self.window() return [turn for turn in self._ledger.transcript if turn.n < start and turn.answer is not None] - def _size(self, turn: Entry) -> int: - """Return the size of ``turn`` as the transcript has it, in characters: its message and its answer.""" - if turn.said is None: - turn.said = self.items.read(turn.item).said - return len(turn.said[1] if turn.said else "") + len(turn.answer or "") - def _latest(self, turns: list[Entry], characters: int) -> list[Entry]: """Return the latest of ``turns`` that ``characters`` holds, in order.""" kept = 0 for turn in reversed(turns): - characters -= self._size(turn) + characters -= _turn_size(self.items, turn) if characters < 0: break kept += 1 @@ -1131,7 +1214,7 @@ def overflowing(self) -> bool: if self.snapshots: return False before, _ = self._upcoming() - return sum(self._size(turn) for turn in before) > self.budget * 4 + return sum(_turn_size(self.items, turn) for turn in before) > self.budget * 4 def to_summarize(self) -> tuple[list[Entry], int] | None: """Return the turns to summarize, and the position of the first turn after them, if any are to be. @@ -1151,8 +1234,9 @@ def history(self) -> list[Item]: """Return the items the model gets, in order: the one place where turns are composed, flat. That is the summary of the earliest turns, if they have one; the turns - before the window, as their prompts and answers, as many of the latest - as ``budget`` holds; then the window's turns, whole. + before the window, as their prompts, the messages steered to them and + their answers, as many of the latest as ``budget`` holds; then the + window's turns, whole. With snapshots, it is the latest snapshot's items, then every item stored after it, whole; before the first snapshot, the whole @@ -1167,7 +1251,9 @@ def history(self) -> list[Item]: if (summary := self.summary) is not None: out.extend(summary_items(self.items, summary)) for turn in self.earlier(): - out.extend((turn.item, self.items.reply(turn.answer or "", turn.agent))) + out.append(turn.item) + out.extend(steered.item for steered in turn.steers) + out.append(self.items.reply(turn.answer or "", turn.agent)) out.extend(entry.item for turn in self._turns for entry in turn.entries) return out @@ -1257,6 +1343,8 @@ async def _transcript(db: DatabaseReader, items: Items, first: int, bounded: tup left = None if bounded is None else bounded[1] * 4 while True: rows = await db.values(sql, [first, before, PAGE if bounded is not None else 2**62]) + # The messages steered to the page's turns: positions between its earliest turn's and the turn after it. + steered = await _steered_in(db, items, int(typing.cast(int, rows[-1][0])), before) if rows else {} for n, prompt, item, status, answer, agent, carried, answered_at, answered_seq, asked_at in rows: at = int(typing.cast(int, n)) entry = Entry( @@ -1271,14 +1359,14 @@ async def _transcript(db: DatabaseReader, items: Items, first: int, bounded: tup answered_at=_int(answered_at), answered_seq=_int(answered_seq), asked_at=_int(asked_at), + steers=steered.get(at, []), ) read.append(entry) if walk is None or left is None or walk.holds(entry.status): continue if entry.answer is not None: - # A turn before the window: the model sees its message and answer, while the budget holds them. - entry.said = items.read(entry.item).said - left -= len(entry.said[1] if entry.said else "") + len(entry.answer or "") + # A turn before the window: the model sees its messages and answer, while the budget holds them. + left -= _turn_size(items, entry) if left < 0: return read[::-1] if bounded is None or len(rows) < PAGE: @@ -1286,6 +1374,67 @@ async def _transcript(db: DatabaseReader, items: Items, first: int, bounded: tup before = read[-1].n +async def _steps(db: DatabaseReader, items: Items, after: int) -> list[Entry]: + """Return the trajectories' entries after position ``after``, in order, each steered message marked as one.""" + stored = await db.values( + f"SELECT s.n, s.item, s.reply, s.agent, st.id FROM {STEPS} s LEFT JOIN {STEERED} st ON st.n = s.n " + "WHERE s.n > ? ORDER BY s.n", + [after], + ) + return [ + Entry( + int(typing.cast(int, n)), + load_item(items, str(item)), + reply=_int(reply), + agent=_str(agent), + steer=_str(steer), + ) + for n, item, reply, agent, steer in stored + ] + + +async def _steered_in(db: DatabaseReader, items: Items, first: int, end: int) -> dict[int, list[Entry]]: + """Return the messages steered to the turns at positions ``first`` (included) to ``end``, by their turn's position. + + It reads only the steered messages: one range of their table's primary key, joined to their items. + """ + rows = await db.values( + f"SELECT st.n, st.turn, st.id, s.item FROM {STEERED} st JOIN {STEPS} s ON s.n = st.n " + "WHERE st.n >= ? AND st.n < ? ORDER BY st.n", + [first, end], + ) + steered: dict[int, list[Entry]] = {} + for n, turn, steer_id, item in rows: + loaded = load_item(items, str(item)) + at = int(typing.cast(int, turn)) + entry = Entry(int(typing.cast(int, n)), loaded, steer=str(steer_id), turn=at, said=items.read(loaded).said) + steered.setdefault(at, []).append(entry) + return steered + + +def _turn_size(items: Items, turn: Entry) -> int: + """Return the size of ``turn`` (a prompt of the transcript) as the model sees it before the window, in characters: + its message, the messages steered to it, and its answer.""" + if turn.said is None: + turn.said = items.read(turn.item).said + said = [turn.said, *(steered.said for steered in turn.steers)] + return sum(len(one[1]) for one in said if one is not None) + len(turn.answer or "") + + +async def steer_stored(db: DatabaseReader, steer_id: str) -> bool: + """Return whether a turn took the steered message ``steer_id``, in a database that may not have the conversation + yet (the agent's ``steered`` asks, which may be the key's first handler).""" + await SCHEMA.ensure(db) + return steer_id in await _steers_stored(db, [steer_id]) + + +async def _steers_stored(db: DatabaseReader, steer_ids: Sequence[str]) -> set[str]: + """Return which of the steered messages ``steer_ids`` a turn took: one read, of their unique ids.""" + marks = ", ".join("?" for _ in steer_ids) + stored = await db.values(f"SELECT id FROM {STEERED} WHERE id IN ({marks})", list(steer_ids)) + return {str(row[0]) for row in stored} + + def _str(value: object) -> str | None: return None if value is None else str(value) @@ -1444,7 +1593,8 @@ def _transcribed(conversation: Conversation, turns: Sequence[Entry]) -> str: lines = [f"Summary of the conversation before:\n{conversation.summary}\n"] if conversation.summary else [] for turn in turns: said = turn.said[1] if turn.said else "" - lines.append(f"User: {said}\nAssistant: {turn.answer or ''}\n") + steered = "".join(f"User: {steer.said[1] if steer.said else ''}\n" for steer in turn.steers) + lines.append(f"User: {said}\n{steered}Assistant: {turn.answer or ''}\n") return "\n".join(lines) diff --git a/python/restate/agents/_messages.py b/python/restate/agents/_messages.py index 303fdd88..c9082f8b 100644 --- a/python/restate/agents/_messages.py +++ b/python/restate/agents/_messages.py @@ -30,7 +30,7 @@ from dataclasses import dataclass, field from restate.agents._api import ChatMessage, Messages, MessagesRequest, Part, Role, ToolPart, text_part -from restate.agents._conversation import SCHEMA, STEPS, TURNS, Items, last_of, load_item, start_of +from restate.agents._conversation import SCHEMA, STEERED, STEPS, TURNS, Items, last_of, load_item, start_of from restate.sessions import DatabaseReader @@ -174,35 +174,51 @@ async def _latest(db: DatabaseReader, items: Items, before: int | None, limit: i } -# --- The summary: per turn, the user's message and the answer ------------------------------------------------------- +# --- The summary: per turn, the user's messages and the answer ------------------------------------------------------ _TURN = f"SELECT n, seq, item, answer, answered_at, answered_seq FROM {TURNS}" +_STEERED = f"SELECT st.n, s.seq, s.item FROM {STEERED} st JOIN {STEPS} s ON s.n = st.n" +"""The messages steered to turns that took them: messages of the user, which the summary shows too.""" -def _summarized(items: Items, rows: typing.Sequence[typing.Sequence[typing.Any]]) -> list[_Message]: - """Render turns' rows (``_TURN``'s columns) as their user's messages and answers, in order. + +def _summarized( + items: Items, + rows: typing.Sequence[typing.Sequence[typing.Any]], + steered: typing.Sequence[typing.Sequence[typing.Any]] = (), +) -> list[_Message]: + """Render turns' rows (``_TURN``'s columns) as their user's messages and answers, with the messages steered to + them (``_STEERED``'s columns), in order. The answer's message has the id it has in the full view (the turn's last - message of the assistant), and the ``seq`` it was stored with. + message of the assistant), and the ``seq`` it was stored with; so does a + steered message. """ - messages: list[_Message] = [] - for n, seq, stored, answer, answered_at, answered_seq in rows: + + def user_message(n: typing.Any, seq: typing.Any, stored: typing.Any) -> _Message: said = items.read(load_item(items, str(stored))).said text = said[1] if said else "" - messages.append(_Message(int(n), "user", [text_part(text)] if text.strip() else [], int(seq))) + return _Message(int(n), "user", [text_part(text)] if text.strip() else [], int(seq)) + + messages: list[_Message] = [] + for n, seq, stored, answer, answered_at, answered_seq in rows: + messages.append(user_message(n, seq, stored)) if answer is not None and answered_at is not None: answered = str(answer) parts: list[Part] = [text_part(answered)] if answered.strip() else [] messages.append(_Message(int(answered_at), "assistant", parts, int(typing.cast(int, answered_seq)))) - return [message for message in messages if message.parts] + messages.extend(user_message(n, seq, stored) for n, seq, stored in steered) + return sorted((message for message in messages if message.parts), key=lambda message: message.id) async def _summary_since(db: DatabaseReader, items: Items, cursor: int, limit: int) -> Messages: """Return the user's messages and answers stored since ``cursor``, in the order they were.""" asked = await db.values(f"{_TURN} WHERE seq > ?", [cursor]) answered = await db.values(f"{_TURN} WHERE answered_seq > ?", [cursor]) + steered = await db.values(f"{_STEERED} WHERE s.seq > ?", [cursor]) rows = {int(typing.cast(int, row[0])): row for row in [*asked, *answered]} - messages = sorted((m for m in _summarized(items, list(rows.values())) if m.seq > cursor), key=lambda m: m.seq) + summarized = _summarized(items, list(rows.values()), steered) + messages = sorted((m for m in summarized if m.seq > cursor), key=lambda m: m.seq) shown = messages[:limit] return { "messages": [_chat(m) for m in shown], @@ -218,9 +234,10 @@ async def _summary_latest(db: DatabaseReader, items: Items, before: int | None, start = await start_of(db) _, now = await last_of(db) cursor = max(start.seq, now) - messages = [m for m in _summarized(items, rows[::-1]) if m.id < end] - first = int(typing.cast(int, rows[-1][0])) if rows else end - [[earlier]] = await db.values(f"SELECT EXISTS (SELECT 1 FROM {TURNS} WHERE n < ?)", [first]) + earliest = int(typing.cast(int, rows[-1][0])) if rows else end + steered = await db.values(f"{_STEERED} WHERE st.n >= ? AND st.n < ? ORDER BY st.n", [earliest, end]) + messages = [m for m in _summarized(items, rows[::-1], steered) if m.id < end] + [[earlier]] = await db.values(f"SELECT EXISTS (SELECT 1 FROM {TURNS} WHERE n < ?)", [earliest]) return { "messages": [_chat(m) for m in _latest_of(messages, limit)], "more": len(messages) > limit or bool(earlier), diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index 9f69e5de..3f06fe0a 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -31,6 +31,8 @@ Resolved, Signal, Status, + Steer, + Steered, Stopped, View, Waiting, @@ -46,7 +48,7 @@ check_window, refusal_of, ) -from restate.agents._conversation import SCHEMA, Items, next_start, restart +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 @@ -171,6 +173,8 @@ class Synced(Messages): "erase", "settle", "compact", + "steer", + "steered", } ) """The names of the built-in handlers and views of every agent, which are reserved.""" @@ -196,6 +200,9 @@ class AgentObject(SessionObject): ``chat()`` turns the object into an agent, and adds: - ``chat``: a handler that takes a message and returns the answer. + - ``steer``: send a message while a turn is answered: the turn takes it + at its next step, or else it is answered as a turn of its own, in the + order it was sent. - ``messages``: read the conversation, one page at a time. - ``clear``: stop the message being answered and forget the conversation, emptying the key's database. @@ -203,8 +210,9 @@ class AgentObject(SessionObject): and the revision, for clients that follow the conversation. These names, ``erase`` (which ``clear`` uses), ``settle`` (which closes - a turn whose invocation is gone) and ``compact`` (which summarizes the - earliest turns), are reserved: + a turn whose invocation is gone), ``compact`` (which summarizes the + earliest turns) and ``steered`` (which answers a steered message that no + turn took), are reserved: registering a handler or view of your own under one of them raises ``ValueError``. You rarely create an ``AgentObject`` yourself: a framework's ``agent()`` creates one and calls ``chat()`` for you. @@ -355,6 +363,7 @@ def chat( summarize: Summarize | None = None, snapshots: bool = False, compact_history: CompactHistory | None = None, + steers: bool = False, ) -> A: ... @typing.overload @@ -368,6 +377,7 @@ def chat( summarize: Summarize | None = None, snapshots: bool = False, compact_history: CompactHistory | None = None, + steers: bool = False, ) -> Callable[[A], A]: ... def chat( @@ -381,6 +391,7 @@ def chat( summarize: Summarize | None = None, snapshots: bool = False, compact_history: CompactHistory | None = None, + steers: bool = False, ) -> Any: """Turn this object into an agent that answers each message with ``answer``. @@ -425,6 +436,10 @@ def chat( ``None``), which is stored as a snapshot. It runs in the private ``compact`` handler, after a turn that asked for it (``Conversation.compact_after``). + steers: Whether ``answer`` takes steered messages at the turn's + steps (``take_steers``), which a framework's integration + says. Without, every message sent to ``steer`` is answered as + a turn of its own, after the turn being answered. Returns: ``answer`` itself. @@ -442,6 +457,7 @@ def decorate( summarize=summarize, snapshots=snapshots, compact_history=compact_history, + steers=steers, ) return decorate @@ -454,6 +470,7 @@ def decorate( summarize=summarize, snapshots=snapshots, compact_history=compact_history, + steers=steers, ) AGENT.on(self, self._settings) answering = answer @@ -463,9 +480,20 @@ async def chat(ctx: SessionContext, text: str) -> Any: chat.__annotations__["return"] = output if output is not None else str + async def steered(ctx: SessionContext, 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 + place in the key's queue: it runs once the turns before it ended. + """ + if await steer_stored(ctx.db, message["id"]): + return + await chat_turn(ctx, message["text"], answering) + self._building = True try: self.handler(name="chat", accept="*/*", input_serde=MessageSerde(), explicit_cancellation=True)(chat) + self.handler(name="steered", ingress_private=True, explicit_cancellation=True)(steered) _followed(self, items) finally: self._building = False @@ -503,7 +531,8 @@ async def _stopped(ctx: ViewContext, step: str) -> Stopped: def _followed(agent: AgentObject, items: Items) -> None: - """Register the views that ``chat()`` adds (``messages``, ``clear``, ``sync``), and the private handlers.""" + """Register the views that ``chat()`` adds (``messages``, ``steer``, ``clear``, ``sync``), and the private + handlers.""" async def messages(ctx: ViewContext, request: MessagesRequest | None = None) -> Messages: """Read the conversation: a page of it, or what changed since a cursor. @@ -514,6 +543,20 @@ 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: + """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, + as for ``chat``. The turn being answered gives it to the model at its + next step (once the model's tool calls have their results), if its + framework takes steered messages and it reaches one; otherwise the + message is answered as a turn of its own, after the turns before it, + in the order it was sent. With no turn being answered, it starts one. + 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}) + async def erase(ctx: SessionContext) -> 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. @@ -581,6 +624,7 @@ async def sync(ctx: ViewContext, request: SyncRequest | None = None) -> Synced: return synced agent.view(messages) + agent.view(accept="*/*", input_serde=MessageSerde())(steer) agent.handler(ingress_private=True)(erase) agent.handler(ingress_private=True)(settle) agent.handler(ingress_private=True)(compact) diff --git a/python/restate/agents/_signals.py b/python/restate/agents/_signals.py index c6731d49..951a077e 100644 --- a/python/restate/agents/_signals.py +++ b/python/restate/agents/_signals.py @@ -13,9 +13,10 @@ 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 -answering, 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. +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. 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 @@ -33,7 +34,19 @@ import restate from restate import ObjectContext, ObjectSharedContext, TerminalError -from restate.agents._api import Answering, Notifications, Resolved, Signal, Status, Waiting, Watch +from restate.agents._api import ( + Answering, + Notifications, + Published, + Resolved, + Signal, + Status, + Steer, + Steered, + Steering, + Waiting, + Watch, +) from restate.sessions import SessionContext, ViewContext from restate.sessions._link import Link from restate.sessions.internal import definition_of @@ -59,6 +72,13 @@ short journals. """ +STEERS = 16 +"""How many messages a turn takes by ``steer``, at most: past them, each is answered as a turn of its own. + +The companion counts each message it holds for the turn, taken or not, until +the turn ends: so a turn that someone steers on and on stays within bounds, +and so does what the companion holds for it.""" + class _WaitRequest(TypedDict): id: str @@ -94,6 +114,24 @@ class Started(Answering, total=False): """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.""" + + +class _Held(TypedDict): + """A message steered to the turn ``turn`` (an invocation id), and when it was sent.""" + + id: str + since: str + turn: str + + +class _Inboxed(_Held, total=False): + """A message steered to a turn, as the companion holds it until the turn ends: with its text until the turn + takes it (``taken``), and then only as one of the turn's ``STEERS``.""" + + text: str def _shown(turn: Started | None) -> Answering | None: @@ -101,6 +139,20 @@ def _shown(turn: Started | None) -> Answering | None: return None if turn is None else {"invocationId": turn["invocationId"], "since": turn["since"]} +async def _inbox(ctx: Context) -> list[_Inboxed]: + """Return the messages steered to the turns not closed yet (the one being answered), taken or not, oldest + first.""" + return await _get(ctx, "inbox", typing.cast(list[_Inboxed], [])) + + +def _waiting(inbox: list[_Inboxed], answering: Started | None) -> list[_Inboxed]: + """Return the messages of ``inbox`` steered to the turn being answered (``answering``) that it has not taken yet, + oldest first: those with their text.""" + if answering is None: + return [] + return [one for one in inbox if one["turn"] == answering["invocationId"] and "text" in one] + + async def _get(ctx: Context, key: str, default: T) -> T: """Return the state ``key``, or ``default`` if it is unset.""" value = await ctx.get(key, type_hint=object) @@ -142,6 +194,10 @@ async def _closed(ctx: ObjectContext, turn: str, gone: bool = False) -> None: 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. + + 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 + answered as one. """ answering = await _get(ctx, "turn", typing.cast(Started | None, None)) is_answered = answering is not None and answering["invocationId"] == turn @@ -150,6 +206,10 @@ async def _closed(ctx: ObjectContext, turn: str, gone: bool = False) -> None: 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) + inbox = await _inbox(ctx) + dropped = any(one["turn"] == turn for one in inbox) + if dropped: + _set_inbox(ctx, [one for one in inbox if one["turn"] != turn]) waiting = await _get(ctx, "waiting", typing.cast(list[_WaitingOn], [])) # A wait's id starts with the invocation id of the turn whose tool waits. over = [wait for wait in waiting if wait["id"].startswith(f"{turn}:")] @@ -157,7 +217,63 @@ async def _closed(ctx: ObjectContext, turn: str, gone: bool = False) -> None: ctx.reject_awakeable(wait["awakeable"], "its turn is over") if over: ctx.set("waiting", [wait for wait in waiting if not wait["id"].startswith(f"{turn}:")]) - if is_answered or over: + if is_answered or over or dropped: + await _publish(ctx) + + +def _set_inbox(ctx: ObjectContext, inbox: list[_Inboxed]) -> None: + if inbox: + ctx.set("inbox", inbox) + else: + 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) @@ -276,11 +392,18 @@ async def await_turn(ctx: ObjectSharedContext, turn: str) -> None: @signals.handler(kind="shared") async def status(ctx: ObjectSharedContext) -> Status: - """Return what the agent is doing: the message it is answering, if any, and its tools' waits.""" - return { - "answering": _shown(await _get(ctx, "turn", typing.cast(Started | None, None))), + """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. @@ -317,9 +440,11 @@ async def ended(ctx: ObjectContext, turn: str) -> None: @signals.handler(ingress_private=True) async def reset(ctx: ObjectContext) -> None: - """Forget all waits and kept signals after the agent's ``clear``, and publish the change.""" + """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) @@ -332,9 +457,22 @@ async def forget(ctx: ObjectContext, wait_id: str) -> None: @signals.handler(ingress_private=True) -async def publish(ctx: ObjectContext) -> Notifications: - """Increase the conversation's revision after a change, and wake up every waiting ``watch``.""" - return await _publish(ctx) +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 at each tool + call's start and end, so it learns of them at no cost of its own. + """ + published: Published = {**await _publish(ctx)} + inbox = await _inbox(ctx) + if any("text" in one for one in inbox): + # Only then is the turn read: a turn that nobody steers costs one read more, of the empty inbox. + answering = await _get(ctx, "turn", typing.cast(Started | None, None)) + steering: list[Steer] = [{"id": one["id"], "text": one.get("text", "")} for one in _waiting(inbox, answering)] + if steering: + published["steering"] = steering + return published @signals.handler(kind="shared") diff --git a/python/restate/agents/_steering.py b/python/restate/agents/_steering.py new file mode 100644 index 00000000..cf02d04f --- /dev/null +++ b/python/restate/agents/_steering.py @@ -0,0 +1,74 @@ +# +# 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 +# +"""Steering: a message sent while a turn is answered joins it, at the turn's next step. + +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``). + +The turn learns of it from its next ``publish``, which it makes at each tool +call's start and end anyway. At its next step, once the model's reply has a +result for each of its tool calls and before the next model request, the +framework integration calls ``take_steers``: the turn stores each message +once, as a message of the user that is a step of the turn, and the +integration puts it in the request that is about to go out. The transaction +that stores them tells the companion they are taken (``taken``), once. + +Once the turn ends, the queued ``steered`` for a message it took does +nothing; one it did not take (it came after the turn's last step, say) is +answered as a turn of its own. So each message is answered once. +""" + +from __future__ import annotations + +from restate.agents import _signals +from restate.agents._conversation import Conversation, Item +from restate.agents._turn import inbox_of, stop_signal_of +from restate.sessions import SessionContext, Transaction + + +async def take_steers(ctx: SessionContext, conversation: Conversation) -> list[Item]: + """Take the messages steered to the turn being answered, at a step: store each once, and return their items. + + Call it right before a model request that follows the model's reply and + its tools' results: the integration then puts the returned items in the + request, after what it already holds. They are messages of the user, as + the framework's ``Items`` makes them (``steered``, or else ``prompt``). + It returns nothing, at no cost, unless the turn runs in a ``chat`` + handler of an agent that takes steered messages, a message it has not + taken is waiting, the turn is not stopping, and each of its tool calls + has a result. + + The messages are steps of the turn from then on: a retried attempt finds + them stored, in ``history``, and never takes them again. They share the + turn's fate: stopped, refused or failed with it, they are not answered + again. + """ + inbox = inbox_of(ctx) + if inbox is None or not inbox.messages: + return [] + stop = stop_signal_of(ctx) + + def going() -> bool: + # A turn stopped while it waited to store them takes none: each is then answered as a turn of its own. + return stop is None or not stop.stopped + + turn = conversation.running() + if not going() or turn is None or turn.open_calls(): + return [] + companion = _signals.signals_of(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]}") + + return await turn.steer(inbox.messages, taken, going) diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py index 51f8337f..597085c2 100644 --- a/python/restate/agents/_turn.py +++ b/python/restate/agents/_turn.py @@ -23,10 +23,10 @@ import time import typing from collections.abc import Awaitable, Callable -from dataclasses import dataclass +from dataclasses import dataclass, field from restate.agents import _signals -from restate.agents._api import Answering +from restate.agents._api import Answering, Published, Steer from restate.agents._calls import settle from restate.agents._common import BUDGET, WINDOW, CompactHistory, Summarize, Window, refusal_of from restate.exceptions import TerminalError @@ -135,6 +135,9 @@ class AgentInfo: ``summarize`` then do not apply.""" compact_history: CompactHistory | None = None """With ``snapshots``, how the agent compacts its history between turns, when a turn asks for it.""" + steers: bool = False + """Whether the agent's framework takes steered messages at a turn's steps (``take_steers``): without, each + message steered to the agent is answered as a turn of its own. A background agent takes none.""" AGENT: Link[AgentInfo] = Link("an agent") @@ -164,6 +167,37 @@ def _unfollowed(ctx: SessionContext) -> bool: return agent is not None and agent.background +@dataclass +class Inbox: + """The messages steered to a ``chat`` turn that it has not taken yet, as the companion last published them. + + The companion returns them from every ``publish``, which a turn makes at + each tool call's start and end: so the turn knows of them without a call + of its own, and finds them at its next step (``take_steers``). The + newest revision wins, as tool calls that run at once publish at once. + """ + + revision: int = -1 + messages: list[Steer] = field(default_factory=list[Steer]) + + def saw(self, published: Published | None) -> None: + """Keep what ``publish`` returned, if it is newer than what is kept.""" + if not isinstance(published, dict): + return + revision = published.get("revision", -1) + if revision > self.revision: + self.revision = revision + self.messages = list(published.get("steering", [])) + + +_INBOX: ScopeKey[Inbox] = ScopeKey("inbox") + + +def inbox_of(ctx: SessionContext) -> Inbox | None: + """Return the inbox of the ``chat`` turn running on ``ctx``, if it takes steered messages: ``None`` outside one.""" + return _INBOX.get(ctx) + + _ANSWERING: ScopeKey[Stored] = ScopeKey("answering") """What the attempt's turn has stored, while it answers.""" @@ -195,7 +229,10 @@ async def publish(ctx: SessionContext) -> None: stored = _ANSWERING.get(ctx) if stored is not None: await _settled(ctx, stored) - await ctx.call(_signals.publish, key=_signals.signals_of(ctx)) + published = await ctx.call(_signals.publish, key=_signals.signals_of(ctx)) + inbox = _INBOX.get(ctx) + if inbox is not None: + inbox.saw(typing.cast(Published | None, published)) async def stepped(ctx: SessionContext) -> None: @@ -285,9 +322,14 @@ async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionCon invocation_id = ctx.request.id companion = _signals.signals_of(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. - started: _signals.Started = {**turn, "settles": True} + started: _signals.Started = {**turn, "settles": True, "steers": steers} await ctx.call(_signals.start, started, key=companion) await stepped(ctx) diff --git a/python/restate/agents/adk/_agent.py b/python/restate/agents/adk/_agent.py index 134b3305..d515e3f9 100644 --- a/python/restate/agents/adk/_agent.py +++ b/python/restate/agents/adk/_agent.py @@ -30,6 +30,23 @@ have (which ADK fails the run for). An agent with a ``retry_config`` of its own is first run again by ADK, within the attempt. - A transfer to a sub-agent holds for the next messages, as ADK's sessions do. +- A message steered to the turn (``steer``) joins it at its next step: as + the model's next request goes out, once its tool calls have their results, + the turn stores it (``take_steers``), and it is added to that request and + to the session, after the results, so that every request after has it too. + It starts an invocation of its own, as a prompt does: after a crash, ADK + resumes the prompt's invocation from that invocation's events. A message + that no request follows (one sent during the model's last reply, or during + a tool after whose result ADK asks the model nothing more: a tool's + ``skip_summarization``, an output schema's ``set_model_response``) is + answered next, as a turn of its own. A tool that runs an agent + (``AgentTool``) is one call: the step comes after its result. Only an + agent whose root is an ``LlmAgent`` takes steered messages, as the + built-in agent does; with a workflow root (a ``SequentialAgent``, say, or + an ``App``'s ``Workflow``), each is answered as a turn of its own. An agent + with ``include_contents='none'`` starts each request at the latest message + of the user: the requests after the one that takes a steered message start + at it, without the prompt and the tools' results before it. - A run of more model calls than its ``RunConfig``'s ``max_llm_calls`` (500 by default) fails its prompt for good: a retry would carry on with a fresh count. - An agent or ``App`` of your own runs as you configured it, with the @@ -61,11 +78,13 @@ from typing import Any, Unpack from google.adk.agents import BaseAgent, LlmAgent, RunConfig +from google.adk.agents.callback_context import CallbackContext from google.adk.agents.invocation_context import LlmCallsLimitExceededError, new_invocation_context_id from google.adk.apps import App, ResumabilityConfig from google.adk.events import Event from google.adk.models import BaseLlm from google.adk.models.llm_request import LlmRequest +from google.adk.models.llm_response import LlmResponse from google.adk.plugins import BasePlugin from google.adk.runners import Runner from google.adk.sessions import BaseSessionService, Session @@ -92,6 +111,7 @@ from restate.agents._conversation import Call, Conversation, Item, Parsed, Result, Run, answer_once from restate.agents._object import AgentObject, AgentOptions from restate.agents._output import answer_with +from restate.agents._steering import take_steers from restate.agents._tools import Tool, current_call, tools_in from restate.agents._turn import tool_call_ended, tool_call_started from restate.exceptions import TerminalError @@ -141,6 +161,12 @@ def prompt(self, text: str) -> Item: content = types.Content(role="user", parts=[types.Part(text=text)]) return Event(author="user", invocation_id=new_invocation_context_id(), content=content) + def steered(self, text: str, steer_id: str) -> Item: + # An invocation of its own too, not the running one, which ADK resumes from its events after a crash: there, + # ADK's runner of a workflow root would take a message of the user last among them for the answer to a tool + # call, which it fails to find on every retry. + return self.prompt(text) + def reply(self, text: str, agent: str | None = None) -> Item: content = types.Content(role="model", parts=[types.Part(text=text)]) return Event(author=agent or self.author, content=content) @@ -231,13 +257,37 @@ async def append_event(self, session: Session, event: Event) -> Event: self._carry(session) return event + async def add_stored(self, session: Session, event: Event) -> None: + """Add ``event``, which the turn stored already (a message steered to it), to ``session``, after its events. + + ADK builds each model request from the session's events: every request after this one has it too, in place. + """ + await super().append_event(session, event) + class _Restate(BasePlugin): - """Publish each tool call as it starts and ends, and mark a tool's error as the tool's.""" + """Publish each tool call as it starts and ends, mark a tool's error as the tool's, and give the model the + messages steered to the turn.""" - def __init__(self, ctx: SessionContext) -> None: + def __init__(self, ctx: SessionContext, sessions: _Sessions) -> None: super().__init__(name="restate") self.ctx = ctx + self.sessions = sessions + + async def before_model_callback( + self, *, callback_context: CallbackContext, llm_request: LlmRequest + ) -> LlmResponse | None: + # The turn's step: a model request is going out, which ADK built from the session's events. Once the + # model's tool calls have their results, the messages steered to the turn join it, after the results. Here + # rather than as the results are stored: ADK makes no request after some (a tool's skip_summarization, an + # output schema's set_model_response), and a message the turn took there would go unread. The agent's own + # before_model_callback runs after this one, on the request with them: a response it gives in place of the + # model's answers them too, as the model would have. + for event in typing.cast(list[Event], await take_steers(self.ctx, self.sessions.conversation)): + await self.sessions.add_stored(callback_context.session, event) + if event.content is not None: + llm_request.contents.append(event.content.model_copy(deep=True)) + return None async def before_tool_callback( self, *, tool: BaseTool, tool_args: dict[str, Any], tool_context: ToolContext @@ -319,14 +369,13 @@ def _model(model: AnyModel | None) -> str | BaseLlm: def _app( - ctx: SessionContext, own: BaseAgent | App | None, model: AnyModel | None, system: str | None, given: Sequence[AnyTool], output: type[BaseModel] | None, ) -> App: - """Return the app that runs the agent: resumable, with this package's plugin. + """Return the app that runs the agent: resumable (each run adds this package's plugin, ``_Restate``). ``output`` is the ``output_schema`` of its root ``LlmAgent``. ADK answers with its JSON, through a tool of its own when the model cannot combine a @@ -347,9 +396,8 @@ def _app( if app.events_compaction_config is not None: _without_compaction() resumable = ResumabilityConfig(is_resumable=True) - update = {"name": _APP, "plugins": [*app.plugins, _Restate(ctx)], "resumability_config": resumable} # The conversation's summary and window stand in for ADK's events compaction (see the module's docstring). - return app.model_copy(update={**update, "events_compaction_config": None}) + return app.model_copy(update={"name": _APP, "resumability_config": resumable, "events_compaction_config": None}) @functools.cache @@ -421,7 +469,11 @@ async def run(conversation: Conversation) -> str: return final turn = conversation.running() invocation = typing.cast(Event, turn.latest_ask().item).invocation_id if turn is not None else None - runner = Runner(app=app, session_service=_Sessions(conversation)) + sessions = _Sessions(conversation) + # This package's plugin, after the app's own: a plugin of the app that answers in the model's place leaves + # the messages steered to the turn to a request that goes out. + plugins = [*app.plugins, _Restate(ctx, sessions)] + runner = Runner(app=app.model_copy(update={"plugins": plugins}), session_service=sessions) answer = "" error: Event | None = None events = runner.run_async(user_id=_USER, session_id=_USER, invocation_id=invocation, run_config=config) @@ -582,7 +634,7 @@ async def _prompt( config: RunConfig | None, ) -> Any: """Answer ``text`` as ``prompt`` does, with its options checked.""" - app = _app(ctx, own, model, system, tools, output) + app = _app(own, model, system, tools, output) def ask(text: str, check: Check | None = None) -> Awaitable[str]: return answer_once(ctx, text, _Items(_root(app).name), _run(ctx, app, config), check=check) @@ -692,10 +744,28 @@ async def answer(ctx: SessionContext, text: str) -> Any: run_config, ) - obj.chat(answer, items=_Items(), output=output, summarize=_summarizer(model) if own is None else None) + obj.chat( + answer, + items=_Items(), + output=output, + summarize=_summarizer(model) if own is None else None, + steers=_steers(own), + ) return obj +def _steers(own: BaseAgent | App | None) -> bool: + """Return whether the agent takes steered messages at its turns' steps: if its root is an ``LlmAgent``, as the + built-in agent's is. + + A workflow root (a ``SequentialAgent``, ``ParallelAgent`` or + ``LoopAgent``, or an ``App``'s ``Workflow``) takes none: each message + steered to it is answered as a turn of its own. + """ + root = own.root_agent if isinstance(own, App) else own + return own is None or isinstance(root, LlmAgent) + + def _summarizer(model: PerSession[AnyModel | None]) -> Summarize: """Return how the built-in agent summarizes its earlier turns: with its own model, in a request of its own.""" diff --git a/python/restate/agents/client.py b/python/restate/agents/client.py index ae58c4ea..9f78f13e 100644 --- a/python/restate/agents/client.py +++ b/python/restate/agents/client.py @@ -43,7 +43,7 @@ "restate.agents.client needs Restate's client: install the client extra, pip install restate-sdk[client]" ) from error -from restate.agents._api import ChatMessage, Messages, MessagesRequest, Signal, Status, View, Waiting +from restate.agents._api import ChatMessage, Messages, MessagesRequest, Signal, Status, Steered, View, Waiting from restate.agents._object import Synced, SyncRequest from restate.client_types import HttpError, RestateClient @@ -119,6 +119,20 @@ async def send(self, text: str, *, idempotency_key: str | None = None) -> str: ) return sent.invocation_id + async def steer(self, text: str, *, idempotency_key: str | None = None) -> Steered: + """Send ``text`` while the agent answers another message: the turn takes it at its next step. + + If the turn being answered reaches a step after it (once its model's + tool calls have their results), the model gets the message there, and + the turn's answer covers it. Otherwise the message is answered as a + turn of its own, after the messages before it, in the order it was + sent; with no turn being answered, it starts one. Either way it is + answered once. Pass an ``idempotency_key`` (made once per message), + so that sending it again (after a lost connection, say) is the same + message; ``follow`` shows it, and its answer, as they come. + """ + return cast(Steered, await self._call("steer", text, idempotency_key=idempotency_key)) + async def messages( self, *, diff --git a/python/restate/agents/internal.py b/python/restate/agents/internal.py index fbdeb788..43493e9e 100644 --- a/python/restate/agents/internal.py +++ b/python/restate/agents/internal.py @@ -33,6 +33,10 @@ 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. +- ``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._common import ( @@ -70,6 +74,7 @@ ) from restate.agents._output import answer_with, asking from restate.agents._signals import Companion, companion_of, deliver, signals +from restate.agents._steering import take_steers from restate.agents._tools import run_tool, wait_for from restate.agents._turn import ( AttemptEnded, @@ -123,6 +128,7 @@ "signals", "stop_signal_of", "stored_answer", + "take_steers", "tool_call_ended", "tool_call_started", "wait_for", diff --git a/python/restate/agents/langgraph/__init__.py b/python/restate/agents/langgraph/__init__.py index 908e71c9..68ba8bc5 100644 --- a/python/restate/agents/langgraph/__init__.py +++ b/python/restate/agents/langgraph/__init__.py @@ -60,6 +60,9 @@ async def get_weather(city: str) -> str: graphs that live long. - Clients see a run step by step, as each step is stored: there is no token streaming. +- A message sent to the agent's ``steer`` view is answered as a turn of its + own, after the one being answered: a graph of your own takes no steered + messages. (The built-in agent's model gets one at the turn's next step.) - Pause a graph with ``interrupt()``, not static breakpoints (``interrupt_before``, ``interrupt_after``), which have nothing to wait for: a compiled graph with them is refused, and a run that stops at a diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 7fc307c3..2d3fe2a5 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -32,15 +32,22 @@ turn that reaches it ends for good, with a terminal error: a retry would carry on from where it stopped, with a fresh limit, and so would never end a model that keeps calling tools. +- A message steered to the turn (the agent's ``steer`` view) joins it at its + next step: once the model's tool calls have their results, the model node + has the turn take it (``take_steers``), which stores it once, and gives it + to the model's next request, after the results. The graph's state keeps it + for the rest of the run, which goes on (within the same recursion limit) + rather than starting again. A retried attempt finds it stored, in the + turn. It is stored, and is in the graph, with the id ``steer:``. - The model sees the earlier turns as their messages and answers, and the latest turns whole (the agent's ``window``), as many as its ``budget`` holds; the built-in agent has its model summarize the earliest once they outgrow it, after a turn's answer (``_conversation.py``). - The model sees the latest messages of a conversation that lives long: as many as fit in ``WINDOW`` of its context window, by LangChain's - ``trim_messages`` (starting with a prompt, counted by - ``count_tokens_approximately``). The conversation is stored whole. That - takes a chat model whose profile has its context window + ``trim_messages`` (starting with a prompt, never with a message steered to + a turn, counted by ``count_tokens_approximately``). The conversation is + stored whole. That takes a chat model whose profile has its context window (``max_input_tokens``), as LangChain's integrations have for the models they know, else the model sees all of it. For a summary of the earlier messages instead, pass a graph of your own, such as ``create_agent`` with @@ -49,7 +56,8 @@ For a graph of your own, ``SqlCheckpointSaver`` stores LangGraph's checkpoints in the object's database instead. Its state is the graph's own: its models read the state, which the agent's ``window``, ``budget`` and summaries do not -reach, and which grows unless the graph trims it (``_graph.py`` shows how). +reach, and which grows unless the graph trims it (``_graph.py`` shows how). It +takes no steered messages: each is answered as a turn of its own. """ from __future__ import annotations @@ -96,6 +104,7 @@ from restate.agents._conversation import Call, Conversation, Item, Parsed, Result, Run, answer_once from restate.agents._object import AgentObject, AgentOptions from restate.agents._output import answer_with, asking +from restate.agents._steering import take_steers from restate.agents._tools import Tool, ToolCall from restate.agents._turn import settings_of, tool_call_ended, tool_call_started from restate.agents.langgraph._context import handler_context @@ -231,6 +240,15 @@ def _text(message: BaseMessage) -> str: return str(message.text) +STEER = "steer:" +"""The prefix of the id that a message steered to a turn has in the graph: ``steer:``.""" + + +def _steered(message: BaseMessage) -> bool: + """Return whether ``message`` is a message steered to a turn (``steer``): a message of the user, but no prompt.""" + return isinstance(message, HumanMessage) and message.id is not None and message.id.startswith(STEER) + + class _Items: """LangChain's messages, as the conversation keeps them: stored as LangChain's dicts, in JSON.""" @@ -244,6 +262,10 @@ def load(self, stored: str) -> Item: def prompt(self, text: str) -> Item: return HumanMessage(text) + def steered(self, text: str, steer_id: str) -> Item: + # Stored with its id, so that every run's state tells it apart from a prompt (``_steered``). + return HumanMessage(text, id=f"{STEER}{steer_id}") + def reply(self, text: str, agent: str | None = None) -> Item: return AIMessage(text) @@ -289,7 +311,10 @@ def _window(model: BaseChatModel) -> int | None: def _latest(messages: Sequence[BaseMessage], window: int | None) -> list[BaseMessage]: """Return the latest of ``messages`` that fit in ``window`` tokens, starting with a prompt (LangChain's ``trim_messages``). - The latest prompt and what followed it are sent whole, even if they do not fit. + The latest prompt and what followed it are sent whole, even if they do not + fit. A message steered to a turn is no prompt: the messages never start at + one, which would leave out its turn's prompt, and the tool calls and + results before it. """ if window is None: return list(messages) @@ -302,10 +327,11 @@ def _latest(messages: Sequence[BaseMessage], window: int | None) -> list[BaseMes token_counter=count_tokens_approximately, start_on="human", ) - if fitting: - return fitting - last_prompt = max((i for i, message in enumerate(messages) if isinstance(message, HumanMessage)), default=0) - return list(messages[last_prompt:]) + # The trim keeps the latest messages: from the first prompt among them, else from the latest prompt, whole. + first = len(messages) - len(fitting) + prompts = [i for i, message in enumerate(messages) if isinstance(message, HumanMessage) and not _steered(message)] + start = next((i for i in prompts if i >= first), prompts[-1] if prompts else 0) + return list(messages[start:]) def _graph( @@ -321,10 +347,13 @@ def _graph( window = _window(model) async def call_model(state: State) -> dict[str, Any]: - messages = _latest(state["messages"], window) + # The turn's step: once the tool calls have their results, the turn takes the messages steered to it. The + # request has them after the results, and the state keeps them for the rest of the run, which goes on. + steered = typing.cast(list[BaseMessage], await take_steers(ctx, conversation)) + messages = _latest([*state["messages"], *steered], window) reply = await bound.ainvoke([SystemMessage(system), *messages] if system else messages) await conversation.turn.append_reply(reply) - return {"messages": [reply]} + return {"messages": [*steered, reply]} async def run_tools(state: State, config: RunnableConfig) -> Any: await tool_call_started(ctx) @@ -359,10 +388,9 @@ async def run(conversation: Conversation) -> str: if (final := conversation.turn.final_reply()) is not None: # The run had stored its final reply, but not its answer: that is the answer, rather than ask again. return final - messages = typing.cast(list[BaseMessage], conversation.history()) graph = _graph(ctx, conversation, chat_model, tools, system) try: - state = typing.cast(State, await graph.ainvoke({"messages": messages})) + state = typing.cast(State, await graph.ainvoke({"messages": conversation.history()})) except Exception as error: if (terminal := terminal_of(error)) is not None: # LangGraph's own error, which a retry would meet again: its recursion limit, say. @@ -532,8 +560,10 @@ def agent( """Create a LangGraph agent: a virtual object with one conversation per key. Its ``chat`` handler answers each message with LangGraph's tool-calling - loop; the conversation is kept in the key's database, a step at a time. - Add handlers of your own with ``@assistant.handler``:: + loop; the conversation is kept in the key's database, a step at a time. A + message sent to its ``steer`` view while it answers joins that turn at the + turn's next step: the model gets it in its next request, after its tool + calls' results. Add handlers of your own with ``@assistant.handler``:: assistant = langgraph.agent("assistant", system="Answer briefly.", tools=[get_weather]) @@ -544,7 +574,8 @@ def agent( Or bring a graph of your own, built on LangGraph's ``MessagesState``: ``chat`` then runs it, with the message as a new ``HumanMessage``, and - answers with its last AI message. Its checkpoints are in the key's + answers with its last AI message. It takes no steered messages: each is + answered as a turn of its own. Its checkpoints are in the key's database, and the prebuilt ``nodes`` (a signal, a workflow, a call, another agent) and the tools of ``tools()`` get the handler's context in it:: @@ -610,7 +641,8 @@ async def answer(ctx: SessionContext, text: str) -> Any: output=output, ) - obj.chat(answer, items=ITEMS, output=output, summarize=_summarizer(model)) + # The model node takes the messages steered to the turn (``take_steers``); a graph of your own takes none. + obj.chat(answer, items=ITEMS, output=output, summarize=_summarizer(model), steers=True) return obj diff --git a/python/restate/agents/langgraph/_graph.py b/python/restate/agents/langgraph/_graph.py index af11aa46..461820e8 100644 --- a/python/restate/agents/langgraph/_graph.py +++ b/python/restate/agents/langgraph/_graph.py @@ -73,6 +73,10 @@ (a model provider refuses such a history), and the note that ends it (such as "Cancelled: ..."), so that the graph's model knows, as the built-in agent's does. So does that run's message, if the graph never stored it. +- A run takes no steered messages: the agent does not say it does + (``chat(..., steers=True)``), as the graph's steps are its own. So a + message sent to the agent's ``steer`` view while a run goes on is answered + as a turn of its own, after it, in the order it was sent. The graph's state is its own. Its models read the state, not the conversation: the agent's ``window`` and ``budget`` do not reach them, and diff --git a/python/restate/agents/openai/__init__.py b/python/restate/agents/openai/__init__.py index 543665ae..9aba9173 100644 --- a/python/restate/agents/openai/__init__.py +++ b/python/restate/agents/openai/__init__.py @@ -36,7 +36,10 @@ async def get_weather(city: str) -> str: desk = openai.agent("desk", agent=triage) ``tools`` gives an ``Agent`` of your own this package's tools. The -conversation is stored one step at a time. +conversation is stored one step at a time. A message sent to an agent's +``steer`` view while it answers joins the turn at its next step, once the +model's tool calls have their outputs: the next model request has it, and +the run goes on. The conversation lives in the key's database, and only its turns write it, so the SDK's ``Session`` stores are not used. Don't use @@ -60,8 +63,8 @@ async def get_weather(city: str) -> str: journaling its model calls with ``ctx.run_typed``. Pick ``restate.ext.openai`` to add durability to an agent you run in your own services; pick this module for a ready-made agent per key, whose conversation lives in a database and which -comes with built-in handlers to chat, read the conversation, watch its -progress and stop a turn. +comes with built-in handlers to chat, steer a turn, read the conversation, +watch its progress and stop a turn. Install the ``openai`` extra (``pip install restate-sdk[openai]``). Requests use OPENAI_API_KEY, unless you pass a model of the SDK of your own. diff --git a/python/restate/agents/openai/_agent.py b/python/restate/agents/openai/_agent.py index 8b7239cd..731b4fe6 100644 --- a/python/restate/agents/openai/_agent.py +++ b/python/restate/agents/openai/_agent.py @@ -28,6 +28,23 @@ resumed call would bypass). A function tool's timeout applies to a resumed call too. A call to a tool the agent does not have (which the SDK fails the run for) is reported back to the model to fix, on the retry. +- A message steered to the turn (the agent's ``steer`` view) joins it at its + next step: once each of the model's tool calls has its output, the input + filter stores the outputs, then the messages waiting (``take_steers``), + and the request about to go out has them after those outputs. The run + goes on in place rather than starting over: its input guardrails run + once, ``max_turns`` counts on, and its tool choice and hooks carry on. As + the SDK makes each request's input afresh, from the run's input and its + own items, each later request of the run gets each message again, right + after the item it came after (if a handoff's ``input_filter`` took that + item out, before the first item after it, as if the filter had kept the + message). A message waiting as a run starts (that an earlier attempt + learned of, say) is stored first, and is in its input; a retried attempt + finds the messages stored, in the history, each once. A message that + comes while the model writes its final answer is answered next, as a turn + of its own. A handoff whose ``input_filter`` takes its own output out of + the next agent's input (as ``remove_all_tools`` does) leaves its call + without an output in the turn: the run takes no message after it. - A run that stops for tool calls that need an approval (``needs_approval``) fails its turn, rather than answer without them, as nothing gives the approval: pause for a person with ``restate.agents``' ``wait_for_signal`` @@ -63,7 +80,9 @@ (the SDK's run context) and the ``run_config`` you give: its ``call_model_input_filter`` runs after this package's, which stores the tools' outputs: the SDK's ``ToolOutputTrimmer`` there trims old tool - outputs in the model's input, not in the conversation. + outputs in the model's input, not in the conversation. An agent used as a + tool (``Agent.as_tool``) runs with the same ``run_config``: its requests + get your filter only, and the turn keeps the call's output, not its run. - The model sees the earlier turns as their messages and answers, and the latest turns whole (the agent's ``window``), as many as its ``budget`` holds; the built-in agent has its model summarize the earliest once they @@ -115,6 +134,7 @@ UserError, ) from agents import Tool as SdkTool +from agents.models.fake_id import FAKE_RESPONSES_ID from agents.run import CallModelData, ModelInputData from agents.tool import invoke_function_tool from agents.tool_context import ToolContext @@ -136,6 +156,7 @@ from restate.agents._conversation import SUMMARY, Call, Conversation, Item, Parsed, Result, Run, answer_once from restate.agents._object import AgentObject, AgentOptions from restate.agents._output import answer_with +from restate.agents._steering import take_steers from restate.agents._tools import Tool, current_call, tools_in from restate.agents._turn import tool_call_ended, tool_call_started from restate.exceptions import TerminalError @@ -285,29 +306,57 @@ def __init__(self, ctx: SessionContext, conversation: Conversation, config: RunC self.requesting = False """Whether a model request is in progress, with the making of its model and client: the only time an error can be its provider refusing.""" + self.steers = _Steers() + """The messages steered to the turn that the run took: each later request gets them again.""" + self.running = 0 + """How many of the run's tool calls are running: a request meanwhile is a nested run's.""" async def model_input(self, data: CallModelData[Any]) -> ModelInputData: - """The run's ``call_model_input_filter``: store the tools' outputs before the next request. + """The run's ``call_model_input_filter``: store the tools' outputs before the next request, and take the + messages steered to the turn, which the request gets after them. + + The SDK makes each request's input afresh, from the run's input and the items the run made, which hold no + steered message: so each message the run took goes back in each later request, in its place. The run's own filter then changes what the model sees, not what is stored. + + An agent used as a tool runs nested in the run's tool call, with the run's ``run_config`` (the SDK passes it + on): its requests come here too, while the call runs. They get the run's own filter only: the turn stores + none of their items (the call's output is what the turn keeps), and they get none of its messages. """ - if outputs := self.conversation.new_results(data.model_data.input): + if self.running: + return await self._filtered(data) + items = data.model_data.input + if outputs := self.conversation.new_results(items): await self.conversation.turn.append(*outputs) - filtered = data.model_data - if self.then is not None: - given = self.then(data) - filtered = await given if inspect.isawaitable(given) else given + # A step, once each of the model's tool calls has its output (else ``take_steers`` takes none): the messages + # go after the input's last item. A last item with nothing to tell it apart in the later requests (a message + # of the user, as a run's first request ends with) takes none: the run took them as it started (``_go``), + # and takes the next at its next step. + after = _key(items[-1]) if items else None + if after is not None and (taken := await take_steers(self.ctx, self.conversation)): + self.steers.took(taken, items, after) + data.model_data.input = self.steers.placed(items) + filtered = await self._filtered(data) # The SDK goes on to make the model and its client (which needs the API key: a refusal if it is missing), and # to request it. (The instructions have run already.) self.requesting = True return filtered + async def _filtered(self, data: CallModelData[Any]) -> ModelInputData: + """Return what the run's own filter makes of ``data``, if it has one.""" + if self.then is None: + return data.model_data + given = self.then(data) + return await given if inspect.isawaitable(given) else given + async def on_llm_end(self, context: RunContextWrapper[Any], agent: Agent[Any], response: ModelResponse) -> None: self.requesting = False reply = typing.cast(list[Item], response.to_input_items()) await self.conversation.turn.append_reply(*reply, agent=agent.name) async def on_tool_start(self, context: RunContextWrapper[Any], agent: Agent[Any], tool: SdkTool) -> None: + self.running += 1 await tool_call_started(self.ctx) async def on_handoff(self, context: RunContextWrapper[Any], from_agent: Agent[Any], to_agent: Agent[Any]) -> None: @@ -315,9 +364,98 @@ async def on_handoff(self, context: RunContextWrapper[Any], from_agent: Agent[An self.conversation.turn.carry(agent=to_agent.name) async def on_tool_end(self, context: RunContextWrapper[Any], agent: Agent[Any], tool: SdkTool, result: Any) -> None: + self.running -= 1 await tool_call_ended(self.ctx) +class _Steers: + """The messages steered to the turn that a run took, which each of the run's later requests gets again, in place. + + The SDK makes each request's input afresh, from the run's input and the + items the run made, which hold no steered message: so ``placed`` puts + each one back, right after the item it came after (``_spliced``). + """ + + def __init__(self) -> None: + self.taken: list[_Steered] = [] + """The messages the run took, in order.""" + self.seen: dict[str, int] = {} + """The keys (``_key``) of the items of the run's requests, once it took a message, in the order it saw them.""" + + def took(self, taken: Sequence[Item], items: Sequence[Item], after: str) -> None: + """Note the messages ``taken`` at a step, after ``items`` (the request's input), whose last item's key is + ``after``.""" + # This request's items came before the messages; an item whose key the run sees later, after them. + self._see(items) + self.taken.extend(_Steered(item, after, len(self.seen)) for item in taken) + + def placed(self, items: list[Item]) -> list[Item]: + """Return ``items`` (a request's input) with each message the run took in its place.""" + if not self.taken: + return items + self._see(items) + return _spliced(items, self.taken, self.seen) + + def _see(self, items: Sequence[Item]) -> None: + """Note the keys of ``items`` that the run had not seen: they come after the messages it took so far.""" + for item in items: + if (key := _key(item)) is not None: + self.seen.setdefault(key, len(self.seen)) + + +@dataclasses.dataclass(frozen=True) +class _Steered: + """A message steered to the turn that the run took, and where it goes in the run's requests.""" + + item: Item + after: str + """The key (``_key``) of the item of the input it came after.""" + seen: int + """How many keys the run had seen as it took it: an item whose key it saw later came after it.""" + + +def _key(item: Item) -> str | None: + """Return what tells ``item`` apart in a run's input, from one request to the next: the id of its tool call, if + it is a call or an output, else its own; ``None`` for an item with neither, such as a message of the user.""" + if not isinstance(item, dict): + return None + fields = typing.cast(dict[str, Any], item) + kind = str(fields.get("type", "message")) + for name in ("call_id", "id"): + value = fields.get(name) + if isinstance(value, str) and value != FAKE_RESPONSES_ID: + return f"{kind}:{name}:{value}" + return None + + +def _spliced(items: list[Any], steered: Sequence[_Steered], seen: dict[str, int]) -> list[Any]: + """Return ``items`` (a request's input) with each steered message in its place. + + That is right after the item it came after; if ``items`` lacks that item (a handoff's ``input_filter`` took it + out, say), right before the first item that came after the message (whose key the run saw later, in ``seen``), + as if the filter had kept the message; else at the end. Messages in the same place keep their order. An item + found more than once counts where it last is, as the SDK keeps an output where its latest copy is. + """ + keys = [_key(item) for item in items] + last = {key: i for i, key in enumerate(keys) if key is not None} + + def place(steer: _Steered) -> int: + if steer.after in last: + return last[steer.after] + 1 + later = (i for i, key in enumerate(keys) if key is not None and seen.get(key, steer.seen) >= steer.seen) + return next(later, len(items)) + + before: dict[int, list[Item]] = {} + for steer in steered: + before.setdefault(place(steer), []).append(steer.item) + spliced: list[Any] = [] + for i, item in enumerate(items): + spliced.extend(before.get(i, ())) + spliced.append(item) + spliced.extend(before.get(len(items), ())) + return spliced + + async def _resume( ctx: SessionContext, conversation: Conversation, agent: Agent[Any], everyone: dict[str, Agent[Any]], context: Any ) -> Agent[Any]: @@ -521,6 +659,9 @@ async def _go( limit: dict[str, Any] = {"max_turns": options.max_turns} if options.max_turns is not None else {} try: agent = await _unwrapped(ctx, _resume(ctx, conversation, current, everyone, options.context)) + # A step: the messages steered to the turn by now (that an earlier attempt learned of, say) are the run's + # input, after what the turn stored. + await take_steers(ctx, conversation) items = typing.cast(Any, _since_compaction(conversation.history())) run = Runner.run(agent, items, context=options.context, hooks=hooks, run_config=filtered, **limit) result = await _unwrapped(ctx, run) @@ -791,8 +932,10 @@ def agent( """Create an OpenAI Agents SDK agent: a virtual object with one conversation per key. Its ``chat`` handler answers each message with the SDK's ``Runner``; the - conversation is kept in the key's database, a step at a time. Add - handlers of your own with ``@assistant.handler``. + conversation is kept in the key's database, a step at a time. A message + sent to its ``steer`` view while it answers joins the turn at its next + step: the run goes on with it, rather than starting over. Add handlers + of your own with ``@assistant.handler``. The built-in agent takes the common options:: @@ -863,7 +1006,9 @@ async def answer(ctx: SessionContext, text: str) -> Any: run, ) - obj.chat(answer, items=ITEMS, output=output, summarize=_summarizer(model) if own is None else None) + # Steered messages: the run's input filter takes them at each step (``_Hooks.model_input``). + summarize = _summarizer(model) if own is None else None + obj.chat(answer, items=ITEMS, output=output, summarize=summarize, steers=True) return obj diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index c63808ba..0f871fb0 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -78,6 +78,16 @@ tool call's writes are done once, however often it runs; ``clear`` empties them. - Usage limits (``usage_limits``) are each run's, so each turn's. +- A message steered to a turn (the agent's ``steer``) joins it at its next + step: once the model's tool calls all have their results stored, before + the next model request (``take_steers``). It is stored as a message of the + user that is a step of the turn, then given to the model with Pydantic + AI's ``RunContext.enqueue``, into that request and the run's history, + after what they hold. The run carries on with it (its usage limits, + retries and capabilities' state too), and a retried attempt, or the next + turn, has the message where the model saw it. A message that comes after + the turn's last step (while the model gives its final answer, say) is + answered next, as a turn of its own. - Each message is answered once per invocation, by the shared conversation of ``restate.agents`` (``_conversation.py``), which also closes a stopped, refused or failed run. @@ -160,6 +170,7 @@ from restate.agents._conversation import Call, Conversation, Item, Parsed, Result, Run, answer_once from restate.agents._object import AgentObject, AgentOptions from restate.agents._output import answer_with +from restate.agents._steering import take_steers from restate.agents._tools import Tool, current_call, tools_in, wait_for from restate.agents._turn import stop_signal_of, tool_call_ended, tool_call_started from restate.agents.pydantic_ai._stores import TurnsHistorySource, calling, conversation_id_of @@ -323,10 +334,16 @@ def read(self, item: Item) -> Parsed: ITEMS = _Items() -def _feedback(message: ModelMessage) -> bool: - """Return whether ``message`` is Pydantic AI's feedback on a reply it could not use (an answer that does not - validate, say), which answers no tool call.""" - return isinstance(message, ModelRequest) and any(isinstance(p, RetryFeedbackPart) for p in message.parts) +def _feedback_of(message: ModelMessage) -> list[str]: + """Return Pydantic AI's feedback on a reply it could not use (an answer that does not validate, say) that + ``message`` holds, which answers no tool call: each ``RetryFeedbackPart``, as JSON. + + A part is its own feedback (it has its own time): a run keeps it as it is when it merges consecutive requests + (a feedback and a message steered after it, say) or resumes from the last one, though the request is new. + """ + if not isinstance(message, ModelRequest): + return [] + return [to_json(p, bytes_mode="base64").decode() for p in message.parts if isinstance(p, RetryFeedbackPart)] def _result_ids(message: ModelMessage) -> set[str]: @@ -351,13 +368,14 @@ def __init__(self, ctx: SessionContext, conversation: Conversation, history: Seq self.conversation = conversation self.model_errors = ErrorSet() """The errors that the run's model requests failed with.""" - self.feedback = {_MESSAGE.dump_json(m) for m in history if _feedback(m)} - """Pydantic AI's feedback requests that the conversation has (``RetryFeedbackPart``), as JSON.""" + self.feedback = {feedback for m in history for feedback in _feedback_of(m)} + """Pydantic AI's feedback that the conversation has (each ``RetryFeedbackPart``), as JSON.""" def _new(self, message: ModelMessage) -> bool: if any(not self.conversation.knows_result(i) for i in _result_ids(message)): return True - return _feedback(message) and _MESSAGE.dump_json(message) not in self.feedback + # Found by its parts, not the whole request: the run's request may be a stored one, merged or resumed. + return any(feedback not in self.feedback for feedback in _feedback_of(message)) async def store(self, messages: Sequence[ModelMessage]) -> None: """Store the requests in ``messages`` (the run's) that the conversation does not have yet. @@ -365,14 +383,17 @@ async def store(self, messages: Sequence[ModelMessage]) -> None: Those are the tools' results, and Pydantic AI's feedback on a reply it could not use (``RetryFeedbackPart``): the rest of what a run adds, its responses, is stored as each arrives. Results are found by their call - ids, not by their place: a history processor (``ProcessHistory``) may - have trimmed the run's messages, and what it makes of them is the - model's to see, not the conversation's. + ids, and feedback by its parts, not by their place or their request: a + history processor (``ProcessHistory``) may have trimmed the run's + messages, and a run merges the consecutive requests it is given (a + message steered after a result, say) and resumes from the last one as + a request of its own. What it makes of them is the model's to see, not + the conversation's. """ new = [m for m in messages if self._new(m)] if new: await self.conversation.turn.append(*new) - self.feedback.update(_MESSAGE.dump_json(m) for m in new if _feedback(m)) + self.feedback.update(feedback for m in new for feedback in _feedback_of(m)) async def before_model_request( self, ctx: RunContext[Any], request_context: ModelRequestContext @@ -654,11 +675,12 @@ def _snapshot_of(messages: Sequence[ModelMessage]) -> list[ModelMessage]: class _Snapshots(AbstractCapability[Any]): """Store the history as a snapshot of the turn when a compaction rewrote it, before the request goes out. - It runs after every other capability's ``before_model_request`` (it is - innermost), so it sees the messages as they are sent. The history was - rewritten, rather than extended, if the messages sent before (or, for a - run's first request, the history the run was given) do not start them: - the next turn, or a retried attempt, then starts from the snapshot. + It runs after every other capability's ``before_model_request`` but + ``_Steering``'s (it is innermost), so it sees the messages as they are + sent. The history was rewritten, rather than extended, if the messages + sent before (or, for a run's first request, the history the run was + given) do not start them: the next turn, or a retried attempt, then + starts from the snapshot. """ def __init__(self, conversation: Conversation, history: Sequence[ModelMessage]) -> None: @@ -680,6 +702,43 @@ async def before_model_request( return request_context +class _Steering(AbstractCapability[Any]): + """Give the model the messages steered to the turn at its next step, in the request about to go out. + + Before each model request, it takes the messages waiting + (``take_steers``), which stores each once as a step of the turn, and + hands them to Pydantic AI's own ``RunContext.enqueue`` (``'asap'``): once + every ``before_model_request`` ran, Pydantic AI adds them to this request + and to the run's history, after what both hold. The run goes on with them + (its usage limits, retries and capabilities' state too); a retried attempt + and the next turns find them in their history, where the model saw them, + and never take them again. + + It is the innermost capability, after ``_Snapshots``: a snapshot that this + request stores covers what was stored before the messages, so every + history after it has them after it. + """ + + def __init__(self, ctx: SessionContext, conversation: Conversation) -> None: + self.ctx = ctx + self.conversation = conversation + + def get_ordering(self) -> CapabilityOrdering: + return CapabilityOrdering(position="innermost", wrapped_by=[_Snapshots]) + + async def before_model_request( + self, ctx: RunContext[Any], request_context: ModelRequestContext + ) -> ModelRequestContext: + messages = request_context.messages + if not messages or isinstance(messages[-1], ModelResponse): + # A response that its provider paused, resumed: Pydantic AI delivers nothing into its request, so a + # message taken now would reach the model after the response, though stored before it. + return request_context + if steered := typing.cast(list[ModelRequest], await take_steers(self.ctx, self.conversation)): + ctx.enqueue(*steered, priority="asap") + return request_context + + def _due(messages: list[ModelMessage], model: Model, context_window: int | None) -> bool: """Return whether ``messages`` are near the default compaction's target for ``model`` (past ``BETWEEN``).""" window = context_window or resolve_context_window(model) or DEFAULT_CONTEXT_WINDOW @@ -733,6 +792,9 @@ async def run(conversation: Conversation) -> str: # The snapshots of an agent that takes its history from them (the agent's ``chat``); a prompt of a session # object that is no Pydantic AI agent takes the turn projection, which never reads one. snapshots = [_Snapshots(conversation, history)] if conversation.snapshots else [] + # The messages steered to a turn of a ``chat`` that takes them, with snapshots or without; any other run takes + # none (``take_steers``). + steering = _Steering(ctx, conversation) spec = _output_spec(output, model) if output is not None else None try: # The tools of ``toolset`` find their call and the handler's context in the run. @@ -741,7 +803,7 @@ async def run(conversation: Conversation) -> str: message_history=history, model=model, toolsets=toolsets, - capabilities=[restate, *chosen, *snapshots], + capabilities=[restate, *chosen, *snapshots, steering], output_type=spec, deps=deps, usage_limits=usage_limits, @@ -971,8 +1033,10 @@ def agent( """Create a Pydantic AI agent: a virtual object with one conversation per key. Its ``chat`` handler answers each message with Pydantic AI; the - conversation is kept in the key's database, a step at a time. Add - handlers of your own with ``@assistant.handler``. + conversation is kept in the key's database, a step at a time. A message + sent to its ``steer`` view while a turn runs joins the turn at its next + step, once the model's tool calls have their results. Add handlers of + your own with ``@assistant.handler``. The built-in agent takes the common options, and gets the harness's compaction (``compaction``) for a conversation that lives long:: @@ -1050,7 +1114,8 @@ async def answer(ctx: SessionContext, text: str) -> Any: ) compact_history = None if any(map(_compacts, _leaves(own))) else _compactor(own, model, context_window) - obj.chat(answer, items=ITEMS, output=output, snapshots=True, compact_history=compact_history) + # Its runs take the messages steered to a turn at the turn's steps (``_Steering``). + obj.chat(answer, items=ITEMS, output=output, snapshots=True, compact_history=compact_history, steers=True) return obj diff --git a/scripts/resurrect.py b/scripts/resurrect.py index 37c7e86a..7851dd26 100644 --- a/scripts/resurrect.py +++ b/scripts/resurrect.py @@ -67,6 +67,7 @@ "tests/agents_langgraph_graph.py", ) OBJECT = ("tests/agents_object.py", "tests/agents_signals.py") +STEERING = ("tests/agents_steering.py", "tests/agents_steering_server.py", "tests/agents_signals.py") SNAPSHOTS = ("tests/agents_snapshots.py", "tests/agents_pydantic_ai_snapshots.py") WAITS = (*EVERY_FRAMEWORK, "tests/agents_snapshots.py") @@ -564,8 +565,8 @@ def kind(self) -> str: " if (final := conversation.turn.final_reply()) is not None:\n" " # The run had stored its final reply, but not its answer: that is the answer, rather than ask again.\n" " return final\n" - " messages = ", - " messages = ", + " graph = ", + " graph = ", LANGGRAPH, ), Mutant( @@ -671,6 +672,111 @@ def kind(self) -> str: "if False:", EVERY_FRAMEWORK, ), + # --- Durability: steering, from no fix --- + Mutant( + "steer-placeholder-runs-taken", + "a steered message the turn took is answered again, as a turn of its own", + f"{AGENTS}/_object.py", + ' if await steer_stored(ctx.db, message["id"]):\n return\n', + " if False:\n return\n", + STEERING, + ), + Mutant( + "steer-latest-ask-counts-steers", + "a steered message counts as the turn's ask: the output check never asks again, ADK resumes another invocation", + f"{AGENTS}/_conversation.py", + '(e for e in reversed(self.items) if e.steer is None and e.said is not None and e.said[0] == "user"),', + '(e for e in reversed(self.items) if e.said is not None and e.said[0] == "user"),', + STEERING, + ), + Mutant( + "steer-history-drops-steers", + "the model forgets the messages steered to a turn once it leaves the window", + f"{AGENTS}/_conversation.py", + " out.extend(steered.item for steered in turn.steers)\n", + "", + STEERING, + ), + Mutant( + "steer-taken-with-open-calls", + "a steered message is taken while a tool call has no result, between a call and its result", + f"{AGENTS}/_steering.py", + " if not going() or turn is None or turn.open_calls():\n", + " if not going() or turn is None:\n", + STEERING, + ), + Mutant( + "steer-publish-without-steering", + "the turn never learns of a steered message: each is answered as a turn of its own", + f"{AGENTS}/_signals.py", + ' if steering:\n published["steering"] = steering\n', + ' if False:\n published["steering"] = steering\n', + STEERING, + ), + Mutant( + "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', + " pass\n", + STEERING, + ), + Mutant( + "steer-drop-keeps-steers", + "a dropped reply leaves the message steered after it marked as taken, so it is never answered", + f"{AGENTS}/_conversation.py", + ' await tx.execute(f"DELETE FROM {STEERED} WHERE n >= ?", [reply])\n', + "", + STEERING, + ), + Mutant( + "steer-cap-counts-waiting", + "a turn takes steered messages without end: the cap counts only those it has not taken", + f"{AGENTS}/_signals.py", + '{"id": one["id"], "since": one["since"], "turn": one["turn"]} if one["id"] in ids else one for one in inbox', + 'one for one in inbox if one["id"] not in ids', + STEERING, + ), + Mutant( + "steer-taken-by-a-stopped-turn", + "a turn that stopped while it waited to store a steered message takes it, so it is never answered", + f"{AGENTS}/_conversation.py", + " if not going():\n return []\n", + "", + STEERING, + ), + Mutant( + "steer-size-misses-the-turns-own", + "the turn just answered leaves its steered messages out of the budget: its summary comes a turn late", + f"{AGENTS}/_conversation.py", + " self.prompt.steers.extend(added)\n", + "", + STEERING, + ), + Mutant( + "steer-langgraph-unmarked", + "LangGraph's steered message is stored as a prompt: a trimmed request may start at it, without its turn", + f"{AGENTS}/langgraph/_agent.py", + ' return HumanMessage(text, id=f"{STEER}{steer_id}")\n', + " return HumanMessage(text)\n", + ("tests/agents_langgraph.py",), + ), + Mutant( + "openai-nested-run-filtered", + "an agent used as a tool gets the turn's steered messages, and the turn stores its tool outputs, with no call", + f"{AGENTS}/openai/_agent.py", + " if self.running:\n return await self._filtered(data)\n", + "", + OPENAI, + ), + Mutant( + "steer-summary-hides-steers", + "the summary view leaves out the messages steered to a turn", + f"{AGENTS}/_messages.py", + " messages.extend(user_message(n, seq, stored) for n, seq, stored in steered)\n", + "", + STEERING, + ), ] diff --git a/scripts/steering_traffic.py b/scripts/steering_traffic.py new file mode 100644 index 00000000..5a51dea3 --- /dev/null +++ b/scripts/steering_traffic.py @@ -0,0 +1,215 @@ +#!/usr/bin/env python3 +# +# 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 +# +"""Measure what an agent's turn costs in Restate traffic, steered or not, against a real Restate server. + +Serves an agent whose framework is the tests' made-up one: each message is +answered by one tool call (which a test gate holds while a message is +steered to the turn), a step, and a reply. For each turn it reads, from +Restate's admin SQL API, the ``chat`` invocation's journal (its commands and +notifications) and the calls it made to the agent's companion, by handler; +and for a steered message, the invocations it costs on its own: the +``steer`` view, the companion's ``steer`` and the agent's ``steered``. + +It measures any checkout of this repository, so that a turn nobody steers +can be compared with one of a version before steering:: + + python scripts/steering_traffic.py --tree . --turns 5 + python scripts/steering_traffic.py --tree ../before --turns 5 # no steering: its plain turns only + +It runs Restate's test harness (Docker, unless the harness is made to run a +local server). +""" + +from __future__ import annotations + +import argparse +import asyncio +import collections +import json +import sys +import typing +import uuid +from pathlib import Path +from typing import Any + + +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() + return typing.cast(list[dict[str, Any]], response.json()["rows"]) + + +async def cost(client: Any, admin: str, invocation: str) -> dict[str, Any]: + """Return the journal of ``invocation`` (commands, notifications) and the companion calls it made, by handler.""" + journal = await query(client, admin, f"SELECT entry_type FROM sys_journal WHERE id = '{invocation}'") + kinds = collections.Counter(row["entry_type"].split(":")[0] for row in journal) + called = await query( + client, + admin, + 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" + ) + return { + "journal": len(journal), + "commands": kinds.get("Command", 0), + "notifications": kinds.get("Notification", 0), + "companion": dict(sorted(companion.items())), + "companion_calls": sum(companion.values()), + } + + +async def measure(args: argparse.Namespace) -> None: + tree = Path(args.tree).resolve() + sys.path[:0] = [str(tree / "python"), str(tree / "tests")] + # pylint: disable=import-outside-toplevel,import-error + import httpx2 + from agents_conversation import ITEMS + from support import RESTATE_IMAGE, Gate + + import restate + from restate import sessions + from restate.agents._conversation import Conversation, answer_once + from restate.agents._object import AgentObject + from restate.agents._turn import stop_signal_of, tool_call_ended, tool_call_started + + try: + from restate.agents._steering import take_steers + except ImportError: + take_steers = None + steering = take_steers is not None + gate = Gate() + + async def answer(ctx: sessions.SessionContext, text: str) -> str: + async def work(conversation: Conversation) -> str: + turn = conversation.turn + if not turn.items: + call = {"id": "call-1", "name": "lookup", "input": {}} + await turn.append_reply({"role": "assistant", "text": "", "calls": [call]}) + await tool_call_started(ctx) + stop = stop_signal_of(ctx) + assert stop is not None + await stop.guard(gate.wait()) + await turn.append({"result": "call-1", "output": "found"}) + await tool_call_ended(ctx) + if take_steers is not None: + await take_steers(ctx, conversation) + reply = f"answered: {text}" + await turn.append_reply({"role": "assistant", "text": reply}) + return reply + + return await answer_once(ctx, text, ITEMS, work) + + agent = AgentObject("measured") + if steering: + agent.chat(answer, items=ITEMS, steers=True) + else: + agent.chat(answer, items=ITEMS) + + text_headers = {"content-type": "text/plain"} + rows: list[dict[str, Any]] = [] + async with restate.create_test_harness(restate.app([agent]), restate_image=RESTATE_IMAGE) as harness: + ingress, admin = harness.ingress_url, harness.admin_api_url + async with httpx2.AsyncClient(timeout=60) as client: + + async def answering(key: str, invocation: str) -> None: + for _ in range(400): + status = (await client.post(f"{ingress}/measured/{key}/status")).json() + if status["answering"] is not None and status["answering"]["invocationId"] == invocation: + return + await asyncio.sleep(0.05) + raise AssertionError("the turn never started") + + async def done(key: str, handler: str) -> list[dict[str, Any]]: + rows = await query( + client, + admin, + "SELECT id, status FROM sys_invocation WHERE target_service_name = 'measured' " + f"AND target_service_key = '{key}' AND target_handler_name = '{handler}'", + ) + return rows + + for n in range(args.turns): + key = f"plain-{n}-{uuid.uuid4().hex[:6]}" + gate.open() + sent = await client.post( + f"{ingress}/measured/{key}/chat/send", content=f"look {n}", headers=text_headers + ) + invocation = sent.json()["invocationId"] + await client.get(f"{ingress}/restate/invocation/{invocation}/attach") + rows.append({"turn": "plain", "n": n, **await cost(client, admin, invocation)}) + + for n in range(args.turns if steering else 0): + key = f"steered-{n}-{uuid.uuid4().hex[:6]}" + gate.close() + sent = await client.post( + f"{ingress}/measured/{key}/chat/send", content=f"look {n}", headers=text_headers + ) + invocation = sent.json()["invocationId"] + await answering(key, invocation) + steered = await client.post( + f"{ingress}/measured/{key}/steer", + content=f"and note {n}", + headers={**text_headers, "idempotency-key": uuid.uuid4().hex}, + ) + assert steered.json()["turn"] == invocation, steered.text + gate.open() + answer_text = (await client.get(f"{ingress}/restate/invocation/{invocation}/attach")).json() + # The placeholder runs once the turn ended: wait for it, then cost the message's own invocations. + for _ in range(400): + placeholders = await done(key, "steered") + if placeholders and all(row["status"] == "completed" for row in placeholders): + break + await asyncio.sleep(0.05) + [view] = await done(key, "steer") + [placeholder] = placeholders + companion_steer = await query( + client, admin, f"SELECT id FROM sys_invocation WHERE invoked_by_id = '{view['id']}'" + ) + rows.append( + { + "turn": "steered", + "n": n, + "answer": answer_text, + **await cost(client, admin, invocation), + "steer_view": (await cost(client, admin, view["id"]))["journal"], + "companion_steer": [ + (await cost(client, admin, row["id"]))["journal"] for row in companion_steer + ], + "placeholder": (await cost(client, admin, placeholder["id"]))["journal"], + } + ) + + for row in rows: + print(json.dumps(row)) + for kind in ("plain", "steered"): + mine = [row for row in rows if row["turn"] == kind] + if not mine: + continue + journals = sorted({row["journal"] for row in mine}) + calls = sorted({json.dumps(row["companion"]) for row in mine}) + print( + f"# {kind:8} turns: {len(mine)}, chat journal entries {journals}, companion calls {calls}", + file=sys.stderr, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + parser.add_argument("--tree", default=".", help="the checkout to measure") + parser.add_argument("--turns", type=int, default=5, help="turns of each kind") + asyncio.run(measure(parser.parse_args())) + + +if __name__ == "__main__": + main() diff --git a/tests/agents_adk.py b/tests/agents_adk.py index 60cf7761..9aad2fd3 100644 --- a/tests/agents_adk.py +++ b/tests/agents_adk.py @@ -25,14 +25,14 @@ pytest.importorskip("google.adk") -from google.adk.agents import LlmAgent, RunConfig # noqa: E402 +from google.adk.agents import LlmAgent, RunConfig, SequentialAgent # noqa: E402 from google.adk.models import BaseLlm # noqa: E402 from google.adk.models.llm_request import LlmRequest # noqa: E402 from google.adk.models.llm_response import LlmResponse # noqa: E402 from google.genai import types # noqa: E402 import restate # noqa: E402 -from support import RESTATE_IMAGE # noqa: E402 # pylint: disable=import-error +from support import RESTATE_IMAGE, Gate # noqa: E402 # pylint: disable=import-error from restate import HarnessEnvironment, sessions # noqa: E402 from restate.agents import adk, tool # noqa: E402 from restate.agents._api import Messages # noqa: E402 @@ -81,7 +81,7 @@ def reply(system: str, contents: list[types.Content]) -> list[types.Part]: call = types.FunctionCall(id=f"call-{rest}", name="transfer_to_agent", args={"agent_name": "refunds"}) return [types.Part(function_call=call)] names = {"weather": "get_weather", "crash": "flaky", "wait": "blocking", "wrong": "get_weather"} - names |= {"forbidden": "forbidden", "unknown": "nonexistent"} + names |= {"forbidden": "forbidden", "unknown": "nonexistent", "hold": "hold"} if word in names and not responses: args: dict[str, Any] = {"town": 3} if word == "wrong" else {"city": rest} return [types.Part(function_call=types.FunctionCall(id=f"call-{rest}", name=names[word], args=args))] @@ -183,11 +183,59 @@ async def generate_content_async( summarizing = adk.agent("adk_summarizing", model=Remembering(), budget=10) """An agent whose model sees 40 characters of the turns before the window: it summarizes the earlier ones.""" +held = Gate() +"""The gate that ``hold``, and the answer to a ``slow`` prompt, wait on until the test opens it.""" +holds: collections.Counter[str] = collections.Counter() + + +@tool +async def hold(city: str) -> str: + """Waits until the test opens the gate.""" + holds[city] += 1 + await held.wait() + return f"held in {city}" + + +seen: collections.defaultdict[str, list[list[types.Content]]] = collections.defaultdict(list) +"""The requests of the steered agents' model, by the key word that ends each test's first prompt.""" +failed: collections.Counter[str] = collections.Counter() + + +class Steered(Scripted): + """The scripted model, keeping each request. Its answer to a prompt that starts with ``slow`` waits on + ``held``; a request whose last message starts with ``crashing`` crashes, once, as a process would.""" + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + contents = list(llm_request.contents) + prompts = _prompts(contents) + seen[prompts[0].split(" ")[-1]].append(contents) + if prompts[-1].startswith("slow"): + await held.wait() + if prompts[-1].startswith("crashing") and not failed[prompts[-1]]: + failed[prompts[-1]] += 1 + raise RuntimeError("the process crashed after the steered message was stored") + async for response in super().generate_content_async(llm_request, stream): + yield response + + +steered = adk.agent("adk_steered", model=Steered(), tools=[hold]) +"""An agent whose turns take the messages steered to them: its root is an LlmAgent (the built-in agent).""" + +pipeline = adk.agent( + "adk_pipeline", + agent=SequentialAgent( + name="pipeline", sub_agents=[LlmAgent(name="step", model=Steered(), tools=adk.tools([hold]))] + ), +) +"""An agent whose root is a workflow agent: its turns take no steered message.""" + @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([assistant, triage, summarizing]), restate_image=RESTATE_IMAGE + restate.app([assistant, triage, summarizing, steered, pipeline]), restate_image=RESTATE_IMAGE ) as harness: yield harness @@ -316,6 +364,166 @@ async def test_a_transfer_to_a_sub_agent_holds_for_the_next_messages(harness: Ha assert (await chat(harness, k, f"weather {k}", "adk_desk")).json() == f"tool said: sunny in {k}" +# --- Steering, against a Restate server ----------------------------------------------------------- + + +async def send(harness: HarnessEnvironment, service: str, key: str, text: str) -> str: + sent = await post(harness, f"{service}/{key}/chat/send", content=text, headers=TEXT) + return typing.cast(str, sent.json()["invocationId"]) + + +async def steer(harness: HarnessEnvironment, service: str, key: str, text: str) -> dict[str, Any]: + # With an idempotency key, as a client sends it: a retried request is the same message. + headers = {**TEXT, "idempotency-key": uuid.uuid4().hex} + steered = await post(harness, f"{service}/{key}/steer", content=text, headers=headers) + assert steered.status_code == 200, steered.text + return typing.cast(dict[str, Any], steered.json()) + + +async def attach(harness: HarnessEnvironment, invocation: str) -> httpx2.Response: + async with httpx2.AsyncClient(timeout=60) as ingress: + return await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation}/attach") + + +async def holding(key: str) -> None: + """Wait until ``hold`` runs for ``key``, its turn's tool call.""" + + async def started() -> bool: + return holds[key] > 0 + + await until(started) + + +async def settled(harness: HarnessEnvironment, service: str, key: str) -> None: + """Wait until the key answered every message queued so far: a chat sent now is answered after them.""" + assert (await chat(harness, key, "done", service)).json() == "echo: done" + + +async def summary(harness: HarnessEnvironment, service: str, key: str) -> list[tuple[str, str]]: + page = typing.cast(Messages, (await post(harness, f"{service}/{key}/messages", json={"view": "summary"})).json()) + return [(message["role"], str(message["parts"][0].get("text"))) for message in page["messages"]] + + +def said(contents: list[types.Content]) -> list[tuple[str, str]]: + """What a request holds, part by part: a message's text, a tool's call, or a tool's result.""" + parts: list[tuple[str, str]] = [] + for content in contents: + for part in content.parts or []: + if part.function_call: + parts.append((str(content.role), f"call {part.function_call.name}")) + elif part.function_response: + parts.append((str(content.role), f"result {part.function_response.name}")) + elif part.text: + parts.append((str(content.role), part.text)) + return parts + + +async def test_a_message_steered_while_a_tool_runs_reaches_the_next_request_once(harness: HarnessEnvironment): + k = key() + held.close() + invocation = await send(harness, "adk_steered", k, f"hold {k}") + await holding(k) + assert (await steer(harness, "adk_steered", k, f"hold Celsius {k}"))["turn"] == invocation + held.open() + # The model's next request has it after the tool's result, once. The model answers it with a call of its own, + # and the request after that one still has it, in place: the turn's answer is the reply to it. + assert (await attach(harness, invocation)).json() == f"tool said: held in Celsius {k}" + steered_request = [ + ("user", f"hold {k}"), + ("model", "call hold"), + ("user", "result hold"), + ("user", f"hold Celsius {k}"), + ] + assert [said(request) for request in seen[k]] == [ + [("user", f"hold {k}")], + steered_request, + [*steered_request, ("model", "call hold"), ("user", "result hold")], + ] + # Stored once, in the turn: the turn of its own that was queued for it does nothing. + await settled(harness, "adk_steered", k) + assert await summary(harness, "adk_steered", k) == [ + ("user", f"hold {k}"), + ("user", f"hold Celsius {k}"), + ("assistant", f"tool said: held in Celsius {k}"), + ("user", "done"), + ("assistant", "echo: done"), + ] + + +async def test_a_crash_after_a_steered_message_is_stored_resumes_the_run_with_it_once(harness: HarnessEnvironment): + k = key() + held.close() + invocation = await send(harness, "adk_steered", k, f"hold {k}") + await holding(k) + await steer(harness, "adk_steered", k, f"crashing {k}") + held.open() + # The request with the message crashed the attempt: the retry resumed the prompt's invocation, from the events + # the turn stored, the message among them. + assert (await attach(harness, invocation)).json() == f"echo: crashing {k}" + after = [("user", f"hold {k}"), ("model", "call hold"), ("user", "result hold"), ("user", f"crashing {k}")] + # Its request is the crashed one: the message once, as the turn stored it, which it did not take again. + assert [said(request) for request in seen[k][1:]] == [after, after] + # Resumed at the step: the tool ran once, and the model was not asked for its call again. + assert holds[k] == 1 and asked[k] == 2 + await settled(harness, "adk_steered", k) + assert await summary(harness, "adk_steered", k) == [ + ("user", f"hold {k}"), + ("user", f"crashing {k}"), + ("assistant", f"echo: crashing {k}"), + ("user", "done"), + ("assistant", "echo: done"), + ] + + +async def test_an_agent_whose_root_is_a_workflow_agent_answers_a_steered_message_as_its_next_turn( + harness: HarnessEnvironment, +): + k = key() + held.close() + invocation = await send(harness, "adk_pipeline", k, f"hold {k}") + await holding(k) + # Its turn takes none: the message is answered as a turn of its own, after it. + assert (await steer(harness, "adk_pipeline", k, f"in Celsius {k}"))["turn"] is None + held.open() + assert (await attach(harness, invocation)).json() == f"tool said: held in {k}" + assert said(seen[k][1]) == [("user", f"hold {k}"), ("model", "call hold"), ("user", "result hold")] + await settled(harness, "adk_pipeline", k) + assert await summary(harness, "adk_pipeline", k) == [ + ("user", f"hold {k}"), + ("assistant", f"tool said: held in {k}"), + ("user", f"in Celsius {k}"), + ("assistant", f"echo: in Celsius {k}"), + ("user", "done"), + ("assistant", "echo: done"), + ] + + +async def test_a_message_steered_during_the_last_model_call_is_answered_as_the_next_turn( + harness: HarnessEnvironment, +): + k = key() + held.close() + invocation = await send(harness, "adk_steered", k, f"slow {k}") + + async def answering() -> bool: + return len(seen[k]) == 1 + + await until(answering) + assert (await steer(harness, "adk_steered", k, f"in Celsius {k}"))["turn"] == invocation + held.open() + # The model's reply was its last: no step comes after it to take the message, which is answered next. + assert (await attach(harness, invocation)).json() == f"echo: slow {k}" + await settled(harness, "adk_steered", k) + assert await summary(harness, "adk_steered", k) == [ + ("user", f"slow {k}"), + ("assistant", f"echo: slow {k}"), + ("user", f"in Celsius {k}"), + ("assistant", f"echo: in Celsius {k}"), + ("user", "done"), + ("assistant", "echo: done"), + ] + + # --- Definitions ---------------------------------------------------------------------------------- @@ -504,7 +712,6 @@ async def test_an_agent_of_your_own_runs_as_configured_its_sub_agents_with_its_m async def test_a_workflows_first_step_stored_before_a_crash_is_not_its_answer(monkeypatch: pytest.MonkeyPatch): - from google.adk.agents import SequentialAgent # pylint: disable=import-outside-toplevel from sessions_object import request # pylint: disable=import-error,import-outside-toplevel from restate.agents.adk import _agent # pylint: disable=import-outside-toplevel @@ -733,3 +940,128 @@ def get_weather(city: str, tool_context: ToolContext) -> str: # pylint: disable assert await asked("hi 3") == "echo: hi 3" # The first turn's events are not loaded any more: its state is, as the latest turn carried it. assert await asked("system x") == "system: City: Lisbon" + + +@pytest.mark.parametrize("last", ["skip_summarization", "set_model_response"]) +async def test_a_step_after_which_adk_asks_the_model_nothing_takes_no_steered_message(last: str): + from google.adk.apps import App # pylint: disable=import-outside-toplevel + from google.adk.plugins import BasePlugin # pylint: disable=import-outside-toplevel + from google.adk.tools.tool_context import ToolContext # pylint: disable=import-outside-toplevel + from pydantic import BaseModel # pylint: disable=import-outside-toplevel + from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel + from sessions_object import Host, request # pylint: disable=import-error,import-outside-toplevel + + from restate.agents._conversation import steer_stored # pylint: disable=import-outside-toplevel + from restate.agents._turn import _INBOX, Inbox # pylint: disable=import-outside-toplevel + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + class Weather(BaseModel): + summary: str + + def quiet(city: str, tool_context: ToolContext) -> str: + """Answers for the model.""" + tool_context.actions.skip_summarization = True + return f"sunny in {city}" + + calls = { + "skip_summarization": ("quiet", {"city": "Rome"}), + "set_model_response": ("set_model_response", {"summary": "sunny"}), + } + name, args = calls[last] + asked_for: list[list[tuple[str, str]]] = [] + + class Ending(BaseLlm): + """Calls the tool after whose result ADK asks it nothing more.""" + + model: str = "ending" + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + asked_for.append(said(list(llm_request.contents))) + call = types.FunctionCall(id="call-1", name=name, args=args) + yield LlmResponse(content=types.Content(role="model", parts=[types.Part(function_call=call)])) + + inbox = Inbox(revision=0) + + class Steerer(BasePlugin): + """A message is steered to the turn while its tool runs, as the companion's publish would return it.""" + + async def before_tool_callback(self, *, tool: Any, tool_args: Any, tool_context: Any) -> None: + inbox.messages.append({"id": "s1", "text": "use Celsius"}) + + output = Weather if last == "set_model_response" else None + # An output schema with tools: the model answers through ADK's set_model_response, whose result is the answer. + mine = LlmAgent(name="mine", model=Ending(), tools=[quiet], output_schema=output) + app = App(name="mine", root_agent=mine, plugins=[Steerer("steerer")]) + asker = sessions.SessionObject("asker") + + @asker.handler + async def ask(ctx: sessions.SessionContext, text: str) -> Any: + # As a chat turn of an agent that takes steered messages sets it. + _INBOX.set(ctx, inbox) + return await adk.prompt(ctx, text, agent=app) + + @asker.handler + 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"}) + 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. + assert asked_for == [[("user", "weather")]] + assert await invoke(asker, "taken", Host(backend, service="asker"), request()) is False + + +async def test_a_steered_message_starts_an_invocation_of_its_own_which_a_retry_resumes_past(): + """After a crash, ADK resumes the prompt's invocation from that invocation's events. Its runner of a root that + is no LlmAgent (a SequentialAgent here, made to take steered messages) would take a message of the user last + among them for the answer to a tool call, which it fails to find on every retry ("No agent to transfer to"). A + steered message starts an invocation of its own instead, as a prompt does.""" + from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel + from sessions_object import Host, request # pylint: disable=import-error,import-outside-toplevel + + 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) + crashed: list[bool] = [] + + def get_weather(city: str) -> str: # pylint: disable=redefined-outer-name + """The weather, while a message is steered to the turn.""" + inbox.messages.append({"id": "s1", "text": "in Celsius"}) + return f"sunny in {city}" + + class Crashing(Scripted): + """Crashes once, on the request with the steered message.""" + + async def generate_content_async( + self, llm_request: LlmRequest, stream: bool = False + ) -> AsyncGenerator[LlmResponse, None]: + if "in Celsius" in _prompts(list(llm_request.contents)) and not crashed: + crashed.append(True) + raise RuntimeError("the process crashed after the steered message was stored") + async for response in super().generate_content_async(llm_request, stream): + yield response + + steps = SequentialAgent(name="steps", sub_agents=[LlmAgent(name="step", model=Crashing(), tools=[get_weather])]) + asker = sessions.SessionObject("asker") + + @asker.handler + async def ask(ctx: sessions.SessionContext, text: str) -> str: + # As a chat turn of an agent that takes steered messages sets it. + _INBOX.set(ctx, inbox) + return await adk.prompt(ctx, text, agent=steps) + + backend, invocation = Backend(), request() + + async def attempt() -> Any: + host = Host(backend, service="asker", answers={"publish": b"null"}) + return await invoke(asker, "ask", host, invocation, "weather Rome") + + with pytest.raises(RuntimeError, match="crashed after the steered message was stored"): + await attempt() + # The retry resumed the prompt's invocation, past the message, and answered it. + assert await attempt() == "echo: in Celsius" diff --git a/tests/agents_langgraph.py b/tests/agents_langgraph.py index 2c8512a9..69cabeb0 100644 --- a/tests/agents_langgraph.py +++ b/tests/agents_langgraph.py @@ -33,7 +33,7 @@ from langchain_core.tools import tool as langchain_tool # noqa: E402 import restate # noqa: E402 -from support import RESTATE_IMAGE # noqa: E402 # pylint: disable=import-error +from support import RESTATE_IMAGE, Gate # noqa: E402 # pylint: disable=import-error from restate import HarnessEnvironment, sessions # noqa: E402 from restate.agents import tool # noqa: E402 from restate.agents.internal import CANCELLED # noqa: E402 @@ -77,13 +77,14 @@ def reply(messages: Sequence[BaseMessage]) -> AIMessage: return AIMessage( "", tool_calls=[{"id": f"call-{rest}-{len(results)}", "name": "get_weather", "args": {"city": rest}}] ) - if word in ("weather", "crash", "wait", "wrong", "unknown", "forbidden") and not results: + if word in ("weather", "crash", "wait", "wrong", "unknown", "forbidden", "hold") and not results: name = { "weather": "get_weather", "crash": "flaky", "wait": "blocking", "wrong": "get_weather", "forbidden": "forbidden", + "hold": "hold", }.get(word, "nonexistent") args: dict[str, Any] = {"town": 3} if word == "wrong" else {"city": rest} return AIMessage("", tool_calls=[{"id": f"call-{rest}", "name": name, "args": args}]) @@ -195,10 +196,58 @@ def _generate( """An agent whose model sees 40 characters of the turns before the window: it summarizes the earlier ones.""" +holding = Gate() +"""Holds the tool ``hold`` until a test opens it: a gate between the test's thread and the served app's.""" +held: set[str] = set() + + +@tool +async def hold(city: str) -> str: + """Waits until the test opens the gate, so that it can steer a message to the turn meanwhile.""" + held.add(city) + await holding.wait() + return f"held in {city}" + + +requests: collections.defaultdict[str, list[list[BaseMessage]]] = collections.defaultdict(list) +"""Each request of the steered agent's model, in order, by the last word of its first message of the user (each +test's own key).""" +booms: set[str] = set() + + +class Recording(Scripted): + """The scripted model, which keeps each request it gets, and crashes the first request whose latest message of + the user is ``boom ...``, as a process that crashed after the turn stored that message.""" + + def _generate( + self, + messages: list[BaseMessage], + stop: list[str] | None = None, + run_manager: CallbackManagerForLLMRun | None = None, + **kwargs: Any, + ) -> ChatResult: + first = next(m for m in messages if isinstance(m, HumanMessage)) + requests[str(first.content).split(" ")[-1]].append(list(messages)) + latest = str(next(m for m in reversed(messages) if isinstance(m, HumanMessage)).content) + if latest.startswith("boom ") and latest not in booms: + booms.add(latest) + raise RuntimeError("the process crashed") + return super()._generate(messages, stop, run_manager, **kwargs) + + +steered_agent = langgraph.agent( + "lg_steered", model=Recording(profile={"max_input_tokens": 4}), tools=[get_weather, hold] +) +"""An agent whose model answers the latest message of the user: a message steered to its turn, once it has one. + +Its window holds no message, so each request is the latest prompt and what followed it, whole: a steered message +taken for a prompt would be the request's start, without the prompt.""" + + @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([assistant, per_user, summarizing]), restate_image=RESTATE_IMAGE + restate.app([assistant, per_user, summarizing, steered_agent]), restate_image=RESTATE_IMAGE ) as harness: yield harness @@ -206,17 +255,21 @@ async def harness(): TEXT = {"content-type": "text/plain"} -async def post(harness: HarnessEnvironment, path: str, **kwargs: Any) -> httpx2.Response: +async def post(harness: HarnessEnvironment, path: str, agent: str = "lg_assistant", **kwargs: Any) -> httpx2.Response: async with httpx2.AsyncClient(timeout=60) as ingress: - return await ingress.post(f"{harness.ingress_url}/lg_assistant/{path}", **kwargs) + return await ingress.post(f"{harness.ingress_url}/{agent}/{path}", **kwargs) async def chat(harness: HarnessEnvironment, key: str, text: str) -> httpx2.Response: return await post(harness, f"{key}/chat", content=text, headers=TEXT) -async def messages(harness: HarnessEnvironment, key: str) -> list[tuple[str, str]]: - page = typing.cast(Messages, (await post(harness, f"{key}/messages", json={})).json()) +async def messages( + harness: HarnessEnvironment, key: str, agent: str = "lg_assistant", view: str | None = None +) -> list[tuple[str, str]]: + """The conversation as the ``messages`` view shows it: the full view, or ``view`` (such as ``"summary"``).""" + wanted = {"view": view} if view else {} + page = typing.cast(Messages, (await post(harness, f"{key}/messages", agent, json=wanted)).json()) return [ (message["role"], " ".join(str(part.get("text", part.get("output", ""))) for part in message["parts"])) for message in page["messages"] @@ -368,6 +421,117 @@ async def test_options_may_be_functions_of_the_context(harness: HarnessEnvironme assert (await ingress.post(url, content=f"weather {k}", headers=TEXT)).json() == f"tool said: sunny in {k}" +# --- Steering ------------------------------------------------------------------------------------- + +STEERED = "lg_steered" + + +async def holding_turn(harness: HarnessEnvironment, k: str) -> str: + """Send ``hold `` to the steered agent, and return its invocation once its tool waits on the closed gate.""" + holding.close() + sent = await post(harness, f"{k}/chat/send", STEERED, content=f"hold {k}", headers=TEXT) + invocation = typing.cast(str, sent.json()["invocationId"]) + + async def in_the_tool() -> bool: + return k in held + + await until(in_the_tool) + return invocation + + +async def steer(harness: HarnessEnvironment, k: str, text: str, turn: str) -> str: + """Steer ``text`` to the turn being answered, invocation ``turn``, and return the message's id.""" + # With an idempotency key, as a client sends it: a retried request is the same message. + headers = {**TEXT, "idempotency-key": uuid.uuid4().hex} + steered = await post(harness, f"{k}/steer", STEERED, content=text, headers=headers) + assert steered.status_code == 200, steered.text + assert steered.json()["turn"] == turn + return typing.cast(str, steered.json()["id"]) + + +async def answered(harness: HarnessEnvironment, invocation: str) -> Any: + """Open the gate, and return the answer of the turn of ``invocation``.""" + holding.open() + async with httpx2.AsyncClient(timeout=60) as ingress: + return (await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation}/attach")).json() + + +def said(request: list[BaseMessage]) -> list[tuple[str, str]]: + return [(type(m).__name__, str(m.content)) for m in request] + + +async def told(harness: HarnessEnvironment, k: str) -> list[str]: + """The messages of the user that the steered agent's conversation stores, in order.""" + return [text for role, text in await messages(harness, k, STEERED) if role == "user"] + + +async def test_a_message_steered_while_a_tool_runs_is_in_the_next_request_after_its_result_once( + harness: HarnessEnvironment, +): + k = key() + invocation = await holding_turn(harness, k) + steered = await steer(harness, k, "use celsius", invocation) + # The model answers the latest message of the user: the turn's answer is the reply to the steered message. + assert await answered(harness, invocation) == "echo: use celsius" + first, second = requests[k] + assert said(first) == [("HumanMessage", f"hold {k}")] + # The very next request has it, once, after the tool's result: as the graph's state holds it, under its id. + assert said(second) == [ + ("HumanMessage", f"hold {k}"), + ("AIMessage", ""), + ("ToolMessage", f"held in {k}"), + ("HumanMessage", "use celsius"), + ] + assert second[-1].id == f"steer:{steered}" + # Stored once, as a step of the turn: the summary has the prompt, the steered message and the answer. + assert await told(harness, k) == [f"hold {k}", "use celsius"] + assert await messages(harness, k, STEERED, "summary") == [ + ("user", f"hold {k}"), + ("user", "use celsius"), + ("assistant", "echo: use celsius"), + ] + + +async def test_a_message_steered_after_a_tool_round_stays_in_the_request_after_the_next_round( + harness: HarnessEnvironment, +): + k = key() + invocation = await holding_turn(harness, k) + # The model acts on it with a tool call of its own: a second round. + await steer(harness, k, f"weather Paris-{k}", invocation) + assert await answered(harness, invocation) == f"tool said: sunny in Paris-{k}" + _, _, after_second = requests[k] + # The run went on with the message in its state, in place: the request after the second round still has it, once. + assert said(after_second) == [ + ("HumanMessage", f"hold {k}"), + ("AIMessage", ""), + ("ToolMessage", f"held in {k}"), + ("HumanMessage", f"weather Paris-{k}"), + ("AIMessage", ""), + ("ToolMessage", f"sunny in Paris-{k}"), + ] + + +async def test_a_crash_after_a_steered_message_is_stored_retries_with_it_stored_once(harness: HarnessEnvironment): + k = key() + invocation = await holding_turn(harness, k) + # The first request with it crashes the attempt, after the turn stored it: Restate retries the turn. + steered = await steer(harness, k, f"boom {k}", invocation) + assert await answered(harness, invocation) == f"echo: boom {k}" + _, crashed, retried = requests[k] + assert said(crashed) == [ + ("HumanMessage", f"hold {k}"), + ("AIMessage", ""), + ("ToolMessage", f"held in {k}"), + ("HumanMessage", f"boom {k}"), + ] + # The retry found it stored, in its place and marked as steered, and did not take it again: its request is the + # crashed one. + assert said(retried) == said(crashed) + assert retried[-1].id == crashed[-1].id == f"steer:{steered}" + assert await told(harness, k) == [f"hold {k}", f"boom {k}"] + + # --- In process ----------------------------------------------------------------------------------- @@ -433,6 +597,26 @@ def test_a_prompt_longer_than_the_window_is_sent_whole(): assert _agent._latest(messages, None) == messages # pyright: ignore[reportPrivateUsage] +def test_a_request_starts_at_a_prompt_never_at_a_steered_message(): + from restate.agents.langgraph import _agent # pylint: disable=import-outside-toplevel + + latest = _agent._latest # pyright: ignore[reportPrivateUsage] + earlier: list[BaseMessage] = [HumanMessage("hi"), AIMessage("hello")] + prompt = HumanMessage("hold on") + call = AIMessage("", tool_calls=[{"id": "c1", "name": "hold", "args": {"city": "x"}}]) + result = ToolMessage("x" * 1000, tool_call_id="c1") + turn: list[BaseMessage] = [prompt, call, result, HumanMessage("use celsius", id="steer:s1")] + # Only the steered message fits, then nothing does: the request is the turn whole, from its prompt. + assert latest([*earlier, *turn], 10) == turn + assert latest([*earlier, *turn], 1) == turn + # Unmarked, LangChain's trim would start at it, without the turn's prompt and the tool's call and result. + unmarked = HumanMessage("use celsius") + assert latest([*earlier, prompt, call, result, unmarked], 10) == [unmarked] + # Nor does a request start at a message steered to an earlier turn: it starts at the next prompt. + before: list[BaseMessage] = [HumanMessage("x" * 400), HumanMessage("in brief", id="steer:s0"), AIMessage("done")] + assert latest([*before, HumanMessage("next")], 20) == [HumanMessage("next")] + + async def test_an_answer_stored_before_a_crash_is_not_asked_for_again(monkeypatch: pytest.MonkeyPatch): """A crash after the model's final reply is stored: the retry answers with it, without asking the model again.""" from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel diff --git a/tests/agents_langgraph_graph.py b/tests/agents_langgraph_graph.py index 37fc214f..1a6a9b22 100644 --- a/tests/agents_langgraph_graph.py +++ b/tests/agents_langgraph_graph.py @@ -532,6 +532,28 @@ async def test_a_prebuilt_node_waits_for_a_signal_in_chat(harness: HarnessEnviro assert (await answering).json() == "approved: yes" +async def test_a_graph_takes_no_steered_message_which_is_answered_as_a_turn_of_its_own(harness: HarnessEnvironment): + key = uuid.uuid4().hex + answering = asyncio.ensure_future(chat(harness, "graph_desk", key, f"refund {key}")) + await until_waiting(harness, key) + headers = {"content-type": "text/plain", "idempotency-key": uuid.uuid4().hex} + steered = await post(harness, f"graph_desk/{key}/steer", content=f"hi {key}", headers=headers) + # The turn being answered takes none: the message waits, as a turn of its own. + assert steered.status_code == 200 and steered.json()["turn"] is None, steered.text + await post(harness, f"graph_desk/{key}/signal", json={"name": "approval", "value": "yes"}) + assert (await answering).json() == "approved: yes" + # Answered next, before the message sent after it. + assert (await chat(harness, "graph_desk", key, "hi again")).json() == "hello again" + assert await messages(harness, "graph_desk", key) == [ + ("user", f"refund {key}"), + ("assistant", "approved: yes"), + ("user", f"hi {key}"), + ("assistant", f"hello {key}"), + ("user", "hi again"), + ("assistant", "hello again"), + ] + + async def test_a_crash_resumes_the_run_and_its_messages_are_copied_once(harness: HarnessEnvironment): key = uuid.uuid4().hex assert (await chat(harness, "graph_desk", key, f"crash {key}")).json() == f"done {key}" diff --git a/tests/agents_openai.py b/tests/agents_openai.py index 4717cd6f..a70bd336 100644 --- a/tests/agents_openai.py +++ b/tests/agents_openai.py @@ -41,6 +41,7 @@ Usage, function_tool, handoff, + input_guardrail, output_guardrail, set_tracing_disabled, tool_input_guardrail, @@ -51,9 +52,9 @@ from openai.types.responses import ResponseFunctionToolCall, ResponseOutputMessage, ResponseOutputText # noqa: E402 import restate # noqa: E402 -from support import RESTATE_IMAGE # noqa: E402 # pylint: disable=import-error +from support import RESTATE_IMAGE, Gate # noqa: E402 # pylint: disable=import-error from restate import HarnessEnvironment, sessions # noqa: E402 -from restate.agents import openai, tool # noqa: E402 +from restate.agents import AgentObject, Tool, openai, tool # noqa: E402 from restate.agents._api import Messages # noqa: E402 from restate.agents.internal import CANCELLED # noqa: E402 @@ -216,11 +217,65 @@ def desk(model: Model) -> Agent[Any]: summarizing = openai.agent("oai_summarizing", model=remembering, budget=10) """An agent whose model sees 40 characters of the turns before the window: it summarizes the earlier ones.""" +STEPPED = ("steer", "twice") +"""The prompts on which the steering tests' model calls tools: ``lookup``, and for ``twice`` then ``get_weather``.""" + + +def steered_reply(system: str | None, items: list[Item]) -> list[Output]: + """The steering tests' model: for ``steer `` and ``twice `` it calls its tools, then answers with the + messages of the user that came after the prompt (those steered to the turn); it echoes any other message.""" + users = [(n, str(i["content"])) for n, i in enumerate(items) if i.get("role") == "user"] + at, prompt = next(((n, text) for n, text in reversed(users) if text.split(" ")[0] in STEPPED), (-1, "")) + turn = items[at + 1 :] + if at < 0 or any(item.get("role") == "assistant" for item in turn): + return [said(f"echo: {users[-1][1]}")] + word, _, rest = prompt.partition(" ") + outputs = [str(item["output"]) for item in turn if item.get("type") == "function_call_output"] + if not outputs: + return [calls(f"lookup-{rest}", "lookup", f'{{"city": "{rest}"}}')] + if word == "twice" and len(outputs) == 1: + return [calls(f"weather-{rest}", "get_weather", f'{{"city": "{rest}"}}')] + heard = [str(item["content"]) for item in turn if item.get("role") == "user"] + return [said(f"heard: {', '.join(heard)}" if heard else f"tool said: {outputs[-1]}")] + + +held = Gate() +"""Holds the steering agent's ``lookup``, and its model's request for ``slow ``, while a test steers.""" +holding: set[str] = set() +"""The keys whose ``lookup``, or request for ``slow``, waits on ``held``.""" + + +class Holding(Scripted): + """The steering tests' model, served: its request for ``slow `` waits on ``held``, as a long answer would.""" + + async def get_response( + self, system_instructions: str | None, input: Any, *args: Any, **kwargs: Any + ) -> ModelResponse: + last = input[-1] if isinstance(input, list) and input else {} + word, _, rest = str(last.get("content", "")).partition(" ") + if last.get("role") == "user" and word == "slow": + holding.add(rest) + await held.wait() + return await super().get_response(system_instructions, input, *args, **kwargs) + + +@tool +async def lookup(city: str) -> str: + """Looks it up, once the test lets it.""" + holding.add(city) + await held.wait() + return f"found {city}" + + +steering = Holding(steered_reply) +steerable = openai.agent("oai_steerable", model=steering, tools=[lookup, get_weather]) +"""An agent that takes the messages steered to its turns, as each agent of this package does.""" + @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([assistant, per_user, triage, forgetful, summarizing]), restate_image=RESTATE_IMAGE + restate.app([assistant, per_user, triage, forgetful, summarizing, steerable]), restate_image=RESTATE_IMAGE ) as harness: yield harness @@ -237,8 +292,10 @@ async def chat(harness: HarnessEnvironment, key: str, text: str, service: str = return await post(harness, f"{service}/{key}/chat", content=text, headers=TEXT) -async def messages(harness: HarnessEnvironment, key: str) -> list[tuple[str, str]]: - page = typing.cast(Messages, (await post(harness, f"oai_assistant/{key}/messages", json={})).json()) +async def messages( + harness: HarnessEnvironment, key: str, service: str = "oai_assistant", view: str = "full" +) -> list[tuple[str, str]]: + page = typing.cast(Messages, (await post(harness, f"{service}/{key}/messages", json={"view": view})).json()) return [ (message["role"], " ".join(str(part.get("text", part.get("output", ""))) for part in message["parts"])) for message in page["messages"] @@ -1159,3 +1216,360 @@ def get_weather(city: str) -> str: # Also on the retry, which runs the stored call again itself. with pytest.raises(ImportError): await ask("weather Rome", invocation) + + +# --- Steering: a message steered to a turn joins it at its next step, in place -------------------- + + +async def send(harness: HarnessEnvironment, key: str, text: str) -> str: + sent = await post(harness, f"oai_steerable/{key}/chat/send", content=text, headers=TEXT) + return typing.cast(str, sent.json()["invocationId"]) + + +async def steer(harness: HarnessEnvironment, key: str, text: str) -> dict[str, Any]: + # With an idempotency key, as a client sends it: a retried request is the same message. + steered = await post( + harness, f"oai_steerable/{key}/steer", content=text, headers={**TEXT, "idempotency-key": uuid.uuid4().hex} + ) + assert steered.status_code == 200, steered.text + return typing.cast(dict[str, Any], steered.json()) + + +async def attach(harness: HarnessEnvironment, invocation: str) -> httpx2.Response: + async with httpx2.AsyncClient(timeout=60) as ingress: + return await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation}/attach") + + +async def waiting(k: str) -> None: + """Wait until the steering agent's tool (or its model's request) for key ``k`` waits on ``held``.""" + + async def holds() -> bool: + return k in holding + + await until(holds) + + +def requests_of(scripted: Scripted, prompt: str) -> list[list[Item]]: + """Return the requests of ``scripted`` for the conversation that starts with ``prompt``, in order.""" + return [items for _, items in scripted.requests if _prompts(items)[:1] == [prompt]] + + +async def test_a_message_steered_while_a_tool_runs_is_in_the_next_request_once(harness: HarnessEnvironment): + k = key() + held.close() + try: + invocation = await send(harness, k, f"steer {k}") + await waiting(k) + assert (await steer(harness, k, "use Celsius"))["turn"] == invocation + finally: + held.open() + # The turn's answer is the model's reply to the message. + assert (await attach(harness, invocation)).json() == "heard: use Celsius" + first, after_the_tool = requests_of(steering, f"steer {k}") + assert _prompts(first) == [f"steer {k}"] + # The request after the tool's output has it, once, right after that output. + assert after_the_tool[-1] == {"role": "user", "content": "use Celsius"} + assert after_the_tool[-2]["type"] == "function_call_output" and after_the_tool[-2]["call_id"] == f"lookup-{k}" + assert _prompts(after_the_tool) == [f"steer {k}", "use Celsius"] + # Its own turn, queued as it was sent, finds it taken: once the key answered a message sent after it, the + # conversation has it once, in the turn, between the prompt and the answer. + assert (await chat(harness, k, f"done {k}", "oai_steerable")).json() == f"echo: done {k}" + assert await messages(harness, k, "oai_steerable", view="summary") == [ + ("user", f"steer {k}"), + ("user", "use Celsius"), + ("assistant", "heard: use Celsius"), + ("user", f"done {k}"), + ("assistant", f"echo: done {k}"), + ] + assert [text for _, text in await messages(harness, k, "oai_steerable")].count("use Celsius") == 1 + + +async def test_a_message_steered_while_the_model_writes_its_answer_is_the_next_turn(harness: HarnessEnvironment): + k = key() + held.close() + try: + invocation = await send(harness, k, f"slow {k}") + await waiting(k) + # The model's request is the turn's last: no step comes after it. + assert (await steer(harness, k, "and also"))["turn"] == invocation + finally: + held.open() + assert (await attach(harness, invocation)).json() == f"echo: slow {k}" + assert (await chat(harness, k, f"done {k}", "oai_steerable")).json() == f"echo: done {k}" + assert await messages(harness, k, "oai_steerable", view="summary") == [ + ("user", f"slow {k}"), + ("assistant", f"echo: slow {k}"), + ("user", "and also"), + ("assistant", "echo: and also"), + ("user", f"done {k}"), + ("assistant", f"echo: done {k}"), + ] + # The turn's one request had no steered message; the message was the next turn's prompt. + assert [_prompts(items) for items in requests_of(steering, f"slow {k}")] == [ + [f"slow {k}"], + [f"slow {k}", "and also"], + [f"slow {k}", "and also", f"done {k}"], + ] + + +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 + messages steered to the turn (``steering``), as the companion does until the turn closes.""" + + def __init__(self) -> None: + super().__init__(start=b"null") + self.revision = 0 + self.sent = 0 + self.steering: list[dict[str, str]] = [] + """The messages steered to the turn being answered. ``taken`` never reaches this companion: each publish + lists them all, those the turn took too.""" + + def __missing__(self, handler: str) -> bytes: + if handler == "close": + # 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}" + self.revision += 1 + return json.dumps({"revision": self.revision, "steering": self.steering}).encode() + + def steer(self, text: str) -> None: + """Steer ``text`` to the turn being answered: its next publish lists it.""" + self.sent += 1 + self.steering.append({"id": f"steer-{self.sent}", "text": text}) + + +def steered_while(companion: Companion, text: str) -> Tool: + """The ``lookup`` tool, which ``text`` is steered to the turn while it runs: the turn learns of it as it ends.""" + + @tool(name="lookup") + async def looking(city: str) -> str: + """Looks it up.""" + companion.steer(text) + return f"found {city}" + + return looking + + +class Chat: + """An agent's ``chat``, invoked in this process with ``companion``: each call is a turn, with an invocation of + its own unless given one (to retry it).""" + + def __init__(self, obj: AgentObject, companion: Companion) -> None: + from sessions_database import Backend # pylint: disable=import-error,import-outside-toplevel + + self.obj = obj + self.companion = companion + self.backend = Backend() + self.hosts: list[Any] = [] + + async def __call__(self, text: str, invocation: Any = None) -> Any: + from sessions_object import Host # pylint: disable=import-error,import-outside-toplevel + + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + host = Host(self.backend, service=self.obj.name, answers=self.companion) + self.hosts.append(host) + return await invoke(self.obj, "chat", host, invocation or request_of(), text) + + async def summary(self) -> list[tuple[str, str]]: + from sessions_object import Host # pylint: disable=import-error,import-outside-toplevel + + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + view = Host(self.backend, service=self.obj.name).view() + page = typing.cast(Messages, await invoke(self.obj, "messages", view, request_of(), {"view": "summary"})) + return [(m["role"], str(m["parts"][0].get("text"))) for m in page["messages"]] + + def taken(self) -> list[Any]: + """Return the companion's ``taken`` that the turns sent, with the transactions that stored the messages.""" + return [send for host in self.hosts for send in host.sends if send.handler == "taken"] + + +CELSIUS = {"role": "user", "content": "use Celsius"} +"""The message steered in the tests in this process, as the model gets it.""" + + +async def test_a_steered_message_keeps_its_place_in_every_later_request_of_the_run(): + companion = Companion() + scripted = Scripted(steered_reply) + chat = Chat( + openai.agent("steered", model=scripted, tools=[steered_while(companion, "use Celsius"), get_weather]), companion + ) + assert await chat("twice Rome") == "heard: use Celsius" + first, second, third = (items for _, items in scripted.requests) + assert first == [{"role": "user", "content": "twice Rome"}] + # Taken after the first round: the request after it has it after the round's output. + assert second[-1] == CELSIUS and second[-2]["call_id"] == "lookup-Rome" + # The run made the request after the second round afresh, from its input and its items: it has the message + # still, once, in its place, before the second round's call and output. + assert third[: len(second)] == second + assert [item.get("call_id") for item in third[len(second) :]] == ["weather-Rome", "weather-Rome"] + assert third.count(CELSIUS) == 1 + # Stored once, though each publish listed it again: the turn has it between its prompt and its answer. + assert await chat.summary() == [ + ("user", "twice Rome"), + ("user", "use Celsius"), + ("assistant", "heard: use Celsius"), + ] + assert len(chat.taken()) == 1 + + +async def test_a_steered_message_stored_before_a_crash_is_in_the_retried_request_once(monkeypatch: pytest.MonkeyPatch): + # The process crashes once the filter stored the message, before the model is asked. + crash_once_after(monkeypatch, "model_input", lambda data: CELSIUS in data.model_data.input) + companion = Companion() + scripted = Scripted(steered_reply) + chat = Chat(openai.agent("steered", model=scripted, tools=[steered_while(companion, "use Celsius")]), companion) + invocation = request_of() + with pytest.raises(RuntimeError, match="crashed"): + await chat("steer Rome", invocation) + assert len(scripted.requests) == 1 + # The retry (the companion lists the message still) runs on from what is stored: the next request has it once, + # after the tool's output, and it is not stored again. + assert await chat("steer Rome", invocation) == "heard: use Celsius" + _, retried = scripted.requests[-1] + assert retried[-1] == CELSIUS and retried.count(CELSIUS) == 1 + assert retried[-2]["call_id"] == "lookup-Rome" + assert await chat.summary() == [ + ("user", "steer Rome"), + ("user", "use Celsius"), + ("assistant", "heard: use Celsius"), + ] + assert len(chat.taken()) == 1 + + +async def test_a_steered_message_runs_no_input_guardrail_again(): + checked: list[Any] = [] + + @input_guardrail + async def counted(context: RunContextWrapper[Any], agent: Agent[Any], input: Any) -> GuardrailFunctionOutput: + checked.append(input) + return GuardrailFunctionOutput(output_info=None, tripwire_triggered=False) + + companion = Companion() + scripted = Scripted(steered_reply) + tools = openai.tools([steered_while(companion, "use Celsius")]) + mine = Agent(name="guarded", model=scripted, tools=tools, input_guardrails=[counted]) + chat = Chat(openai.agent("steered", agent=mine), companion) + assert await chat("steer Rome") == "heard: use Celsius" + # The run took the message in place, and went on: its input guardrail ran once, on its input. + assert len(scripted.requests) == 2 and CELSIUS in scripted.requests[-1][1] + assert len(checked) == 1 + + +async def test_a_message_steered_before_the_runs_first_request_is_in_it(): + companion = Companion() + scripted = Scripted(steered_reply) + chat = Chat(openai.agent("steered", model=scripted), companion) + # The turn's first publish, as it starts, lists the message: the run starts with it, after the prompt. + companion.steer("use Celsius") + assert await chat("hi there") == "echo: use Celsius" + assert [items for _, items in scripted.requests] == [[{"role": "user", "content": "hi there"}, CELSIUS]] + assert await chat.summary() == [("user", "hi there"), ("user", "use Celsius"), ("assistant", "echo: use Celsius")] + + +def shown(items: list[Item]) -> list[str]: + """Return ``items`` as text, to compare: each message's text, and each call's and output's kind and id.""" + return [str(item["content"]) if "role" in item else f"{item['type']} {item['call_id']}" for item in items] + + +async def test_a_steered_message_keeps_its_place_after_a_handoff_whose_input_filter_took_out_its_item(): + def desk_script(system: str | None, items: list[Item]) -> list[Output]: + outputs = {item["call_id"] for item in items if item.get("type") == "function_call_output"} + if system == "You triage.": + if "call-l1" in outputs: + return [calls("call-h", "transfer_to_refunds", "{}")] + return [calls("call-l1", "lookup", '{"city": "Rome"}')] + if "call-l2" not in outputs: + return [calls("call-l2", "lookup", '{"city": "Oslo"}')] + return [said("refunded")] + + companion = Companion() + scripted = Scripted(desk_script) + refunds_tools = openai.tools([steered_while(companion, "and Fahrenheit")]) + refunds = Agent(name="refunds", instructions="You handle refunds.", model=scripted, tools=refunds_tools) + triage = Agent( + name="triage", + instructions="You triage.", + model=scripted, + tools=openai.tools([steered_while(companion, "use Celsius")]), + handoffs=[handoff(refunds, input_filter=handoff_filters.remove_all_tools)], + ) + chat = Chat(openai.agent("steered", agent=triage), companion) + assert await chat("handoff x") == "refunded" + assert [shown(items) for _, items in scripted.requests] == [ + ["handoff x"], + ["handoff x", "function_call call-l1", "function_call_output call-l1", "use Celsius"], + # The filter took the tools' items out of the refunds agent's input, the output the message came after + # among them: it stays where the filter would have kept it, after the prompt and before what came after. + ["handoff x", "use Celsius"], + ["handoff x", "use Celsius", "function_call call-l2", "function_call_output call-l2"], + ] + # The filter took the handoff's own output out too, so the turn has no output for its call: the run takes no + # message after it (a turn of its own answers the one steered during the refunds' lookup). + assert companion.sent == 2 + assert await chat.summary() == [("user", "handoff x"), ("user", "use Celsius"), ("assistant", "refunded")] + + +async def test_an_agent_used_as_a_tool_gets_none_of_the_turns_messages_and_stores_none_of_its_items(): + def nested_script(system: str | None, items: list[Item]) -> list[Output]: + outputs = {item["call_id"] for item in items if item.get("type") == "function_call_output"} + if system == "You look things up.": + return [said("sunny") if "call-w" in outputs else calls("call-w", "forecast", '{"city": "Rome"}')] + if "call-l" not in outputs: + return [calls("call-l", "lookup", '{"city": "Rome"}')] + if "call-s" not in outputs: + return [calls("call-s", "ask_helper", '{"input": "Rome?"}')] + return [said("done")] + + @function_tool + def forecast(city: str) -> str: + """The forecast for a city.""" + return f"sunny in {city}" + + companion = Companion() + scripted = Scripted(nested_script) + helper = Agent(name="helper", instructions="You look things up.", model=scripted, tools=[forecast]) + tools = [*openai.tools([steered_while(companion, "use Celsius")]), helper.as_tool("ask_helper", "Asks the helper.")] + chat = Chat(openai.agent("steered", agent=Agent(name="desk", model=scripted, tools=tools)), companion) + assert await chat("nest x") == "done" + # The helper's run gets the turn's run_config, its input filter with it: its requests are its own still. + assert [shown(items) for _, items in scripted.requests] == [ + ["nest x"], + ["nest x", "function_call call-l", "function_call_output call-l", "use Celsius"], + ["Rome?"], + ["Rome?", "function_call call-w", "function_call_output call-w"], + [ + "nest x", + "function_call call-l", + "function_call_output call-l", + "use Celsius", + "function_call call-s", + "function_call_output call-s", + ], + ] + # The turn stored its own items, none of the helper's: the next turn's request has no output without its call. + assert await chat("thanks") == "done" + assert "function_call_output call-w" not in shown(scripted.requests[-1][1]) + + +def test_steered_messages_keep_their_order_and_go_after_their_item_or_before_what_came_after_them(): + from restate.agents.openai import _agent # pylint: disable=import-outside-toplevel + + spliced, steered = _agent._spliced, _agent._Steered # pyright: ignore[reportPrivateUsage] + output = {"type": "function_call_output", "call_id": "c1", "output": "found"} + later = {"type": "function_call", "call_id": "c2", "name": "lookup", "arguments": "{}"} + hi, first, second = ({"role": "user", "content": text} for text in ("hi", "first", "second")) + # Both taken after the output, once the run had seen one key; the call's key it saw after. + seen = {"function_call_output:call_id:c1": 0, "function_call:call_id:c2": 1} + taken = [ + steered(first, "function_call_output:call_id:c1", 1), + steered(second, "function_call_output:call_id:c1", 1), + ] + assert spliced([hi, output, later], taken, seen) == [hi, output, first, second, later] + # The item twice: after its last copy, the one the SDK keeps. + assert spliced([output, later, output], taken, seen) == [output, later, output, first, second] + # The item they came after taken out: before the first item that came after them, in order; else at the end. + assert spliced([hi, later], taken, seen) == [hi, first, second, later] + assert spliced([hi], taken, seen) == [hi, first, second] diff --git a/tests/agents_pydantic_ai.py b/tests/agents_pydantic_ai.py index 0e869f63..737c4bb1 100644 --- a/tests/agents_pydantic_ai.py +++ b/tests/agents_pydantic_ai.py @@ -14,6 +14,7 @@ import asyncio import collections +import json import typing import uuid from dataclasses import dataclass @@ -41,7 +42,7 @@ from pydantic_ai.models.function import AgentInfo, FunctionModel # noqa: E402 import restate # noqa: E402 -from support import RESTATE_IMAGE # noqa: E402 # pylint: disable=import-error +from support import RESTATE_IMAGE, Gate # noqa: E402 # pylint: disable=import-error from restate import HarnessEnvironment, sessions # noqa: E402 from restate.agents import tool # noqa: E402 from restate.agents.internal import CANCELLED # noqa: E402 @@ -456,11 +457,57 @@ def trimmed_agent() -> Any: trimmed = pydantic_ai.agent("pai_trimmed", agent=trimmed_agent()) +steer_gate = Gate() +"""Holds the steered agent's ``hold`` tool, and its model's answer to ``slow``, until a test opens it.""" + +held: set[str] = set() +"""The keys whose tool call, or model request, waits on ``steer_gate``.""" + +steered_requests: collections.defaultdict[str, list[list[ModelMessage]]] = collections.defaultdict(list) +"""What the steered agent's model was given, request by request, by the key that ends the conversation's first +prompt.""" + +crashing: set[str] = set() +"""The keys whose model crashes once, as it answers a message steered to the turn.""" + + +async def steerable(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + """The scripted model of the steered agent: it answers the latest message of the user, after looking it up + (``hold``) for ``look ``, and taking its time (``steer_gate``) for ``slow ``.""" + prompts = _prompts(messages) + k = prompts[0].split(" ")[-1] + steered_requests[k].append(list(messages)) + latest = prompts[-1] + word = latest.split(" ")[0] + last = messages[-1] + returned = isinstance(last, ModelRequest) and any(isinstance(p, ToolReturnPart) for p in last.parts) + if word == "look" and not returned: + return ModelResponse(parts=[ToolCallPart("hold", {"key": k}, tool_call_id=f"call-{k}")]) + if word == "slow": + held.add(k) + await steer_gate.wait() + if k in crashing and word not in ("look", "slow"): + crashing.discard(k) + raise RuntimeError("the process crashed while the model answered") + return ModelResponse(parts=[TextPart(f"answer to: {latest}")]) + + +@tool +async def hold(key: str) -> str: + """Looks something up, until the test opens the gate.""" + held.add(key) + await steer_gate.wait() + return f"found for {key}" + + +steered = pydantic_ai.agent("pai_steered", model=FunctionModel(steerable), tools=[hold]) + @pytest.fixture(scope="session") async def harness(): async with restate.create_test_harness( - restate.app([assistant, per_user, trimmed, shop, desk, remembering, memories]), restate_image=RESTATE_IMAGE + restate.app([assistant, per_user, trimmed, shop, desk, remembering, memories, steered]), + restate_image=RESTATE_IMAGE, ) as harness: yield harness @@ -1054,3 +1101,169 @@ def test_a_tool_call_asked_for_again_is_read_as_an_error(): [again, found] = ITEMS.read(ModelRequest(parts=[retried, returned])).results assert (again.id, again.output, again.failed) == ("c1", '{"error":"not like that"}', True) assert (found.id, found.output, found.failed) == ("c2", "found", False) + + +async def steered_post(harness: HarnessEnvironment, path: str, **kwargs: Any) -> httpx2.Response: + async with httpx2.AsyncClient(timeout=60) as ingress: + return await ingress.post(f"{harness.ingress_url}/pai_steered/{path}", **kwargs) + + +async def steered_send(harness: HarnessEnvironment, k: str, text: str) -> str: + sent = await steered_post(harness, f"{k}/chat/send", content=text, headers=TEXT) + return typing.cast(str, sent.json()["invocationId"]) + + +async def steer(harness: HarnessEnvironment, k: str, text: str) -> dict[str, Any]: + # With an idempotency key, as a client sends it: a retried request is the same message. + headers = {**TEXT, "idempotency-key": uuid.uuid4().hex} + steered_to = await steered_post(harness, f"{k}/steer", content=text, headers=headers) + assert steered_to.status_code == 200, steered_to.text + return typing.cast(dict[str, Any], steered_to.json()) + + +async def attach(harness: HarnessEnvironment, invocation: str) -> httpx2.Response: + async with httpx2.AsyncClient(timeout=60) as ingress: + return await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation}/attach") + + +async def steered_thread(harness: HarnessEnvironment, k: str, view: str | None = None) -> list[tuple[str, str]]: + """The steered agent's conversation, as the ``messages`` view shows it: whole, or its ``summary``.""" + wanted = {"view": view} if view is not None else {} + page = typing.cast(Messages, (await steered_post(harness, f"{k}/messages", json=wanted)).json()) + return [ + (message["role"], " ".join(str(part.get("text", part.get("output", ""))) for part in message["parts"])) + for message in page["messages"] + ] + + +async def until_held(k: str) -> None: + async def holding() -> bool: + return k in held + + await until(holding) + + +async def test_a_message_steered_while_a_tool_runs_reaches_the_next_request_after_the_tools_result( + harness: HarnessEnvironment, +): + k = key() + steer_gate.close() + invocation = await steered_send(harness, k, f"look {k}") + await until_held(k) + assert (await steer(harness, k, "use Celsius"))["turn"] == invocation + steer_gate.open() + # The turn's answer is the model's reply to it. + assert (await attach(harness, invocation)).json() == "answer to: use Celsius" + # The very next request had it, after the tool's result, once. + [first, following] = steered_requests[k] + assert _prompts(first) == [f"look {k}"] + assert _prompts(following) == [f"look {k}", "use Celsius"] + last = following[-1] + assert isinstance(last, ModelRequest) and [type(p) for p in last.parts] == [ToolReturnPart, UserPromptPart] + # Stored once, as a step of the turn; its queued turn of its own then did nothing. + assert (await steered_post(harness, f"{k}/chat", content=f"thanks {k}", headers=TEXT)).status_code == 200 + thread = await steered_thread(harness, k) + assert [text for role, text in thread if role == "user"] == [f"look {k}", "use Celsius", f"thanks {k}"] + assert await steered_thread(harness, k, "summary") == [ + ("user", f"look {k}"), + ("user", "use Celsius"), + ("assistant", "answer to: use Celsius"), + ("user", f"thanks {k}"), + ("assistant", f"answer to: thanks {k}"), + ] + + +async def test_a_crash_after_a_steered_message_is_stored_gives_the_retried_request_it_once( + harness: HarnessEnvironment, +): + k = key() + crashing.add(k) + steer_gate.close() + invocation = await steered_send(harness, k, f"look {k}") + await until_held(k) + await steer(harness, k, "use Celsius") + steer_gate.open() + assert (await attach(harness, invocation)).json() == "answer to: use Celsius" + # The attempt crashed as the model answered the stored message: the retried attempt carries on from the + # history, which has it, and takes it no more. + assert k not in crashing + [_, crashed, retried] = steered_requests[k] + assert _prompts(crashed) == _prompts(retried) == [f"look {k}", "use Celsius"] + thread = await steered_thread(harness, k) + assert [text for role, text in thread if role == "user"] == [f"look {k}", "use Celsius"] + + +async def test_a_message_steered_during_the_last_model_request_is_answered_as_the_next_turn( + harness: HarnessEnvironment, +): + k = key() + steer_gate.close() + invocation = await steered_send(harness, k, f"slow {k}") + await until_held(k) + assert (await steer(harness, k, "use Celsius"))["turn"] == invocation + steer_gate.open() + # The turn reaches no step after it: the message is answered next, as a turn of its own. + assert (await attach(harness, invocation)).json() == f"answer to: slow {k}" + done = await steered_post(harness, f"{k}/chat", content=f"done {k}", headers=TEXT) + assert done.json() == f"answer to: done {k}" + assert await steered_thread(harness, k, "summary") == [ + ("user", f"slow {k}"), + ("assistant", f"answer to: slow {k}"), + ("user", "use Celsius"), + ("assistant", "answer to: use Celsius"), + ("user", f"done {k}"), + ("assistant", f"answer to: done {k}"), + ] + + +async def test_a_chat_without_snapshots_that_answers_with_prompt_takes_steered_messages_in_order(): + """A ``chat`` of one's own that takes steered messages and answers with ``prompt``: the turn projection, with no + snapshot (nor ``_Snapshots``). Two messages steered during a tool call reach the next request, in order.""" + from agents_pydantic_ai_snapshots import ANSWERS, Chat # pylint: disable=import-outside-toplevel,import-error + from sessions_object import Host, request # pylint: disable=import-outside-toplevel,import-error + + from restate.agents import AgentObject # pylint: disable=import-outside-toplevel + from restate.agents.pydantic_ai._agent import ITEMS # pylint: disable=import-outside-toplevel + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + answers = {**ANSWERS} + requests: list[list[ModelMessage]] = [] + + def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + requests.append(list(messages)) + latest = _prompts(messages)[-1] + last = messages[-1] + returned = isinstance(last, ModelRequest) and any(isinstance(p, ToolReturnPart) for p in last.parts) + if latest.startswith("look") and not returned: + return ModelResponse(parts=[ToolCallPart("lookup", {"about": "weather"}, tool_call_id="call-look")]) + return ModelResponse(parts=[TextPart(f"answer to: {latest}")]) + + 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() + return f"found {about}" + + mine: Any = Agent(FunctionModel(model), tools=[lookup]) + obj = AgentObject("pai_steered_turns", window=0) + + @obj.chat(items=ITEMS, steers=True) + async def answer(ctx: sessions.SessionContext, text: str) -> str: + return await pydantic_ai.prompt(ctx, text, agent=mine) + + chat = Chat(obj) + host = Host(chat.backend, service=chat.name, answers=answers) + assert await invoke(obj, "chat", host, request(), "look it up") == "answer to: and be brief" + # Both reached the request after the tool's result, in order, each once; one transaction took both. + assert _prompts(requests[-1]) == ["look it up", "use Celsius", "and be brief"] + assert [send.handler for send in host.sends].count("taken") == 1 + # The next turn has them between the earlier turn's message and its answer (a window of 0: no trajectory). + assert await chat("and then?") == "answer to: and then?" + assert _prompts(requests[-1]) == ["look it up", "use Celsius", "and be brief", "and then?"] + thread = await chat.messages() + assert [text for role, text in thread if role == "user"] == [ + "look it up", + "use Celsius", + "and be brief", + "and then?", + ] diff --git a/tests/agents_pydantic_ai_snapshots.py b/tests/agents_pydantic_ai_snapshots.py index 686337f7..766deac4 100644 --- a/tests/agents_pydantic_ai_snapshots.py +++ b/tests/agents_pydantic_ai_snapshots.py @@ -16,6 +16,7 @@ # pylint: disable=missing-function-docstring,missing-class-docstring,import-error +import dataclasses import json import warnings from typing import Any @@ -24,8 +25,8 @@ pytest.importorskip("pydantic_ai") -from pydantic_ai import Agent # noqa: E402 -from pydantic_ai.capabilities import ReinjectSystemPrompt # noqa: E402 +from pydantic_ai import Agent, RunContext # noqa: E402 +from pydantic_ai.capabilities import AbstractCapability, ReinjectSystemPrompt # noqa: E402 from pydantic_ai.messages import ( # noqa: E402 CompactionPart, ModelMessage, @@ -34,6 +35,7 @@ SystemPromptPart, TextPart, ToolCallPart, + ToolReturnPart, UserPromptPart, ) from pydantic_ai.models.function import AgentInfo, FunctionModel # noqa: E402 @@ -513,3 +515,131 @@ async def own(ctx: sessions.SessionContext, text: str) -> str: await chat("q1 hi") await invoke(obj, "own", chat.host(), request(), "q2 hi") 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.""" + + +def _returned(message: ModelMessage) -> bool: + return isinstance(message, ModelRequest) and any(isinstance(p, ToolReturnPart) for p in message.parts) + + +def _said(messages: list[ModelMessage], text: str) -> int: + """How often the model is given ``text`` as a message of the user in ``messages``.""" + return _prompts(messages).count(text) + + +class CompactAtTheStep(AbstractCapability[Any]): + """A compaction of one's own, which compacts at a step: at the request after a tool's result, the turns before the + one being answered become a summary. A history that starts with the summary is left as it is.""" + + def compact(self, messages: list[ModelMessage]) -> list[ModelMessage]: + # The turn's prompt, its tool call and the result stay. + return [ModelRequest(parts=[CompactionPart(SUMMARY)]), *messages[-3:]] + + async def before_model_request(self, ctx: RunContext[Any], request_context: Any) -> Any: + messages = request_context.messages + if not _returned(messages[-1]) or _has_summary(messages[0]): + return request_context + compacted = self.compact(messages) + ctx.messages[:] = compacted + return dataclasses.replace(request_context, messages=compacted) + + +async def test_a_message_steered_at_the_request_that_stores_a_snapshot_comes_after_it(): + """The snapshot that a compaction stores at a step stands for what was stored before the message steered to the + turn at that step: a retried attempt, and the next turn, give the model the message after the snapshot's items, + once.""" + answers = {**ANSWERS} + requests: list[list[ModelMessage]] = [] + crashed: list[bool] = [] + + def model(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + requests.append(list(messages)) + latest = _prompts(messages)[-1] + if latest.startswith("look") and not _returned(messages[-1]): + return ModelResponse(parts=[ToolCallPart("lookup", {"about": "weather"}, tool_call_id="call-look")]) + if latest == "use Celsius" and not crashed: + crashed.append(True) + raise RuntimeError("the process crashed while the model answered") + return ModelResponse(parts=[TextPart(f"answer to: {latest}")]) + + def lookup(about: str) -> str: + """Look something up; meanwhile, a client steers the turn.""" + answers["publish"] = STEERING + return f"found {about}" + + mine: Any = Agent(FunctionModel(model), tools=[lookup], capabilities=[CompactAtTheStep()]) + chat = Chat(pydantic_ai.agent("pai_steered_snapshot", agent=mine)) + await chat("hi") + hosts = [Host(chat.backend, service=chat.name, answers=answers) for _ in range(2)] + invocation = request() + with pytest.raises(RuntimeError, match="crashed"): + await invoke(chat.obj, "chat", hosts[0], invocation, "look it up") + # The step after the tool's result stored the snapshot, then the steered message: the snapshot stands for what + # was stored before it, and does not hold it. + stored = await chat.read() + assert stored["snapshots"] == 1 + assert all("use Celsius" not in item for item in stored["items"]) + steps = [_agent.ITEMS.load(step) for step in stored["steps"]] + assert _said(steps, "use Celsius") == 1 and _returned(steps[-2]) + # The request had the summary, the turn, then the message after the tool's result. + sent = requests[-1] + assert _has_summary(sent[0]) and _said(sent, "use Celsius") == 1 + assert [type(p) for p in sent[-1].parts] == [ToolReturnPart, UserPromptPart] + # The retried attempt (the companion still lists the message) carries on from the snapshot and the message, once. + assert await invoke(chat.obj, "chat", hosts[1], invocation, "look it up") == "answer to: use Celsius" + retried = requests[-1] + assert _has_summary(retried[0]) and _said(retried, "use Celsius") == 1 + assert _sent(retried) == _sent(sent) + assert [send.handler for host in hosts for send in host.sends].count("taken") == 1 + # The next turn's history is the snapshot's items, then the message, the answer, and the turn's own prompt. + assert await chat("and then?") == "answer to: and then?" + following = requests[-1] + assert _has_summary(following[0]) and _said(following, "use Celsius") == 1 + assert _prompts(following)[-2:] == ["use Celsius", "and then?"] + stored = await chat.read() + assert stored["snapshots"] == 1 + assert _said([_agent.ITEMS.load(step) for step in stored["steps"]], "use Celsius") == 1 + + +async def test_a_feedback_and_the_message_steered_after_it_are_stored_once_through_a_retry_and_the_next_turn(): + """A retried attempt resumes from its last request, Pydantic AI's feedback on a reply it could not use, as a + request of its own; a later run merges it with the message steered after it. Either is found by its parts: + stored once, and the model gets each once.""" + from pydantic import BaseModel # pylint: disable=import-outside-toplevel + + class Plan(BaseModel): + title: str + + requests: list[list[ModelMessage]] = [] + crashed: list[bool] = [] + + def planner(messages: list[ModelMessage], info: AgentInfo) -> ModelResponse: + requests.append(list(messages)) + # The model gets the feedback as a message of the user, which says it is automated. + feedback = any("automated feedback" in prompt for prompt in _prompts([messages[-1]])) + if feedback and not crashed: + crashed.append(True) + raise RuntimeError("the process crashed while the model answered") + latest = [prompt for prompt in _prompts(messages) if "automated feedback" not in prompt][-1] + if latest.startswith("plan") and not feedback: + return ModelResponse(parts=[TextPart("not json")]) + return ModelResponse(parts=[TextPart(json.dumps({"title": latest}))]) + + chat = Chat(pydantic_ai.agent("pai_feedback_steered", model=FunctionModel(planner), output=Plan)) + invocation = request() + 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}) + 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. + feedbacks = [prompt for prompt in _prompts(requests[-1]) if "automated feedback" in prompt] + assert len(feedbacks) == 2 and _said(requests[-1], "use Celsius") == 1 + steps = [_agent.ITEMS.load(step) for step in (await chat.read())["steps"]] + stored = [p for m in steps if isinstance(m, ModelRequest) for p in m.parts] + assert [type(p).__name__ for p in stored] == ["RetryFeedbackPart", "UserPromptPart", "RetryFeedbackPart"] + assert [text for role, text in await chat.messages() if role == "user"] == ["plan one", "use Celsius", "plan two"] diff --git a/tests/agents_signals.py b/tests/agents_signals.py index 54e2c374..99980d83 100644 --- a/tests/agents_signals.py +++ b/tests/agents_signals.py @@ -24,8 +24,8 @@ from support import RESTATE_IMAGE # pylint: disable=import-error from restate import HarnessEnvironment from restate.agents import _signals -from restate.agents._api import Signal, Status, Watch -from restate.agents._signals import Started, WaitRequest +from restate.agents._api import Published, Signal, Status, Steer, Steered, Watch +from restate.agents._signals import STEERS, Started, WaitRequest @pytest.fixture(scope="session") @@ -44,6 +44,8 @@ async def start_turn(ctx: restate.Context, request: dict[str, str]) -> 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 await ctx.object_call(_signals.start, key=request["key"], arg=turn) @@ -52,9 +54,33 @@ async def close_turn(ctx: restate.Context, request: dict[str, str]) -> None: await ctx.object_call(_signals.close, key=request["key"], arg=request["invocation"]) +@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)) + + +@driver.handler() +async def publish_to(ctx: restate.Context, key: str) -> Published: + return await ctx.object_call(_signals.publish, key=key, arg=None) + + +@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"]) + + +@driver.handler() +async def reset_of(ctx: restate.Context, key: str) -> None: + await ctx.object_call(_signals.reset, key=key, arg=None) + + settled: list[tuple[str, str]] = [] """The turns the companion had the agent settle: its key, and the turn's invocation.""" +placeholders: list[tuple[str, dict[str, str]]] = [] +"""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.""" @@ -65,6 +91,11 @@ async def settle(ctx: restate.ObjectContext, invocation_id: str) -> bool: return True +@assistant.handler() +async def steered(ctx: restate.ObjectContext, message: dict[str, str]) -> None: + placeholders.append((ctx.key(), message)) + + failing = restate.Service("failing") @@ -210,3 +241,88 @@ async def sent() -> bool: 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}) + + +async def published(harness: HarnessEnvironment, k: str) -> Published: + return await harness.client.service_call(publish_to, arg=k) + + +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) + + +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"}) + 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)) + 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"]}) + assert "steering" not in await published(harness, k) + assert "steering" not in await harness.client.object_call(_signals.status, key=k, arg=None) + # 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 + + +async def test_a_turn_that_does_not_steer_holds_nothing_and_a_closed_turn_forgets_what_it_left( + harness: HarnessEnvironment, +): + 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"}) + 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"}) + + # 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"}) + 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"}) + 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) + 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"}) + 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. + assert await steer(harness, k, "past", "one more") == {"id": "past", "turn": None} + assert len((await published(harness, k)).get("steering", [])) == STEERS + # The same message sent again is held once. + 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)]}) + 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) diff --git a/tests/agents_steering.py b/tests/agents_steering.py new file mode 100644 index 00000000..1851e4cf --- /dev/null +++ b/tests/agents_steering.py @@ -0,0 +1,338 @@ +# +# 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 +# +"""Steering, in process: a message steered to a turn is stored once, as a step of the turn, and the model keeps +seeing it. The companion is tested against a server in ``agents_signals.py``, and the agent's ``steered`` in +``agents_steering_server.py``.""" + +# pylint: disable=missing-function-docstring,import-error + +import asyncio +import json +import typing +from collections.abc import Awaitable, Callable +from typing import Any + +import pytest +from agents_conversation import ITEMS, PUBLISHED +from sessions_database import Backend +from sessions_object import Host, Send, request + +from restate import sessions +from restate.agents import _messages +from restate.agents._api import Steer +from restate.agents._common import Check +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 +from restate.sessions.internal import Schema, invoke + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + +Run = Callable[[sessions.SessionContext, Conversation], Awaitable[str]] + + +class Asker: + """An object whose handler answers each prompt with ``run``, steered with what ``inbox`` holds.""" + + def __init__(self, run: Run, *, window: int | None = None, check: Check | None = None): + self.obj = sessions.SessionObject("asker") + 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))) + + async def ran(conversation: Conversation) -> str: + return await run(ctx, conversation) + + return await answer_once(ctx, text, ITEMS, ran, window, check=check) + + @self.obj.view + async def messages(ctx: sessions.ViewContext, wanted: Any = None) -> Any: + return await _messages.page(ctx.db, ITEMS, wanted) + + async def ask(self, text: str, invocation: Any = None) -> Any: + host = Host(self.backend, service="asker", answers=PUBLISHED) + self.hosts.append(host) + return await invoke(self.obj, "ask", host, invocation or request(), text) + + async def messages(self, wanted: Any = None) -> Any: + return await invoke(self.obj, "messages", Host(self.backend, service="asker").view(), request(), wanted) + + def taken(self) -> list[Send]: + return [send for host in self.hosts for send in host.sends if send.handler == "taken"] + + +async def tool_round(conversation: Conversation, call_id: str = "c1") -> None: + """The model's reply with one tool call, and the call's result: the turn's next step is a boundary.""" + call = {"id": call_id, "name": "lookup", "input": {}} + await conversation.turn.append_reply({"role": "assistant", "text": "", "calls": [call]}) + await conversation.turn.append({"result": call_id, "output": "found"}) + + +def texts(items: list[Any]) -> list[str]: + return [item["text"] for item in items] + + +async def test_a_steered_message_is_stored_once_as_a_step_of_the_turn_and_taken_with_its_store(): + took: list[list[str]] = [] + seen: list[list[Any]] = [] + + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + await tool_round(conversation) + took.append(texts(await take_steers(ctx, conversation))) + # A later boundary, with the same message still listed: it is not taken again. + took.append(texts(await take_steers(ctx, conversation))) + seen.append(conversation.history()) + turn = conversation.turn + assert [entry.steer for entry in turn.steered()] == ["s1"] + # The transcript's entry of the turn has it too, as the next turn loads it. + assert [entry.steer for entry in turn.prompt.steers] == ["s1"] + # A steered message asks for no answer of its own: the turn goes on from its prompt. + assert turn.latest_ask() is turn.prompt + await turn.append_reply({"role": "assistant", "text": "done, in Celsius"}) + return "done, in Celsius" + + asker = Asker(run) + asker.inbox = [{"id": "s1", "text": "use Celsius"}] + assert await asker.ask("the weather?") == "done, in Celsius" + assert took == [["use Celsius"], []] + # The model's next request has the message, after the tool's result. + assert seen[0][-1] == {"role": "user", "text": "use Celsius"} + assert seen[0][-2] == {"result": "c1", "output": "found"} + # 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", + ["s1"], + "asker:alice", + "taken:s1", + ) + + +async def test_no_message_is_taken_while_a_tool_call_has_no_result_nor_outside_a_chat_turn(): + took: list[list[Any]] = [] + + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + call = {"id": "c1", "name": "lookup", "input": {}} + await conversation.turn.append_reply({"role": "assistant", "text": "", "calls": [call]}) + took.append(await take_steers(ctx, conversation)) + await conversation.turn.append({"result": "c1", "output": "found"}) + _INBOX.delete(ctx) + took.append(await take_steers(ctx, conversation)) + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + asker = Asker(run) + asker.inbox = [{"id": "s1", "text": "use Celsius"}] + assert await asker.ask("the weather?") == "done" + assert took == [[], []] + assert asker.taken() == [] + + +async def test_a_retried_attempt_finds_the_steered_message_stored_and_never_takes_it_again(): + attempts: list[list[str]] = [] + + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + if not conversation.turn.items: + await tool_round(conversation) + took = await take_steers(ctx, conversation) + attempts.append([entry.steer or "" for entry in conversation.turn.steered()]) + if len(attempts) == 1: + assert texts(took) == ["use Celsius"] + raise RuntimeError("the process crashed after storing the message") + assert took == [] + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + asker, invocation = Asker(run), request() + asker.inbox = [{"id": "s1", "text": "use Celsius"}] + with pytest.raises(RuntimeError, match="crashed"): + await asker.ask("the weather?", invocation) + assert await asker.ask("the weather?", invocation) == "done" + # The retry loaded it as stored, marked as steered. + assert attempts == [["s1"], ["s1"]] + assert len(asker.taken()) == 1 + + +async def test_an_earlier_turn_reaches_the_model_as_its_prompt_its_steered_messages_and_its_answer(): + seen: list[list[Any]] = [] + + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + seen.append(conversation.history()) + if conversation.turn.prompt.item["text"] == "first": + await tool_round(conversation) + await take_steers(ctx, conversation) + await conversation.turn.append_reply({"role": "assistant", "text": "ok"}) + return "ok" + + # A window of 0: every earlier turn is its messages and answer. + asker = Asker(run, window=0) + asker.inbox = [{"id": "s1", "text": "and then?"}] + await asker.ask("first") + asker.inbox = [] + await asker.ask("second") + assert seen[1] == [ + {"role": "user", "text": "first"}, + {"role": "user", "text": "and then?"}, + {"role": "assistant", "text": "ok"}, + {"role": "user", "text": "second"}, + ] + + +async def test_a_turn_counts_its_steered_messages_toward_the_budget_as_soon_as_it_is_answered(): + loaded: list[Conversation] = [] + + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + loaded.append(conversation) + await tool_round(conversation) + await take_steers(ctx, conversation) + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + # A window of 0: the turn just answered is before the next turn's window. + asker = Asker(run, window=0) + asker.inbox = [{"id": "s1", "text": "use Celsius"}] + await asker.ask("the weather?") + [conversation] = loaded + # Its message and answer (16 characters) fit in 20; with the steered message (11 more), they do not: the + # summary due now is asked for now, rather than after the next turn. + conversation.budget = 5 + assert conversation.overflowing() + + +async def test_a_turn_stopped_while_it_waits_to_store_steered_messages_takes_none(): + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + await tool_round(conversation) + stop = Stop() + _STOP.set(ctx, stop) + # Another write holds the ledger: the turn waits for it to store the message, and stops meanwhile. + async with conversation._ledger.lock: # pylint: disable=protected-access + taking = asyncio.ensure_future(take_steers(ctx, conversation)) + for _ in range(10): + await asyncio.sleep(0) + stop.stop(RuntimeError("stopped")) + assert await taking == [] + # Not the turn's: its agent's ``steered`` answers it as a turn of its own. + assert not await steer_stored(ctx.db, "s1") + assert conversation.turn.steered() == [] + _STOP.delete(ctx) + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + asker = Asker(run) + asker.inbox = [{"id": "s1", "text": "use Celsius"}] + assert await asker.ask("the weather?") == "done" + assert asker.taken() == [] + + +async def test_the_output_check_asks_again_once_after_a_steered_message(): + replies = iter(["prose", "still prose"]) + asked: list[str] = [] + + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + turn = conversation.turn + if not turn.items: + await tool_round(conversation) + await take_steers(ctx, conversation) + asked.append(turn.latest_ask().item["text"]) + reply = next(replies) + await turn.append_reply({"role": "assistant", "text": reply}) + return reply + + asker = Asker(run, check=lambda answer, _: "Reply again, with JSON only." if answer == "prose" else None) + asker.inbox = [{"id": "s1", "text": "be brief"}] + assert await asker.ask("plan") == "still prose" + # The steered message did not count as the ask: the check asked again, once. + assert asked == ["plan", "Reply again, with JSON only."] + + +async def test_a_dropped_reply_drops_the_message_steered_after_it_which_is_then_free_to_be_answered(): + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + turn = conversation.turn + await tool_round(conversation) + await take_steers(ctx, conversation) + assert await steer_stored(ctx.db, "s1") + [reply] = turn.replies() + await turn.drop(reply) + assert turn.steered() == [] + assert turn.prompt.steers == [] + assert not await steer_stored(ctx.db, "s1") + [[rows]] = await ctx.db.values(f"SELECT COUNT(*) FROM {STEERED}") + assert rows == 0 + await turn.append_reply({"role": "assistant", "text": "again"}) + return "again" + + asker = Asker(run) + asker.inbox = [{"id": "s1", "text": "use Celsius"}] + assert await asker.ask("the weather?") == "again" + + +async def test_the_summary_view_shows_a_steered_message_in_its_turn_once_from_a_cursor(): + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + await tool_round(conversation) + await take_steers(ctx, conversation) + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + asker = Asker(run) + asker.inbox = [{"id": "s1", "text": "use Celsius"}] + await asker.ask("the weather?") + page = await asker.messages({"view": "summary"}) + shown = [(m["role"], m["parts"][0]["text"]) for m in page["messages"]] + assert shown == [("user", "the weather?"), ("user", "use Celsius"), ("assistant", "done")] + # The full view shows it within the turn, between the tool's call and the answer. + full = await asker.messages() + assert [m["role"] for m in full["messages"]] == ["user", "assistant", "user", "assistant"] + # From a cursor before it, it comes once; from the cursor after it, not again. + since = await asker.messages({"view": "summary", "cursor": 1}) + assert ("user", "use Celsius") in [(m["role"], m["parts"][0]["text"]) for m in since["messages"]] + after = await asker.messages({"view": "summary", "cursor": since["cursor"]}) + assert after["messages"] == [] + + +async def test_a_database_of_the_schema_before_steering_takes_its_table_as_it_opens(): + before = Schema(SCHEMA.name, list(SCHEMA.steps[:-1])) + obj = sessions.SessionObject("asker") + backend = Backend() + + @obj.handler + async def old(ctx: sessions.SessionContext) -> None: + await before.ensure(ctx.db) + + @obj.handler + async def new(ctx: sessions.SessionContext) -> bool: + return await steer_stored(ctx.db, "s1") + + await invoke(obj, "old", Host(backend, service="asker"), request()) + assert await invoke(obj, "new", Host(backend, service="asker"), request()) is False + + +async def test_steer_dedupes_by_id_within_the_messages_given(): + async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: + await tool_round(conversation) + took = await take_steers(ctx, conversation) + assert texts(took) == ["once"] + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + asker = Asker(run) + asker.inbox = typing.cast(list[Steer], [{"id": "s1", "text": "once"}, {"id": "s1", "text": "once"}]) + assert await asker.ask("hi") == "done" diff --git a/tests/agents_steering_server.py b/tests/agents_steering_server.py new file mode 100644 index 00000000..6ace06d5 --- /dev/null +++ b/tests/agents_steering_server.py @@ -0,0 +1,260 @@ +# +# 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 +# +"""Steering an agent's turn, with a scripted framework, against a real Restate server: a message sent to ``steer`` +joins the turn at its next step, or is answered as a turn of its own, once, in the order it was sent.""" + +# pylint: disable=missing-function-docstring,import-error + +import asyncio +import typing +import uuid +from collections.abc import Awaitable, Callable + +import httpx2 +import pytest +from agents_conversation import ITEMS +from support import RESTATE_IMAGE, Gate + +import restate +from restate import HarnessEnvironment, sessions +from restate.agents._api import Messages, Status +from restate.agents._conversation import Conversation, answer_once +from restate.agents._object import AgentObject +from restate.agents._steering import take_steers +from restate.agents._turn import stop_signal_of, tool_call_ended, tool_call_started + + +@pytest.fixture(scope="session") +def anyio_backend(): + return "asyncio" + + +pytestmark = [pytest.mark.anyio] + +gate = Gate() +requests: dict[str, list[list[str]]] = {} +"""What each key's model was asked, request by request: the texts of the messages of the user.""" + + +def framework(name: str, steers: bool) -> AgentObject: + """An agent whose scripted model calls one tool for ``look ...``, which waits on ``gate``, then answers with + the messages of the user it was given since the prompt; ``wait`` waits on ``gate`` and echoes.""" + + async def answer(ctx: sessions.SessionContext, text: str) -> str: + async def work(conversation: Conversation) -> str: + turn = conversation.turn + if text.startswith("look ") and not turn.items: + call = {"id": "call-1", "name": "lookup", "input": {}} + await turn.append_reply({"role": "assistant", "text": "", "calls": [call]}) + await tool_call_started(ctx) + stop = stop_signal_of(ctx) + assert stop is not None + await stop.guard(gate.wait()) + await turn.append({"result": "call-1", "output": "found"}) + await tool_call_ended(ctx) + elif text == "wait": + stop = stop_signal_of(ctx) + assert stop is not None + await stop.guard(gate.wait()) + # The step before the model's next request: it takes the messages steered to the turn. + await take_steers(ctx, conversation) + history = conversation.history() + asked = [item["text"] for item in history if item.get("role") == "user"] + requests.setdefault(ctx.key, []).append(asked) + heard = [entry.item["text"] for entry in turn.steered()] + reply = f"{text}: {', '.join(heard)}" if heard else f"echo: {text}" + await turn.append_reply({"role": "assistant", "text": reply}) + return reply + + return await answer_once(ctx, text, ITEMS, work) + + agent = AgentObject(name) + agent.chat(answer, items=ITEMS, steers=steers) + return agent + + +steerer = framework("steerer", steers=True) +plain = framework("plain", steers=False) + + +@pytest.fixture(scope="session") +async def harness(): + async with restate.create_test_harness(restate.app([steerer, plain]), restate_image=RESTATE_IMAGE) as harness: + yield harness + + +TEXT = {"content-type": "text/plain"} + + +async def post(harness: HarnessEnvironment, path: str, body: object = None, **headers: str) -> httpx2.Response: + async with httpx2.AsyncClient(timeout=60) as ingress: + if isinstance(body, str): + return await ingress.post(f"{harness.ingress_url}/{path}", content=body, headers={**TEXT, **headers}) + return await ingress.post(f"{harness.ingress_url}/{path}", json=body, headers=headers) + + +async def send(harness: HarnessEnvironment, agent: str, key: str, text: str) -> str: + sent = await post(harness, f"{agent}/{key}/chat/send", text) + return typing.cast(str, sent.json()["invocationId"]) + + +async def steer(harness: HarnessEnvironment, agent: str, key: str, text: str) -> dict[str, typing.Any]: + # With an idempotency key, as a client sends it: a retried request is the same message. + steered = await post(harness, f"{agent}/{key}/steer", text, **{"idempotency-key": str(uuid.uuid4())}) + assert steered.status_code == 200, steered.text + return typing.cast(dict[str, typing.Any], steered.json()) + + +async def attach(harness: HarnessEnvironment, invocation_id: str) -> httpx2.Response: + async with httpx2.AsyncClient(timeout=60) as ingress: + return await ingress.get(f"{harness.ingress_url}/restate/invocation/{invocation_id}/attach") + + +async def status(harness: HarnessEnvironment, agent: str, key: str) -> Status: + return typing.cast(Status, (await post(harness, f"{agent}/{key}/status")).json()) + + +async def until(check: Callable[[], Awaitable[bool]]) -> None: + for _ in range(300): + if await check(): + return + await asyncio.sleep(0.05) + raise AssertionError("never") + + +async def answering(harness: HarnessEnvironment, agent: str, key: str, invocation: str) -> None: + async def started() -> bool: + turn = (await status(harness, agent, key))["answering"] + return turn is not None and turn["invocationId"] == invocation + + await until(started) + + +async def summary(harness: HarnessEnvironment, agent: str, key: str) -> list[tuple[str, str]]: + page = typing.cast(Messages, (await post(harness, f"{agent}/{key}/messages", {"view": "summary"})).json()) + return [(m["role"], typing.cast(str, m["parts"][0].get("text"))) for m in page["messages"]] + + +async def settled(harness: HarnessEnvironment, agent: str, key: str) -> None: + """Wait until the key answered every message queued so far: a chat sent now is answered after them.""" + async with httpx2.AsyncClient(timeout=60) as ingress: + done = await ingress.post(f"{harness.ingress_url}/{agent}/{key}/chat", content="done", headers=TEXT) + assert done.status_code == 200, done.text + + +async def test_a_message_steered_during_a_tool_call_joins_the_turn_at_its_next_step(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + gate.close() + invocation = await send(harness, "steerer", key, "look up the weather") + await answering(harness, "steerer", key, invocation) + steered = await steer(harness, "steerer", key, "use Celsius") + assert steered["turn"] == invocation + assert [one["id"] for one in (await status(harness, "steerer", key)).get("steering", [])] == [steered["id"]] + gate.open() + answer = await attach(harness, invocation) + assert answer.json() == "look up the weather: use Celsius" + # The model's request after the tool's result had the message. + assert requests[key][-1] == ["look up the weather", "use Celsius"] + + async def taken() -> bool: + return "steering" not in await status(harness, "steerer", key) + + await until(taken) + # Its queued turn of its own did nothing: the turn took it. The summary shows it in the turn. + await settled(harness, "steerer", key) + assert await summary(harness, "steerer", key) == [ + ("user", "look up the weather"), + ("user", "use Celsius"), + ("assistant", "look up the weather: use Celsius"), + ("user", "done"), + ("assistant", "echo: done"), + ] + + +async def test_a_message_steered_after_the_last_step_is_answered_next_in_the_order_it_was_sent( + harness: HarnessEnvironment, +): + key = str(uuid.uuid4()) + gate.close() + invocation = await send(harness, "steerer", key, "wait") + await answering(harness, "steerer", key, invocation) + # The turn reaches no step after them: the steered message is answered next, then the chat sent after it. + assert (await steer(harness, "steerer", key, "and also"))["turn"] == invocation + after = await send(harness, "steerer", key, "then this") + gate.open() + assert (await attach(harness, invocation)).json() == "echo: wait" + assert (await attach(harness, after)).json() == "echo: then this" + assert await summary(harness, "steerer", key) == [ + ("user", "wait"), + ("assistant", "echo: wait"), + ("user", "and also"), + ("assistant", "echo: and also"), + ("user", "then this"), + ("assistant", "echo: then this"), + ] + + +async def test_a_message_steered_with_no_turn_being_answered_starts_one(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + gate.open() + assert (await steer(harness, "steerer", key, "hello"))["turn"] is None + await settled(harness, "steerer", key) + assert await summary(harness, "steerer", key) == [ + ("user", "hello"), + ("assistant", "echo: hello"), + ("user", "done"), + ("assistant", "echo: done"), + ] + + +async def test_an_agent_that_takes_no_steered_messages_answers_each_as_a_turn_of_its_own( + harness: HarnessEnvironment, +): + key = str(uuid.uuid4()) + gate.close() + invocation = await send(harness, "plain", key, "look up the weather") + await answering(harness, "plain", key, invocation) + assert (await steer(harness, "plain", key, "use Celsius"))["turn"] is None + gate.open() + assert (await attach(harness, invocation)).json() == "echo: look up the weather" + await settled(harness, "plain", key) + assert await summary(harness, "plain", key) == [ + ("user", "look up the weather"), + ("assistant", "echo: look up the weather"), + ("user", "use Celsius"), + ("assistant", "echo: use Celsius"), + ("user", "done"), + ("assistant", "echo: done"), + ] + + +async def test_a_message_steered_to_a_turn_killed_before_its_next_step_is_answered_once(harness: HarnessEnvironment): + key = str(uuid.uuid4()) + gate.close() + invocation = await send(harness, "steerer", key, "look twice") + await answering(harness, "steerer", key, invocation) + await steer(harness, "steerer", key, "first steer") + gate.open() + assert (await attach(harness, invocation)).json() == "look twice: first steer" + # A second turn is killed while its tool runs, before the step at which it would take the message. + gate.close() + killed = await send(harness, "steerer", key, "look again") + await answering(harness, "steerer", key, killed) + await steer(harness, "steerer", key, "second steer") + async with httpx2.AsyncClient(timeout=60) as admin: + response = await admin.patch(f"{harness.admin_api_url}/invocations/{killed}/kill") + assert response.status_code < 300, response.text + gate.open() + await settled(harness, "steerer", key) + shown = await summary(harness, "steerer", key) + # The killed turn never took the message: it is answered once, as a turn of its own. + assert [text for role, text in shown if role == "user"].count("second steer") == 1 + assert ("assistant", "echo: second steer") in shown