From 8da9e9dd5956718a1791e21176d7bb65b54f3174 Mon Sep 17 00:00:00 2001 From: worldmozara Date: Tue, 14 Jul 2026 22:11:25 +0800 Subject: [PATCH] core/agent: streaming loop with text deltas and tool-call reconstruction --- src/plyngent/agent/loop.py | 98 +++++++++++++++++++++++++++++++------- 1 file changed, 81 insertions(+), 17 deletions(-) diff --git a/src/plyngent/agent/loop.py b/src/plyngent/agent/loop.py index 3f81bcc..c23ea26 100644 --- a/src/plyngent/agent/loop.py +++ b/src/plyngent/agent/loop.py @@ -1,6 +1,6 @@ from __future__ import annotations -from typing import TYPE_CHECKING +from typing import TYPE_CHECKING, cast from msgspec import UNSET @@ -11,10 +11,12 @@ from plyngent.lmproto.openai_compatible.model import ( ChatCompletionsParam, ToolChatMessage, ) +from plyngent.typedef import Unset # noqa: TC001 from .events import ( AgentEvent, AssistantMessageEvent, + ErrorEvent, MaxRoundsEvent, TextDeltaEvent, ToolCallEvent, @@ -43,7 +45,11 @@ async def _execute_tool_calls( for call in tool_calls: yield ToolCallEvent(tool_call=call) if isinstance(call, AssistantFunctionToolCall): - result_text = await tools.execute(call.function.name, call.function.arguments) + try: + result_text = await tools.execute(call.function.name, call.function.arguments) + except Exception as exc: # noqa: BLE001 + result_text = f"error: tool {call.function.name!r} failed: {exc}" + yield ErrorEvent(message=result_text) tool_msg = ToolChatMessage(content=result_text, tool_call_id=call.id) else: tool_msg = ToolChatMessage( @@ -54,7 +60,61 @@ async def _execute_tool_calls( yield ToolResultEvent(message=tool_msg) -async def run_chat_loop( # noqa: PLR0913 +async def _stream_and_build_assistant( # noqa: C901 + 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 = [] + 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] + raw_lines: list[bytes] = [] + content_parts: list[str] = [] + finish_reason: str | None = None + text_events = [] + 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 + + full_content = "".join(content_parts) or "" + tool_calls: list[AnyAssistantToolCall] | Unset = UNSET + + 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) + + assistant = AssistantChatMessage(content=full_content or None, tool_calls=tool_calls) + return assistant, text_events + + +async def run_chat_loop( # noqa: PLR0913, C901 client: ChatClient, messages: list[AnyChatMessage], *, @@ -63,12 +123,14 @@ async def run_chat_loop( # noqa: PLR0913 max_rounds: int = DEFAULT_MAX_ROUNDS, temperature: float | None = None, on_limit: LimitContinueHook | None = None, + stream: bool = True, ) -> 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. + Continues until the model returns no tool calls, or ``max_rounds`` is hit. - If ``on_limit`` is set and returns True when the cap is reached, another - batch of ``max_rounds`` is granted and the loop continues. """ tool_items: Sequence[AnyToolItem] | None = None if tools is not None and len(tools) > 0: @@ -86,17 +148,23 @@ async def run_chat_loop( # noqa: PLR0913 temperature=temperature if temperature is not None else UNSET, tools=list(tool_items) if tool_items is not None else UNSET, ) - 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 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) + messages.append(assistant) yield AssistantMessageEvent(message=assistant) - if isinstance(assistant.content, str) and assistant.content: - yield TextDeltaEvent(content=assistant.content) - tool_calls = assistant.tool_calls if tool_calls is UNSET or not tool_calls: return @@ -112,7 +180,3 @@ async def run_chat_loop( # noqa: PLR0913 continue yield MaxRoundsEvent(rounds=allowance, continued=False) return - - -def collect_assistant_messages(events: Sequence[AgentEvent]) -> list[AssistantChatMessage]: - return [e.message for e in events if isinstance(e, AssistantMessageEvent)]