From bd1e55369456e002b91645c2ac1193e8df810221 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 08:08:56 +0000 Subject: [PATCH 01/13] Agents: steering, a message joins the turn being answered at its next step A key answers one message at a time: a message sent while a turn runs waits in Restate's queue and becomes the next turn, and to change course a client stops the turn and sends again. This adds a second way in, which leaves `chat` as it is: the agent's new `steer` view. The turn being answered takes a steered message at its next step (once the model's tool calls have their results, before its next request) and the model's next request has it; a message that comes after the turn's last step is answered next, as a turn of its own, in the order it was sent. Delivery, in the companion (`signals`) - `steer({id, text})` sends the agent's new private `steered` handler the message at once (idempotency key `steer:`): the message takes its place in the key's queue as it is sent, so it is answered in order with the chats sent after it, and even if the turn dies. - If the turn being answered takes steered messages (`Started.steers`, a field the turn's `start` adds, like `settles`), the companion also holds the message in its `inbox`, up to 16 per turn (`STEERS`). - `publish`, which a turn calls at each tool call's start and end, returns that turn's inbox (`steering`, absent when empty). So a turn learns of a message at no cost of its own: nobody steering costs nothing. - `taken(ids)` forgets messages a turn stored; `_closed` forgets those a closed turn left; `reset` (clear) forgets them all. `status` shows the inbox (`steering`). In the turn - `take_steers(ctx, conversation)` (in `restate.agents.internal`), called by an integration at a step, stores the messages the last `publish` listed that are not stored, as messages of the user that are steps of the turn, and returns them for the model's next request. It does nothing outside a `chat` turn, for an agent whose framework does not say it takes them (`chat(..., steers=True)`), for a background agent, while a tool call has no result, or once the turn stops. - `Turn.steer` stores them in one transaction, each once by its id (in the new table `restate_agent_steers`, an appended migration), with the companion's `taken` sent from the transaction's outbox: once, with the commit. The store is shielded from a stop cancelling it mid-commit. - `steered` answers the message as a turn of its own, unless the database has it (a turn took it): so each message is answered once. A dropped reply drops the messages steered after it, which `steered` then answers. - `latest_ask` skips steered messages: the output check still asks again, ADK still resumes the prompt's invocation, and a graph keeps its run. - The model keeps seeing them: an earlier turn is its prompt, the messages steered to it, and its answer, in `history`, the budget and the summaries. The summary view shows them under their ids. Also `AgentClient.steer` and the terminal chat's `/steer`. Tests: the core in process (stored once, through a crash, after a drop, with the output check, in the history and the summary view, the migration), the companion on a server (the inbox, the placeholder, the cap, a closed turn, clear), and an agent with a scripted framework on a server (a message steered during a tool call reaches the next request; one after the last step, or to an agent that takes none, or with no turn, is answered as a turn of its own in order; a turn killed before its step leaves it to its own turn). 569 passed; the 3 failures are the usual host.docker.internal ones. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/__main__.py | 10 +- python/restate/agents/_api.py | 47 +++- python/restate/agents/_conversation.py | 179 ++++++++++++--- python/restate/agents/_messages.py | 33 ++- python/restate/agents/_object.py | 51 ++++- python/restate/agents/_signals.py | 148 ++++++++++++- python/restate/agents/_steering.py | 71 ++++++ python/restate/agents/_turn.py | 50 ++++- python/restate/agents/client.py | 16 +- python/restate/agents/internal.py | 5 + tests/agents_signals.py | 109 ++++++++- tests/agents_steering.py | 293 +++++++++++++++++++++++++ tests/agents_steering_server.py | 260 ++++++++++++++++++++++ 13 files changed, 1211 insertions(+), 61 deletions(-) create mode 100644 python/restate/agents/_steering.py create mode 100644 tests/agents_steering.py create mode 100644 tests/agents_steering_server.py 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..4cd79e6b 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): @@ -624,16 +641,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 +690,60 @@ async def ask_again(self, text: str) -> None: """ await self.append(self._ledger.items.prompt(text)) + async def steer( + self, messages: Sequence[Steer], sent: Callable[[Transaction, list[str]], None] | None = None + ) -> list[Entry]: + """Store the steered ``messages`` that are not stored yet as messages of the user in the turn, and return them. + + Each is stored once, by its id, as a step of the turn after those + before: the model reads it in its next request, and in every turn's + history after, between the turn's message and its answer. ``sent`` + adds what goes out with the transaction that stores them (the + companion's ``taken``), which commits once. + + The store finishes even if the turn stops meanwhile: the turn then + closes after it, at the positions after its messages. + """ + ledger = self._ledger + + async def store() -> list[Entry]: + async with ledger.lock: + self._check_writable() + wanted = list({message["id"]: message for message in messages}.values()) + if not wanted: + return [] + ids = [message["id"] for message in wanted] + marks = ", ".join("?" for _ in ids) + stored = await ledger.db.values(f"SELECT id FROM {STEERED} WHERE id IN ({marks})", ids) + known = {str(row[0]) for row in stored} + new = [message for message in wanted if message["id"] not in known] + if not new: + return [] + n = ledger.next + added = [ + Entry(n + i, ledger.items.prompt(message["text"]), 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( + f"INSERT INTO {STEPS} (n, turn, seq, item, reply, agent) VALUES (?, ?, ?, ?, NULL, NULL)", + [entry.n, self.n, ledger.seq + i, ledger.items.dump(entry.item)], + ) + await tx.execute( + f"INSERT INTO {STEERED} (n, turn, id) VALUES (?, ?, ?)", [entry.n, self.n, entry.steer] + ) + if sent is not None: + sent(tx, [typing.cast(str, entry.steer) for entry in added]) + ledger.seq += len(added) + ledger.next = n + len(added) + self._add(added) + return added + + # 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(store()) + 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 +780,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,6 +795,7 @@ 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) @@ -960,13 +1042,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 +1064,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 @@ -1094,10 +1164,11 @@ def _before(self) -> list[Entry]: 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.""" + """Return the size of ``turn`` as the transcript has it, in characters: its message, the messages steered + to it, 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 "") + return len(turn.said[1] if turn.said else "") + _steered_size(turn) + len(turn.answer or "") def _latest(self, turns: list[Entry], characters: int) -> list[Entry]: """Return the latest of ``turns`` that ``characters`` holds, in order.""" @@ -1151,8 +1222,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 +1239,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 +1331,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,6 +1347,7 @@ 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): @@ -1278,7 +1355,7 @@ async def _transcript(db: DatabaseReader, items: Items, first: int, bounded: tup 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 "") + left -= len(entry.said[1] if entry.said else "") + _steered_size(entry) + len(entry.answer or "") if left < 0: return read[::-1] if bounded is None or len(rows) < PAGE: @@ -1286,6 +1363,55 @@ 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 n, item, reply, agent FROM {STEPS} WHERE n > ? ORDER BY n", [after]) + steered = { + int(typing.cast(int, n)): str(i) + for n, i in await db.values(f"SELECT n, id FROM {STEERED} WHERE n > ?", [after]) + } + return [ + Entry( + int(typing.cast(int, n)), + load_item(items, str(item)), + reply=_int(reply), + agent=_str(agent), + steer=steered.get(int(typing.cast(int, n))), + ) + for n, item, reply, agent 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 _steered_size(turn: Entry) -> int: + """Return the size, in characters, of the messages steered to ``turn`` (a prompt of the transcript).""" + return sum(len(steered.said[1]) if steered.said else 0 for steered in turn.steers) + + +async def steer_stored(db: DatabaseReader, steer_id: str) -> bool: + """Return whether a turn took the steered message ``steer_id``: one read, of its unique id.""" + await SCHEMA.ensure(db) + return await db.get(f"SELECT 1 FROM {STEERED} WHERE id = ?", [steer_id]) is not None + + def _str(value: object) -> str | None: return None if value is None else str(value) @@ -1444,7 +1570,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..a3f70ba3 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,16 +174,25 @@ 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: @@ -194,15 +203,21 @@ def _summarized(items: Items, rows: typing.Sequence[typing.Sequence[typing.Any]] 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] + for n, seq, stored in steered: + 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 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,7 +233,9 @@ 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] + 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] 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]) return { diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index 9f69e5de..6b9d2acc 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -31,6 +31,7 @@ Resolved, Signal, Status, + Steered, Stopped, View, Waiting, @@ -46,7 +47,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 +172,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 +199,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 +209,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 +362,7 @@ def chat( summarize: Summarize | None = None, snapshots: bool = False, compact_history: CompactHistory | None = None, + steers: bool = False, ) -> A: ... @typing.overload @@ -368,6 +376,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 +390,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 +435,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 +456,7 @@ def decorate( summarize=summarize, snapshots=snapshots, compact_history=compact_history, + steers=steers, ) return decorate @@ -454,6 +469,7 @@ def decorate( summarize=summarize, snapshots=snapshots, compact_history=compact_history, + steers=steers, ) AGENT.on(self, self._settings) answering = answer @@ -463,9 +479,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: _signals.SteerRequest) -> 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 +530,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 +542,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 +623,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..f499f5e6 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,12 @@ 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. + +Each one the turn takes is one more model request in it, so they keep a turn +that someone steers on and on within bounds.""" + class _WaitRequest(TypedDict): id: str @@ -94,6 +113,25 @@ 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 SteerRequest(TypedDict): + """A message to steer to the turn being answered: its id (the ``steer`` request's) and its text.""" + + id: str + text: str + + +class _Inboxed(TypedDict): + """A steered message the turn ``turn`` (an invocation id) has not taken yet, and when it was sent.""" + + id: str + text: str + since: str + turn: str def _shown(turn: Started | None) -> Answering | None: @@ -101,6 +139,19 @@ 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 steered messages that their turns have not taken yet, oldest first.""" + return await _get(ctx, "inbox", typing.cast(list[_Inboxed], [])) + + +async def _steering(ctx: Context) -> list[_Inboxed]: + """Return the steered messages that the turn being answered has not taken yet, oldest first.""" + answering = await _get(ctx, "turn", typing.cast(Started | None, None)) + if answering is None: + return [] + return [one for one in await _inbox(ctx) if one["turn"] == answering["invocationId"]] + + 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 +193,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 that it did not take are forgotten here: + each 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 +205,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,10 +216,61 @@ 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, request: SteerRequest) -> 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 = request["id"], request["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} + 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: + """Forget the steered messages ``ids``, which their turn has stored, and publish the change. + + The turn sends it with the transaction that stores them, so it comes once + they are stored, and once only. + """ + inbox = await _inbox(ctx) + left = [one for one in inbox if one["id"] not in ids] + if len(left) != len(inbox): + _set_inbox(ctx, left) + await _publish(ctx) + + @signals.handler(kind="shared", inactivity_timeout=PARKED) async def wait(ctx: ObjectSharedContext, request: WaitRequest) -> Any: """Wait for the signal ``name`` as wait ``id``, and return its value.""" @@ -276,11 +386,16 @@ 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 { + """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.""" + shown: Status = { "answering": _shown(await _get(ctx, "turn", typing.cast(Started | None, None))), "waiting": [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))], } + steering: list[Steering] = [{"id": one["id"], "since": one["since"]} for one in await _steering(ctx)] + 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 +432,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 not taken 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 +449,18 @@ 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)} + steering: list[Steer] = [{"id": one["id"], "text": one["text"]} for one in await _steering(ctx)] + 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..a4bde9a3 --- /dev/null +++ b/python/restate/agents/_steering.py @@ -0,0 +1,71 @@ +# +# 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, Entry +from restate.agents._turn import steering_of, stop_signal_of +from restate.sessions import SessionContext, Transaction + + +async def take_steers(ctx: SessionContext, conversation: Conversation) -> list[Entry]: + """Take the messages steered to the turn being answered, at a step: store each once, and return the new ones. + + Call it right before a model request that follows the model's reply and + its tools' results: the integration then gives the model the returned + entries' items (messages of the user), after what the request already + holds. It returns nothing, at no cost, unless the turn runs in a ``chat`` + handler of an agent that takes steered messages, a message 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. + """ + steering = steering_of(ctx) + if steering is None or not steering.messages: + return [] + stop = stop_signal_of(ctx) + if stop is not None and stop.stopped: + return [] + turn = conversation.running() + if turn is None or turn.open_calls(): + return [] + taken = {entry.steer for entry in turn.steered()} + fresh = [message for message in steering.messages if message["id"] not in taken] + if not fresh: + return [] + companion = _signals.signals_of(ctx) + + def sent(tx: Transaction, ids: list[str]) -> None: + # With the transaction that stores them: so they are forgotten once stored, and once. + tx.send(_signals.taken, ids, key=companion, idempotency_key=f"taken:{ids[0]}") + + return await turn.steer(fresh, sent) diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py index 51f8337f..7b1b466f 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 Steering: + """The messages steered to a ``chat`` turn that it has not taken yet, as the companion last told it. + + 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", [])) + + +_STEERING: ScopeKey[Steering] = ScopeKey("steering") + + +def steering_of(ctx: SessionContext) -> Steering | None: + """Return the steered messages of the ``chat`` turn running on ``ctx``, if it takes them: ``None`` outside one.""" + return _STEERING.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)) + steering = _STEERING.get(ctx) + if steering is not None: + steering.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. + _STEERING.set(ctx, Steering()) # 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/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..0d924332 100644 --- a/python/restate/agents/internal.py +++ b/python/restate/agents/internal.py @@ -33,6 +33,9 @@ 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 messages steered to it, which + the next model request gets (an integration says it calls it with + ``chat(..., steers=True)``). """ from restate.agents._common import ( @@ -70,6 +73,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 +127,7 @@ "signals", "stop_signal_of", "stored_answer", + "take_steers", "tool_call_ended", "tool_call_started", "wait_for", diff --git a/tests/agents_signals.py b/tests/agents_signals.py index 54e2c374..cc8d2737 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, 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(_signals.SteerRequest, 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,77 @@ 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) + + +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 diff --git a/tests/agents_steering.py b/tests/agents_steering.py new file mode 100644 index 00000000..6738ce66 --- /dev/null +++ b/tests/agents_steering.py @@ -0,0 +1,293 @@ +# +# 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 and the agent's ``steered`` are tested against a server in ``agents_object.py``.""" + +# pylint: disable=missing-function-docstring,import-error + +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._conversation import ( + SCHEMA, + STEERED, + Conversation, + Entry, + answer_once, + steer_stored, +) +from restate.agents._steering import take_steers +from restate.agents._turn import _STEERING, Steering +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 ``steering`` holds.""" + + def __init__(self, run: Run, *, window: int | None = None, check: Callable[[str], str | None] | None = None): + self.obj = sessions.SessionObject("asker") + self.backend = Backend() + self.steering: 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. + _STEERING.set(ctx, Steering(revision=0, messages=list(self.steering))) + + 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(entries: list[Entry]) -> list[str]: + return [entry.item["text"] for entry in entries] + + +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"] + # 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.steering = [{"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[Entry]] = [] + + 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"}) + _STEERING.delete(ctx) + took.append(await take_steers(ctx, conversation)) + await conversation.turn.append_reply({"role": "assistant", "text": "done"}) + return "done" + + asker = Asker(run) + asker.steering = [{"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.steering = [{"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.steering = [{"id": "s1", "text": "and then?"}] + await asker.ask("first") + asker.steering = [] + 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_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.steering = [{"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 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.steering = [{"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.steering = [{"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.steering = 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 From 86644f358fa4b6eabd1f46459d01e22c72840e9f Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 08:16:04 +0000 Subject: [PATCH 02/13] Agents: steering's docs, its traffic measured, and the mutants that guard it Docs (the reference pages) - turns.md: a Steering section (the steer view, the step, where a message waits, how the turn learns of it from publish, taken once, answered as a turn of its own otherwise, a taken message is its turn's, what the model and clients see, limits), the words "step" and "steered message", the status and summary views, the revisions, Started.steers alongside settles, and the guarantees: a turn is joined only by a steered message, at a step, and each is answered once. - conversation.md: restate_agent_steers (the first appended migration), a steered message as a step, earlier turns as prompt, steered messages and answer in history, budgets and summaries. - README.md: the bullet, the view, the module. - learnings.md: D6 to D10, why a view of its own, the placeholder queued as the message is sent, the inbox on publish, in place not a restart, and a late message as the next turn. - porting.md: the TypeScript step; access-patterns.md: the traffic. Traffic, measured (scripts/steering_traffic.py, against a real server, any checkout): a turn of one tool call that nobody steers costs what it costs before steering, 47 journal entries and 5 companion calls (start, three publishes, close), on this branch and on main alike. A turn that takes a steered message: 57 entries and 6 calls, the sixth 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). Mutants (scripts/resurrect.py, durability, from no fix): the placeholder answering a taken message, latest_ask counting steered messages, history dropping them, a message taken while a call is open, publish without the inbox, taken never sent, a drop keeping their rows, and the summary view hiding them. Co-Authored-By: Claude Opus 5.5 --- docs/agents/README.md | 5 + docs/agents/access-patterns.md | 17 +++ docs/agents/conversation.md | 40 ++++-- docs/agents/learnings.md | 5 + docs/agents/porting.md | 11 ++ docs/agents/turns.md | 104 ++++++++++++++-- scripts/resurrect.py | 66 ++++++++++ scripts/steering_traffic.py | 215 +++++++++++++++++++++++++++++++++ 8 files changed, 444 insertions(+), 19 deletions(-) create mode 100644 scripts/steering_traffic.py diff --git a/docs/agents/README.md b/docs/agents/README.md index 5e9737ca..82ce7c6d 100644 --- a/docs/agents/README.md +++ b/docs/agents/README.md @@ -17,6 +17,10 @@ building it, and how to port it. - **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 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. @@ -77,6 +81,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..d02b08d2 100644 --- a/docs/agents/access-patterns.md +++ b/docs/agents/access-patterns.md @@ -290,6 +290,23 @@ 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, 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..52c97417 100644 --- a/docs/agents/conversation.md +++ b/docs/agents/conversation.md @@ -21,8 +21,9 @@ records). 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)). - **Every row belongs to one turn,** and only the latest turn, while it has @@ -32,14 +33,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 +125,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 +161,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.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 +316,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 +364,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 +554,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/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..6e2aff77 100644 --- a/docs/agents/porting.md +++ b/docs/agents/porting.md @@ -108,6 +108,17 @@ 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`), `taken`, `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, 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/turns.md b/docs/agents/turns.md index 80fcc716..9720be97 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 @@ -59,6 +61,10 @@ 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, which the + turn being answered takes at its next step ([Steering a turn](#steering-a-turn)). ## One message at a time @@ -66,7 +72,8 @@ 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 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 `chat` never joins the turn being answered; one sent +to `steer` does, at the turn's next step ([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 +87,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 @@ -368,6 +376,77 @@ sequenceDiagram before `clear` is answered first, then forgotten too. Message ids and cursors carry on. +## Steering a turn + +A message sent to the agent's `steer` view, rather than to `chat`, joins 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), and the model's next request has it, after those results. +The turn's answer then answers the turn's messages so far. + +```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 +``` + +- **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`). +- **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. +- **Or answered as a turn of its own.** Once the turn ended, `steered` runs. + It does nothing if the turn stored the message, and otherwise answers it + as a turn of its own: a message that came after the turn's last step + (while the model wrote its final answer), one sent to a turn that takes no + steered messages, or past the cap, or one sent with no turn being answered + (it starts one, as `chat/send` does). So each message is answered once, in + the order it was sent with the messages after it. +- **A taken message is its turn's.** A turn killed after it took a message + ends `ended`, and the message with it. 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 (`STEERS`); past them, each is a turn of its own. A LangGraph + graph (`graph=`) and a background agent take none. + ## The turn model A turn is the unit of the conversation, in storage, in code and in views. @@ -451,9 +530,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 +545,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 +563,18 @@ 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**: in the turn that took it, or as a + turn of its own, in the order it was sent with the messages after it. - **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/scripts/resurrect.py b/scripts/resurrect.py index 37c7e86a..dd54865b 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") @@ -671,6 +672,71 @@ 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 turn is None or turn.open_calls():\n", + " if 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-summary-hides-steers", + "the summary view leaves out the messages steered to a turn", + f"{AGENTS}/_messages.py", + " for n, seq, stored in steered:\n", + " for n, seq, stored in ():\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() From c0c7c36345e50e9ac0be709bb841a6672c2080ba Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 08:17:17 +0000 Subject: [PATCH 03/13] Agents: testing.md lists the steering tests The in-process tier gains agents_steering, the server tier agents_steering_server. Co-Authored-By: Claude Opus 5.5 --- docs/agents/testing.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 From 96f1e40766530ba42b1d55626d7574ff8fe8f635 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 08:33:41 +0000 Subject: [PATCH 04/13] Agents, LangGraph: the built-in agent takes steered messages at each step, in place The built-in LangGraph agent (`langgraph.agent(model=...)`) now takes the messages steered to the turn being answered (the `steer` view) at the turn's next step, and says so: `chat(..., steers=True)`. At the step - The model node (`call_model`) calls `take_steers` before its request. The graph reaches that node after the tools node, which stores every tool call's result and then publishes (`tool_call_ended`): the publish's response is how the turn learns of the waiting messages. `take_steers` stores each once, as a message of the user that is a step of the turn, and returns the new ones; it returns nothing, at no cost, when nothing waits, a call has no result, the turn stops, or it is no `chat` turn (a handler of your own that calls `langgraph.prompt`). At the run's start the node calls it too: a message the turn already knows of then goes into the first request. - The request gets them after what the state holds, so after the tool results, and the node returns them with the reply (`{"messages": [*steered, reply]}`), in the order the conversation stored them: the graph's state keeps them for the rounds after. Why in place - Running the graph again with the message would reset LangGraph's recursion limit, which is what ends a turn whose model keeps calling tools. Returned in the model node's update, the messages are appended by the add_messages reducer and the run goes on, within its limit. Marked as steered - In the graph, a steered message has the id `steer:`, which no other message has. The conversation stores it without one, so `_history` marks the stored ones that the history holds (a retried attempt's own, and earlier turns'), by the entries whose items `history()` returns (`Turn.steered()`, and the transcript's `steers`). - `_latest` trims each request to the model's window with LangChain's `trim_messages(start_on="human")`, and sends the latest prompt whole when nothing fits. A steered message is a `HumanMessage`, so either could have started the request at it, without the turn's prompt and the tool calls and results before it. Now a request starts at the first prompt among the messages that fit, never at a steered message, else at the latest prompt, whole. Without steered messages, requests are as before. `graph=` agents take none - A graph of your own takes no steered messages in this version: its steps are its own, and `_graph_agent` does not pass `steers=True`. So `steer` answers that no turn takes the message, and it is answered as a turn of its own, after the one being answered, in the order it was sent. Docs: the steering point in `_agent.py`'s docstring and `agent()`'s, the `graph=` notes in `_graph.py` and the package's, and LangGraph's section of docs/agents/frameworks.md. Tests, in tests/agents_langgraph.py, on a server: an agent whose scripted model records each request and answers the latest message of the user, with a window that holds no message, so that each request is the latest prompt and what followed it (a steered message taken for a prompt would show). - A message steered while a tool waits on a gate is in the very next request, once, after the tool's result, under `steer:`. The turn's answer is the reply to it; it is stored once, and the summary view shows the prompt, the message and the answer. - A message steered after the first tool round, which has the model call a tool again, is still in the request after the second round, in place. - A crash after the message is stored (the model's request with it fails the attempt): the retried attempt's request is the crashed one, with the message once and marked, and it is not stored again. - `_latest` (in process): a request starts at the turn's prompt, never at a steered message, when only that fits and when nothing does, nor at one of an earlier turn. - tests/agents_langgraph_graph.py: a `graph=` agent's `steer` answers that no turn takes the message, which is answered next, as a turn of its own. Each of these mutations fails them: no state update, no marks on the history, no steered messages in the request, no `steers=True`, `_latest` ignoring the marks. tests/agents_langgraph*.py, agents_steering.py and agents_steering_server.py: 110 passed. All agents tests: 393 passed, and the usual host.docker.internal failure of examples_agents.py. Co-Authored-By: Claude Opus 5.5 --- docs/agents/frameworks.md | 8 + python/restate/agents/langgraph/__init__.py | 3 + python/restate/agents/langgraph/_agent.py | 85 +++++++-- python/restate/agents/langgraph/_graph.py | 4 + tests/agents_langgraph.py | 198 +++++++++++++++++++- tests/agents_langgraph_graph.py | 22 +++ 6 files changed, 295 insertions(+), 25 deletions(-) diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index ce9c5f19..bd8f7db1 100644 --- a/docs/agents/frameworks.md +++ b/docs/agents/frameworks.md @@ -357,6 +357,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 +385,12 @@ 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:` (the conversation stores it without one, so a retried attempt + marks the ones it loads), and 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/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..6d7a0f70 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. In the graph its id is ``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 @@ -93,9 +101,10 @@ model_of, provider_of, ) -from restate.agents._conversation import Call, Conversation, Item, Parsed, Result, Run, answer_once +from restate.agents._conversation import Call, Conversation, Entry, 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 @@ -286,10 +295,43 @@ def _window(model: BaseChatModel) -> int | None: return int(limit * WINDOW) if limit is not None and limit > 0 else None +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) + + +def _as_steered(entry: Entry) -> BaseMessage: + """Return the message steered to a turn that ``entry`` is, as the graph holds it: under its ``steer:`` id.""" + return typing.cast(BaseMessage, entry.item).model_copy(update={"id": f"{STEER}{entry.steer}"}) + + +def _history(conversation: Conversation) -> list[BaseMessage]: + """Return what the model gets (``Conversation.history``), each message steered to a turn under its ``steer:`` id. + + The conversation stores a steered message as a message of the user, + without that id. So the messages that a retried attempt finds in its turn, + and those of the turns before, are marked here, as the run marks the + messages it takes (``_as_steered``). + """ + steered = {id(entry.item): entry for turn in conversation.transcript for entry in turn.steers} + steered |= {id(entry.item): entry for turn in conversation.turns for entry in turn.steered()} + return [ + _as_steered(entry) if (entry := steered.get(id(item))) is not None else typing.cast(BaseMessage, item) + for item in conversation.history() + ] + + 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 +344,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 +364,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 = [_as_steered(entry) for entry in 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 +405,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": _history(conversation)})) 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 +577,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 +591,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 +658,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/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}" From 8c587a7d71ea1a4402e273d3cb8ea7c9e04cfdfe Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 08:35:07 +0000 Subject: [PATCH 05/13] Agents: ADK takes steered messages, at the step before each model request An ADK agent whose root is an LlmAgent (the built-in agent, or an agent or App of your own with one) now says it takes steered messages (`chat(..., steers=True)`): a message sent to `steer` while a turn runs joins the turn at its next step, and the model's next request has it. The run goes on in place; it is not restarted. At the step - `_Restate`, this package's plugin, takes the messages in its `before_model_callback`, as a model request goes out that ADK built from the session's events. `take_steers` stores those the last `publish` listed, only once each of the model's tool calls has a result, and each once. Each is appended to that request's contents and, through `_Sessions.steered`, to the in-memory session after its events (not stored again), so that every later request of the run has it, in place. - The plugin is now built per run, with that run's `_Sessions` (it was built with the app, which no longer takes `ctx`): it reaches the conversation without a private accessor of ADK's. It still comes after the app's own plugins. Why at the request, rather than as the tools' results are stored - Taking them in `_Sessions.append_event`, after a function response, was the first plan. ADK makes no request after some function responses: a tool's `skip_summarization` (its result is the answer), an output schema answered through `set_model_response` (ADK makes the final event itself), and, from ADK's code, an auth or confirmation request that ends the invocation. In process, a message steered while such a tool ran was stored by the turn and never shown to the model. At the request, a step that ADK ends takes nothing, and the core answers the message next, as a turn of its own. An invocation of its own - A steered message is an event of the user that starts an invocation of its own (`_Items.prompt`, which the core makes it with), never the running one. After a crash, ADK resumes the prompt's invocation from that invocation's events (`latest_ask` skips steered messages, so `_final` and `_run` keep resuming the prompt's, or the output re-ask's). ADK's classic runner, which runs a root that is no LlmAgent and resumes its LlmAgents from their agent states, takes a message of the user last among those events for the answer to a tool call, and raises `ValueError: No agent to transfer to for resuming agent from function response step` on every retry: with the running invocation's id, a crash after the store fails the turn for good. On ADK 2.14 an LlmAgent root runs on the node runner, which loads no agent states, so a same-id message would resume there; a test that forces a SequentialAgent root to take steered messages shows the failure. - With the default `include_contents`, ADK puts a user event of another invocation in the request (`_should_include_event_in_context` does not filter by invocation): it is in every request after it, in its place. Only LlmAgent roots, for now - `_steers(own)`: the built-in agent, or an agent or App whose root is an LlmAgent. A workflow root (a SequentialAgent, ParallelAgent or LoopAgent, an App's Workflow) passes `steers=False`: each message steered to it is answered as a turn of its own. Limits, documented in the module's docstring - A message sent during the model's last reply, or while a tool runs after whose result ADK asks the model nothing more, is answered next, 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. - An AgentTool is one call: ADK gives its runner our plugin, but the parent's call is open while its agent runs, so its requests take nothing; the step is after its result. - An agent's own before_model_callback runs after the plugins: one that answers in the model's place gets the request with the message in it. Tests, in tests/agents_adk.py (against a Restate server unless in process) - A message steered while a tool runs reaches the next request, after the function response, once; the model answers it with a call of its own and the request after that still has it in place; it is stored once (its own queued turn does nothing), and the summary view shows prompt, steered message, answer. - A crash on the request with the message: the retry resumes the prompt's invocation, its request has the message once, the tool ran once, the model was not asked for the call again, and nothing is stored twice. - A SequentialAgent root takes no steered message (`turn` is None), which is answered as the next turn. - A message steered during the last model call is answered as the next turn. - In process: after a tool's skip_summarization, and after an output schema's set_model_response, no request follows and the turn stores nothing. - In process: on ADK's classic runner (a SequentialAgent root made to take steered messages), the retry after a crash resumes past the message. Each fails without what it covers (taking at the function response, steers for no root or for every root, no append to the request, no append to the session, the running invocation's id). tests/agents_adk.py 34 passed (27 before), tests/agents_steering.py 9, tests/agents_steering_server.py 5; ruff, pyright (0 errors) and mypy clean. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/adk/_agent.py | 92 +++++++- tests/agents_adk.py | 342 +++++++++++++++++++++++++++- 2 files changed, 419 insertions(+), 15 deletions(-) diff --git a/python/restate/agents/adk/_agent.py b/python/restate/agents/adk/_agent.py index 134b3305..44fec11e 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 @@ -172,8 +192,9 @@ def read(self, item: Item) -> Parsed: class _Sessions(BaseSessionService): """ADK's sessions, of which there is one: the conversation, an event per item.""" - def __init__(self, conversation: Conversation) -> None: + def __init__(self, ctx: SessionContext, conversation: Conversation) -> None: super().__init__() + self.ctx = ctx self.conversation = conversation def session(self) -> Session: @@ -231,13 +252,44 @@ async def append_event(self, session: Session, event: Event) -> Event: self._carry(session) return event + async def steered(self, session: Session) -> list[Event]: + """Take the messages steered to the turn (``take_steers``), and add them to ``session`` after its events. + + ADK builds each model request from the session's events: every + request after this one has them too, in place. Each starts an + invocation of its own, as a prompt does (``_Items.prompt``), and 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. + """ + taken = [typing.cast(Event, entry.item) for entry in await take_steers(self.ctx, self.conversation)] + for event in taken: + # Not stored again: the turn stored them as it took them. + await super().append_event(session, event) + return taken + 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. + for event in await self.sessions.steered(callback_context.session): + 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 +371,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 +398,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 +471,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(ctx, 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 +636,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 +746,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/tests/agents_adk.py b/tests/agents_adk.py index 60cf7761..25a97fb6 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 _STEERING, Steering # 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)])) + + steering = Steering(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: + steering.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. + _STEERING.set(ctx, steering) + 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 _STEERING, Steering # pylint: disable=import-outside-toplevel + from restate.sessions.internal import invoke # pylint: disable=import-outside-toplevel + + steering = Steering(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.""" + steering.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. + _STEERING.set(ctx, steering) + 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" From a60e9c6260e15662d211b243fe55dc98e200c1ac Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 08:37:40 +0000 Subject: [PATCH 06/13] Agents: Pydantic AI takes steered messages at a turn's steps, in place The Pydantic AI agent now says it takes steered messages (`steers=True` in `agent()`'s `chat`), and every run takes them at each step, with a capability of its own (`_Steering`). At the boundary - Before each model request, once the model's tool calls all have their results stored (by `_Restate`, earlier in the same before-chain), it calls `take_steers`: the messages that wait are stored once, as messages of the user that are steps of the turn (`ModelRequest(parts=[UserPromptPart])`), and returned. - It then enqueues those same stored requests with Pydantic AI's own `RunContext.enqueue(..., priority="asap")`. Content enqueued from a `before_model_request` hook is delivered once the whole before-chain ran (`_agent_graph._apply_before_model_request`), into the request being prepared and into the run's history (`ctx.messages`), after what both hold: right after the tools' results. So what the model gets is what is stored, where it is stored, and a retried attempt or the next turn rebuilds the same history from the database. - In place: the run carries on with the message. Starting `agent.run` again from the stored history would reset the run's usage limits, its retries and its capabilities' state. - A whole `ModelRequest` enqueued outside a tool call is never spliced into the request that carries the tools' results (Pydantic AI splices only parts enqueued during a tool call, and `ctx.tool_call_id` is None in `before_model_request`): it is a request of its own, which `_Restate` never stores again (it holds no result and no feedback). - It takes nothing for a resumed suspended response: that request ends with the paused `ModelResponse`, into which Pydantic AI delivers nothing, so a message taken then would reach the model after the response though stored before it. The next step takes it. - Outside a `chat` turn that takes steered messages (a `prompt` in a handler of one's own, a background agent), `take_steers` returns nothing, at no cost. Snapshots - `_Steering` is innermost, ordered after `_Snapshots` (`CapabilityOrdering(position="innermost", wrapped_by=[_Snapshots])`). A snapshot covers every position stored so far (`ledger.next - 1`) and holds the request's messages as the before-chain left them, before the delivery: a message stored before a snapshot of the same request would be covered by it but not in it, and lost from the next turn's history and from a retried attempt. Stored after it, it comes after the snapshot's items. - Without snapshots (a `chat` of one's own that answers with `prompt`), the ordering has nothing to follow, and the turn projection has the messages in their turn. Feedback found by its parts (`_Restate._new`) - A run merges the consecutive requests it is given, and resumes from the last one as a new `ModelRequest` with the same parts. `_Restate` found Pydantic AI's feedback (`RetryFeedbackPart`) by the whole request's JSON, so a merged or resumed feedback request was taken as new and stored again, then again in every later turn. That was so before steering: a retry after a stored feedback stored it twice. Steering makes consecutive requests common (a result, or a feedback, then the steered message): after a retry that took a message right after a stored feedback, the next turn stored the feedback and the message again, and the views showed the message twice. Each `RetryFeedbackPart` is now found by its own JSON (it has its own time), which merges and resumes keep. Limitations - A turn learns of a message only from the companion's `publish`, which a tool call's start and end make: a step that ran no tool (Pydantic AI's feedback on an answer that does not validate, say) takes only what the last publish listed. A message that comes during the last model request is answered next, as a turn of its own. - Not changed here: what other capabilities enqueue as requests of their own (the harness's Planning, SystemReminders, DetectLoops, WarnNearLimits, Memory, ...) reaches the model and the run's history, but `_Restate` stores only results and feedback, so a retried attempt and the next turn do not have it. Tests - On a server, with a scripted model and a tool held on a `support.Gate`: a message steered during the tool call reaches the very next request, after the tool's result, once; the turn's answer is the reply to it; it is stored once, the summary view shows the prompt, the message and the answer, and its queued turn of its own does nothing. A crash while the model answers it: the retried request has it once, and it is not stored again. A message steered during the last model request is answered as the next turn. - In process: a compaction (of one's own, at the step) and a steered message in the same request: the snapshot does not hold it, and a retried attempt and the next turn get it after the snapshot's items, once (this fails with `_Steering` ordered before `_Snapshots`). A `chat` of one's own that answers with `prompt`, without snapshots, takes two messages, in order. A feedback, a crash, a retry that takes a message right after the feedback, and the next turn: each stored once, and given to the model once (this fails with the former `_new`). - tests/agents_pydantic_ai.py, agents_pydantic_ai_snapshots.py, agents_pydantic_ai_stores.py, agents_steering.py and agents_steering_server.py: 84 passed (78 before). ruff, pyright and mypy are clean. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/pydantic_ai/_agent.py | 110 ++++++++-- tests/agents_pydantic_ai.py | 217 +++++++++++++++++++- tests/agents_pydantic_ai_snapshots.py | 134 +++++++++++- 3 files changed, 437 insertions(+), 24 deletions(-) diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index c63808ba..bc0ae513 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,48 @@ 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, once the model's tool calls all have their + results stored, it takes the messages that wait (``take_steers``): each + is stored once, as a message of the user that is a step of the turn. It + then enqueues the stored messages with Pydantic AI's own + ``RunContext.enqueue`` (``'asap'``), which delivers them once every + ``before_model_request`` ran: into this request and into the run's + history, after what both hold. So the model gets each message as it is + stored and where it is stored, and the run goes on with it, its usage + limits, retries and capabilities' state with it, rather than starting + again; the turn's next requests, a retried attempt and the next turns + have it in their history, in place, and never take it again. + + It is the innermost capability, after ``_Snapshots``: a snapshot that this + request stores stands for the positions stored before the messages, so + the history after it (the next request's, a retried attempt's, the next + turn's) 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 + steered = await take_steers(self.ctx, self.conversation) + if steered: + ctx.enqueue(*(typing.cast(ModelRequest, entry.item) for entry in 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 +797,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 +808,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 +1038,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 +1119,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/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"] From 95f753f3779fdde195cee633c027985edc1738e4 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 08:39:43 +0000 Subject: [PATCH 07/13] Agents: OpenAI takes steered messages at each step, in place An OpenAI agent's turn now takes the messages steered to it (the agent's steer view) at its next step, and its run goes on with them: agent() says so to chat (steers=True), and the run's call_model_input_filter, which already stores the tools' outputs right before each model request, takes them there. At a step - The input filter (_Hooks.model_input) stores the new tool outputs, then calls take_steers: once each of the model's tool calls has its output, the turn stores the waiting messages, once each, as messages of the user, and the request about to go out has them right after those outputs. - A run also takes the messages waiting as it starts (_go, after _resume closed a cut-off reply): they are stored before Runner.run, so they are in its input. A request whose last item cannot be told apart in later requests (a message of the user, as a run's first request ends with) takes none: its messages wait for the next step. In place, not a restart - A new Runner.run from the stored history would run the input guardrails again, start max_turns over, and reset the tool-choice tracking, the agent-start hooks and the trace. The filter changes only the request about to go out, and the run goes on as it was. The splice - The SDK makes each request's input afresh (the run's input plus its generated items, copied for the filter), so a message put in one request is gone from the next. The run keeps the messages it took, each with the key of the item it came after (a tool call's or output's call_id, else the item's own id: _key), and puts each back in every later request right after that item. Messages taken together keep their order; an item found twice counts where its last copy is, the one the SDK's dedup keeps. User messages have no id, so that dedup never drops them. - If a handoff's input_filter took that item out, the message goes right before the first item that came after it (one whose key the run first saw after taking it), as if the filter had kept the message; else at the end. - The run's own call_model_input_filter (a ToolOutputTrimmer, say) still runs after, on the input with the messages. On a retry - Stored messages come back in conversation.history() at their places, after the outputs they followed, so a resumed run's input has each once; take_steers never returns one twice, and the new run splices only what it took itself. A stored message after the outputs leaves _resume no open call, and makes run() ask the model rather than take a final reply. Limitations - A message is taken only at a step: one that comes while the model writes its final answer is answered next, as a turn of its own (the agent's steered handler); one sent during a long tool call waits for it. - A handoff whose input_filter takes its own output out of the next agent's input (handoff_filters.remove_all_tools does) leaves the handoff call without an output in the turn, so take_steers takes nothing for the rest of that run (by design: the turn has an open call). - The fallback placement counts an item with no key (a Chat Completions message, whose id is the SDK's fake one) as having come before the message. - The SDK's ToolOutputTrimmer counts user messages as turns, so a steered message makes the outputs before it older to it, as on a resumed run. Found, not changed (independent of steering): a handoff whose input_filter removes the handoff's own output never has that output stored, so every later request of the conversation carries the transfer call without an output, which OpenAI's APIs reject. The same on the base commit. Tests (tests/agents_openai.py) - Against a server: a message steered while a tool waits on a gate is in the next request, once, right after the tool's output; the turn's answer is the reply to it; the summary view shows the prompt, the message and the answer; its queued turn of its own finds it taken. A message steered while the model writes its final answer is the next turn. - In process, with a companion whose publish lists the messages: a message taken after the first of two tool rounds is in the request after the second, in its place, once, stored once and taken once; after a crash once the filter stored it, the retry's request has it once and it is not stored again; an input guardrail runs once; a message waiting as the run starts is in its first request; after a remove_all_tools handoff the message keeps its place and the run takes no more; the splice itself (order, a repeated item, the fallback). - A mutation of each fails a test: no steers=True, the current request only, no take as the run starts, always at the end, the fallback at the end, the first copy of an item, a restart instead of in place (the guardrail test). tests/agents_openai.py 54 passed (46 + 8), agents_steering.py 9 and agents_steering_server.py 5; the whole suite 577 passed, with the 3 usual host.docker.internal failures. ruff, pyright and mypy are clean. Co-Authored-By: Claude Opus 5.5 --- python/restate/agents/openai/__init__.py | 9 +- python/restate/agents/openai/_agent.py | 116 ++++++- tests/agents_openai.py | 382 ++++++++++++++++++++++- 3 files changed, 494 insertions(+), 13 deletions(-) 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..61692e68 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`` @@ -115,6 +132,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 +154,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,14 +304,35 @@ 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.steered: list[_Steered] = [] + """The messages steered to the turn that the run took, in order: each later request gets them again.""" + 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.""" 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. """ - if outputs := self.conversation.new_results(data.model_data.input): + items = data.model_data.input + if outputs := self.conversation.new_results(items): await self.conversation.turn.append(*outputs) + # 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)): + # This request's items came before the messages; an item whose key the run sees later, after them. + self._see(items) + self.steered.extend(_Steered(entry.item, after, len(self.seen)) for entry in taken) + if self.steered: + self._see(items) + data.model_data.input = _spliced(items, self.steered, self.seen) filtered = data.model_data if self.then is not None: given = self.then(data) @@ -317,6 +357,65 @@ async def on_handoff(self, context: RunContextWrapper[Any], from_agent: Agent[An async def on_tool_end(self, context: RunContextWrapper[Any], agent: Agent[Any], tool: SdkTool, result: Any) -> None: await tool_call_ended(self.ctx) + 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 @@ -521,6 +620,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 +893,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 +967,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/tests/agents_openai.py b/tests/agents_openai.py index 4717cd6f..0310aa42 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,318 @@ 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")] + + +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] From 870be69fbed159b420d3477086a8008235d0143a Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 08:47:29 +0000 Subject: [PATCH 08/13] Docs: how each framework takes steered messages, and their limits frameworks.md gains a row for steered messages and a note per framework: OpenAI splices them back into each later request, Pydantic AI enqueues them after any snapshot, ADK adds them as a request goes out, for an LlmAgent root. It also fixes the note on which capability is innermost, and records two gaps found on the way: an OpenAI handoff filter that drops the handoff's output, and what other Pydantic AI capabilities enqueue, which is not stored. Co-Authored-By: Claude Opus 5.5 --- docs/agents/frameworks.md | 56 ++++++++++++++++++++++++++++++++++++--- docs/agents/turns.md | 9 +++++-- 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index bd8f7db1..75343bec 100644 --- a/docs/agents/frameworks.md +++ b/docs/agents/frameworks.md @@ -59,6 +59,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 +110,20 @@ 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. +- **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 +267,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 +362,20 @@ 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. +- **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 diff --git a/docs/agents/turns.md b/docs/agents/turns.md index 9720be97..9de9337b 100644 --- a/docs/agents/turns.md +++ b/docs/agents/turns.md @@ -444,8 +444,13 @@ sequenceDiagram - **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 (`STEERS`); past them, each is a turn of its own. A LangGraph - graph (`graph=`) and a background agent take none. + messages (`STEERS`); 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 From 8b73fd660022a36b70b7f4dfcdff152ec68f27d3 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 09:26:46 +0000 Subject: [PATCH 09/13] Agents, steering: count a turn's every message toward its cap, and the fixes of a self review - The cap of 16 counted only the messages a turn had not taken yet, so a turn steered on and on took them without end. `taken` now marks a message (without its text) rather than forgetting it, until the turn ends: the cap counts them all, and a message sent again is still the turn's. - `taken` publishes only a change: one for a turn already closed wakes no follower. - A turn that stopped while it waited for the ledger to store a message takes none, which is then answered as a turn of its own. A message the turn did take shares its fate (stopped, refused, failed), as documented. - The turn just answered counts its steered messages toward the budget in the same attempt (`prompt.steers`), so a summary due is asked for now rather than a turn later; a drop trims them. - `publish` reads the companion's turn only when a message waits, and `status` reads it once; `_steps` reads the steered marks with one join; `Turn.steer` uses `_INSERT_STEP`. - ADK: the agent's own before_model_callback runs after the plugin's, on the request with the message, so a response it gives answers it too. Three more mutants in scripts/resurrect.py guard the cap, the stop and the budget; the publish mutant follows the code it mutates. Co-Authored-By: Claude Opus 5.5 --- docs/agents/access-patterns.md | 5 +- docs/agents/frameworks.md | 4 +- docs/agents/porting.md | 15 +++--- docs/agents/turns.md | 13 +++-- python/restate/agents/_conversation.py | 42 ++++++++------- python/restate/agents/_messages.py | 3 +- python/restate/agents/_signals.py | 71 ++++++++++++++++---------- python/restate/agents/_steering.py | 12 +++-- python/restate/agents/adk/_agent.py | 4 +- scripts/resurrect.py | 28 +++++++++- tests/agents_signals.py | 11 ++++ tests/agents_steering.py | 52 ++++++++++++++++++- 12 files changed, 194 insertions(+), 66 deletions(-) diff --git a/docs/agents/access-patterns.md b/docs/agents/access-patterns.md index d02b08d2..38340e2f 100644 --- a/docs/agents/access-patterns.md +++ b/docs/agents/access-patterns.md @@ -292,8 +292,9 @@ In order of value for effort: 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, 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 + 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 diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index 75343bec..e64eb237 100644 --- a/docs/agents/frameworks.md +++ b/docs/agents/frameworks.md @@ -368,7 +368,9 @@ sequenceDiagram `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. + 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 diff --git a/docs/agents/porting.md b/docs/agents/porting.md index 6e2aff77..02d455dc 100644 --- a/docs/agents/porting.md +++ b/docs/agents/porting.md @@ -112,12 +112,15 @@ Each step stands on its own, with its tests, and the earlier ones are small. 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`), `taken`, `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, and - each framework putting the messages in its next request in place: pi's + `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/turns.md b/docs/agents/turns.md index 9de9337b..fab03dde 100644 --- a/docs/agents/turns.md +++ b/docs/agents/turns.md @@ -434,9 +434,13 @@ sequenceDiagram steered messages, or past the cap, or one sent with no turn being answered (it starts one, as `chat/send` does). So each message is answered once, in the order it was sent with the messages after it. -- **A taken message is its turn's.** A turn killed after it took a message - ends `ended`, and the message with it. A reply dropped to ask the model - again drops the messages taken after it, which `steered` then answers. +- **A taken message is its turn's.** A turn that ends otherwise than + answered after it took a message (stopped, refused, failed, or killed and + so `ended`) ends so for the message too, which is not answered again, as + the turn's own message is not. A turn that stops while it waits to store + messages takes none: each is then a turn of its own. 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 @@ -444,7 +448,8 @@ sequenceDiagram - **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 (`STEERS`); past them, each is a turn of its own. A message is + messages in all (`STEERS`; the companion counts those it took until the + turn ends); 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 diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index 4cd79e6b..03567745 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -691,7 +691,10 @@ async def ask_again(self, text: str) -> None: await self.append(self._ledger.items.prompt(text)) async def steer( - self, messages: Sequence[Steer], sent: Callable[[Transaction, list[str]], None] | None = None + self, + messages: Sequence[Steer], + sent: Callable[[Transaction, list[str]], None] | None = None, + going: Callable[[], bool] | None = None, ) -> list[Entry]: """Store the steered ``messages`` that are not stored yet as messages of the user in the turn, and return them. @@ -701,19 +704,25 @@ async def steer( adds what goes out with the transaction that stores them (the companion's ``taken``), which commits once. - The store finishes even if the turn stops meanwhile: the turn then - closes after it, at the positions after its messages. + ``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). """ ledger = self._ledger async def store() -> list[Entry]: async with ledger.lock: + if going is not None and not going(): + return [] self._check_writable() wanted = list({message["id"]: message for message in messages}.values()) if not wanted: return [] ids = [message["id"] for message in wanted] marks = ", ".join("?" for _ in ids) + # The database, not the turn: an id is stored once in the conversation, by whichever turn took it. stored = await ledger.db.values(f"SELECT id FROM {STEERED} WHERE id IN ({marks})", ids) known = {str(row[0]) for row in stored} new = [message for message in wanted if message["id"] not in known] @@ -726,10 +735,8 @@ async def store() -> list[Entry]: ] async with ledger.db.transaction() as tx: for i, entry in enumerate(added, 1): - await tx.execute( - f"INSERT INTO {STEPS} (n, turn, seq, item, reply, agent) VALUES (?, ?, ?, ?, NULL, NULL)", - [entry.n, self.n, ledger.seq + i, ledger.items.dump(entry.item)], - ) + dumped = ledger.items.dump(entry.item) + await tx.execute(_INSERT_STEP, [entry.n, self.n, ledger.seq + i, dumped, None, None]) await tx.execute( f"INSERT INTO {STEERED} (n, turn, id) VALUES (?, ?, ?)", [entry.n, self.n, entry.steer] ) @@ -738,6 +745,8 @@ async def store() -> list[Entry]: 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 added # Shielded: a stop cancels the run, but a commit that is under way still lands, and the turn's own record of @@ -799,6 +808,7 @@ async def drop(self, reply: int) -> None: 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. @@ -1365,20 +1375,16 @@ async def _transcript(db: DatabaseReader, items: Items, first: int, bounded: tup 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 n, item, reply, agent FROM {STEPS} WHERE n > ? ORDER BY n", [after]) - steered = { - int(typing.cast(int, n)): str(i) - for n, i in await db.values(f"SELECT n, id FROM {STEERED} WHERE n > ?", [after]) - } + 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=steered.get(int(typing.cast(int, n))), + int(typing.cast(int, n)), load_item(items, str(item)), reply=_int(reply), agent=_str(agent), steer=_str(i) ) - for n, item, reply, agent in stored + for n, item, reply, agent, i in stored ] diff --git a/python/restate/agents/_messages.py b/python/restate/agents/_messages.py index a3f70ba3..e2f530d8 100644 --- a/python/restate/agents/_messages.py +++ b/python/restate/agents/_messages.py @@ -236,8 +236,7 @@ async def _summary_latest(db: DatabaseReader, items: Items, before: int | None, 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] - 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]) + [[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/_signals.py b/python/restate/agents/_signals.py index f499f5e6..a1b47759 100644 --- a/python/restate/agents/_signals.py +++ b/python/restate/agents/_signals.py @@ -75,8 +75,9 @@ STEERS = 16 """How many messages a turn takes by ``steer``, at most: past them, each is answered as a turn of its own. -Each one the turn takes is one more model request in it, so they keep a turn -that someone steers on and on within bounds.""" +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): @@ -125,31 +126,38 @@ class SteerRequest(TypedDict): text: str -class _Inboxed(TypedDict): - """A steered message the turn ``turn`` (an invocation id) has not taken yet, and when it was sent.""" +class _Held(TypedDict): + """A message steered to the turn ``turn`` (an invocation id), and when it was sent.""" id: str - text: 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: """Return ``turn`` as ``status`` shows it: without what only the companion reads.""" return None if turn is None else {"invocationId": turn["invocationId"], "since": turn["since"]} async def _inbox(ctx: Context) -> list[_Inboxed]: - """Return the steered messages that their turns have not taken yet, oldest first.""" + """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], [])) -async def _steering(ctx: Context) -> list[_Inboxed]: - """Return the steered messages that the turn being answered has not taken yet, oldest first.""" - answering = await _get(ctx, "turn", typing.cast(Started | None, None)) +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 await _inbox(ctx) if one["turn"] == answering["invocationId"]] + 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: @@ -194,9 +202,9 @@ async def _closed(ctx: ObjectContext, turn: str, gone: bool = False) -> None: 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 that it did not take are forgotten here: - each was queued as a turn of its own as it was sent (``steer``), and - is answered as one. + 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 @@ -249,6 +257,7 @@ async def steer(ctx: ObjectContext, request: SteerRequest) -> Steered: 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()) @@ -259,16 +268,20 @@ async def steer(ctx: ObjectContext, request: SteerRequest) -> Steered: @signals.handler(ingress_private=True) async def taken(ctx: ObjectContext, ids: list[str]) -> None: - """Forget the steered messages ``ids``, which their turn has stored, and publish the change. + """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. + 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) - left = [one for one in inbox if one["id"] not in ids] - if len(left) != len(inbox): - _set_inbox(ctx, left) - await _publish(ctx) + marked: list[_Inboxed] = [ + {"id": one["id"], "since": one["since"], "turn": one["turn"]} if one["id"] in ids else one for one in inbox + ] + if marked != inbox: + _set_inbox(ctx, marked) + await _publish(ctx) @signals.handler(kind="shared", inactivity_timeout=PARKED) @@ -388,11 +401,13 @@ async def await_turn(ctx: ObjectSharedContext, turn: str) -> None: async def status(ctx: ObjectSharedContext) -> Status: """Return what the agent is doing: the message it is answering, if any, its tools' waits, and the messages steered to it that it has not taken yet.""" + answering = await _get(ctx, "turn", typing.cast(Started | None, None)) shown: Status = { - "answering": _shown(await _get(ctx, "turn", typing.cast(Started | None, None))), + "answering": _shown(answering), "waiting": [_public(wait) for wait in await _get(ctx, "waiting", typing.cast(list[_WaitingOn], []))], } - steering: list[Steering] = [{"id": one["id"], "since": one["since"]} for one in await _steering(ctx)] + 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 @@ -432,8 +447,8 @@ async def ended(ctx: ObjectContext, turn: str) -> None: @signals.handler(ingress_private=True) async def reset(ctx: ObjectContext) -> None: - """Forget all waits, kept signals and steered messages not taken 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``.""" + """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") @@ -457,9 +472,13 @@ async def publish(ctx: ObjectContext) -> Published: call's start and end, so it learns of them at no cost of its own. """ published: Published = {**await _publish(ctx)} - steering: list[Steer] = [{"id": one["id"], "text": one["text"]} for one in await _steering(ctx)] - if steering: - published["steering"] = steering + 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 diff --git a/python/restate/agents/_steering.py b/python/restate/agents/_steering.py index a4bde9a3..c1891696 100644 --- a/python/restate/agents/_steering.py +++ b/python/restate/agents/_steering.py @@ -47,7 +47,9 @@ async def take_steers(ctx: SessionContext, conversation: Conversation) -> list[E 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. + 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. """ steering = steering_of(ctx) if steering is None or not steering.messages: @@ -65,7 +67,11 @@ async def take_steers(ctx: SessionContext, conversation: Conversation) -> list[E companion = _signals.signals_of(ctx) def sent(tx: Transaction, ids: list[str]) -> None: - # With the transaction that stores them: so they are forgotten once stored, and once. + # 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(fresh, sent) + 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 + + return await turn.steer(fresh, sent, going) diff --git a/python/restate/agents/adk/_agent.py b/python/restate/agents/adk/_agent.py index 44fec11e..1544d2f0 100644 --- a/python/restate/agents/adk/_agent.py +++ b/python/restate/agents/adk/_agent.py @@ -285,7 +285,9 @@ async def before_model_callback( # 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. + # 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 await self.sessions.steered(callback_context.session): if event.content is not None: llm_request.contents.append(event.content.model_copy(deep=True)) diff --git a/scripts/resurrect.py b/scripts/resurrect.py index dd54865b..08b42f64 100644 --- a/scripts/resurrect.py +++ b/scripts/resurrect.py @@ -709,8 +709,8 @@ def kind(self) -> str: "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', + ' if steering:\n published["steering"] = steering\n', + ' if False:\n published["steering"] = steering\n', STEERING, ), Mutant( @@ -729,6 +729,30 @@ def kind(self) -> str: "", 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 going is not None and 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-summary-hides-steers", "the summary view leaves out the messages steered to a turn", diff --git a/tests/agents_signals.py b/tests/agents_signals.py index cc8d2737..94ebac23 100644 --- a/tests/agents_signals.py +++ b/tests/agents_signals.py @@ -276,6 +276,12 @@ async def test_a_steered_message_is_queued_as_a_turn_and_held_for_a_turn_that_st 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( @@ -315,3 +321,8 @@ async def test_a_turn_takes_at_most_its_share_of_steered_messages(harness: Harne # 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 index 6738ce66..0644d880 100644 --- a/tests/agents_steering.py +++ b/tests/agents_steering.py @@ -13,6 +13,7 @@ # pylint: disable=missing-function-docstring,import-error +import asyncio import json import typing from collections.abc import Awaitable, Callable @@ -35,7 +36,7 @@ steer_stored, ) from restate.agents._steering import take_steers -from restate.agents._turn import _STEERING, Steering +from restate.agents._turn import _STEERING, _STOP, Steering, Stop from restate.sessions.internal import Schema, invoke @@ -107,6 +108,8 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: 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"}) @@ -199,6 +202,52 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: ] +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.steering = [{"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.steering = [{"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] = [] @@ -229,6 +278,7 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: [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 From 36874854c7c4d16f46dcd37ec0d7ef154958f409 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 10:25:05 +0000 Subject: [PATCH 10/13] Docs: steering in the turn semantics, with its cases A turn is now its prompt, the messages steered to it that it took, and what answers them: turns.md says so where it defines a turn, its life (a step takes the messages its publish lists), how it ends (the messages it took go as it goes, those it did not are answered after it), its attempts, stopping and clear, the turn model and the guarantees. "One message at a time" is now "One turn at a time". Steering a turn opens with what happens to a steered message, then three cases as sequence diagrams: taken at the next step, too late for the turn, and a turn that stops. README and conversation.md follow. Co-Authored-By: Claude Opus 5.5 --- docs/agents/README.md | 17 +-- docs/agents/conversation.md | 5 +- docs/agents/turns.md | 235 ++++++++++++++++++++++++++---------- 3 files changed, 186 insertions(+), 71 deletions(-) diff --git a/docs/agents/README.md b/docs/agents/README.md index 82ce7c6d..d98ba8b4 100644 --- a/docs/agents/README.md +++ b/docs/agents/README.md @@ -14,9 +14,10 @@ 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 @@ -40,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"] @@ -48,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 | diff --git a/docs/agents/conversation.md b/docs/agents/conversation.md index 52c97417..36c04b72 100644 --- a/docs/agents/conversation.md +++ b/docs/agents/conversation.md @@ -16,7 +16,8 @@ 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*; @@ -25,7 +26,7 @@ records). 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,** diff --git a/docs/agents/turns.md b/docs/agents/turns.md index fab03dde..7c3b9af6 100644 --- a/docs/agents/turns.md +++ b/docs/agents/turns.md @@ -33,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. @@ -63,17 +67,23 @@ stored and loaded, and the `settle` and `compact` handlers, are this SDK's ([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, which the - turn being answered takes at its next step ([Steering a turn](#steering-a-turn)). +- **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. A message sent to `chat` never joins the turn being answered; one sent -to `steer` does, at the turn's next step ([Steering a turn](#steering-a-turn)). +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 @@ -112,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 @@ -140,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` @@ -151,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 @@ -172,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 @@ -316,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 @@ -371,18 +400,39 @@ 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`, joins the -turn being answered: the turn takes it at its next **step** (once the +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), and the model's next request has it, after those results. -The turn's answer then answers the turn's messages so far. +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 @@ -411,11 +461,65 @@ sequenceDiagram 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`). + 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 @@ -426,21 +530,18 @@ sequenceDiagram 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. -- **Or answered as a turn of its own.** Once the turn ended, `steered` runs. - It does nothing if the turn stored the message, and otherwise answers it - as a turn of its own: a message that came after the turn's last step - (while the model wrote its final answer), one sent to a turn that takes no - steered messages, or past the cap, or one sent with no turn being answered - (it starts one, as `chat/send` does). So each message is answered once, in - the order it was sent with the messages after it. -- **A taken message is its turn's.** A turn that ends otherwise than - answered after it took a message (stopped, refused, failed, or killed and - so `ended`) ends so for the message too, which is not answered again, as - the turn's own message is not. A turn that stops while it waits to store - messages takes none: each is then a turn of its own. A reply dropped to - ask the model again drops the messages taken after it, which `steered` - then answers. + 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 @@ -448,13 +549,12 @@ sequenceDiagram - **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 (`STEERS`; the companion counts those it took until the - turn ends); 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 + 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 @@ -471,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)): @@ -493,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, sent, going) take steered messages, at a step (take_steers) answered(answer), close(ctx, status, note) its end: once ``` @@ -522,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 @@ -583,8 +692,12 @@ turns of their own. A TypeScript companion, which has no `steer`, answers - **One turn at a time per key**, in the order messages came. A turn is joined only by a message sent to `steer`, at a step. -- **A steered message is answered once**: in the turn that took it, or as a - turn of its own, in the order it was sent with the messages after it. +- **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, From 529a0215b183c36cec55aca7aa87222ee47b9c3b Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 11:08:36 +0000 Subject: [PATCH 11/13] Agents: steered messages come from the conversation as the framework's items The conversation now makes and stores a steered message as the framework's own item, so no integration works out which messages were steered: - `Items.steered(text, steer_id)`, optional like `summary`, makes the item of a steered message (by default a prompt's). LangGraph's carries the id `steer:` and is stored with it, so its history, a retried attempt's and the next turns', tells steered messages from prompts with no pass that marks them again (`_history` and `_as_steered` are gone). ADK's says why it starts an invocation of its own. - `take_steers` returns the items the request gets, not entries: no integration reads an `Entry` for them any more. - `Turn.steer` skips the messages the turn already holds, and takes `taken` and `going` as required callbacks. Its stored-id query is shared with `steer_stored`. Readability: - ADK's plugin takes the messages at its step. `_Sessions` only adds the stored events to the session (`add_stored`), and no longer needs `ctx`. - OpenAI's in-run state for the messages a run took, which each later request gets again, is one small class (`_Steers`), out of the run hooks. - The turn's copy of its waiting messages is `Inbox` (it was `Steering`, a name `_api` also gives another type), and the companion's `steer` takes `_api.Steer`, which `SteerRequest` duplicated. - One `_turn_size` for the budget and the transcript's walk, and one helper for the summary view's messages of the user. Mutants: the steering ones follow the code, and `steer-langgraph-unmarked` is new. `langgraph-reask-stored-reply` matched no code since the steering change, and now does. Co-Authored-By: Claude Opus 5.5 --- docs/agents/conversation.md | 12 +- docs/agents/frameworks.md | 11 +- docs/agents/turns.md | 2 +- python/restate/agents/_conversation.py | 159 +++++++++++--------- python/restate/agents/_messages.py | 15 +- python/restate/agents/_object.py | 3 +- python/restate/agents/_signals.py | 11 +- python/restate/agents/_steering.py | 45 +++--- python/restate/agents/_turn.py | 20 +-- python/restate/agents/adk/_agent.py | 32 ++-- python/restate/agents/internal.py | 7 +- python/restate/agents/langgraph/_agent.py | 51 +++---- python/restate/agents/openai/_agent.py | 43 ++++-- python/restate/agents/pydantic_ai/_agent.py | 33 ++-- scripts/resurrect.py | 24 ++- tests/agents_adk.py | 16 +- tests/agents_signals.py | 4 +- tests/agents_steering.py | 52 +++---- 18 files changed, 276 insertions(+), 264 deletions(-) diff --git a/docs/agents/conversation.md b/docs/agents/conversation.md index 36c04b72..3249c7fc 100644 --- a/docs/agents/conversation.md +++ b/docs/agents/conversation.md @@ -163,12 +163,12 @@ CREATE TABLE restate_agent_steers ( -- the messages steered to a turn t 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.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. + 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 diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index e64eb237..9462ac02 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. @@ -439,8 +441,9 @@ sequenceDiagram (`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:` (the conversation stores it without one, so a retried attempt - marks the ones it loads), and a trimmed request never starts at one, which - would leave out the turn's prompt and its tool results. + `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/turns.md b/docs/agents/turns.md index 7c3b9af6..e1919343 100644 --- a/docs/agents/turns.md +++ b/docs/agents/turns.md @@ -605,7 +605,7 @@ Turn one prompt, the messages steered to it, and wh open_calls(), replies(), last_reply(), final_reply(), stored() its own state, never another turn's 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, sent, going) take steered messages, at a step (take_steers) + steer(messages, taken, going) take steered messages, at a step (take_steers) answered(answer), close(ctx, status, note) its end: once ``` diff --git a/python/restate/agents/_conversation.py b/python/restate/agents/_conversation.py index 03567745..d0cf3d3b 100644 --- a/python/restate/agents/_conversation.py +++ b/python/restate/agents/_conversation.py @@ -269,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: @@ -327,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.""" @@ -523,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 = ?" ) @@ -691,18 +704,16 @@ async def ask_again(self, text: str) -> None: await self.append(self._ledger.items.prompt(text)) async def steer( - self, - messages: Sequence[Steer], - sent: Callable[[Transaction, list[str]], None] | None = None, - going: Callable[[], bool] | None = None, - ) -> list[Entry]: - """Store the steered ``messages`` that are not stored yet as messages of the user in the turn, and return them. - - Each is stored once, by its id, as a step of the turn after those - before: the model reads it in its next request, and in every turn's - history after, between the turn's message and its answer. ``sent`` - adds what goes out with the transaction that stores them (the - companion's ``taken``), which commits once. + 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 @@ -710,48 +721,46 @@ async def steer( turn stops: the turn then closes after it, at the positions after its messages, and the messages are its own (stopped with it). """ - ledger = self._ledger - - async def store() -> list[Entry]: - async with ledger.lock: - if going is not None and not going(): - return [] - self._check_writable() - wanted = list({message["id"]: message for message in messages}.values()) - if not wanted: - return [] - ids = [message["id"] for message in wanted] - marks = ", ".join("?" for _ in ids) - # The database, not the turn: an id is stored once in the conversation, by whichever turn took it. - stored = await ledger.db.values(f"SELECT id FROM {STEERED} WHERE id IN ({marks})", ids) - known = {str(row[0]) for row in stored} - new = [message for message in wanted if message["id"] not in known] - if not new: - return [] - n = ledger.next - added = [ - Entry(n + i, ledger.items.prompt(message["text"]), steer=message["id"]) - for i, message in enumerate(new) - ] - async with ledger.db.transaction() as tx: - for i, entry in enumerate(added, 1): - dumped = ledger.items.dump(entry.item) - await tx.execute(_INSERT_STEP, [entry.n, self.n, ledger.seq + i, dumped, None, None]) - await tx.execute( - f"INSERT INTO {STEERED} (n, turn, id) VALUES (?, ?, ?)", [entry.n, self.n, entry.steer] - ) - if sent is not None: - sent(tx, [typing.cast(str, entry.steer) for entry in added]) - 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 added - + 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(store()) + 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. @@ -1173,18 +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, the messages steered - to it, 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 "") + _steered_size(turn) + 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 @@ -1212,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. @@ -1363,9 +1365,8 @@ async def _transcript(db: DatabaseReader, items: Items, first: int, bounded: tup 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 "") + _steered_size(entry) + 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: @@ -1382,9 +1383,13 @@ async def _steps(db: DatabaseReader, items: Items, after: int) -> list[Entry]: ) return [ Entry( - int(typing.cast(int, n)), load_item(items, str(item)), reply=_int(reply), agent=_str(agent), steer=_str(i) + int(typing.cast(int, n)), + load_item(items, str(item)), + reply=_int(reply), + agent=_str(agent), + steer=_str(steer), ) - for n, item, reply, agent, i in stored + for n, item, reply, agent, steer in stored ] @@ -1407,15 +1412,27 @@ async def _steered_in(db: DatabaseReader, items: Items, first: int, end: int) -> return steered -def _steered_size(turn: Entry) -> int: - """Return the size, in characters, of the messages steered to ``turn`` (a prompt of the transcript).""" - return sum(len(steered.said[1]) if steered.said else 0 for steered in turn.steers) +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``: one read, of its unique id.""" + """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 await db.get(f"SELECT 1 FROM {STEERED} WHERE id = ?", [steer_id]) is not None + 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: diff --git a/python/restate/agents/_messages.py b/python/restate/agents/_messages.py index e2f530d8..c9082f8b 100644 --- a/python/restate/agents/_messages.py +++ b/python/restate/agents/_messages.py @@ -194,19 +194,20 @@ def _summarized( 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)))) - for n, seq, stored in steered: - 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))) + 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) diff --git a/python/restate/agents/_object.py b/python/restate/agents/_object.py index 6b9d2acc..3f06fe0a 100644 --- a/python/restate/agents/_object.py +++ b/python/restate/agents/_object.py @@ -31,6 +31,7 @@ Resolved, Signal, Status, + Steer, Steered, Stopped, View, @@ -479,7 +480,7 @@ 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: _signals.SteerRequest) -> None: + 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 diff --git a/python/restate/agents/_signals.py b/python/restate/agents/_signals.py index a1b47759..951a077e 100644 --- a/python/restate/agents/_signals.py +++ b/python/restate/agents/_signals.py @@ -119,13 +119,6 @@ class Started(Answering, total=False): SDK, or of an agent whose framework cannot take them) gets none: each is answered as a turn of its own.""" -class SteerRequest(TypedDict): - """A message to steer to the turn being answered: its id (the ``steer`` request's) and its text.""" - - id: str - text: str - - class _Held(TypedDict): """A message steered to the turn ``turn`` (an invocation id), and when it was sent.""" @@ -236,7 +229,7 @@ def _set_inbox(ctx: ObjectContext, inbox: list[_Inboxed]) -> None: @signals.handler(ingress_private=True) -async def steer(ctx: ObjectContext, request: SteerRequest) -> Steered: +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 @@ -247,7 +240,7 @@ async def steer(ctx: ObjectContext, request: SteerRequest) -> Steered: at its next step, through ``publish``, until it takes it (``taken``) or ends. """ - message_id, text = request["id"], request["text"] + 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}") diff --git a/python/restate/agents/_steering.py b/python/restate/agents/_steering.py index c1891696..cf02d04f 100644 --- a/python/restate/agents/_steering.py +++ b/python/restate/agents/_steering.py @@ -31,47 +31,44 @@ from __future__ import annotations from restate.agents import _signals -from restate.agents._conversation import Conversation, Entry -from restate.agents._turn import steering_of, stop_signal_of +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[Entry]: - """Take the messages steered to the turn being answered, at a step: store each once, and return the new ones. +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 gives the model the returned - entries' items (messages of the user), after what the request already - holds. It returns nothing, at no cost, unless the turn runs in a ``chat`` - handler of an agent that takes steered messages, a message is waiting, - the turn is not stopping, and each of its tool calls has a result. + 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. """ - steering = steering_of(ctx) - if steering is None or not steering.messages: + inbox = inbox_of(ctx) + if inbox is None or not inbox.messages: return [] stop = stop_signal_of(ctx) - if stop is not None and stop.stopped: - return [] + + 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 turn is None or turn.open_calls(): - return [] - taken = {entry.steer for entry in turn.steered()} - fresh = [message for message in steering.messages if message["id"] not in taken] - if not fresh: + if not going() or turn is None or turn.open_calls(): return [] companion = _signals.signals_of(ctx) - def sent(tx: Transaction, ids: list[str]) -> None: + 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]}") - 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 - - return await turn.steer(fresh, sent, going) + return await turn.steer(inbox.messages, taken, going) diff --git a/python/restate/agents/_turn.py b/python/restate/agents/_turn.py index 7b1b466f..597085c2 100644 --- a/python/restate/agents/_turn.py +++ b/python/restate/agents/_turn.py @@ -168,8 +168,8 @@ def _unfollowed(ctx: SessionContext) -> bool: @dataclass -class Steering: - """The messages steered to a ``chat`` turn that it has not taken yet, as the companion last told it. +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 @@ -190,12 +190,12 @@ def saw(self, published: Published | None) -> None: self.messages = list(published.get("steering", [])) -_STEERING: ScopeKey[Steering] = ScopeKey("steering") +_INBOX: ScopeKey[Inbox] = ScopeKey("inbox") -def steering_of(ctx: SessionContext) -> Steering | None: - """Return the steered messages of the ``chat`` turn running on ``ctx``, if it takes them: ``None`` outside one.""" - return _STEERING.get(ctx) +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") @@ -230,9 +230,9 @@ async def publish(ctx: SessionContext) -> None: if stored is not None: await _settled(ctx, stored) published = await ctx.call(_signals.publish, key=_signals.signals_of(ctx)) - steering = _STEERING.get(ctx) - if steering is not None: - steering.saw(typing.cast(Published | None, published)) + inbox = _INBOX.get(ctx) + if inbox is not None: + inbox.saw(typing.cast(Published | None, published)) async def stepped(ctx: SessionContext) -> None: @@ -326,7 +326,7 @@ async def chat_turn(ctx: SessionContext, text: str, answer: Callable[[SessionCon steers = settings.steers and not settings.background if steers: # Before the first publish, which may return messages steered to the turn already. - _STEERING.set(ctx, Steering()) + _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, "steers": steers} diff --git a/python/restate/agents/adk/_agent.py b/python/restate/agents/adk/_agent.py index 1544d2f0..d515e3f9 100644 --- a/python/restate/agents/adk/_agent.py +++ b/python/restate/agents/adk/_agent.py @@ -161,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) @@ -192,9 +198,8 @@ def read(self, item: Item) -> Parsed: class _Sessions(BaseSessionService): """ADK's sessions, of which there is one: the conversation, an event per item.""" - def __init__(self, ctx: SessionContext, conversation: Conversation) -> None: + def __init__(self, conversation: Conversation) -> None: super().__init__() - self.ctx = ctx self.conversation = conversation def session(self) -> Session: @@ -252,22 +257,12 @@ async def append_event(self, session: Session, event: Event) -> Event: self._carry(session) return event - async def steered(self, session: Session) -> list[Event]: - """Take the messages steered to the turn (``take_steers``), and add them to ``session`` after its events. + 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 them too, in place. Each starts an - invocation of its own, as a prompt does (``_Items.prompt``), and 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. + ADK builds each model request from the session's events: every request after this one has it too, in place. """ - taken = [typing.cast(Event, entry.item) for entry in await take_steers(self.ctx, self.conversation)] - for event in taken: - # Not stored again: the turn stored them as it took them. - await super().append_event(session, event) - return taken + await super().append_event(session, event) class _Restate(BasePlugin): @@ -288,7 +283,8 @@ async def before_model_callback( # 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 await self.sessions.steered(callback_context.session): + 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 @@ -473,7 +469,7 @@ 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 - sessions = _Sessions(ctx, 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)] diff --git a/python/restate/agents/internal.py b/python/restate/agents/internal.py index 0d924332..43493e9e 100644 --- a/python/restate/agents/internal.py +++ b/python/restate/agents/internal.py @@ -33,9 +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 messages steered to it, which - the next model request gets (an integration says it calls it with - ``chat(..., steers=True)``). +- ``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 ( diff --git a/python/restate/agents/langgraph/_agent.py b/python/restate/agents/langgraph/_agent.py index 6d7a0f70..2d3fe2a5 100644 --- a/python/restate/agents/langgraph/_agent.py +++ b/python/restate/agents/langgraph/_agent.py @@ -38,7 +38,7 @@ 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. In the graph its id is ``steer:``. + 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 @@ -101,7 +101,7 @@ model_of, provider_of, ) -from restate.agents._conversation import Call, Conversation, Entry, Item, Parsed, Result, Run, answer_once +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 @@ -240,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.""" @@ -253,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) @@ -295,36 +308,6 @@ def _window(model: BaseChatModel) -> int | None: return int(limit * WINDOW) if limit is not None and limit > 0 else None -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) - - -def _as_steered(entry: Entry) -> BaseMessage: - """Return the message steered to a turn that ``entry`` is, as the graph holds it: under its ``steer:`` id.""" - return typing.cast(BaseMessage, entry.item).model_copy(update={"id": f"{STEER}{entry.steer}"}) - - -def _history(conversation: Conversation) -> list[BaseMessage]: - """Return what the model gets (``Conversation.history``), each message steered to a turn under its ``steer:`` id. - - The conversation stores a steered message as a message of the user, - without that id. So the messages that a retried attempt finds in its turn, - and those of the turns before, are marked here, as the run marks the - messages it takes (``_as_steered``). - """ - steered = {id(entry.item): entry for turn in conversation.transcript for entry in turn.steers} - steered |= {id(entry.item): entry for turn in conversation.turns for entry in turn.steered()} - return [ - _as_steered(entry) if (entry := steered.get(id(item))) is not None else typing.cast(BaseMessage, item) - for item in conversation.history() - ] - - 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``). @@ -366,7 +349,7 @@ def _graph( async def call_model(state: State) -> dict[str, Any]: # 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 = [_as_steered(entry) for entry in await take_steers(ctx, conversation)] + 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) @@ -407,7 +390,7 @@ async def run(conversation: Conversation) -> str: return final graph = _graph(ctx, conversation, chat_model, tools, system) try: - state = typing.cast(State, await graph.ainvoke({"messages": _history(conversation)})) + 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. diff --git a/python/restate/agents/openai/_agent.py b/python/restate/agents/openai/_agent.py index 61692e68..22cd0722 100644 --- a/python/restate/agents/openai/_agent.py +++ b/python/restate/agents/openai/_agent.py @@ -304,10 +304,8 @@ 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.steered: list[_Steered] = [] - """The messages steered to the turn that the run took, in order: each later request gets them again.""" - 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.""" + self.steers = _Steers() + """The messages steered to the turn that the run took: each later request gets them again.""" async def model_input(self, data: CallModelData[Any]) -> ModelInputData: """The run's ``call_model_input_filter``: store the tools' outputs before the next request, and take the @@ -327,12 +325,8 @@ async def model_input(self, data: CallModelData[Any]) -> ModelInputData: # 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)): - # This request's items came before the messages; an item whose key the run sees later, after them. - self._see(items) - self.steered.extend(_Steered(entry.item, after, len(self.seen)) for entry in taken) - if self.steered: - self._see(items) - data.model_data.input = _spliced(items, self.steered, self.seen) + self.steers.took(taken, items, after) + data.model_data.input = self.steers.placed(items) filtered = data.model_data if self.then is not None: given = self.then(data) @@ -357,6 +351,35 @@ async def on_handoff(self, context: RunContextWrapper[Any], from_agent: Agent[An async def on_tool_end(self, context: RunContextWrapper[Any], agent: Agent[Any], tool: SdkTool, result: Any) -> None: 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: diff --git a/python/restate/agents/pydantic_ai/_agent.py b/python/restate/agents/pydantic_ai/_agent.py index bc0ae513..0f871fb0 100644 --- a/python/restate/agents/pydantic_ai/_agent.py +++ b/python/restate/agents/pydantic_ai/_agent.py @@ -703,24 +703,20 @@ async def before_model_request( 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, once the model's tool calls all have their - results stored, it takes the messages that wait (``take_steers``): each - is stored once, as a message of the user that is a step of the turn. It - then enqueues the stored messages with Pydantic AI's own - ``RunContext.enqueue`` (``'asap'``), which delivers them once every - ``before_model_request`` ran: into this request and into the run's - history, after what both hold. So the model gets each message as it is - stored and where it is stored, and the run goes on with it, its usage - limits, retries and capabilities' state with it, rather than starting - again; the turn's next requests, a retried attempt and the next turns - have it in their history, in place, and never take it again. + """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 stands for the positions stored before the messages, so - the history after it (the next request's, a retried attempt's, the next - turn's) has them after it. + 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: @@ -738,9 +734,8 @@ async def before_model_request( # 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 - steered = await take_steers(self.ctx, self.conversation) - if steered: - ctx.enqueue(*(typing.cast(ModelRequest, entry.item) for entry in steered), priority="asap") + if steered := typing.cast(list[ModelRequest], await take_steers(self.ctx, self.conversation)): + ctx.enqueue(*steered, priority="asap") return request_context diff --git a/scripts/resurrect.py b/scripts/resurrect.py index 08b42f64..e4a728dc 100644 --- a/scripts/resurrect.py +++ b/scripts/resurrect.py @@ -565,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( @@ -701,8 +701,8 @@ def kind(self) -> str: "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 turn is None or turn.open_calls():\n", - " if turn is None:\n", + " if not going() or turn is None or turn.open_calls():\n", + " if not going() or turn is None:\n", STEERING, ), Mutant( @@ -741,7 +741,7 @@ def kind(self) -> str: "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 going is not None and not going():\n return []\n", + " if not going():\n return []\n", "", STEERING, ), @@ -749,16 +749,24 @@ def kind(self) -> str: "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", + " 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( "steer-summary-hides-steers", "the summary view leaves out the messages steered to a turn", f"{AGENTS}/_messages.py", - " for n, seq, stored in steered:\n", - " for n, seq, stored in ():\n", + " messages.extend(user_message(n, seq, stored) for n, seq, stored in steered)\n", + "", STEERING, ), ] diff --git a/tests/agents_adk.py b/tests/agents_adk.py index 25a97fb6..9aad2fd3 100644 --- a/tests/agents_adk.py +++ b/tests/agents_adk.py @@ -952,7 +952,7 @@ async def test_a_step_after_which_adk_asks_the_model_nothing_takes_no_steered_me 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 _STEERING, Steering # 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): @@ -982,13 +982,13 @@ async def generate_content_async( call = types.FunctionCall(id="call-1", name=name, args=args) yield LlmResponse(content=types.Content(role="model", parts=[types.Part(function_call=call)])) - steering = Steering(revision=0) + 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: - steering.messages.append({"id": "s1", "text": "use Celsius"}) + 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. @@ -999,7 +999,7 @@ async def before_tool_callback(self, *, tool: Any, tool_args: Any, tool_context: @asker.handler async def ask(ctx: sessions.SessionContext, text: str) -> Any: # As a chat turn of an agent that takes steered messages sets it. - _STEERING.set(ctx, steering) + _INBOX.set(ctx, inbox) return await adk.prompt(ctx, text, agent=app) @asker.handler @@ -1023,15 +1023,15 @@ async def test_a_steered_message_starts_an_invocation_of_its_own_which_a_retry_r 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 _STEERING, Steering # 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 - steering = Steering(revision=0) + 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.""" - steering.messages.append({"id": "s1", "text": "in Celsius"}) + inbox.messages.append({"id": "s1", "text": "in Celsius"}) return f"sunny in {city}" class Crashing(Scripted): @@ -1052,7 +1052,7 @@ async def generate_content_async( @asker.handler async def ask(ctx: sessions.SessionContext, text: str) -> str: # As a chat turn of an agent that takes steered messages sets it. - _STEERING.set(ctx, steering) + _INBOX.set(ctx, inbox) return await adk.prompt(ctx, text, agent=steps) backend, invocation = Backend(), request() diff --git a/tests/agents_signals.py b/tests/agents_signals.py index 94ebac23..99980d83 100644 --- a/tests/agents_signals.py +++ b/tests/agents_signals.py @@ -24,7 +24,7 @@ from support import RESTATE_IMAGE # pylint: disable=import-error from restate import HarnessEnvironment from restate.agents import _signals -from restate.agents._api import Published, Signal, Status, Steered, Watch +from restate.agents._api import Published, Signal, Status, Steer, Steered, Watch from restate.agents._signals import STEERS, Started, WaitRequest @@ -57,7 +57,7 @@ async def close_turn(ctx: restate.Context, request: dict[str, str]) -> None: @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(_signals.SteerRequest, message)) + return await ctx.object_call(_signals.steer, key=request["key"], arg=typing.cast(Steer, message)) @driver.handler() diff --git a/tests/agents_steering.py b/tests/agents_steering.py index 0644d880..b4ac17c9 100644 --- a/tests/agents_steering.py +++ b/tests/agents_steering.py @@ -27,16 +27,10 @@ from restate import sessions from restate.agents import _messages from restate.agents._api import Steer -from restate.agents._conversation import ( - SCHEMA, - STEERED, - Conversation, - Entry, - answer_once, - steer_stored, -) +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 _STEERING, _STOP, Steering, Stop +from restate.agents._turn import _INBOX, _STOP, Inbox, Stop from restate.sessions.internal import Schema, invoke @@ -51,18 +45,18 @@ def anyio_backend(): class Asker: - """An object whose handler answers each prompt with ``run``, steered with what ``steering`` holds.""" + """An object whose handler answers each prompt with ``run``, steered with what ``inbox`` holds.""" - def __init__(self, run: Run, *, window: int | None = None, check: Callable[[str], str | None] | None = None): + def __init__(self, run: Run, *, window: int | None = None, check: Check | None = None): self.obj = sessions.SessionObject("asker") self.backend = Backend() - self.steering: list[Steer] = [] + 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. - _STEERING.set(ctx, Steering(revision=0, messages=list(self.steering))) + _INBOX.set(ctx, Inbox(revision=0, messages=list(self.inbox))) async def ran(conversation: Conversation) -> str: return await run(ctx, conversation) @@ -92,8 +86,8 @@ async def tool_round(conversation: Conversation, call_id: str = "c1") -> None: await conversation.turn.append({"result": call_id, "output": "found"}) -def texts(entries: list[Entry]) -> list[str]: - return [entry.item["text"] for entry in entries] +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(): @@ -116,7 +110,7 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: return "done, in Celsius" asker = Asker(run) - asker.steering = [{"id": "s1", "text": "use Celsius"}] + 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. @@ -133,20 +127,20 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: async def test_no_message_is_taken_while_a_tool_call_has_no_result_nor_outside_a_chat_turn(): - took: list[list[Entry]] = [] + 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"}) - _STEERING.delete(ctx) + _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.steering = [{"id": "s1", "text": "use Celsius"}] + asker.inbox = [{"id": "s1", "text": "use Celsius"}] assert await asker.ask("the weather?") == "done" assert took == [[], []] assert asker.taken() == [] @@ -168,7 +162,7 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: return "done" asker, invocation = Asker(run), request() - asker.steering = [{"id": "s1", "text": "use Celsius"}] + 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" @@ -190,9 +184,9 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: # A window of 0: every earlier turn is its messages and answer. asker = Asker(run, window=0) - asker.steering = [{"id": "s1", "text": "and then?"}] + asker.inbox = [{"id": "s1", "text": "and then?"}] await asker.ask("first") - asker.steering = [] + asker.inbox = [] await asker.ask("second") assert seen[1] == [ {"role": "user", "text": "first"}, @@ -214,7 +208,7 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: # A window of 0: the turn just answered is before the next turn's window. asker = Asker(run, window=0) - asker.steering = [{"id": "s1", "text": "use Celsius"}] + 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 @@ -243,7 +237,7 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: return "done" asker = Asker(run) - asker.steering = [{"id": "s1", "text": "use Celsius"}] + asker.inbox = [{"id": "s1", "text": "use Celsius"}] assert await asker.ask("the weather?") == "done" assert asker.taken() == [] @@ -262,8 +256,8 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: 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.steering = [{"id": "s1", "text": "be brief"}] + 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."] @@ -286,7 +280,7 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: return "again" asker = Asker(run) - asker.steering = [{"id": "s1", "text": "use Celsius"}] + asker.inbox = [{"id": "s1", "text": "use Celsius"}] assert await asker.ask("the weather?") == "again" @@ -298,7 +292,7 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: return "done" asker = Asker(run) - asker.steering = [{"id": "s1", "text": "use Celsius"}] + 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"]] @@ -339,5 +333,5 @@ async def run(ctx: sessions.SessionContext, conversation: Conversation) -> str: return "done" asker = Asker(run) - asker.steering = typing.cast(list[Steer], [{"id": "s1", "text": "once"}, {"id": "s1", "text": "once"}]) + asker.inbox = typing.cast(list[Steer], [{"id": "s1", "text": "once"}, {"id": "s1", "text": "once"}]) assert await asker.ask("hi") == "done" From a3d0b1cca91876836205a4ba9c0898847bc41e45 Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 11:09:35 +0000 Subject: [PATCH 12/13] Agents: the steering tests' docstring names the files that test against a server agents_object.py has no steering tests: the companion's are in agents_signals.py, and the agent's steered placeholder's in agents_steering_server.py (Pullfrog's review). Co-Authored-By: Claude Opus 5.5 --- tests/agents_steering.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/tests/agents_steering.py b/tests/agents_steering.py index b4ac17c9..1851e4cf 100644 --- a/tests/agents_steering.py +++ b/tests/agents_steering.py @@ -9,7 +9,8 @@ # 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 and the agent's ``steered`` are tested against a server in ``agents_object.py``.""" +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 From fe793715188ec2095d6738e9b69af8844b5ff37f Mon Sep 17 00:00:00 2001 From: Claude Date: Mon, 9 Nov 2026 11:25:12 +0000 Subject: [PATCH 13/13] Agents: an OpenAI agent used as a tool runs apart from the turn `Agent.as_tool` runs its agent with the outer run's `run_config` when given none, so its requests went through the turn's input filter: they got the messages steered to the turn, and the turn stored the nested run's tool outputs, without their calls (which breaks the next turn's request). A request made while one of the run's tool calls runs is a nested run's (the SDK makes the run's own next request once all of its calls end): it gets the user's own filter only. Test and mutant `openai-nested-run-filtered`. Co-Authored-By: Claude Opus 5.5 --- docs/agents/frameworks.md | 6 ++++ python/restate/agents/openai/_agent.py | 26 +++++++++++++--- scripts/resurrect.py | 8 +++++ tests/agents_openai.py | 42 ++++++++++++++++++++++++++ 4 files changed, 77 insertions(+), 5 deletions(-) diff --git a/docs/agents/frameworks.md b/docs/agents/frameworks.md index 9462ac02..a0cfc039 100644 --- a/docs/agents/frameworks.md +++ b/docs/agents/frameworks.md @@ -121,6 +121,12 @@ In every framework, a tool can wait for a person: `wait_for_signal`. 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 diff --git a/python/restate/agents/openai/_agent.py b/python/restate/agents/openai/_agent.py index 22cd0722..731b4fe6 100644 --- a/python/restate/agents/openai/_agent.py +++ b/python/restate/agents/openai/_agent.py @@ -80,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 @@ -306,6 +308,8 @@ def __init__(self, ctx: SessionContext, conversation: Conversation, config: RunC 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, and take the @@ -315,7 +319,13 @@ async def model_input(self, data: CallModelData[Any]) -> ModelInputData: 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 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) @@ -327,21 +337,26 @@ async def model_input(self, data: CallModelData[Any]) -> ModelInputData: 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 = data.model_data - if self.then is not None: - given = self.then(data) - filtered = await given if inspect.isawaitable(given) else given + 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: @@ -349,6 +364,7 @@ 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) diff --git a/scripts/resurrect.py b/scripts/resurrect.py index e4a728dc..7851dd26 100644 --- a/scripts/resurrect.py +++ b/scripts/resurrect.py @@ -761,6 +761,14 @@ def kind(self) -> str: " 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", diff --git a/tests/agents_openai.py b/tests/agents_openai.py index 0310aa42..a70bd336 100644 --- a/tests/agents_openai.py +++ b/tests/agents_openai.py @@ -1512,6 +1512,48 @@ def desk_script(system: str | None, items: list[Item]) -> list[Output]: 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