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`.
- **`@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/`)
+1
View File
@@ -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 = [
"",
+2 -1
View File
@@ -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,
*,
+55 -47
View File
@@ -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."""
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: list[TextDeltaEvent] = []
if isinstance(assistant.content, str) and assistant.content:
text_events.append(TextDeltaEvent(content=assistant.content))
return assistant, text_events
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:
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 = []
if isinstance(assistant.content, str) and assistant.content:
text_events.append(TextDeltaEvent(content=assistant.content))
return assistant, text_events
return await _non_stream_assistant(client, param)
stream_iter = await raw_lines_attr(param) # type: ignore[misc]
# 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:
try:
chunk = stream_decoder.decode(raw)
except Exception: # noqa: BLE001
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
if stream_decoder is None:
continue
try:
chunk = stream_decoder.decode(raw)
except Exception: # noqa: BLE001 — skip malformed SSE payloads
continue
if not chunk.choices:
continue
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))
full_content = "".join(content_parts) or ""
full_content = "".join(content_parts)
tool_calls: list[AnyAssistantToolCall] | Unset = UNSET
# 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
if finish_reason in ("tool_calls", "function_call"):
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)
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)
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)
assistant, text_events = await _non_stream_assistant(client, param)
for event in text_events:
yield event
messages.append(assistant)
yield AssistantMessageEvent(message=assistant)
+2 -2
View File
@@ -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
View File
@@ -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):
+1
View File
@@ -29,6 +29,7 @@ SLASH_COMMANDS: tuple[str, ...] = (
"/tools",
"/rounds",
"/retry",
"/status",
)
_TOOLS_ARGS: tuple[str, ...] = ("on", "off")
+18 -2
View File
@@ -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.")
+1
View File
@@ -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,
)
async for line in resp.iter_lines():
if not line or line == b"data: [DONE]":
continue
if line.startswith(b"data: "):
yield line[6:]
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):
+1 -1
View File
@@ -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,
+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]"
def _format_result( # noqa: PLR0913
def _format_result(
*,
returncode: int | None,
workdir_display: str,