mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
232 lines
7.2 KiB
Python
232 lines
7.2 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING, Literal, overload
|
|
|
|
import pytest
|
|
|
|
from plyngent.agent import ChatAgent
|
|
from plyngent.cli.retry import (
|
|
DEFAULT_MAX_AUTO_RETRIES,
|
|
DEFAULT_RETRY_DELAYS_SECONDS,
|
|
default_retry_delays,
|
|
retry_pending_with_retries,
|
|
run_turn_with_retries,
|
|
sleep_cancellable,
|
|
)
|
|
from plyngent.config.models import DatabaseConfig
|
|
from plyngent.lmproto.openai_compatible.model import (
|
|
AssistantChatMessage,
|
|
ChatCompletionChoice,
|
|
ChatCompletionChunk,
|
|
ChatCompletionResponse,
|
|
ChatCompletionsParam,
|
|
ChunkChoice,
|
|
DeltaMessage,
|
|
UserChatMessage,
|
|
)
|
|
from plyngent.memory import MemoryStore
|
|
|
|
if TYPE_CHECKING:
|
|
from collections.abc import AsyncIterator
|
|
|
|
|
|
def test_default_retry_delays_schedule() -> None:
|
|
assert DEFAULT_MAX_AUTO_RETRIES == 10
|
|
assert default_retry_delays(0) == ()
|
|
assert default_retry_delays(4) == (5.0, 10.0, 15.0, 20.0)
|
|
delays = default_retry_delays(10)
|
|
assert delays[:4] == (5.0, 10.0, 15.0, 20.0)
|
|
assert delays[4:] == (30.0, 40.0, 50.0, 60.0, 70.0, 80.0)
|
|
assert delays == DEFAULT_RETRY_DELAYS_SECONDS
|
|
assert len(DEFAULT_RETRY_DELAYS_SECONDS) == 10
|
|
|
|
|
|
def _response(message: AssistantChatMessage) -> ChatCompletionResponse:
|
|
return ChatCompletionResponse(
|
|
id="1",
|
|
object="chat.completion",
|
|
created=0,
|
|
model="test",
|
|
choices=[
|
|
ChatCompletionChoice(
|
|
index=0,
|
|
message=message,
|
|
logprobs={},
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
system_fingerprint="",
|
|
usage={},
|
|
)
|
|
|
|
|
|
async def _as_stream(response: ChatCompletionResponse) -> AsyncIterator[ChatCompletionChunk]:
|
|
message = response.choices[0].message
|
|
content = message.content if isinstance(message.content, str) else ""
|
|
yield ChatCompletionChunk(
|
|
id=response.id,
|
|
object="chat.completion.chunk",
|
|
created=response.created,
|
|
model=response.model,
|
|
choices=[
|
|
ChunkChoice(
|
|
index=0,
|
|
delta=DeltaMessage(content=content),
|
|
finish_reason="stop",
|
|
)
|
|
],
|
|
)
|
|
|
|
|
|
class FlakyClient:
|
|
calls: int
|
|
fail_times: int
|
|
|
|
def __init__(self, fail_times: int) -> None:
|
|
self.calls = 0
|
|
self.fail_times = fail_times
|
|
|
|
@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 <= self.fail_times:
|
|
msg = f"fail-{self.calls}"
|
|
raise RuntimeError(msg)
|
|
response = _response(AssistantChatMessage(content="ok"))
|
|
if stream:
|
|
return _as_stream(response)
|
|
return response
|
|
|
|
|
|
async def test_sleep_cancellable_completes() -> None:
|
|
assert await sleep_cancellable(0.01) is True
|
|
|
|
|
|
async def test_auto_retry_eventually_succeeds(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
async def no_wait(_seconds: float) -> bool:
|
|
return True
|
|
|
|
monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", no_wait)
|
|
store = await MemoryStore.open(DatabaseConfig())
|
|
session = await store.create_session(name="t")
|
|
client = FlakyClient(fail_times=2)
|
|
agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)
|
|
|
|
ok = await run_turn_with_retries(
|
|
agent,
|
|
starter=lambda: agent.run("hello"),
|
|
delays=(0.01, 0.01, 0.01),
|
|
)
|
|
assert ok is True
|
|
assert client.calls == 3
|
|
loaded = await store.list_messages(session.sid)
|
|
assert len(loaded) == 2
|
|
assert isinstance(loaded[0], UserChatMessage)
|
|
assert loaded[0].content == "hello"
|
|
await store.close()
|
|
|
|
|
|
async def test_run_cancellable_cancels_task() -> None:
|
|
import asyncio
|
|
|
|
from plyngent.cli.retry import run_cancellable
|
|
|
|
started = asyncio.Event()
|
|
|
|
async def hang() -> None:
|
|
started.set()
|
|
await asyncio.sleep(60)
|
|
|
|
task = asyncio.create_task(run_cancellable(hang()))
|
|
_ = await started.wait()
|
|
_ = task.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
async def test_cancel_turn_rolls_back_and_sets_pending() -> None:
|
|
import asyncio
|
|
|
|
from plyngent.cli.display import render_events
|
|
from plyngent.cli.retry import run_cancellable
|
|
|
|
store = await MemoryStore.open(DatabaseConfig())
|
|
session = await store.create_session(name="t")
|
|
|
|
class HangClient:
|
|
@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
|
|
await asyncio.sleep(60)
|
|
return _response(AssistantChatMessage(content="never"))
|
|
|
|
agent = ChatAgent(HangClient(), model="m", memory=store, session_id=session.sid)
|
|
turn = asyncio.create_task(run_cancellable(render_events(agent.run("cancel-me"))))
|
|
await asyncio.sleep(0.05)
|
|
_ = turn.cancel()
|
|
with pytest.raises(asyncio.CancelledError):
|
|
await turn
|
|
|
|
assert agent.pending_retry_text == "cancel-me"
|
|
assert len(agent.messages) == 1
|
|
assert isinstance(agent.messages[0], UserChatMessage)
|
|
loaded = await store.list_messages(session.sid)
|
|
assert len(loaded) == 1
|
|
assert isinstance(loaded[0], UserChatMessage)
|
|
await store.close()
|
|
|
|
|
|
async def test_manual_retry_after_exhausted(monkeypatch: pytest.MonkeyPatch) -> None:
|
|
async def no_wait(_seconds: float) -> bool:
|
|
return True
|
|
|
|
monkeypatch.setattr("plyngent.cli.retry.sleep_cancellable", no_wait)
|
|
store = await MemoryStore.open(DatabaseConfig())
|
|
session = await store.create_session(name="t")
|
|
client = FlakyClient(fail_times=5)
|
|
agent = ChatAgent(client, model="m", memory=store, session_id=session.sid)
|
|
|
|
ok = await run_turn_with_retries(
|
|
agent,
|
|
starter=lambda: agent.run("hold-me"),
|
|
delays=(0.01,), # one auto-retry only → 2 attempts total, both fail
|
|
)
|
|
assert ok is False
|
|
assert agent.pending_retry_text == "hold-me"
|
|
assert len(await store.list_messages(session.sid)) == 1
|
|
|
|
# Now succeed on manual /retry path (no second user message)
|
|
client.fail_times = client.calls # next call succeeds
|
|
ok2 = await retry_pending_with_retries(agent)
|
|
assert ok2 is True
|
|
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 == "hold-me"
|
|
assert sum(1 for m in loaded if isinstance(m, UserChatMessage)) == 1
|
|
await store.close()
|