mirror of
https://github.com/NCBM/plyngent.git
synced 2026-07-23 05:55:16 +08:00
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:
+42
-21
@@ -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,
|
||||||
|
|||||||
@@ -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:
|
||||||
|
|||||||
Reference in New Issue
Block a user