Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions python/restate/agents/_common.py
Original file line number Diff line number Diff line change
Expand Up @@ -220,6 +220,16 @@ def check_window(window: Window) -> Window:
CANCELLATION = 409
"""The status code of Restate's cancellation error, which a stopped turn fails with."""


def is_cancellation(error: BaseException) -> bool:
"""Return whether ``error`` is Restate's cancellation error, which a stopped turn fails with too.

Its status code alone does not tell: a call fails with 409 for other reasons, such as a workflow run a second
time, which fail the turn rather than stop it.
"""
return isinstance(error, TerminalError) and error.status_code == CANCELLATION and error.message == "cancelled"


CANCELLED = "Cancelled: the turn was stopped before it finished."
"""The answer of a cancelled turn, and the result of each tool call it cut off, which the model sees later in the conversation."""

Expand Down
4 changes: 2 additions & 2 deletions python/restate/agents/_conversation.py
Original file line number Diff line number Diff line change
Expand Up @@ -94,7 +94,6 @@
from restate.agents._calls import outcomes
from restate.agents._common import (
BUDGET,
CANCELLATION,
CANCELLED,
ENDED,
SUMMARIZE,
Expand All @@ -103,6 +102,7 @@
Window,
check_window,
cut_off_result,
is_cancellation,
next_prompt_id,
refusal_of,
stored_answer,
Expand Down Expand Up @@ -1402,7 +1402,7 @@ async def answering() -> str:
# The attempt ended: stop here, and let the next attempt carry on the run.
raise
except TerminalError as error:
if error.status_code == CANCELLATION:
if is_cancellation(error):
# Cancelled: close the turn, so that the next prompt starts afresh.
await turn.close(ctx, "cancelled", CANCELLED)
await conversation.seal()
Expand Down
15 changes: 11 additions & 4 deletions python/restate/agents/pydantic_ai/_agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -144,7 +144,6 @@
)
from pydantic_ai import Tool as PydanticTool
from restate.agents._common import (
CANCELLATION,
Check,
CompactHistory,
ErrorSet,
Expand All @@ -154,6 +153,7 @@
check_own,
configured,
from_tool,
is_cancellation,
model_of,
provider_of,
)
Expand Down Expand Up @@ -222,6 +222,9 @@ async def run(run_ctx: RunContext[Any], **args: Any) -> str:
_APPROVED = frozenset({"approved", "approve", "yes", "ok", "allow", "allowed"})
"""The texts that approve a tool call, whatever their case; any other text denies it, with that reason."""

_DENIED = frozenset({"", "denied", "deny", "no", "rejected", "reject", "declined", "decline", "refused", "refuse"})
"""The texts that deny a tool call with no reason of their own, whatever their case: the model is told it was denied."""


def approval_of(value: Any) -> ToolApproved | ToolDenied:
"""Return what the value of an ``approval`` signal decides about its tool call.
Expand All @@ -230,12 +233,16 @@ def approval_of(value: Any) -> ToolApproved | ToolDenied:
``{"approved": true, "args": {...}}`` approves it with other arguments.
``false`` or ``{"approved": false, "reason": "..."}`` denies it, and so
does any other text, which is the denial's reason for the model: a value
that does not say "approved" never runs the tool.
that does not say "approved" never runs the tool. A text that only says
"no" (``"denied"``, ``"rejected"``...) tells the model the call was denied.
"""
if value is True:
return ToolApproved()
if isinstance(value, str):
return ToolApproved() if value.strip().lower() in _APPROVED else ToolDenied(value)
said = value.strip().lower()
if said in _APPROVED:
return ToolApproved()
return ToolDenied() if said in _DENIED else ToolDenied(value)
if isinstance(value, dict):
decision = typing.cast(dict[str, Any], value)
if decision.get("approved") is True:
Expand Down Expand Up @@ -435,7 +442,7 @@ async def on_tool_execute_error(
stop = stop_signal_of(self.ctx)
if stop is not None and stop.reason is not None:
raise stop.reason from error
if isinstance(error, TerminalError) and error.status_code != CANCELLATION:
if isinstance(error, TerminalError) and not is_cancellation(error):
await tool_call_ended(self.ctx)
return f"Error: {error.message}"
# Mark the error as the tool's, so that it is never mistaken for the model provider refusing.
Expand Down
24 changes: 22 additions & 2 deletions tests/agents_conversation.py
Original file line number Diff line number Diff line change
Expand Up @@ -174,15 +174,35 @@ async def run(conversation: Conversation) -> str:
]


async def test_a_409_that_is_no_cancellation_fails_the_turn_rather_than_cancelling_it():
"""A call that fails with 409 for another reason (a workflow run a second time) fails the turn with its error."""
seen: list[list[Item]] = []

async def run(conversation: Conversation) -> str:
seen.append([entry.item for entry in conversation.entries])
if len(seen) == 1:
raise restate.TerminalError("the workflow method was already invoked", status_code=409)
return "fine"

ask = asking(run)
with pytest.raises(restate.TerminalError, match="already invoked"):
await ask("first")
assert await ask("second") == "fine"
assert seen[1][-2:] == [
{"role": "assistant", "text": "Error: the workflow method was already invoked"},
{"role": "user", "text": "second"},
]


async def test_a_stopped_run_stores_nothing_more():
kept: list[Conversation] = []

async def run(conversation: Conversation) -> str:
kept.append(conversation)
raise restate.TerminalError("stopped", status_code=409)
raise restate.TerminalError("cancelled", status_code=409)

ask = asking(run)
with pytest.raises(restate.TerminalError, match="stopped"):
with pytest.raises(restate.TerminalError, match="cancelled"):
await ask("first")
[conversation] = kept
turn = conversation.turn
Expand Down
4 changes: 3 additions & 1 deletion tests/agents_pydantic_ai.py
Original file line number Diff line number Diff line change
Expand Up @@ -835,7 +835,7 @@ class StopBeforeTheNextRequest(AbstractCapability[Any]):

async def before_model_request(self, ctx: RunContext[Any], request_context: Any) -> Any:
if any(isinstance(p, ToolReturnPart) for m in ctx.messages if isinstance(m, ModelRequest) for p in m.parts):
raise restate.TerminalError("stopped", status_code=409)
raise restate.TerminalError("cancelled", status_code=409)
return request_context


Expand Down Expand Up @@ -975,6 +975,8 @@ def test_an_approvals_signal_approves_only_what_says_so():
assert approval_of(approving) == ToolApproved()
assert approval_of({"approved": True, "args": {"city": "Oslo"}}) == ToolApproved(override_args={"city": "Oslo"})
assert approval_of(False) == ToolDenied()
for denying in ("no", "Denied", " rejected ", ""): # no reason of their own: the model is told it was denied
assert approval_of(denying) == ToolDenied()
assert approval_of("not today") == ToolDenied("not today")
assert approval_of({"approved": False, "reason": "too much"}) == ToolDenied("too much")
for anything in (None, 1, ["yes"], {"approved": "yes"}):
Expand Down
2 changes: 1 addition & 1 deletion tests/agents_snapshots.py
Original file line number Diff line number Diff line change
Expand Up @@ -459,7 +459,7 @@ async def run(conversation: Conversation) -> str:
await conversation.turn.append_reply(_calling("c1"))
await wait(contexts[-1])
# The turn is stopped before the call's result is stored.
raise TerminalError("stopped", status_code=409)
raise TerminalError("cancelled", status_code=409)

@obj.handler
async def ask(ctx: sessions.SessionContext, text: str) -> str:
Expand Down
Loading