mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/agent+cli: defer DB write until turn success; auto/manual retry
Failed turns roll back memory and keep pending_retry_text; persist only after success. CLI auto-retries 10s/20s/30s (cancellable) and /retry.
This commit is contained in:
@@ -55,7 +55,7 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init, default local user, se
|
|||||||
- **`ChatClient`** Protocol for `chat_completions`.
|
- **`ChatClient`** Protocol for `chat_completions`.
|
||||||
- **`@tool` / `ToolRegistry`**: decorator infers JSON Schema from type hints; execute tools by name.
|
- **`@tool` / `ToolRegistry`**: decorator infers JSON Schema from type hints; execute tools by name.
|
||||||
- **`run_chat_loop`**: multi-round tool loop, yields `AgentEvent` stream; optional `on_limit` to continue past max rounds.
|
- **`run_chat_loop`**: multi-round tool loop, yields `AgentEvent` stream; optional `on_limit` to continue past max rounds.
|
||||||
- **`ChatAgent`**: wrapper with optional `MemoryStore` bind (load/append messages).
|
- **`ChatAgent`**: wrapper with optional `MemoryStore` bind (load/append messages on success only); `pending_retry_text` + `retry()` after failed turns.
|
||||||
|
|
||||||
### Tools (`tools/`)
|
### Tools (`tools/`)
|
||||||
|
|
||||||
@@ -73,7 +73,8 @@ Module-level `@tool` handlers. Call `set_workspace_root()` before use.
|
|||||||
Click app + readline REPL. Entry: `plyngent` / `python -m plyngent`.
|
Click app + readline REPL. Entry: `plyngent` / `python -m plyngent`.
|
||||||
|
|
||||||
- **`plyngent chat`**: provider/model selection (flags or interactive), SQLite sessions via config `[database]` (file DB under user data if unset/`:memory:`), resumes latest session by default (`--new` / `--session`).
|
- **`plyngent chat`**: provider/model selection (flags or interactive), SQLite sessions via config `[database]` (file DB under user data if unset/`:memory:`), resumes latest session by default (`--new` / `--session`).
|
||||||
- Slash: `/history`, `/sessions`, `/resume`, `/rounds`, …
|
- Slash: `/history`, `/sessions`, `/resume`, `/rounds`, `/retry`, …
|
||||||
|
- Failed turns: not written to DB; auto-retry 10s/20s/30s (Ctrl+C cancels wait); `/retry` manual.
|
||||||
- **`plyngent providers`**: list config providers.
|
- **`plyngent providers`**: list config providers.
|
||||||
- **`plyngent config path|edit`**: print or open config in `$EDITOR` (`shlex`-split, e.g. `codium --wait`).
|
- **`plyngent config path|edit`**: print or open config in `$EDITOR` (`shlex`-split, e.g. `codium --wait`).
|
||||||
- If no providers and `$EDITOR` is set, chat/providers prompt to edit config then reload.
|
- If no providers and `$EDITOR` is set, chat/providers prompt to edit config then reload.
|
||||||
|
|||||||
+43
-16
@@ -31,6 +31,7 @@ class ChatAgent:
|
|||||||
temperature: float | None
|
temperature: float | None
|
||||||
on_limit: LimitContinueHook | None
|
on_limit: LimitContinueHook | None
|
||||||
messages: list[AnyChatMessage]
|
messages: list[AnyChatMessage]
|
||||||
|
pending_retry_text: str | None
|
||||||
|
|
||||||
def __init__( # noqa: PLR0913
|
def __init__( # noqa: PLR0913
|
||||||
self,
|
self,
|
||||||
@@ -54,6 +55,7 @@ class ChatAgent:
|
|||||||
self.temperature = temperature
|
self.temperature = temperature
|
||||||
self.on_limit = on_limit
|
self.on_limit = on_limit
|
||||||
self.messages = list(messages) if messages is not None else []
|
self.messages = list(messages) if messages is not None else []
|
||||||
|
self.pending_retry_text = None
|
||||||
|
|
||||||
async def load_history(self) -> None:
|
async def load_history(self) -> None:
|
||||||
"""Replace in-memory messages from the bound memory session."""
|
"""Replace in-memory messages from the bound memory session."""
|
||||||
@@ -61,6 +63,7 @@ class ChatAgent:
|
|||||||
msg = "load_history requires memory and session_id"
|
msg = "load_history requires memory and session_id"
|
||||||
raise RuntimeError(msg)
|
raise RuntimeError(msg)
|
||||||
self.messages = await self.memory.list_messages(self.session_id)
|
self.messages = await self.memory.list_messages(self.session_id)
|
||||||
|
self.pending_retry_text = None
|
||||||
|
|
||||||
async def bind_session(self, session_id: int, *, load: bool = True) -> None:
|
async def bind_session(self, session_id: int, *, load: bool = True) -> None:
|
||||||
"""Attach a memory session id; optionally load existing messages."""
|
"""Attach a memory session id; optionally load existing messages."""
|
||||||
@@ -75,24 +78,48 @@ class ChatAgent:
|
|||||||
if self.memory is not None and self.session_id is not None:
|
if self.memory is not None and self.session_id is not None:
|
||||||
_ = await self.memory.append_message(self.session_id, message)
|
_ = await self.memory.append_message(self.session_id, message)
|
||||||
|
|
||||||
|
async def _run_from_user_message(self, user_msg: UserChatMessage) -> AsyncIterator[AgentEvent]:
|
||||||
|
"""Run the tool loop for an already-appended user message; persist only on success."""
|
||||||
|
pre_len = len(self.messages) - 1
|
||||||
|
if pre_len < 0 or self.messages[pre_len] is not user_msg:
|
||||||
|
pre_len = len(self.messages)
|
||||||
|
self.messages.append(user_msg)
|
||||||
|
|
||||||
|
try:
|
||||||
|
async for event in run_chat_loop(
|
||||||
|
self.client,
|
||||||
|
self.messages,
|
||||||
|
model=self.model,
|
||||||
|
tools=self.tools,
|
||||||
|
max_rounds=self.max_rounds,
|
||||||
|
temperature=self.temperature,
|
||||||
|
on_limit=self.on_limit,
|
||||||
|
):
|
||||||
|
yield event
|
||||||
|
except Exception:
|
||||||
|
# Roll back the whole turn so a failed user message is not left half-applied.
|
||||||
|
del self.messages[pre_len:]
|
||||||
|
self.pending_retry_text = user_msg.content
|
||||||
|
raise
|
||||||
|
|
||||||
|
for message in self.messages[pre_len:]:
|
||||||
|
await self._persist(message)
|
||||||
|
self.pending_retry_text = None
|
||||||
|
|
||||||
async def run(self, user_text: str) -> AsyncIterator[AgentEvent]:
|
async def run(self, user_text: str) -> AsyncIterator[AgentEvent]:
|
||||||
"""Append a user message, run the tool loop, yield events, persist new messages."""
|
"""Append a user message, run the tool loop, yield events, persist only on success."""
|
||||||
user_msg = UserChatMessage(content=user_text)
|
user_msg = UserChatMessage(content=user_text)
|
||||||
self.messages.append(user_msg)
|
self.messages.append(user_msg)
|
||||||
await self._persist(user_msg)
|
async for event in self._run_from_user_message(user_msg):
|
||||||
|
|
||||||
start_len = len(self.messages)
|
|
||||||
async for event in run_chat_loop(
|
|
||||||
self.client,
|
|
||||||
self.messages,
|
|
||||||
model=self.model,
|
|
||||||
tools=self.tools,
|
|
||||||
max_rounds=self.max_rounds,
|
|
||||||
temperature=self.temperature,
|
|
||||||
on_limit=self.on_limit,
|
|
||||||
):
|
|
||||||
yield event
|
yield event
|
||||||
|
|
||||||
# Persist messages appended by the loop after the user message.
|
async def retry(self) -> AsyncIterator[AgentEvent]:
|
||||||
for message in self.messages[start_len:]:
|
"""Re-run the last failed user turn (no duplicate user message in history/DB)."""
|
||||||
await self._persist(message)
|
text = self.pending_retry_text
|
||||||
|
if text is None:
|
||||||
|
msg = "nothing to retry"
|
||||||
|
raise RuntimeError(msg)
|
||||||
|
user_msg = UserChatMessage(content=text)
|
||||||
|
self.messages.append(user_msg)
|
||||||
|
async for event in self._run_from_user_message(user_msg):
|
||||||
|
yield event
|
||||||
|
|||||||
@@ -28,6 +28,7 @@ SLASH_COMMANDS: tuple[str, ...] = (
|
|||||||
"/model",
|
"/model",
|
||||||
"/tools",
|
"/tools",
|
||||||
"/rounds",
|
"/rounds",
|
||||||
|
"/retry",
|
||||||
)
|
)
|
||||||
|
|
||||||
_TOOLS_ARGS: tuple[str, ...] = ("on", "off")
|
_TOOLS_ARGS: tuple[str, ...] = ("on", "off")
|
||||||
|
|||||||
@@ -7,8 +7,8 @@ from typing import TYPE_CHECKING
|
|||||||
import click
|
import click
|
||||||
from msgspec import UNSET
|
from msgspec import UNSET
|
||||||
|
|
||||||
from plyngent.cli.display import render_events
|
|
||||||
from plyngent.cli.readline_setup import setup_readline
|
from plyngent.cli.readline_setup import setup_readline
|
||||||
|
from plyngent.cli.retry import retry_pending_with_retries, run_user_text_with_retries
|
||||||
from plyngent.cli.selection import select_model, select_provider
|
from plyngent.cli.selection import select_model, select_provider
|
||||||
from plyngent.lmproto.openai_compatible.model import (
|
from plyngent.lmproto.openai_compatible.model import (
|
||||||
AssistantChatMessage,
|
AssistantChatMessage,
|
||||||
@@ -35,6 +35,10 @@ Commands:
|
|||||||
/model [id] Show or switch model
|
/model [id] Show or switch model
|
||||||
/tools [on|off] Show or toggle tools
|
/tools [on|off] Show or toggle tools
|
||||||
/rounds [n] Show or set max tool-loop rounds
|
/rounds [n] Show or set max tool-loop rounds
|
||||||
|
/retry Retry the last failed user turn (after errors)
|
||||||
|
|
||||||
|
On network/API errors, the user turn is not saved to the DB. Auto-retry
|
||||||
|
waits 10s, 20s, then 30s (Ctrl+C cancels waits; use /retry later).
|
||||||
|
|
||||||
Tab completes slash commands and some arguments (provider, model, tools).
|
Tab completes slash commands and some arguments (provider, model, tools).
|
||||||
Use --session ID or /resume to continue a prior chat after restart.
|
Use --session ID or /resume to continue a prior chat after restart.
|
||||||
@@ -141,6 +145,7 @@ def _cmd_rounds(state: ReplState, arg: str) -> None:
|
|||||||
|
|
||||||
def _cmd_clear(state: ReplState) -> None:
|
def _cmd_clear(state: ReplState) -> None:
|
||||||
state.agent.messages.clear()
|
state.agent.messages.clear()
|
||||||
|
state.agent.pending_retry_text = None
|
||||||
click.echo("conversation cleared (in-memory only; DB history kept)")
|
click.echo("conversation cleared (in-memory only; DB history kept)")
|
||||||
|
|
||||||
|
|
||||||
@@ -197,6 +202,15 @@ def _cmd_history(state: ReplState, arg: str) -> None:
|
|||||||
click.echo(f"session={state.session_id} messages={len(messages)} showing={len(messages) - start}")
|
click.echo(f"session={state.session_id} messages={len(messages)} showing={len(messages) - start}")
|
||||||
for offset, message in enumerate(messages[start:]):
|
for offset, message in enumerate(messages[start:]):
|
||||||
click.echo(_format_history_message(start + offset, message))
|
click.echo(_format_history_message(start + offset, message))
|
||||||
|
if state.agent.pending_retry_text is not None:
|
||||||
|
click.secho(
|
||||||
|
f"(pending retry) user: {_preview_content(state.agent.pending_retry_text)}",
|
||||||
|
fg="yellow",
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
async def _cmd_retry(state: ReplState) -> None:
|
||||||
|
_ = await retry_pending_with_retries(state.agent)
|
||||||
|
|
||||||
|
|
||||||
async def _dispatch_slash(state: ReplState, command: str, arg: str) -> bool:
|
async def _dispatch_slash(state: ReplState, command: str, arg: str) -> bool:
|
||||||
@@ -214,6 +228,7 @@ async def _dispatch_slash(state: ReplState, command: str, arg: str) -> bool:
|
|||||||
"model": lambda: _cmd_model(state, arg),
|
"model": lambda: _cmd_model(state, arg),
|
||||||
"tools": lambda: _cmd_tools(state, arg),
|
"tools": lambda: _cmd_tools(state, arg),
|
||||||
"rounds": lambda: _cmd_rounds(state, arg),
|
"rounds": lambda: _cmd_rounds(state, arg),
|
||||||
|
"retry": lambda: _cmd_retry(state),
|
||||||
}
|
}
|
||||||
handler = handlers.get(command)
|
handler = handlers.get(command)
|
||||||
if handler is None:
|
if handler is None:
|
||||||
@@ -269,8 +284,4 @@ async def run_repl(state: ReplState) -> None:
|
|||||||
|
|
||||||
click.secho("user: ", fg="green", nl=False)
|
click.secho("user: ", fg="green", nl=False)
|
||||||
click.echo(line)
|
click.echo(line)
|
||||||
try:
|
_ = await run_user_text_with_retries(state.agent, line)
|
||||||
await render_events(state.agent.run(line))
|
|
||||||
except Exception as exc: # noqa: BLE001 — show API/runtime errors in REPL
|
|
||||||
click.secho(f"error: {exc}", fg="red")
|
|
||||||
click.echo() # match render_events spacing before next prompt
|
|
||||||
|
|||||||
@@ -0,0 +1,102 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
import asyncio
|
||||||
|
import time
|
||||||
|
from typing import TYPE_CHECKING
|
||||||
|
|
||||||
|
import click
|
||||||
|
|
||||||
|
from plyngent.cli.display import render_events
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import AsyncIterator, Callable
|
||||||
|
|
||||||
|
from plyngent.agent import AgentEvent
|
||||||
|
from plyngent.agent.chat import ChatAgent
|
||||||
|
|
||||||
|
# Wait before retry attempt 1, 2, and 3 (after the first failure).
|
||||||
|
DEFAULT_RETRY_DELAYS_SECONDS: tuple[float, ...] = (10.0, 20.0, 30.0)
|
||||||
|
_PREVIEW_LEN = 80
|
||||||
|
|
||||||
|
|
||||||
|
async def sleep_cancellable(seconds: float) -> bool:
|
||||||
|
"""Sleep in short steps so Ctrl+C can cancel. Returns False if interrupted."""
|
||||||
|
deadline = time.monotonic() + seconds
|
||||||
|
try:
|
||||||
|
while True:
|
||||||
|
remaining = deadline - time.monotonic()
|
||||||
|
if remaining <= 0:
|
||||||
|
return True
|
||||||
|
await asyncio.sleep(min(0.5, remaining))
|
||||||
|
except asyncio.CancelledError:
|
||||||
|
return False
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
async def run_turn_with_retries(
|
||||||
|
agent: ChatAgent,
|
||||||
|
*,
|
||||||
|
starter: Callable[[], AsyncIterator[AgentEvent]],
|
||||||
|
delays: tuple[float, ...] = DEFAULT_RETRY_DELAYS_SECONDS,
|
||||||
|
) -> bool:
|
||||||
|
"""Run a chat turn with automatic retries on failure.
|
||||||
|
|
||||||
|
``starter`` produces the event stream for the current attempt (``agent.run``
|
||||||
|
or ``agent.retry``). Returns True if the turn completed successfully.
|
||||||
|
"""
|
||||||
|
max_retries = len(delays)
|
||||||
|
attempt = 0
|
||||||
|
while True:
|
||||||
|
try:
|
||||||
|
await render_events(starter())
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
click.echo()
|
||||||
|
click.secho("interrupted", fg="yellow")
|
||||||
|
return False
|
||||||
|
except Exception as exc: # noqa: BLE001 — surface and optionally retry
|
||||||
|
click.secho(f"error: {exc}", fg="red")
|
||||||
|
if attempt >= max_retries:
|
||||||
|
if agent.pending_retry_text is not None:
|
||||||
|
click.secho(
|
||||||
|
"auto-retry exhausted; use /retry to try again, or send a new message",
|
||||||
|
fg="yellow",
|
||||||
|
)
|
||||||
|
click.echo()
|
||||||
|
return False
|
||||||
|
|
||||||
|
wait = delays[attempt]
|
||||||
|
attempt += 1
|
||||||
|
click.secho(
|
||||||
|
f"auto-retry {attempt}/{max_retries} in {wait:g}s "
|
||||||
|
f"(Ctrl+C to cancel; then /retry later)",
|
||||||
|
fg="yellow",
|
||||||
|
)
|
||||||
|
try:
|
||||||
|
ok = await sleep_cancellable(wait)
|
||||||
|
except KeyboardInterrupt:
|
||||||
|
ok = False
|
||||||
|
if not ok:
|
||||||
|
click.secho("auto-retry cancelled; use /retry to try again", fg="yellow")
|
||||||
|
click.echo()
|
||||||
|
return False
|
||||||
|
click.secho(f"retrying ({attempt}/{max_retries})…", fg="yellow")
|
||||||
|
else:
|
||||||
|
return True
|
||||||
|
|
||||||
|
|
||||||
|
async def run_user_text_with_retries(agent: ChatAgent, text: str) -> bool:
|
||||||
|
"""Send a new user message with auto-retry."""
|
||||||
|
return await run_turn_with_retries(agent, starter=lambda: agent.run(text))
|
||||||
|
|
||||||
|
|
||||||
|
async def retry_pending_with_retries(agent: ChatAgent) -> bool:
|
||||||
|
"""Retry the last failed user turn with auto-retry."""
|
||||||
|
if agent.pending_retry_text is None:
|
||||||
|
click.echo("nothing to retry")
|
||||||
|
return False
|
||||||
|
preview = agent.pending_retry_text
|
||||||
|
if len(preview) > _PREVIEW_LEN:
|
||||||
|
preview = preview[:_PREVIEW_LEN] + "…"
|
||||||
|
click.echo(f"retrying: {preview}")
|
||||||
|
return await run_turn_with_retries(agent, starter=agent.retry)
|
||||||
@@ -66,6 +66,7 @@ class ReplState:
|
|||||||
session = await self.memory.create_session(name=name)
|
session = await self.memory.create_session(name=name)
|
||||||
self.session_id = session.sid
|
self.session_id = session.sid
|
||||||
self.agent = self._make_agent()
|
self.agent = self._make_agent()
|
||||||
|
self.agent.pending_retry_text = None
|
||||||
|
|
||||||
async def resume_session(self, session_id: int) -> None:
|
async def resume_session(self, session_id: int) -> None:
|
||||||
row = await self.memory.get_session(session_id)
|
row = await self.memory.get_session(session_id)
|
||||||
|
|||||||
@@ -2,6 +2,8 @@ from __future__ import annotations
|
|||||||
|
|
||||||
from typing import TYPE_CHECKING, Literal, overload
|
from typing import TYPE_CHECKING, Literal, overload
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
from plyngent.agent import (
|
from plyngent.agent import (
|
||||||
AssistantMessageEvent,
|
AssistantMessageEvent,
|
||||||
ChatAgent,
|
ChatAgent,
|
||||||
@@ -225,3 +227,91 @@ async def test_chat_agent_memory_roundtrip() -> None:
|
|||||||
await agent2.load_history()
|
await agent2.load_history()
|
||||||
assert len(agent2.messages) == 2 # noqa: PLR2004
|
assert len(agent2.messages) == 2 # noqa: PLR2004
|
||||||
await store.close()
|
await store.close()
|
||||||
|
|
||||||
|
|
||||||
|
async def test_chat_agent_failed_turn_not_persisted() -> 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 agent.messages == []
|
||||||
|
assert agent.pending_retry_text == "hello"
|
||||||
|
loaded = await store.list_messages(session.sid)
|
||||||
|
assert loaded == []
|
||||||
|
await store.close()
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
if stream:
|
||||||
|
|
||||||
|
async def empty() -> AsyncIterator[ChatCompletionChunk]:
|
||||||
|
empty_chunks: list[ChatCompletionChunk] = []
|
||||||
|
for chunk in empty_chunks:
|
||||||
|
yield chunk
|
||||||
|
|
||||||
|
return empty()
|
||||||
|
return _response(AssistantChatMessage(content="recovered"))
|
||||||
|
|
||||||
|
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"
|
||||||
|
assert await store.list_messages(session.sid) == []
|
||||||
|
|
||||||
|
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 # noqa: PLR2004
|
||||||
|
assert isinstance(loaded[0], UserChatMessage)
|
||||||
|
assert loaded[0].content == "ping"
|
||||||
|
await store.close()
|
||||||
|
|||||||
@@ -0,0 +1,138 @@
|
|||||||
|
from __future__ import annotations
|
||||||
|
|
||||||
|
from typing import TYPE_CHECKING, Literal, overload
|
||||||
|
|
||||||
|
from plyngent.agent import ChatAgent
|
||||||
|
from plyngent.cli.retry import 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,
|
||||||
|
UserChatMessage,
|
||||||
|
)
|
||||||
|
from plyngent.memory import MemoryStore
|
||||||
|
|
||||||
|
if TYPE_CHECKING:
|
||||||
|
from collections.abc import AsyncIterator
|
||||||
|
|
||||||
|
import pytest
|
||||||
|
|
||||||
|
|
||||||
|
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={},
|
||||||
|
)
|
||||||
|
|
||||||
|
|
||||||
|
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)
|
||||||
|
if stream:
|
||||||
|
|
||||||
|
async def empty() -> AsyncIterator[ChatCompletionChunk]:
|
||||||
|
empty_chunks: list[ChatCompletionChunk] = []
|
||||||
|
for chunk in empty_chunks:
|
||||||
|
yield chunk
|
||||||
|
|
||||||
|
return empty()
|
||||||
|
return _response(AssistantChatMessage(content="ok"))
|
||||||
|
|
||||||
|
|
||||||
|
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 # noqa: PLR2004
|
||||||
|
loaded = await store.list_messages(session.sid)
|
||||||
|
assert len(loaded) == 2 # noqa: PLR2004
|
||||||
|
assert isinstance(loaded[0], UserChatMessage)
|
||||||
|
assert loaded[0].content == "hello"
|
||||||
|
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 await store.list_messages(session.sid) == []
|
||||||
|
|
||||||
|
# Now succeed on manual /retry path
|
||||||
|
client.fail_times = client.calls # next call succeeds
|
||||||
|
ok2 = await retry_pending_with_retries(agent)
|
||||||
|
# retry_pending uses default delays; force zero-wait already patched
|
||||||
|
# But fail_times set so first attempt of manual retry succeeds immediately
|
||||||
|
assert ok2 is True
|
||||||
|
assert agent.pending_retry_text is None
|
||||||
|
loaded = await store.list_messages(session.sid)
|
||||||
|
assert len(loaded) == 2 # noqa: PLR2004
|
||||||
|
assert isinstance(loaded[0], UserChatMessage)
|
||||||
|
assert loaded[0].content == "hold-me"
|
||||||
|
await store.close()
|
||||||
Reference in New Issue
Block a user