Spaces:
Paused
Paused
| # -*- 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 | |