wehub-resource-sync
c3bf08ac8d
K8s Workspace Integration Tests / k8s-workspace-tests (push) Waiting to run
Pre-commit / run (ubuntu-latest) (push) Waiting to run
Python Unittest Coverage / test (macos-15, 3.11) (push) Waiting to run
Python Unittest Coverage / test (ubuntu-latest, 3.11) (push) Waiting to run
Python Unittest Coverage / test (windows-latest, 3.11) (push) Waiting to run
Web UI / check (push) Waiting to run
105 lines
3.8 KiB
Python
105 lines
3.8 KiB
Python
# -*- coding: utf-8 -*-
|
|
"""Shared utility helpers for model-examples scripts."""
|
|
import base64 as _b64
|
|
from typing import AsyncGenerator
|
|
|
|
from agentscope.message import (
|
|
TextBlock,
|
|
ThinkingBlock,
|
|
DataBlock,
|
|
Base64Source,
|
|
)
|
|
from agentscope.model import ChatResponse
|
|
|
|
|
|
async def stream_and_collect(
|
|
gen: AsyncGenerator[ChatResponse, None],
|
|
) -> ChatResponse:
|
|
"""Stream delta chunks printing text in real-time; return the final chunk.
|
|
|
|
Only delta chunks (is_last=False) are printed. The final accumulated
|
|
chunk (is_last=True) is returned so callers can read ToolCallBlock objects
|
|
from it without re-printing the entire content. Text from the final chunk
|
|
is printed only when no text was streamed in any delta chunk (e.g. some
|
|
models batch the answer in the last chunk).
|
|
|
|
Streaming ``DataBlock`` chunks (e.g. omni audio output) are reported as
|
|
per-chunk size summaries while the stream is being consumed, so callers
|
|
can see the audio arriving incrementally rather than only in the final
|
|
cumulative chunk.
|
|
"""
|
|
final: ChatResponse | None = None
|
|
in_thinking = False
|
|
text_was_streamed = False
|
|
# Track per-block audio progress: block_id -> (media_type, chunk_count,
|
|
# total_bytes).
|
|
audio_progress: dict[str, list] = {}
|
|
async for chunk in gen:
|
|
if chunk.is_last:
|
|
final = chunk
|
|
continue # Skip printing; full content is in the final chunk
|
|
for block in chunk.content:
|
|
if isinstance(block, ThinkingBlock):
|
|
if not in_thinking:
|
|
print("[Thinking] ", end="", flush=True)
|
|
in_thinking = True
|
|
print(block.thinking, end="", flush=True)
|
|
elif isinstance(block, TextBlock):
|
|
if in_thinking:
|
|
print() # Newline after thinking content
|
|
print("--- Answer ---")
|
|
in_thinking = False
|
|
print(block.text, end="", flush=True)
|
|
text_was_streamed = True
|
|
elif isinstance(block, DataBlock) and isinstance(
|
|
block.source,
|
|
Base64Source,
|
|
):
|
|
# Streaming binary delta (e.g. omni audio output).
|
|
if in_thinking:
|
|
print()
|
|
print("--- Answer ---")
|
|
in_thinking = False
|
|
delta_bytes = len(_b64.b64decode(block.source.data))
|
|
state = audio_progress.setdefault(
|
|
block.id,
|
|
[block.source.media_type, 0, 0],
|
|
)
|
|
state[1] += 1
|
|
state[2] += delta_bytes
|
|
print(
|
|
f"\n[Audio chunk #{state[1]} ({state[0]}): "
|
|
f"+{delta_bytes}B, total={state[2]}B]",
|
|
flush=True,
|
|
)
|
|
# If text was not streamed in any delta, print it from the final chunk now.
|
|
if not text_was_streamed and final is not None:
|
|
final_text = "".join(
|
|
b.text
|
|
for b in final.content
|
|
if isinstance(b, TextBlock) and b.text
|
|
)
|
|
if final_text:
|
|
if in_thinking:
|
|
print()
|
|
print("--- Answer ---")
|
|
in_thinking = False
|
|
print(final_text)
|
|
if in_thinking:
|
|
print()
|
|
# Report cumulative audio output from the final chunk.
|
|
if final is not None:
|
|
for block in final.content:
|
|
if isinstance(block, DataBlock) and isinstance(
|
|
block.source,
|
|
Base64Source,
|
|
):
|
|
media = block.source.media_type
|
|
byte_size = len(_b64.b64decode(block.source.data))
|
|
print(
|
|
f"[Audio output (cumulative): {media}, {byte_size} bytes]",
|
|
)
|
|
print()
|
|
assert final is not None
|
|
return final
|