core/agent: yield stream text deltas as chunks arrive

Stop buffering the full SSE response before emitting TextDeltaEvent so
CLI/UI can render assistant output token-by-token.
This commit is contained in:
2026-07-15 10:27:49 +08:00
parent 53ac1c8d44
commit 4741341888
2 changed files with 90 additions and 21 deletions
+42 -21
View File
@@ -125,19 +125,17 @@ async def _non_stream_assistant(
return assistant, text_events return assistant, text_events
async def _stream_and_build_assistant( async def _stream_round(
client: ChatClient, client: ChatClient,
param: ChatCompletionsParam, param: ChatCompletionsParam,
) -> tuple[AssistantChatMessage, list[TextDeltaEvent]]: ) -> AsyncIterator[AgentEvent]:
"""Stream one completion via the normal client API. """Stream one completion; yield text deltas as chunks arrive, then assistant.
Pattern: ``stream = await client.chat_completions(..., stream=True)`` then Pattern: ``stream = await client.chat_completions(..., stream=True)`` then
``async for chunk in stream``. That return type (async function → async ``async for chunk in stream``. Tool-call deltas are merged after the stream.
iterator) is accepted as the library interface.
""" """
stream = await client.chat_completions(param, stream=True) stream = await client.chat_completions(param, stream=True)
content_parts: list[str] = [] content_parts: list[str] = []
text_events: list[TextDeltaEvent] = []
tool_deltas: list[StreamToolCallDelta] = [] tool_deltas: list[StreamToolCallDelta] = []
async for chunk in stream: async for chunk in stream:
@@ -147,7 +145,7 @@ async def _stream_and_build_assistant(
delta = choice.delta delta = choice.delta
if isinstance(delta.content, str) and delta.content: if isinstance(delta.content, str) and delta.content:
content_parts.append(delta.content) content_parts.append(delta.content)
text_events.append(TextDeltaEvent(content=delta.content)) yield TextDeltaEvent(content=delta.content)
if delta.tool_calls is not UNSET and delta.tool_calls: if delta.tool_calls is not UNSET and delta.tool_calls:
tool_deltas.extend(delta.tool_calls) tool_deltas.extend(delta.tool_calls)
@@ -162,7 +160,39 @@ async def _stream_and_build_assistant(
content=full_content or None, content=full_content or None,
tool_calls=tool_calls, tool_calls=tool_calls,
) )
return assistant, text_events yield AssistantMessageEvent(message=assistant)
async def _assistant_round(
client: ChatClient,
param: ChatCompletionsParam,
messages: list[AnyChatMessage],
*,
stream: bool,
) -> AsyncIterator[AgentEvent]:
"""One model turn: yield events and append the assistant message to ``messages``."""
if stream:
async for event in _stream_round(client, param):
if isinstance(event, AssistantMessageEvent):
messages.append(event.message)
yield event
return
assistant, text_events = await _non_stream_assistant(client, param)
for event in text_events:
yield event
messages.append(assistant)
yield AssistantMessageEvent(message=assistant)
def _last_assistant(messages: list[AnyChatMessage], pre_len: int) -> AssistantChatMessage:
if len(messages) <= pre_len:
msg = "model round produced no assistant message"
raise RuntimeError(msg)
last = messages[-1]
if not isinstance(last, AssistantChatMessage):
msg = "model round did not end with an assistant message"
raise TypeError(msg)
return last
async def run_chat_loop( async def run_chat_loop(
@@ -207,21 +237,12 @@ async def run_chat_loop(
tools=list(tool_items) if tool_items is not None else UNSET, tools=list(tool_items) if tool_items is not None else UNSET,
) )
if stream: pre_len = len(messages)
assistant, text_events = await _stream_and_build_assistant(client, param) async for event in _assistant_round(client, param, messages, stream=stream):
else:
assistant, text_events = await _non_stream_assistant(client, param)
for event in text_events:
yield event yield event
assistant = _last_assistant(messages, pre_len)
messages.append(assistant)
yield AssistantMessageEvent(message=assistant)
tool_calls = assistant.tool_calls tool_calls = assistant.tool_calls
if tool_calls is UNSET or not tool_calls: if tool_calls is UNSET or not tool_calls or tools is None:
return
if tools is None:
return return
async for event in _execute_tool_calls( async for event in _execute_tool_calls(
tools, tools,
+48
View File
@@ -180,6 +180,54 @@ async def test_run_chat_loop_text_only() -> None:
assert len(client.calls) == 1 assert len(client.calls) == 1
async def test_stream_yields_deltas_incrementally() -> None:
"""Text deltas are yielded as chunks arrive, not only after the full stream."""
class ChunkClient:
@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
if not stream:
return _response(AssistantChatMessage(content="ab"))
async def chunks() -> AsyncIterator[ChatCompletionChunk]:
for part in ("a", "b"):
yield ChatCompletionChunk(
id="1",
object="chat.completion.chunk",
created=0,
model="t",
choices=[
ChunkChoice(
index=0,
delta=DeltaMessage(content=part),
finish_reason=None,
)
],
)
return chunks()
messages: list[AnyChatMessage] = [UserChatMessage(content="hi")]
events = [e async for e in run_chat_loop(ChunkClient(), messages, model="m", stream=True)]
deltas = [e for e in events if isinstance(e, TextDeltaEvent)]
assert [d.content for d in deltas] == ["a", "b"]
assert any(isinstance(e, AssistantMessageEvent) for e in events)
assert isinstance(messages[-1], AssistantChatMessage)
assert messages[-1].content == "ab"
async def test_run_chat_loop_with_tools() -> None: async def test_run_chat_loop_with_tools() -> None:
@tool @tool
def add(a: int, b: int) -> int: def add(a: int, b: int) -> int: