mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
core/agent: fix Phase A streaming await bug and polish
Do not await async generators for raw SSE; close stream on cancel; merge tool calls without relying only on finish_reason; quieter tool output; /status slash command.
This commit is contained in:
@@ -54,8 +54,9 @@ Async SQLAlchemy + aiosqlite. `MemoryStore`: schema init, default local user, se
|
||||
|
||||
- **`ChatClient`** Protocol for `chat_completions`.
|
||||
- **`@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.
|
||||
- **`ChatAgent`**: wrapper with optional `MemoryStore` bind (load/append messages on success only); `pending_retry_text` + `retry()` after failed turns.
|
||||
- **`run_chat_loop`**: multi-round tool loop; default **streaming** text deltas + raw-SSE tool-call merge; optional `on_limit`.
|
||||
- **`ChatAgent`**: optional `MemoryStore` (persist on success only); `stream` flag; `pending_retry_text` + `retry()`.
|
||||
- Events: text_delta, assistant_message, tool_call/result, max_rounds, **error**, **cancelled**.
|
||||
|
||||
### Tools (`tools/`)
|
||||
|
||||
|
||||
@@ -107,6 +107,7 @@ ignore = [
|
||||
"COM812", # missing-trailing-comma
|
||||
"PLC0415", # import-outside-top-level
|
||||
"TID252",
|
||||
"PLR0913", # too-many-arguments (common on constructors/composed functions)
|
||||
]
|
||||
allowed-confusables = [
|
||||
",",
|
||||
|
||||
@@ -30,10 +30,11 @@ class ChatAgent:
|
||||
max_rounds: int
|
||||
temperature: float | None
|
||||
on_limit: LimitContinueHook | None
|
||||
stream: bool
|
||||
messages: list[AnyChatMessage]
|
||||
pending_retry_text: str | None
|
||||
|
||||
def __init__( # noqa: PLR0913
|
||||
def __init__(
|
||||
self,
|
||||
client: ChatClient,
|
||||
*,
|
||||
|
||||
+37
-29
@@ -60,61 +60,74 @@ async def _execute_tool_calls(
|
||||
yield ToolResultEvent(message=tool_msg)
|
||||
|
||||
|
||||
async def _stream_and_build_assistant( # noqa: C901
|
||||
async def _non_stream_assistant(
|
||||
client: ChatClient,
|
||||
param: ChatCompletionsParam,
|
||||
) -> tuple[AssistantChatMessage, list[TextDeltaEvent]]:
|
||||
"""Stream a round, return the assistant message and any text deltas."""
|
||||
raw_lines_attr = getattr(client, "chat_completions_raw_lines", None)
|
||||
if raw_lines_attr is None:
|
||||
response = await client.chat_completions(param, stream=False)
|
||||
if not response.choices:
|
||||
msg = "chat completion response contained no choices"
|
||||
raise RuntimeError(msg)
|
||||
assistant = response.choices[0].message
|
||||
text_events = []
|
||||
text_events: list[TextDeltaEvent] = []
|
||||
if isinstance(assistant.content, str) and assistant.content:
|
||||
text_events.append(TextDeltaEvent(content=assistant.content))
|
||||
return assistant, text_events
|
||||
|
||||
stream_iter = await raw_lines_attr(param) # type: ignore[misc]
|
||||
|
||||
async def _stream_and_build_assistant(
|
||||
client: ChatClient,
|
||||
param: ChatCompletionsParam,
|
||||
) -> tuple[AssistantChatMessage, list[TextDeltaEvent]]:
|
||||
"""Stream a round, return the assistant message and any text deltas.
|
||||
|
||||
Requires ``client.chat_completions_raw_lines`` (async generator of raw SSE
|
||||
payload bytes). Falls back to non-streaming if unavailable.
|
||||
"""
|
||||
raw_lines_attr = getattr(client, "chat_completions_raw_lines", None)
|
||||
if raw_lines_attr is None:
|
||||
return await _non_stream_assistant(client, param)
|
||||
|
||||
# Async generators are not awaitable — call to get the iterator.
|
||||
stream_iter = raw_lines_attr(param)
|
||||
raw_lines: list[bytes] = []
|
||||
content_parts: list[str] = []
|
||||
finish_reason: str | None = None
|
||||
text_events = []
|
||||
text_events: list[TextDeltaEvent] = []
|
||||
stream_decoder = getattr(client, "stream_decoder", None)
|
||||
|
||||
async for raw in stream_iter:
|
||||
raw_lines.append(raw)
|
||||
if stream_decoder is not None:
|
||||
if stream_decoder is None:
|
||||
continue
|
||||
try:
|
||||
chunk = stream_decoder.decode(raw)
|
||||
except Exception: # noqa: BLE001
|
||||
except Exception: # noqa: BLE001 — skip malformed SSE payloads
|
||||
continue
|
||||
if not chunk.choices:
|
||||
continue
|
||||
if chunk.choices:
|
||||
delta_text = chunk.choices[0].delta.content
|
||||
if isinstance(delta_text, str) and delta_text:
|
||||
content_parts.append(delta_text)
|
||||
text_events.append(TextDeltaEvent(content=delta_text))
|
||||
fr = chunk.choices[0].finish_reason
|
||||
if isinstance(fr, str):
|
||||
finish_reason = fr
|
||||
|
||||
full_content = "".join(content_parts) or ""
|
||||
full_content = "".join(content_parts)
|
||||
tool_calls: list[AnyAssistantToolCall] | Unset = UNSET
|
||||
|
||||
if finish_reason in ("tool_calls", "function_call"):
|
||||
# Reconstruct tool calls from raw lines whenever present (do not rely solely
|
||||
# on finish_reason — some providers omit it on intermediate chunks).
|
||||
from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls
|
||||
|
||||
calls = merge_stream_tool_calls(raw_lines)
|
||||
if calls:
|
||||
tool_calls = cast("list[AnyAssistantToolCall]", calls)
|
||||
|
||||
assistant = AssistantChatMessage(content=full_content or None, tool_calls=tool_calls)
|
||||
assistant = AssistantChatMessage(
|
||||
content=full_content or None,
|
||||
tool_calls=tool_calls,
|
||||
)
|
||||
return assistant, text_events
|
||||
|
||||
|
||||
async def run_chat_loop( # noqa: PLR0913, C901
|
||||
async def run_chat_loop(
|
||||
client: ChatClient,
|
||||
messages: list[AnyChatMessage],
|
||||
*,
|
||||
@@ -127,8 +140,8 @@ async def run_chat_loop( # noqa: PLR0913, C901
|
||||
) -> AsyncIterator[AgentEvent]:
|
||||
"""Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
|
||||
|
||||
When ``stream=True``, text tokens yield as they arrive; for tool-calling
|
||||
rounds, the full response is reconstructed from stream deltas.
|
||||
When ``stream=True``, text tokens yield as they arrive; tool calls are
|
||||
reconstructed from raw SSE payloads when the client supports raw streaming.
|
||||
|
||||
Continues until the model returns no tool calls, or ``max_rounds`` is hit.
|
||||
"""
|
||||
@@ -151,16 +164,11 @@ async def run_chat_loop( # noqa: PLR0913, C901
|
||||
|
||||
if stream:
|
||||
assistant, text_events = await _stream_and_build_assistant(client, param)
|
||||
else:
|
||||
assistant, text_events = await _non_stream_assistant(client, param)
|
||||
|
||||
for event in text_events:
|
||||
yield event
|
||||
else:
|
||||
response = await client.chat_completions(param, stream=False)
|
||||
if not response.choices:
|
||||
msg = "chat completion response contained no choices"
|
||||
raise RuntimeError(msg)
|
||||
assistant = response.choices[0].message
|
||||
if isinstance(assistant.content, str) and assistant.content:
|
||||
yield TextDeltaEvent(content=assistant.content)
|
||||
|
||||
messages.append(assistant)
|
||||
yield AssistantMessageEvent(message=assistant)
|
||||
|
||||
@@ -43,7 +43,7 @@ def _database_config(store: ConfigStore) -> DatabaseConfig:
|
||||
return msgspec.convert(raw, DatabaseConfig)
|
||||
|
||||
|
||||
async def _run_chat( # noqa: PLR0913
|
||||
async def _run_chat(
|
||||
*,
|
||||
config_path: Path | None,
|
||||
provider_name: str | None,
|
||||
@@ -135,7 +135,7 @@ def main() -> None:
|
||||
show_default=True,
|
||||
help="Max tool-loop rounds per user turn.",
|
||||
)
|
||||
def chat_cmd( # noqa: PLR0913
|
||||
def chat_cmd(
|
||||
config_path: Path | None,
|
||||
provider_name: str | None,
|
||||
model: str | None,
|
||||
|
||||
+15
-13
@@ -19,7 +19,14 @@ if TYPE_CHECKING:
|
||||
|
||||
from plyngent.agent import AgentEvent
|
||||
|
||||
_TOOL_RESULT_PREVIEW = 200
|
||||
_TOOL_RESULT_PREVIEW = 120
|
||||
_TOOL_ARGS_PREVIEW = 80
|
||||
|
||||
|
||||
def _preview(text: str, limit: int) -> str:
|
||||
if len(text) <= limit:
|
||||
return text
|
||||
return text[:limit] + "…"
|
||||
|
||||
|
||||
async def render_events(events: AsyncIterator[AgentEvent]) -> None: # noqa: C901, PLR0912
|
||||
@@ -35,20 +42,15 @@ async def render_events(events: AsyncIterator[AgentEvent]) -> None: # noqa: C90
|
||||
elif isinstance(event, ToolCallEvent):
|
||||
call = event.tool_call
|
||||
if isinstance(call, AssistantFunctionToolCall):
|
||||
click.secho(
|
||||
f"\n[tool call] {call.function.name}({call.function.arguments})",
|
||||
fg="yellow",
|
||||
)
|
||||
args = _preview(call.function.arguments, _TOOL_ARGS_PREVIEW)
|
||||
click.secho(f"\n[tool] {call.function.name}({args})", fg="yellow")
|
||||
else:
|
||||
click.secho(f"\n[tool call] custom id={call.id}", fg="yellow")
|
||||
click.secho(f"\n[tool] custom id={call.id}", fg="yellow")
|
||||
elif isinstance(event, ToolResultEvent):
|
||||
content = event.message.content
|
||||
preview = (
|
||||
content
|
||||
if len(content) <= _TOOL_RESULT_PREVIEW
|
||||
else content[:_TOOL_RESULT_PREVIEW] + "…"
|
||||
)
|
||||
click.secho(f"[tool result] {preview}", fg="magenta")
|
||||
preview = _preview(event.message.content, _TOOL_RESULT_PREVIEW)
|
||||
# Single-line summary; full result stays in the message history for the model.
|
||||
one_line = preview.replace("\n", " ")
|
||||
click.secho(f"[tool ok] {one_line}", fg="magenta")
|
||||
elif isinstance(event, ErrorEvent):
|
||||
click.secho(f"\n[error] {event.message}", fg="bright_red")
|
||||
elif isinstance(event, CancelledEvent):
|
||||
|
||||
@@ -29,6 +29,7 @@ SLASH_COMMANDS: tuple[str, ...] = (
|
||||
"/tools",
|
||||
"/rounds",
|
||||
"/retry",
|
||||
"/status",
|
||||
)
|
||||
|
||||
_TOOLS_ARGS: tuple[str, ...] = ("on", "off")
|
||||
|
||||
@@ -36,10 +36,11 @@ Commands:
|
||||
/tools [on|off] Show or toggle tools
|
||||
/rounds [n] Show or set max tool-loop rounds
|
||||
/retry Retry the last failed user turn (after errors)
|
||||
/status Show session/provider/tools/rounds status
|
||||
|
||||
On network/API errors or Ctrl+C during a turn, the turn is not saved to
|
||||
the DB. Auto-retry waits 10s, 20s, then 30s (Ctrl+C cancels waits or the
|
||||
in-flight turn; use /retry later).
|
||||
in-flight turn; use /retry later). Streaming is on by default.
|
||||
|
||||
Tab completes slash commands and some arguments (provider, model, tools).
|
||||
Use --session ID or /resume to continue a prior chat after restart.
|
||||
@@ -51,6 +52,19 @@ _DEFAULT_HISTORY_LINES = 20
|
||||
_CONTENT_PREVIEW = 200
|
||||
|
||||
|
||||
def _cmd_status(state: ReplState) -> None:
|
||||
pending = state.agent.pending_retry_text
|
||||
pending_disp = "yes" if pending else "no"
|
||||
click.echo(
|
||||
f"provider={state.provider_name} model={state.model}\n"
|
||||
f"session={state.session_id} messages={len(state.agent.messages)} "
|
||||
f"pending_retry={pending_disp}\n"
|
||||
f"tools={'on' if state.tools_enabled else 'off'} "
|
||||
f"rounds={state.max_rounds} stream={'on' if state.agent.stream else 'off'}\n"
|
||||
f"workspace={state.workspace}"
|
||||
)
|
||||
|
||||
|
||||
async def _cmd_sessions(state: ReplState) -> None:
|
||||
sessions = await state.memory.list_sessions()
|
||||
if not sessions:
|
||||
@@ -230,6 +244,7 @@ async def _dispatch_slash(state: ReplState, command: str, arg: str) -> bool:
|
||||
"tools": lambda: _cmd_tools(state, arg),
|
||||
"rounds": lambda: _cmd_rounds(state, arg),
|
||||
"retry": lambda: _cmd_retry(state),
|
||||
"status": lambda: _cmd_status(state),
|
||||
}
|
||||
handler = handlers.get(command)
|
||||
if handler is None:
|
||||
@@ -261,7 +276,8 @@ async def run_repl(state: ReplState) -> None:
|
||||
click.echo(
|
||||
f"plyngent chat provider={state.provider_name} model={state.model} "
|
||||
f"session={state.session_id} tools={'on' if state.tools_enabled else 'off'} "
|
||||
f"rounds={state.max_rounds} messages={len(state.agent.messages)}"
|
||||
f"rounds={state.max_rounds} messages={len(state.agent.messages)} "
|
||||
f"stream={'on' if state.agent.stream else 'off'}"
|
||||
)
|
||||
click.echo("Type /help for commands. Empty line is ignored.")
|
||||
|
||||
|
||||
@@ -93,6 +93,7 @@ async def run_turn_with_retries(
|
||||
try:
|
||||
await run_cancellable(render_events(starter()))
|
||||
except asyncio.CancelledError:
|
||||
# Do not auto-retry cancelled turns — user intent was stop, not retry.
|
||||
click.echo()
|
||||
click.secho(
|
||||
"cancelled (turn not saved); use /retry to try again",
|
||||
|
||||
@@ -46,19 +46,12 @@ class BaseOpenAIClient:
|
||||
if line.startswith(b"data: "):
|
||||
yield self.chunk_decoder.decode(line[6:])
|
||||
|
||||
async def _parse_sse_stream(self, resp: AsyncResponse) -> AsyncIterator[StreamChatCompletionChunk]:
|
||||
"""Parse SSE using the tolerant ::class:`StreamChatCompletionChunk` decoder."""
|
||||
lines = resp.iter_lines()
|
||||
async for line in lines:
|
||||
if not line or line == b"data: [DONE]":
|
||||
continue
|
||||
if line.startswith(b"data: "):
|
||||
yield self.stream_decoder.decode(line[6:])
|
||||
|
||||
async def chat_completions_raw_lines(self, param: ChatCompletionsParam) -> AsyncIterator[bytes]:
|
||||
"""Yield raw SSE ``data: `` payload bytes for manual accumulation.
|
||||
|
||||
Each yielded bytes object is a complete JSON object (no ``data: `` prefix).
|
||||
The HTTP response is closed when the iterator finishes or is closed early
|
||||
(e.g. task cancellation during streaming).
|
||||
"""
|
||||
data = self.encoder.encode(msgspec.structs.replace(param, stream=True))
|
||||
resp = await self.session.post(
|
||||
@@ -67,11 +60,22 @@ class BaseOpenAIClient:
|
||||
headers={"Content-Type": "application/json"},
|
||||
stream=True,
|
||||
)
|
||||
try:
|
||||
async for line in resp.iter_lines():
|
||||
if not line or line == b"data: [DONE]":
|
||||
continue
|
||||
if line.startswith(b"data: "):
|
||||
yield line[6:]
|
||||
finally:
|
||||
# Best-effort abort of the in-flight stream (level-2 cancel toward HTTP).
|
||||
aclose = getattr(resp, "aclose", None)
|
||||
if callable(aclose):
|
||||
await aclose() # pyright: ignore[reportGeneralTypeIssues]
|
||||
else:
|
||||
close = getattr(resp, "close", None)
|
||||
if callable(close):
|
||||
_ = close()
|
||||
|
||||
|
||||
|
||||
class OpenAIClient(BaseOpenAIClient):
|
||||
|
||||
@@ -46,7 +46,7 @@ def _prepare_dest(source: Path, dest: Path, root: Path, *, overwrite: bool, dst_
|
||||
return dest
|
||||
|
||||
|
||||
def _copy_or_move_validated( # noqa: PLR0913
|
||||
def _copy_or_move_validated(
|
||||
source: Path,
|
||||
dest: Path,
|
||||
root: Path,
|
||||
|
||||
@@ -26,7 +26,7 @@ def _truncate(text: str, label: str) -> str:
|
||||
return text[:DEFAULT_MAX_OUTPUT_CHARS] + f"\n...[{label} truncated]"
|
||||
|
||||
|
||||
def _format_result( # noqa: PLR0913
|
||||
def _format_result(
|
||||
*,
|
||||
returncode: int | None,
|
||||
workdir_display: str,
|
||||
|
||||
Reference in New Issue
Block a user