mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 07:17:54 +08:00
core/agent: keep executed tool results when a failed turn is rolled back
A tool batch counted as committed only once every call had its result, so a batch cut short was dropped whole and the retry ran tools that had already had effects. Keep the calls that did report a result: their assistant row stays as the pairing carrier (a request must answer every tool_calls id), without the text of the unfinished turn, and the calls without a result are absent so the retry re-issues them instead of assuming they ran. The kept rows go to the store right away, and a failed rollback write is not allowed to mask the failure being rolled back.
This commit is contained in:
@@ -1,9 +1,11 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import contextlib
|
||||
from contextvars import ContextVar, Token
|
||||
from pathlib import Path
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
import msgspec
|
||||
from msgspec import UNSET
|
||||
|
||||
from plyngent.lmproto.openai_compatible.model import (
|
||||
@@ -37,7 +39,7 @@ from .usage import TokenUsage
|
||||
if TYPE_CHECKING:
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable, Sequence
|
||||
|
||||
from plyngent.lmproto.openai_compatible.model import AnyChatMessage
|
||||
from plyngent.lmproto.openai_compatible.model import AnyAssistantToolCall, AnyChatMessage
|
||||
from plyngent.memory import MemoryStore
|
||||
|
||||
from .events import AgentEvent
|
||||
@@ -206,6 +208,75 @@ def committed_prefix_end(messages: Sequence[AnyChatMessage], user_index: int) ->
|
||||
return end
|
||||
|
||||
|
||||
def rollback_tail(messages: Sequence[AnyChatMessage], start: int) -> list[AnyChatMessage]:
|
||||
"""Rows to keep from *start* when a failed turn is rolled back.
|
||||
|
||||
The trailing tool batch is committed only once every call has its result
|
||||
(:func:`committed_prefix_end`), so a batch cut short — cancelled before the
|
||||
last result landed, or ids that do not pair up exactly — would be dropped
|
||||
whole and the retry would run calls that already had external effects. Keep
|
||||
the executed ones: their assistant message stays as the carrier for the
|
||||
pairing (API messages must answer every ``tool_calls`` id), without the text
|
||||
of the unfinished turn, and the calls that never reported a result are
|
||||
absent so the retry re-issues them instead of assuming they ran.
|
||||
|
||||
Developer rows (host notices, directive checkpoints) are already durable, so
|
||||
they stay in place; unfinished assistant text and stray rows go.
|
||||
"""
|
||||
body = list(messages[start:])
|
||||
kept: list[AnyChatMessage] = []
|
||||
index = 0
|
||||
while index < len(body):
|
||||
message = body[index]
|
||||
if isinstance(message, DeveloperChatMessage):
|
||||
kept.append(message)
|
||||
index += 1
|
||||
continue
|
||||
if isinstance(message, AssistantChatMessage) and message.tool_calls is not UNSET and message.tool_calls:
|
||||
results: list[ToolChatMessage] = []
|
||||
follow = index + 1
|
||||
while follow < len(body):
|
||||
following = body[follow]
|
||||
if not isinstance(following, ToolChatMessage):
|
||||
break
|
||||
results.append(following)
|
||||
follow += 1
|
||||
kept.extend(_executed_tool_pair(message, message.tool_calls, results))
|
||||
index = follow
|
||||
continue
|
||||
index += 1
|
||||
return kept
|
||||
|
||||
|
||||
def _executed_tool_pair(
|
||||
assistant: AssistantChatMessage,
|
||||
calls: Sequence[AnyAssistantToolCall],
|
||||
results: Sequence[ToolChatMessage],
|
||||
) -> list[AnyChatMessage]:
|
||||
"""Assistant call record plus the results of the calls that produced one.
|
||||
|
||||
Pairs by call id (first result per id), so the history keeps the API's
|
||||
one-result-per-call shape even when a model repeats an id. Text and
|
||||
reasoning go with the rest of the unfinished turn.
|
||||
"""
|
||||
by_id: dict[str, ToolChatMessage] = {}
|
||||
for result in results:
|
||||
_ = by_id.setdefault(result.tool_call_id, result)
|
||||
executed: list[AnyAssistantToolCall] = []
|
||||
paired: list[AnyChatMessage] = []
|
||||
seen: set[str] = set()
|
||||
for call in calls:
|
||||
if call.id in seen or call.id not in by_id:
|
||||
continue
|
||||
seen.add(call.id)
|
||||
executed.append(call)
|
||||
paired.append(by_id[call.id])
|
||||
if not executed:
|
||||
return []
|
||||
record = msgspec.structs.replace(assistant, content=UNSET, reasoning_content=UNSET, tool_calls=executed)
|
||||
return [record, *paired]
|
||||
|
||||
|
||||
class ChatAgent:
|
||||
"""Thin wrapper: chat client + optional tools + optional memory bind."""
|
||||
|
||||
@@ -475,10 +546,20 @@ class ChatAgent:
|
||||
msg = "nothing to retry"
|
||||
raise RuntimeError(msg)
|
||||
|
||||
def _rollback_uncommitted(self, user_index: int) -> None:
|
||||
"""Drop incomplete suffix; keep committed tool rounds after the user."""
|
||||
async def _rollback_uncommitted(self, user_index: int) -> None:
|
||||
"""Drop the unfinished suffix; keep committed rounds and executed tools.
|
||||
|
||||
Rows before the checkpoint cursor are already stored, so only the kept
|
||||
suffix is written: a restart has to continue from the same tool results
|
||||
the model is about to see, or the retry re-runs their effects.
|
||||
"""
|
||||
end = committed_prefix_end(self.messages, user_index)
|
||||
del self.messages[end:]
|
||||
self.messages[end:] = rollback_tail(self.messages, end)
|
||||
if len(self.messages) > self._persist_from:
|
||||
# A store error must not mask the failure being rolled back; the
|
||||
# rows stay in memory, and retry() writes what is still unstored.
|
||||
with contextlib.suppress(Exception):
|
||||
await self._persist_range(self._persist_from, len(self.messages))
|
||||
|
||||
def _developer_tail_end(self, start: int) -> int:
|
||||
"""Extend *start* through trailing developer checkpoint messages."""
|
||||
@@ -492,7 +573,8 @@ class ChatAgent:
|
||||
|
||||
After each completed tool batch, commits assistant+tool messages to the
|
||||
DB and keeps them on failure so :meth:`retry` continues without redoing
|
||||
side-effecting tools. Unfinished assistant/stream suffix is rolled back.
|
||||
side-effecting tools. Without such a batch, only the unfinished text is
|
||||
dropped: calls that did report a result survive (:func:`rollback_tail`).
|
||||
"""
|
||||
user_index = self._user_index(user_msg)
|
||||
if self.todo_stack is not None:
|
||||
@@ -550,7 +632,7 @@ class ChatAgent:
|
||||
completed = True
|
||||
except BaseException:
|
||||
if not completed:
|
||||
self._rollback_uncommitted(user_index)
|
||||
await self._rollback_uncommitted(user_index)
|
||||
raise
|
||||
|
||||
self.last_turn_usage = turn_usage
|
||||
@@ -687,7 +769,8 @@ class ChatAgent:
|
||||
user_msg = self.messages[user_index]
|
||||
assert isinstance(user_msg, UserChatMessage)
|
||||
self._ensure_system_prompt()
|
||||
# Already-committed messages stay; only new rounds are persisted.
|
||||
self._persist_from = len(self.messages)
|
||||
# Already-committed rows stay; write any the failed rollback could not.
|
||||
if self._persist_from < len(self.messages):
|
||||
await self._persist_range(self._persist_from, len(self.messages))
|
||||
async for event in self._run_from_user_message(user_msg):
|
||||
yield event
|
||||
|
||||
@@ -639,6 +639,77 @@ async def test_chat_agent_failed_turn_keeps_user_in_db() -> None:
|
||||
await store.close()
|
||||
|
||||
|
||||
async def test_rollback_keeps_executed_tool_result_in_store() -> None:
|
||||
"""Rolling back a cut-short batch keeps its executed result, in memory and on disk."""
|
||||
store = await MemoryStore.open(DatabaseConfig())
|
||||
session = await store.create_session(name="t")
|
||||
ran: list[str] = []
|
||||
|
||||
@tool(register=False)
|
||||
def first() -> str:
|
||||
ran.append("first")
|
||||
return "first ran"
|
||||
|
||||
client = ScriptedClient([_response(AssistantChatMessage(content="finished"))])
|
||||
agent = ChatAgent(
|
||||
client,
|
||||
model="m",
|
||||
tools=ToolRegistry([first]),
|
||||
memory=store,
|
||||
session_id=session.sid,
|
||||
stream=False,
|
||||
)
|
||||
user = UserChatMessage(content="do both")
|
||||
await store.append_message(session.sid, user)
|
||||
agent.replace_messages(
|
||||
[
|
||||
user,
|
||||
AssistantChatMessage(
|
||||
content="running both",
|
||||
tool_calls=[
|
||||
AssistantFunctionToolCall(
|
||||
id="call_1", function=AssistantFunctionTool(name="first", arguments="{}")
|
||||
),
|
||||
AssistantFunctionToolCall(
|
||||
id="call_2", function=AssistantFunctionTool(name="other", arguments="{}")
|
||||
),
|
||||
],
|
||||
),
|
||||
ToolChatMessage(tool_call_id="call_1", content="first ran"),
|
||||
],
|
||||
persist_from=1,
|
||||
)
|
||||
|
||||
await agent._rollback_uncommitted(0) # the failure path itself
|
||||
|
||||
kept = [message for message in agent.messages if isinstance(message, ToolChatMessage)]
|
||||
assert [message.content for message in kept] == ["first ran"]
|
||||
assert agent.pending_retry_text == "do both"
|
||||
# What the model continues from is on disk too, or a restart re-runs it.
|
||||
assert agent.persist_from == len(agent.messages)
|
||||
|
||||
stored = await store.list_messages(session.sid)
|
||||
assert [type(message).__name__ for message in stored] == [
|
||||
"UserChatMessage",
|
||||
"AssistantChatMessage",
|
||||
"ToolChatMessage",
|
||||
]
|
||||
record = stored[1]
|
||||
assert isinstance(record, AssistantChatMessage)
|
||||
assert record.content is UNSET
|
||||
calls = record.tool_calls
|
||||
assert calls is not UNSET
|
||||
assert [call.id for call in calls] == ["call_1"]
|
||||
|
||||
# The retry answers from the stored result instead of running the tool again.
|
||||
events = [event async for event in agent.retry()]
|
||||
assert ran == []
|
||||
assert any(isinstance(event, TextDeltaEvent) and event.content == "finished" for event in events)
|
||||
sent = client.calls[-1].messages
|
||||
assert any(isinstance(message, ToolChatMessage) and message.content == "first ran" for message in sent)
|
||||
await store.close()
|
||||
|
||||
|
||||
async def test_chat_agent_retry_after_failure() -> None:
|
||||
store = await MemoryStore.open(DatabaseConfig())
|
||||
session = await store.create_session(name="t")
|
||||
|
||||
@@ -0,0 +1,99 @@
|
||||
"""Rollback of a failed turn keeps the tool work that already ran.
|
||||
|
||||
``committed_prefix_end`` commits a batch only once every call has its result, so
|
||||
``rollback_tail`` keeps the executed part of a batch that was cut short: the
|
||||
retry then continues from real results instead of running those tools again.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from msgspec import UNSET
|
||||
|
||||
from plyngent.agent.chat import rollback_tail
|
||||
from plyngent.lmproto.openai_compatible.model import (
|
||||
AnyChatMessage,
|
||||
AssistantChatMessage,
|
||||
AssistantFunctionTool,
|
||||
AssistantFunctionToolCall,
|
||||
DeveloperChatMessage,
|
||||
ToolChatMessage,
|
||||
UserChatMessage,
|
||||
)
|
||||
|
||||
|
||||
def _call(call_id: str, name: str = "some_tool") -> AssistantFunctionToolCall:
|
||||
return AssistantFunctionToolCall(id=call_id, function=AssistantFunctionTool(name=name, arguments="{}"))
|
||||
|
||||
|
||||
def test_rollback_tail_keeps_the_executed_call_of_a_cut_short_batch() -> None:
|
||||
messages: list[AnyChatMessage] = [
|
||||
UserChatMessage(content="do both"),
|
||||
AssistantChatMessage(content="first, then the other", tool_calls=[_call("call_a"), _call("call_b")]),
|
||||
ToolChatMessage(tool_call_id="call_a", content="a ran"),
|
||||
]
|
||||
|
||||
tail = rollback_tail(messages, 1)
|
||||
|
||||
assert len(tail) == 2
|
||||
record, result = tail
|
||||
assert isinstance(record, AssistantChatMessage)
|
||||
# The assistant row survives only as the carrier of the pairing.
|
||||
assert record.content is UNSET
|
||||
assert record.reasoning_content is UNSET
|
||||
calls = record.tool_calls
|
||||
assert calls is not UNSET
|
||||
assert [call.id for call in calls] == ["call_a"]
|
||||
assert isinstance(result, ToolChatMessage)
|
||||
assert result.content == "a ran"
|
||||
|
||||
|
||||
def test_rollback_tail_drops_a_batch_that_reported_nothing() -> None:
|
||||
messages: list[AnyChatMessage] = [
|
||||
UserChatMessage(content="go"),
|
||||
AssistantChatMessage(content="about to call", tool_calls=[_call("call_a")]),
|
||||
]
|
||||
|
||||
assert rollback_tail(messages, 1) == []
|
||||
|
||||
|
||||
def test_rollback_tail_drops_unfinished_text_and_keeps_developer_rows() -> None:
|
||||
messages: list[AnyChatMessage] = [
|
||||
UserChatMessage(content="go"),
|
||||
AssistantChatMessage(content=UNSET, tool_calls=[_call("call_a")]),
|
||||
ToolChatMessage(tool_call_id="call_a", content="a ran"),
|
||||
DeveloperChatMessage(content="checkpoint"),
|
||||
AssistantChatMessage(content="half-written answer"),
|
||||
]
|
||||
|
||||
tail = rollback_tail(messages, 1)
|
||||
|
||||
assert [type(message) for message in tail] == [AssistantChatMessage, ToolChatMessage, DeveloperChatMessage]
|
||||
|
||||
|
||||
def test_rollback_tail_pairs_a_repeated_call_id_once() -> None:
|
||||
messages: list[AnyChatMessage] = [
|
||||
UserChatMessage(content="go"),
|
||||
AssistantChatMessage(content=UNSET, tool_calls=[_call("call_a"), _call("call_a")]),
|
||||
ToolChatMessage(tool_call_id="call_a", content="only once"),
|
||||
]
|
||||
|
||||
tail = rollback_tail(messages, 1)
|
||||
|
||||
assert len(tail) == 2
|
||||
record, result = tail
|
||||
assert isinstance(record, AssistantChatMessage)
|
||||
calls = record.tool_calls
|
||||
assert calls is not UNSET
|
||||
assert [call.id for call in calls] == ["call_a"]
|
||||
assert isinstance(result, ToolChatMessage)
|
||||
assert result.content == "only once"
|
||||
|
||||
|
||||
def test_rollback_tail_ignores_rows_before_start() -> None:
|
||||
messages: list[AnyChatMessage] = [
|
||||
UserChatMessage(content="go"),
|
||||
AssistantChatMessage(content=UNSET, tool_calls=[_call("call_a")]),
|
||||
ToolChatMessage(tool_call_id="call_a", content="a ran"),
|
||||
]
|
||||
|
||||
assert rollback_tail(messages, len(messages)) == []
|
||||
Reference in New Issue
Block a user