mirror of
https://github.com/NCBM/plyngent.git
synced 2026-11-07 15:37:55 +08:00
NCBM
7d9b979604
A caller that goes away mid-batch (cancel, display error) left the results of tools that had already run out of history, so the retry ran them again. Append the collected results first, then emit the events.
1082 lines
38 KiB
Python
1082 lines
38 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Literal, overload
|
|
|
|
import pytest
|
|
from msgspec import UNSET
|
|
|
|
from plyngent.agent import (
|
|
AssistantMessageEvent,
|
|
ChatAgent,
|
|
MaxRoundsEvent,
|
|
ReasoningDeltaEvent,
|
|
TextDeltaEvent,
|
|
ToolCallEvent,
|
|
ToolRegistry,
|
|
ToolResultEvent,
|
|
UsageEvent,
|
|
run_chat_loop,
|
|
tool,
|
|
)
|
|
from plyngent.config.models import DatabaseConfig
|
|
from plyngent.lmproto.openai_compatible.model import (
|
|
AnyChatMessage,
|
|
AssistantChatMessage,
|
|
AssistantFunctionTool,
|
|
AssistantFunctionToolCall,
|
|
ChatCompletionChoice,
|
|
ChatCompletionChunk,
|
|
ChatCompletionResponse,
|
|
ChatCompletionsParam,
|
|
ChunkChoice,
|
|
DeltaMessage,
|
|
FinishReason,
|
|
StreamFunctionDelta,
|
|
StreamToolCallDelta,
|
|
ToolChatMessage,
|
|
UserChatMessage,
|
|
)
|
|
from plyngent.memory import MemoryStore
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
|
|
|
|
def _chunks_from_response(response: ChatCompletionResponse) -> list[ChatCompletionChunk]:
|
|
"""Turn a full response into stream chunks (library-style stream=True path)."""
|
|
message = response.choices[0].message
|
|
chunks: list[ChatCompletionChunk] = []
|
|
if isinstance(message.content, str) and message.content:
|
|
chunks.append(
|
|
ChatCompletionChunk(
|
|
id=response.id,
|
|
object="chat.completion.chunk",
|
|
created=response.created,
|
|
model=response.model,
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(content=message.content),
|
|
finish_reason=None,
|
|
)
|
|
],
|
|
)
|
|
)
|
|
# Emit terminal finish_reason when content was streamed without one.
|
|
fr = response.choices[0].finish_reason
|
|
if (
|
|
fr
|
|
and (isinstance(message.content, str) and message.content)
|
|
and (message.tool_calls is UNSET or not message.tool_calls)
|
|
):
|
|
chunks.append(
|
|
ChatCompletionChunk(
|
|
id=response.id,
|
|
object="chat.completion.chunk",
|
|
created=response.created,
|
|
model=response.model,
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(),
|
|
finish_reason=fr,
|
|
)
|
|
],
|
|
)
|
|
)
|
|
tool_calls = message.tool_calls
|
|
if tool_calls is not UNSET and tool_calls:
|
|
deltas: list[StreamToolCallDelta] = []
|
|
for i, call in enumerate(tool_calls):
|
|
if isinstance(call, AssistantFunctionToolCall):
|
|
deltas.append(
|
|
StreamToolCallDelta(
|
|
index=i,
|
|
id=call.id,
|
|
type="function",
|
|
function=StreamFunctionDelta(
|
|
name=call.function.name,
|
|
arguments=call.function.arguments,
|
|
),
|
|
)
|
|
)
|
|
if deltas:
|
|
chunks.append(
|
|
ChatCompletionChunk(
|
|
id=response.id,
|
|
object="chat.completion.chunk",
|
|
created=response.created,
|
|
model=response.model,
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(tool_calls=deltas),
|
|
finish_reason="tool_calls",
|
|
)
|
|
],
|
|
)
|
|
)
|
|
if not chunks:
|
|
chunks.append(
|
|
ChatCompletionChunk(
|
|
id=response.id,
|
|
object="chat.completion.chunk",
|
|
created=response.created,
|
|
model=response.model,
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(),
|
|
finish_reason=response.choices[0].finish_reason or "stop",
|
|
)
|
|
],
|
|
)
|
|
)
|
|
return chunks
|
|
|
|
|
|
class ScriptedClient:
|
|
"""Scripted chat completions; supports stream=True via chunked responses."""
|
|
|
|
_responses: list[ChatCompletionResponse]
|
|
calls: list[ChatCompletionsParam]
|
|
|
|
def __init__(self, responses: list[ChatCompletionResponse]) -> None:
|
|
self._responses = list(responses)
|
|
self.calls = []
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
self.calls.append(param)
|
|
if not self._responses:
|
|
msg = "no more scripted responses"
|
|
raise RuntimeError(msg)
|
|
response = self._responses.pop(0)
|
|
if stream:
|
|
return self._as_stream(response)
|
|
return response
|
|
|
|
async def _as_stream(self, response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]:
|
|
for chunk in _chunks_from_response(response):
|
|
yield chunk
|
|
|
|
|
|
def _response(
|
|
message: AssistantChatMessage,
|
|
*,
|
|
usage: dict[str, object] | None = None,
|
|
finish_reason: FinishReason | None = "stop",
|
|
) -> ChatCompletionResponse:
|
|
return ChatCompletionResponse(
|
|
id="1",
|
|
object="chat.completion",
|
|
created=0,
|
|
model="test",
|
|
choices=[
|
|
ChatCompletionChoice(
|
|
index=0,
|
|
message=message,
|
|
logprobs={},
|
|
finish_reason=finish_reason,
|
|
)
|
|
],
|
|
system_fingerprint="",
|
|
usage=usage if usage is not None else {},
|
|
)
|
|
|
|
|
|
async def test_run_chat_loop_text_only() -> None:
|
|
client = ScriptedClient(
|
|
[
|
|
_response(AssistantChatMessage(content="hello")),
|
|
]
|
|
)
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
|
|
assert isinstance(events[0], TextDeltaEvent)
|
|
assert events[0].content == "hello"
|
|
assert isinstance(events[1], AssistantMessageEvent)
|
|
assert len(messages) == 2
|
|
assert len(client.calls) == 1
|
|
|
|
|
|
async def test_non_stream_emits_usage() -> None:
|
|
client = ScriptedClient(
|
|
[
|
|
_response(
|
|
AssistantChatMessage(content="hi"),
|
|
usage={"prompt_tokens": 9, "completion_tokens": 2, "total_tokens": 11},
|
|
),
|
|
]
|
|
)
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="x")]
|
|
events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
|
|
usages = [e for e in events if isinstance(e, UsageEvent)]
|
|
assert len(usages) == 1
|
|
assert usages[0].usage.prompt_tokens == 9
|
|
assert usages[0].usage.completion_tokens == 2
|
|
assert usages[0].usage.total_tokens == 11
|
|
assert usages[0].usage.source == "api"
|
|
|
|
|
|
async def test_non_stream_estimates_usage_when_missing() -> None:
|
|
client = ScriptedClient([_response(AssistantChatMessage(content="hello"))])
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
|
|
usages = [e for e in events if isinstance(e, UsageEvent)]
|
|
assert len(usages) == 1
|
|
assert usages[0].usage.source == "estimate"
|
|
assert usages[0].usage.total_tokens > 0
|
|
|
|
|
|
async def test_chat_agent_accumulates_session_usage() -> None:
|
|
client = ScriptedClient(
|
|
[
|
|
_response(
|
|
AssistantChatMessage(content="a"),
|
|
usage={"prompt_tokens": 5, "completion_tokens": 1, "total_tokens": 6},
|
|
),
|
|
_response(
|
|
AssistantChatMessage(content="b"),
|
|
usage={"prompt_tokens": 7, "completion_tokens": 3, "total_tokens": 10},
|
|
),
|
|
]
|
|
)
|
|
agent = ChatAgent(client, model="m", stream=False)
|
|
_ = [e async for e in agent.run("one")]
|
|
assert agent.last_turn_usage.total_tokens == 6
|
|
assert agent.last_request_usage.total_tokens == 6
|
|
assert agent.last_turn_rounds == 1
|
|
assert agent.session_usage.total_tokens == 6
|
|
_ = [e async for e in agent.run("two")]
|
|
assert agent.last_turn_usage.total_tokens == 10
|
|
assert agent.last_request_usage.total_tokens == 10
|
|
assert agent.session_usage.total_tokens == 16
|
|
|
|
|
|
async def test_chat_agent_turn_usage_sums_tool_rounds() -> None:
|
|
"""Multi-round tool loop: turn usage is billing sum; last_request is final call."""
|
|
|
|
@tool(register=False)
|
|
def ping() -> str:
|
|
return "pong"
|
|
|
|
registry = ToolRegistry([ping])
|
|
client = ScriptedClient(
|
|
[
|
|
_response(
|
|
AssistantChatMessage(
|
|
content="",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="1",
|
|
function=AssistantFunctionTool(name="ping", arguments="{}"),
|
|
)
|
|
],
|
|
),
|
|
usage={"prompt_tokens": 100, "completion_tokens": 5, "total_tokens": 105},
|
|
),
|
|
_response(
|
|
AssistantChatMessage(content="done"),
|
|
usage={"prompt_tokens": 200, "completion_tokens": 3, "total_tokens": 203},
|
|
),
|
|
]
|
|
)
|
|
agent = ChatAgent(client, model="m", tools=registry, stream=False)
|
|
_ = [e async for e in agent.run("go")]
|
|
assert agent.last_turn_rounds == 2
|
|
assert agent.last_request_usage.prompt_tokens == 200
|
|
assert agent.last_turn_usage.prompt_tokens == 300
|
|
assert agent.last_turn_usage.total_tokens == 308
|
|
# Context size is last request prompt, not billed sum
|
|
assert agent.context_tokens == 200
|
|
assert agent.context_tokens_source == "api"
|
|
|
|
|
|
async def test_context_tokens_falls_back_to_message_estimate() -> None:
|
|
agent = ChatAgent(
|
|
ScriptedClient([]),
|
|
model="m",
|
|
stream=False,
|
|
messages=[UserChatMessage(content="12345678")],
|
|
)
|
|
assert agent.last_request_usage.is_zero()
|
|
assert agent.context_tokens_source == "estimate"
|
|
assert agent.context_tokens >= 1
|
|
|
|
|
|
async def test_stream_yields_deltas_incrementally() -> None:
|
|
"""Text deltas are yielded as chunks arrive, not only after the full stream."""
|
|
|
|
class ChunkClient:
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del param
|
|
if not stream:
|
|
return _response(AssistantChatMessage(content="ab"))
|
|
|
|
async def chunks() -> AsyncIterator[ChatCompletionChunk]:
|
|
for part in ("a", "b"):
|
|
yield ChatCompletionChunk(
|
|
id="1",
|
|
object="chat.completion.chunk",
|
|
created=0,
|
|
model="t",
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(content=part),
|
|
finish_reason=None,
|
|
)
|
|
],
|
|
)
|
|
|
|
return chunks()
|
|
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
events = [e async for e in run_chat_loop(ChunkClient(), messages, model="m", stream=True)]
|
|
deltas = [e for e in events if isinstance(e, TextDeltaEvent)]
|
|
assert [d.content for d in deltas] == ["a", "b"]
|
|
assert any(isinstance(e, AssistantMessageEvent) for e in events)
|
|
assert isinstance(messages[-1], AssistantChatMessage)
|
|
assert messages[-1].content == "ab"
|
|
|
|
|
|
async def test_stream_yields_reasoning_deltas() -> None:
|
|
class ReasoningClient:
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del param
|
|
if not stream:
|
|
return _response(AssistantChatMessage(content="ans", reasoning_content="think"))
|
|
|
|
async def chunks() -> AsyncIterator[ChatCompletionChunk]:
|
|
for part in ("th", "ink"):
|
|
yield ChatCompletionChunk(
|
|
id="1",
|
|
object="chat.completion.chunk",
|
|
created=0,
|
|
model="t",
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(reasoning_content=part),
|
|
finish_reason=None,
|
|
)
|
|
],
|
|
)
|
|
yield ChatCompletionChunk(
|
|
id="1",
|
|
object="chat.completion.chunk",
|
|
created=0,
|
|
model="t",
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(content="ans"),
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
)
|
|
|
|
return chunks()
|
|
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
events = [e async for e in run_chat_loop(ReasoningClient(), messages, model="m", stream=True)]
|
|
reasoning = [e for e in events if isinstance(e, ReasoningDeltaEvent)]
|
|
assert [r.content for r in reasoning] == ["th", "ink"]
|
|
deltas = [e for e in events if isinstance(e, TextDeltaEvent)]
|
|
assert [d.content for d in deltas] == ["ans"]
|
|
assert isinstance(messages[-1], AssistantChatMessage)
|
|
assert messages[-1].content == "ans"
|
|
assert messages[-1].reasoning_content == "think"
|
|
|
|
|
|
async def test_non_stream_yields_reasoning() -> None:
|
|
client = ScriptedClient([_response(AssistantChatMessage(content="ok", reasoning_content="plan"))])
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
events = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
|
|
reasoning = [e for e in events if isinstance(e, ReasoningDeltaEvent)]
|
|
assert len(reasoning) == 1
|
|
assert reasoning[0].content == "plan"
|
|
assert isinstance(messages[-1], AssistantChatMessage)
|
|
assert messages[-1].reasoning_content == "plan"
|
|
|
|
|
|
async def test_run_chat_loop_with_tools() -> None:
|
|
@tool(register=False)
|
|
def add(a: int, b: int) -> int:
|
|
return a + b
|
|
|
|
registry = ToolRegistry([add])
|
|
client = ScriptedClient(
|
|
[
|
|
_response(
|
|
AssistantChatMessage(
|
|
content="",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="c1",
|
|
function=AssistantFunctionTool(name="add", arguments='{"a": 1, "b": 2}'),
|
|
)
|
|
],
|
|
)
|
|
),
|
|
_response(AssistantChatMessage(content="3")),
|
|
]
|
|
)
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="1+2")]
|
|
events = [e async for e in run_chat_loop(client, messages, model="m", tools=registry)]
|
|
types = [type(e) for e in events]
|
|
assert ToolCallEvent in types
|
|
assert ToolResultEvent in types
|
|
assert any(isinstance(e, TextDeltaEvent) and e.content == "3" for e in events)
|
|
assert len(client.calls) == 2
|
|
# second call includes tool result message
|
|
assert any(getattr(m, "tool_call_id", None) == "c1" for m in client.calls[1].messages)
|
|
|
|
|
|
async def test_max_rounds() -> None:
|
|
@tool(register=False)
|
|
def ping() -> str:
|
|
return "pong"
|
|
|
|
registry = ToolRegistry([ping])
|
|
forever = _response(
|
|
AssistantChatMessage(
|
|
content="",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="c",
|
|
function=AssistantFunctionTool(name="ping", arguments="{}"),
|
|
)
|
|
],
|
|
)
|
|
)
|
|
client = ScriptedClient([forever, forever, forever])
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="x")]
|
|
events = [e async for e in run_chat_loop(client, messages, model="m", tools=registry, max_rounds=2)]
|
|
assert any(isinstance(e, MaxRoundsEvent) and e.rounds == 2 and not e.continued for e in events)
|
|
assert len(client.calls) == 2
|
|
|
|
|
|
async def test_max_rounds_continue_hook() -> None:
|
|
@tool(register=False)
|
|
def ping() -> str:
|
|
return "pong"
|
|
|
|
registry = ToolRegistry([ping])
|
|
forever = _response(
|
|
AssistantChatMessage(
|
|
content="",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="c",
|
|
function=AssistantFunctionTool(name="ping", arguments="{}"),
|
|
)
|
|
],
|
|
)
|
|
)
|
|
final = _response(AssistantChatMessage(content="done"))
|
|
client = ScriptedClient([forever, forever, final])
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="x")]
|
|
asks: list[str] = []
|
|
|
|
def on_limit(reason: str) -> bool:
|
|
asks.append(reason)
|
|
return len(asks) == 1
|
|
|
|
events = [
|
|
e
|
|
async for e in run_chat_loop(
|
|
client,
|
|
messages,
|
|
model="m",
|
|
tools=registry,
|
|
max_rounds=2,
|
|
on_limit=on_limit,
|
|
)
|
|
]
|
|
assert len(asks) == 1
|
|
assert any(isinstance(e, MaxRoundsEvent) and e.continued for e in events)
|
|
assert any(isinstance(e, TextDeltaEvent) and e.content == "done" for e in events)
|
|
assert len(client.calls) == 3
|
|
|
|
|
|
async def test_max_rounds_async_continue_hook() -> None:
|
|
@tool(register=False)
|
|
def ping() -> str:
|
|
return "pong"
|
|
|
|
registry = ToolRegistry([ping])
|
|
forever = _response(
|
|
AssistantChatMessage(
|
|
content="",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="c",
|
|
function=AssistantFunctionTool(name="ping", arguments="{}"),
|
|
)
|
|
],
|
|
)
|
|
)
|
|
final = _response(AssistantChatMessage(content="done"))
|
|
client = ScriptedClient([forever, forever, final])
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="x")]
|
|
asks: list[str] = []
|
|
|
|
async def on_limit(reason: str) -> bool:
|
|
asks.append(reason)
|
|
return True
|
|
|
|
events = [
|
|
e
|
|
async for e in run_chat_loop(
|
|
client,
|
|
messages,
|
|
model="m",
|
|
tools=registry,
|
|
max_rounds=2,
|
|
on_limit=on_limit,
|
|
)
|
|
]
|
|
assert len(asks) == 1
|
|
assert any(isinstance(e, MaxRoundsEvent) and e.continued for e in events)
|
|
assert any(isinstance(e, TextDeltaEvent) and e.content == "done" for e in events)
|
|
|
|
|
|
async def test_default_max_rounds_is_generous() -> None:
|
|
from plyngent.agent.loop import DEFAULT_MAX_ROUNDS
|
|
|
|
assert DEFAULT_MAX_ROUNDS >= 16
|
|
|
|
|
|
async def test_chat_agent_memory_roundtrip() -> None:
|
|
store = await MemoryStore.open(DatabaseConfig())
|
|
session = await store.create_session(name="t")
|
|
client = ScriptedClient([_response(AssistantChatMessage(content="yo"))])
|
|
agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)
|
|
events = [e async for e in agent.run("hi")]
|
|
assert any(isinstance(e, TextDeltaEvent) and e.content == "yo" for e in events)
|
|
|
|
loaded = await store.list_messages(session.sid)
|
|
assert len(loaded) == 2
|
|
assert isinstance(loaded[0], UserChatMessage)
|
|
assert loaded[0].content == "hi"
|
|
|
|
agent2 = ChatAgent(client, model="m", memory=store, session_id=session.sid)
|
|
await agent2.load_history()
|
|
assert len(agent2.messages) == 2
|
|
await store.close()
|
|
|
|
|
|
async def test_chat_agent_failed_turn_keeps_user_in_db() -> None:
|
|
store = await MemoryStore.open(DatabaseConfig())
|
|
session = await store.create_session(name="t")
|
|
|
|
class BoomClient:
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del param, stream
|
|
msg = "network down"
|
|
raise RuntimeError(msg)
|
|
|
|
agent = ChatAgent(BoomClient(), model="m", memory=store, session_id=session.sid)
|
|
with pytest.raises(RuntimeError, match="network down"):
|
|
_ = [e async for e in agent.run("hello")]
|
|
|
|
assert len(agent.messages) == 1
|
|
assert isinstance(agent.messages[0], UserChatMessage)
|
|
assert agent.pending_retry_text == "hello"
|
|
loaded = await store.list_messages(session.sid)
|
|
assert len(loaded) == 1
|
|
assert isinstance(loaded[0], UserChatMessage)
|
|
assert loaded[0].content == "hello"
|
|
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_cancelled_batch_records_every_computed_result() -> None:
|
|
"""Leaving mid-batch keeps the results that were already computed."""
|
|
ran: list[str] = []
|
|
|
|
@tool(register=False)
|
|
def first() -> str:
|
|
ran.append("first")
|
|
return "first ran"
|
|
|
|
@tool(register=False)
|
|
def second() -> str:
|
|
ran.append("second")
|
|
return "second ran"
|
|
|
|
client = ScriptedClient(
|
|
[
|
|
_response(
|
|
AssistantChatMessage(
|
|
content="running both",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="call_1", function=AssistantFunctionTool(name="first", arguments="{}")
|
|
),
|
|
AssistantFunctionToolCall(
|
|
id="call_2", function=AssistantFunctionTool(name="second", arguments="{}")
|
|
),
|
|
],
|
|
)
|
|
)
|
|
]
|
|
)
|
|
agent = ChatAgent(client, model="m", tools=ToolRegistry([first, second]), stream=False)
|
|
|
|
stream = agent.run("do both")
|
|
async for event in stream:
|
|
if isinstance(event, ToolResultEvent):
|
|
break
|
|
await stream.aclose()
|
|
|
|
results = [message for message in agent.messages if isinstance(message, ToolChatMessage)]
|
|
assert [message.content for message in results] == ["first ran", "second ran"]
|
|
assert sorted(ran) == ["first", "second"]
|
|
|
|
|
|
async def test_chat_agent_retry_after_failure() -> None:
|
|
store = await MemoryStore.open(DatabaseConfig())
|
|
session = await store.create_session(name="t")
|
|
|
|
class FlakyClient:
|
|
calls: int
|
|
|
|
def __init__(self) -> None:
|
|
self.calls = 0
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del param
|
|
self.calls += 1
|
|
if self.calls == 1:
|
|
msg = "temporary"
|
|
raise RuntimeError(msg)
|
|
response = _response(AssistantChatMessage(content="recovered"))
|
|
if stream:
|
|
|
|
async def as_stream() -> AsyncIterator[ChatCompletionChunk]:
|
|
for chunk in _chunks_from_response(response):
|
|
yield chunk
|
|
|
|
return as_stream()
|
|
return response
|
|
|
|
client = FlakyClient()
|
|
agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)
|
|
with pytest.raises(RuntimeError, match="temporary"):
|
|
_ = [e async for e in agent.run("ping")]
|
|
assert agent.pending_retry_text == "ping"
|
|
# User already in DB after first attempt
|
|
assert len(await store.list_messages(session.sid)) == 1
|
|
|
|
events = [e async for e in agent.retry()]
|
|
assert any(isinstance(e, TextDeltaEvent) and e.content == "recovered" for e in events)
|
|
assert agent.pending_retry_text is None
|
|
loaded = await store.list_messages(session.sid)
|
|
assert len(loaded) == 2
|
|
assert isinstance(loaded[0], UserChatMessage)
|
|
assert loaded[0].content == "ping"
|
|
# Single user message (no duplicate on retry)
|
|
assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1
|
|
await store.close()
|
|
|
|
|
|
async def test_retry_after_resume_orphan_user() -> None:
|
|
store = await MemoryStore.open(DatabaseConfig())
|
|
session = await store.create_session(name="t")
|
|
_ = await store.append_message(session.sid, UserChatMessage(content="left hanging"))
|
|
|
|
client = ScriptedClient([_response(AssistantChatMessage(content="ok now"))])
|
|
agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)
|
|
await agent.load_history()
|
|
assert agent.pending_retry_text == "left hanging"
|
|
|
|
events = [e async for e in agent.retry()]
|
|
assert any(isinstance(e, TextDeltaEvent) and e.content == "ok now" for e in events)
|
|
loaded = await store.list_messages(session.sid)
|
|
assert len(loaded) == 2
|
|
assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1
|
|
await store.close()
|
|
|
|
|
|
async def test_retry_keeps_committed_tools_after_second_round_fails() -> None:
|
|
"""Tools that already ran stay in history; retry continues without re-calling them."""
|
|
store = await MemoryStore.open(DatabaseConfig())
|
|
session = await store.create_session(name="t")
|
|
calls: list[str] = []
|
|
|
|
@tool(register=False)
|
|
def side_effect() -> str:
|
|
calls.append("ran")
|
|
return "done-once"
|
|
|
|
registry = ToolRegistry([side_effect])
|
|
|
|
class ToolThenBoom:
|
|
n: int
|
|
|
|
def __init__(self) -> None:
|
|
self.n = 0
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del stream
|
|
self.n += 1
|
|
if self.n == 1:
|
|
return _response(
|
|
AssistantChatMessage(
|
|
content="",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="c1",
|
|
function=AssistantFunctionTool(
|
|
name="side_effect",
|
|
arguments="{}",
|
|
),
|
|
)
|
|
],
|
|
)
|
|
)
|
|
if self.n == 2:
|
|
msg = "api down after tools"
|
|
raise RuntimeError(msg)
|
|
# Retry continues: third call is next model round with tools in history.
|
|
assert any(getattr(m, "tool_call_id", None) == "c1" for m in param.messages)
|
|
return _response(AssistantChatMessage(content="finished"))
|
|
|
|
client = ToolThenBoom()
|
|
agent = ChatAgent(
|
|
client,
|
|
model="m",
|
|
tools=registry,
|
|
memory=store,
|
|
session_id=session.sid,
|
|
stream=False,
|
|
)
|
|
with pytest.raises(RuntimeError, match="api down after tools"):
|
|
_ = [e async for e in agent.run("do it")]
|
|
|
|
assert calls == ["ran"]
|
|
assert agent.pending_retry_text == "do it"
|
|
# User + assistant(tool_calls) + tool result kept; incomplete second assistant dropped.
|
|
assert isinstance(agent.messages[-1], ToolChatMessage)
|
|
loaded = await store.list_messages(session.sid)
|
|
assert len(loaded) == 3 # user, assistant, tool — committed before boom
|
|
|
|
events = [e async for e in agent.retry()]
|
|
assert any(isinstance(e, TextDeltaEvent) and e.content == "finished" for e in events)
|
|
assert calls == ["ran"] # tool not re-executed
|
|
assert agent.pending_retry_text is None
|
|
await store.close()
|
|
|
|
|
|
async def test_chat_agent_system_prompt_prepended() -> None:
|
|
from plyngent.lmproto.openai_compatible.model import SystemChatMessage
|
|
|
|
client = ScriptedClient([_response(AssistantChatMessage(content="ok"))])
|
|
agent = ChatAgent(client, model="m", system_prompt="Be brief.", stream=False)
|
|
_ = [e async for e in agent.run("hi")]
|
|
assert isinstance(agent.messages[0], SystemChatMessage)
|
|
assert agent.messages[0].content == "Be brief."
|
|
assert isinstance(client.calls[0].messages[0], SystemChatMessage)
|
|
|
|
|
|
async def test_tool_result_char_budget() -> None:
|
|
@tool(register=False)
|
|
def big() -> str:
|
|
return "x" * 100
|
|
|
|
registry = ToolRegistry([big])
|
|
client = ScriptedClient(
|
|
[
|
|
_response(
|
|
AssistantChatMessage(
|
|
content="",
|
|
tool_calls=[
|
|
AssistantFunctionToolCall(
|
|
id="1",
|
|
function=AssistantFunctionTool(name="big", arguments="{}"),
|
|
)
|
|
],
|
|
)
|
|
),
|
|
_response(AssistantChatMessage(content="done")),
|
|
]
|
|
)
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="go")]
|
|
_ = [
|
|
e
|
|
async for e in run_chat_loop(
|
|
client,
|
|
messages,
|
|
model="m",
|
|
tools=registry,
|
|
stream=False,
|
|
max_tool_result_chars=20,
|
|
parallel_tools=False,
|
|
)
|
|
]
|
|
from plyngent.lmproto.openai_compatible.model import ToolChatMessage
|
|
|
|
tool_msgs = [m for m in messages if isinstance(m, ToolChatMessage)]
|
|
assert len(tool_msgs) == 1
|
|
assert tool_msgs[0].content.startswith("x" * 20)
|
|
assert "Truncated" in tool_msgs[0].content
|
|
|
|
|
|
async def test_empty_completion_raises() -> None:
|
|
client = ScriptedClient([_response(AssistantChatMessage(content=None))])
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
with pytest.raises(RuntimeError, match="empty model completion"):
|
|
_ = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
|
|
|
|
|
|
async def test_length_finish_raises() -> None:
|
|
client = ScriptedClient(
|
|
[
|
|
_response(
|
|
AssistantChatMessage(content="partial"),
|
|
finish_reason="length",
|
|
)
|
|
]
|
|
)
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
with pytest.raises(RuntimeError, match="truncated"):
|
|
_ = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
|
|
|
|
|
|
async def test_content_filter_finish_raises() -> None:
|
|
client = ScriptedClient(
|
|
[
|
|
_response(
|
|
AssistantChatMessage(content=""),
|
|
finish_reason="content_filter",
|
|
)
|
|
]
|
|
)
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
with pytest.raises(RuntimeError, match="content filter"):
|
|
_ = [e async for e in run_chat_loop(client, messages, model="m", stream=False)]
|
|
|
|
|
|
async def test_stream_missing_terminal_empty_raises() -> None:
|
|
class EmptyStreamClient:
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del param
|
|
if not stream:
|
|
return _response(AssistantChatMessage(content="x"))
|
|
|
|
async def chunks() -> AsyncIterator[ChatCompletionChunk]:
|
|
# Usage-only style chunk with no finish_reason and no content.
|
|
yield ChatCompletionChunk(
|
|
id="1",
|
|
object="chat.completion.chunk",
|
|
created=0,
|
|
model="t",
|
|
choices=[],
|
|
usage={"prompt_tokens": 1, "completion_tokens": 0, "total_tokens": 1},
|
|
)
|
|
|
|
return chunks()
|
|
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
with pytest.raises(RuntimeError, match="stream ended without a terminal"):
|
|
_ = [e async for e in run_chat_loop(EmptyStreamClient(), messages, model="m", stream=True)]
|
|
|
|
|
|
async def test_stream_payload_without_finish_reason_ok() -> None:
|
|
"""Some providers omit finish_reason; non-empty text still counts as terminal."""
|
|
|
|
class PayloadNoFinishClient:
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[False] = False
|
|
) -> ChatCompletionResponse: ...
|
|
|
|
@overload
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: Literal[True]
|
|
) -> AsyncIterator[ChatCompletionChunk]: ...
|
|
|
|
async def chat_completions(
|
|
self, param: ChatCompletionsParam, *, stream: bool = False
|
|
) -> ChatCompletionResponse | AsyncIterator[ChatCompletionChunk]:
|
|
del param
|
|
if not stream:
|
|
return _response(AssistantChatMessage(content="ok"))
|
|
|
|
async def chunks() -> AsyncIterator[ChatCompletionChunk]:
|
|
yield ChatCompletionChunk(
|
|
id="1",
|
|
object="chat.completion.chunk",
|
|
created=0,
|
|
model="t",
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(content="ok"),
|
|
finish_reason=None,
|
|
)
|
|
],
|
|
)
|
|
|
|
return chunks()
|
|
|
|
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
|
|
events = [e async for e in run_chat_loop(PayloadNoFinishClient(), messages, model="m", stream=True)]
|
|
assert any(isinstance(e, TextDeltaEvent) and e.content == "ok" for e in events)
|