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:
2026-07-14 22:22:14 +08:00
parent 1b582e82d3
commit e063efda20
12 changed files with 118 additions and 83 deletions
+3 -2
View File
@@ -54,8 +54,9 @@ 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; default **streaming** text deltas + raw-SSE tool-call merge; optional `on_limit`.
- **`ChatAgent`**: wrapper with optional `MemoryStore` bind (load/append messages on success only); `pending_retry_text` + `retry()` after failed turns. - **`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/`) ### Tools (`tools/`)
+1
View File
@@ -107,6 +107,7 @@ ignore = [
"COM812", # missing-trailing-comma "COM812", # missing-trailing-comma
"PLC0415", # import-outside-top-level "PLC0415", # import-outside-top-level
"TID252", "TID252",
"PLR0913", # too-many-arguments (common on constructors/composed functions)
] ]
allowed-confusables = [ allowed-confusables = [
"", "",
+2 -1
View File
@@ -30,10 +30,11 @@ class ChatAgent:
max_rounds: int max_rounds: int
temperature: float | None temperature: float | None
on_limit: LimitContinueHook | None on_limit: LimitContinueHook | None
stream: bool
messages: list[AnyChatMessage] messages: list[AnyChatMessage]
pending_retry_text: str | None pending_retry_text: str | None
def __init__( # noqa: PLR0913 def __init__(
self, self,
client: ChatClient, client: ChatClient,
*, *,
+37 -29
View File
@@ -60,61 +60,74 @@ async def _execute_tool_calls(
yield ToolResultEvent(message=tool_msg) yield ToolResultEvent(message=tool_msg)
async def _stream_and_build_assistant( # noqa: C901 async def _non_stream_assistant(
client: ChatClient, client: ChatClient,
param: ChatCompletionsParam, param: ChatCompletionsParam,
) -> tuple[AssistantChatMessage, list[TextDeltaEvent]]: ) -> 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) response = await client.chat_completions(param, stream=False)
if not response.choices: if not response.choices:
msg = "chat completion response contained no choices" msg = "chat completion response contained no choices"
raise RuntimeError(msg) raise RuntimeError(msg)
assistant = response.choices[0].message assistant = response.choices[0].message
text_events = [] text_events: list[TextDeltaEvent] = []
if isinstance(assistant.content, str) and assistant.content: if isinstance(assistant.content, str) and assistant.content:
text_events.append(TextDeltaEvent(content=assistant.content)) text_events.append(TextDeltaEvent(content=assistant.content))
return assistant, text_events 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] = [] raw_lines: list[bytes] = []
content_parts: list[str] = [] content_parts: list[str] = []
finish_reason: str | None = None text_events: list[TextDeltaEvent] = []
text_events = []
stream_decoder = getattr(client, "stream_decoder", None) stream_decoder = getattr(client, "stream_decoder", None)
async for raw in stream_iter: async for raw in stream_iter:
raw_lines.append(raw) raw_lines.append(raw)
if stream_decoder is not None: if stream_decoder is None:
continue
try: try:
chunk = stream_decoder.decode(raw) chunk = stream_decoder.decode(raw)
except Exception: # noqa: BLE001 except Exception: # noqa: BLE001 — skip malformed SSE payloads
continue
if not chunk.choices:
continue continue
if chunk.choices:
delta_text = chunk.choices[0].delta.content delta_text = chunk.choices[0].delta.content
if isinstance(delta_text, str) and delta_text: if isinstance(delta_text, str) and delta_text:
content_parts.append(delta_text) content_parts.append(delta_text)
text_events.append(TextDeltaEvent(content=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 tool_calls: list[AnyAssistantToolCall] | Unset = UNSET
# Reconstruct tool calls from raw lines whenever present (do not rely solely
if finish_reason in ("tool_calls", "function_call"): # on finish_reason — some providers omit it on intermediate chunks).
from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls from plyngent.lmproto.openai_compatible.client import merge_stream_tool_calls
calls = merge_stream_tool_calls(raw_lines) calls = merge_stream_tool_calls(raw_lines)
if calls: if calls:
tool_calls = cast("list[AnyAssistantToolCall]", 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 return assistant, text_events
async def run_chat_loop( # noqa: PLR0913, C901 async def run_chat_loop(
client: ChatClient, client: ChatClient,
messages: list[AnyChatMessage], messages: list[AnyChatMessage],
*, *,
@@ -127,8 +140,8 @@ async def run_chat_loop( # noqa: PLR0913, C901
) -> AsyncIterator[AgentEvent]: ) -> AsyncIterator[AgentEvent]:
"""Multi-round chat/tool loop; mutates ``messages`` in place and yields events. """Multi-round chat/tool loop; mutates ``messages`` in place and yields events.
When ``stream=True``, text tokens yield as they arrive; for tool-calling When ``stream=True``, text tokens yield as they arrive; tool calls are
rounds, the full response is reconstructed from stream deltas. reconstructed from raw SSE payloads when the client supports raw streaming.
Continues until the model returns no tool calls, or ``max_rounds`` is hit. 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: if stream:
assistant, text_events = await _stream_and_build_assistant(client, param) 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: for event in text_events:
yield event 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) messages.append(assistant)
yield AssistantMessageEvent(message=assistant) yield AssistantMessageEvent(message=assistant)
+2 -2
View File
@@ -43,7 +43,7 @@ def _database_config(store: ConfigStore) -> DatabaseConfig:
return msgspec.convert(raw, DatabaseConfig) return msgspec.convert(raw, DatabaseConfig)
async def _run_chat( # noqa: PLR0913 async def _run_chat(
*, *,
config_path: Path | None, config_path: Path | None,
provider_name: str | None, provider_name: str | None,
@@ -135,7 +135,7 @@ def main() -> None:
show_default=True, show_default=True,
help="Max tool-loop rounds per user turn.", help="Max tool-loop rounds per user turn.",
) )
def chat_cmd( # noqa: PLR0913 def chat_cmd(
config_path: Path | None, config_path: Path | None,
provider_name: str | None, provider_name: str | None,
model: str | None, model: str | None,
+15 -13
View File
@@ -19,7 +19,14 @@ if TYPE_CHECKING:
from plyngent.agent import AgentEvent 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 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): elif isinstance(event, ToolCallEvent):
call = event.tool_call call = event.tool_call
if isinstance(call, AssistantFunctionToolCall): if isinstance(call, AssistantFunctionToolCall):
click.secho( args = _preview(call.function.arguments, _TOOL_ARGS_PREVIEW)
f"\n[tool call] {call.function.name}({call.function.arguments})", click.secho(f"\n[tool] {call.function.name}({args})", fg="yellow")
fg="yellow",
)
else: 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): elif isinstance(event, ToolResultEvent):
content = event.message.content preview = _preview(event.message.content, _TOOL_RESULT_PREVIEW)
preview = ( # Single-line summary; full result stays in the message history for the model.
content one_line = preview.replace("\n", " ")
if len(content) <= _TOOL_RESULT_PREVIEW click.secho(f"[tool ok] {one_line}", fg="magenta")
else content[:_TOOL_RESULT_PREVIEW] + ""
)
click.secho(f"[tool result] {preview}", fg="magenta")
elif isinstance(event, ErrorEvent): elif isinstance(event, ErrorEvent):
click.secho(f"\n[error] {event.message}", fg="bright_red") click.secho(f"\n[error] {event.message}", fg="bright_red")
elif isinstance(event, CancelledEvent): elif isinstance(event, CancelledEvent):
+1
View File
@@ -29,6 +29,7 @@ SLASH_COMMANDS: tuple[str, ...] = (
"/tools", "/tools",
"/rounds", "/rounds",
"/retry", "/retry",
"/status",
) )
_TOOLS_ARGS: tuple[str, ...] = ("on", "off") _TOOLS_ARGS: tuple[str, ...] = ("on", "off")
+18 -2
View File
@@ -36,10 +36,11 @@ Commands:
/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) /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 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 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). 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.
@@ -51,6 +52,19 @@ _DEFAULT_HISTORY_LINES = 20
_CONTENT_PREVIEW = 200 _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: async def _cmd_sessions(state: ReplState) -> None:
sessions = await state.memory.list_sessions() sessions = await state.memory.list_sessions()
if not 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), "tools": lambda: _cmd_tools(state, arg),
"rounds": lambda: _cmd_rounds(state, arg), "rounds": lambda: _cmd_rounds(state, arg),
"retry": lambda: _cmd_retry(state), "retry": lambda: _cmd_retry(state),
"status": lambda: _cmd_status(state),
} }
handler = handlers.get(command) handler = handlers.get(command)
if handler is None: if handler is None:
@@ -261,7 +276,8 @@ async def run_repl(state: ReplState) -> None:
click.echo( click.echo(
f"plyngent chat provider={state.provider_name} model={state.model} " f"plyngent chat provider={state.provider_name} model={state.model} "
f"session={state.session_id} tools={'on' if state.tools_enabled else 'off'} " 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.") click.echo("Type /help for commands. Empty line is ignored.")
+1
View File
@@ -93,6 +93,7 @@ async def run_turn_with_retries(
try: try:
await run_cancellable(render_events(starter())) await run_cancellable(render_events(starter()))
except asyncio.CancelledError: except asyncio.CancelledError:
# Do not auto-retry cancelled turns — user intent was stop, not retry.
click.echo() click.echo()
click.secho( click.secho(
"cancelled (turn not saved); use /retry to try again", "cancelled (turn not saved); use /retry to try again",
@@ -46,19 +46,12 @@ class BaseOpenAIClient:
if line.startswith(b"data: "): if line.startswith(b"data: "):
yield self.chunk_decoder.decode(line[6:]) 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]: async def chat_completions_raw_lines(self, param: ChatCompletionsParam) -> AsyncIterator[bytes]:
"""Yield raw SSE ``data: `` payload bytes for manual accumulation. """Yield raw SSE ``data: `` payload bytes for manual accumulation.
Each yielded bytes object is a complete JSON object (no ``data: `` prefix). 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)) data = self.encoder.encode(msgspec.structs.replace(param, stream=True))
resp = await self.session.post( resp = await self.session.post(
@@ -67,11 +60,22 @@ class BaseOpenAIClient:
headers={"Content-Type": "application/json"}, headers={"Content-Type": "application/json"},
stream=True, stream=True,
) )
try:
async for line in resp.iter_lines(): async for line in resp.iter_lines():
if not line or line == b"data: [DONE]": if not line or line == b"data: [DONE]":
continue continue
if line.startswith(b"data: "): if line.startswith(b"data: "):
yield line[6:] 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): class OpenAIClient(BaseOpenAIClient):
+1 -1
View File
@@ -46,7 +46,7 @@ def _prepare_dest(source: Path, dest: Path, root: Path, *, overwrite: bool, dst_
return dest return dest
def _copy_or_move_validated( # noqa: PLR0913 def _copy_or_move_validated(
source: Path, source: Path,
dest: Path, dest: Path,
root: Path, root: Path,
+1 -1
View File
@@ -26,7 +26,7 @@ def _truncate(text: str, label: str) -> str:
return text[:DEFAULT_MAX_OUTPUT_CHARS] + f"\n...[{label} truncated]" return text[:DEFAULT_MAX_OUTPUT_CHARS] + f"\n...[{label} truncated]"
def _format_result( # noqa: PLR0913 def _format_result(
*, *,
returncode: int | None, returncode: int | None,
workdir_display: str, workdir_display: str,