mirror of
https://github.com/agentscope-ai/ReMe.git
synced 2026-09-22 00:32:49 +00:00
105 lines
3.9 KiB
Python
105 lines
3.9 KiB
Python
"""Common utilities: hashing and async stream task execution."""
|
|
|
|
import asyncio
|
|
import hashlib
|
|
from collections.abc import AsyncGenerator
|
|
from typing import Any, Literal
|
|
|
|
from .logger_utils import get_logger
|
|
from ..enumeration import ChunkEnum
|
|
from ..schema import StreamChunk
|
|
|
|
|
|
def hash_text(text: str, encoding: str = "utf-8") -> str:
|
|
"""Return SHA-256 hex digest of text."""
|
|
return hashlib.sha256(text.encode(encoding)).hexdigest()
|
|
|
|
|
|
def _format_chunk(
|
|
chunk: StreamChunk,
|
|
output_format: Literal["str", "bytes", "chunk"],
|
|
) -> str | bytes | StreamChunk:
|
|
"""Render a StreamChunk in the requested transport format."""
|
|
if output_format == "chunk":
|
|
return chunk
|
|
data = "data:[DONE]\n\n" if chunk.done else f"data:{chunk.model_dump_json()}\n\n"
|
|
return data.encode() if output_format == "bytes" else data
|
|
|
|
|
|
async def execute_stream_task(
|
|
stream_queue: asyncio.Queue[StreamChunk],
|
|
task: asyncio.Task[Any],
|
|
task_name: str | None = None,
|
|
output_format: Literal["str", "bytes", "chunk"] = "str",
|
|
) -> AsyncGenerator[str | bytes | StreamChunk, None]:
|
|
"""Yield chunks from stream_queue while monitoring task; cancels task on exit.
|
|
|
|
output_format: "str"/"bytes" emit SSE frames, "chunk" emits raw StreamChunk.
|
|
"""
|
|
logger = get_logger()
|
|
consumer: asyncio.Task[StreamChunk] | None = None
|
|
try:
|
|
while True:
|
|
consumer = get_chunk = asyncio.create_task(stream_queue.get())
|
|
done, _pending = await asyncio.wait({get_chunk, task}, return_when=asyncio.FIRST_COMPLETED)
|
|
|
|
# Producer still running — relay the next chunk and continue.
|
|
if task not in done:
|
|
chunk = get_chunk.result()
|
|
yield _format_chunk(chunk, output_format)
|
|
if chunk.done:
|
|
return
|
|
continue
|
|
|
|
# Producer finished. Capture any pending chunk, then stop the consumer wait
|
|
# so we can inspect task state safely.
|
|
pending_chunk: StreamChunk | None = None
|
|
if get_chunk in done:
|
|
pending_chunk = get_chunk.result()
|
|
else:
|
|
get_chunk.cancel()
|
|
try:
|
|
await get_chunk
|
|
except asyncio.CancelledError:
|
|
pass
|
|
|
|
# Surface task failure first — an exception trumps trailing data.
|
|
if task.cancelled():
|
|
msg = f"Task cancelled: {task_name}" if task_name else "Task cancelled"
|
|
raise asyncio.CancelledError(msg)
|
|
exc = task.exception()
|
|
if exc is not None:
|
|
log_msg = f"Task error in {task_name}: {exc}" if task_name else f"Task error: {exc}"
|
|
logger.error(log_msg, exc_info=exc)
|
|
raise exc
|
|
|
|
# Producer ended cleanly — flush pending + drain queue so no chunk is lost,
|
|
# then emit the terminal sentinel.
|
|
if pending_chunk is not None:
|
|
yield _format_chunk(pending_chunk, output_format)
|
|
if pending_chunk.done:
|
|
return
|
|
while not stream_queue.empty():
|
|
chunk = stream_queue.get_nowait()
|
|
yield _format_chunk(chunk, output_format)
|
|
if chunk.done:
|
|
return
|
|
|
|
yield _format_chunk(StreamChunk(chunk_type=ChunkEnum.DONE, chunk="", done=True), output_format)
|
|
return
|
|
|
|
finally:
|
|
# Cancel consumer wait if still pending (e.g. on consumer aclose).
|
|
if consumer is not None and not consumer.done():
|
|
consumer.cancel()
|
|
try:
|
|
await consumer
|
|
except asyncio.CancelledError:
|
|
pass
|
|
# Cancel producer task if still running to avoid resource leaks.
|
|
if not task.done():
|
|
task.cancel()
|
|
try:
|
|
await task
|
|
except asyncio.CancelledError:
|
|
pass
|