wehub-resource-sync
139 lines
4.1 KiB
Python
139 lines
4.1 KiB
Python
"""Composable stream transforms for stream blocks.
|
|
|
|
Each transform is an async generator function that wraps a block
|
|
stream. Compose with ``pipe()``:
|
|
|
|
stream = pipe(
|
|
block_stream.stream(session, text),
|
|
skip_blocks(ReasoningBlock),
|
|
skip_intermediate_ends(),
|
|
)
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import AsyncIterator
|
|
|
|
from ._blocks import (
|
|
AnyBlock,
|
|
ResponseEndBlock,
|
|
TextDone,
|
|
)
|
|
|
|
|
|
def pipe(stream: AsyncIterator[AnyBlock], *transforms: object) -> AsyncIterator[AnyBlock]:
|
|
"""Compose transforms left-to-right.
|
|
|
|
``pipe(stream, t1, t2)`` is equivalent to ``t2(t1(stream))``.
|
|
"""
|
|
for t in transforms:
|
|
stream = t(stream)
|
|
return stream
|
|
|
|
|
|
def skip_blocks(*types: type) -> object:
|
|
"""Drop blocks of the given types.
|
|
|
|
Usage::
|
|
|
|
stream = skip_blocks(ReasoningBlock)(renderer.stream(...))
|
|
"""
|
|
|
|
def _transform(stream: AsyncIterator[AnyBlock]) -> AsyncIterator[AnyBlock]:
|
|
async def _inner() -> AsyncIterator[AnyBlock]:
|
|
async for block in stream:
|
|
if not isinstance(block, tuple(types)):
|
|
yield block
|
|
|
|
return _inner()
|
|
|
|
return _transform
|
|
|
|
|
|
def skip_intermediate_ends() -> object:
|
|
"""Suppress ``ResponseEndBlock`` events from tool loop iterations.
|
|
|
|
Only yields the final ``ResponseEndBlock`` — the one not followed
|
|
by another block from the same turn.
|
|
"""
|
|
|
|
def _transform(stream: AsyncIterator[AnyBlock]) -> AsyncIterator[AnyBlock]:
|
|
async def _inner() -> AsyncIterator[AnyBlock]:
|
|
pending_end: ResponseEndBlock | None = None
|
|
async for block in stream:
|
|
if isinstance(block, ResponseEndBlock):
|
|
# Buffer it — might not be the last one.
|
|
pending_end = block
|
|
else:
|
|
# A non-end block arrived, so the buffered end
|
|
# was intermediate. Drop it.
|
|
pending_end = None
|
|
yield block
|
|
# Stream ended. The buffered end IS the final one.
|
|
if pending_end is not None:
|
|
yield pending_end
|
|
|
|
return _inner()
|
|
|
|
return _transform
|
|
|
|
|
|
def merge_text_across_iterations() -> object:
|
|
"""Merge ``TextDone`` blocks across tool loop iterations.
|
|
|
|
When the tool loop runs N iterations, each may produce a
|
|
``TextDone``. This merges them into a single ``TextDone`` at
|
|
the end. ``TextChunk`` blocks pass through unchanged.
|
|
"""
|
|
|
|
def _transform(stream: AsyncIterator[AnyBlock]) -> AsyncIterator[AnyBlock]:
|
|
async def _inner() -> AsyncIterator[AnyBlock]:
|
|
accumulated = ""
|
|
async for block in stream:
|
|
if isinstance(block, TextDone):
|
|
accumulated += block.full_text
|
|
elif isinstance(block, ResponseEndBlock):
|
|
if accumulated:
|
|
yield TextDone(
|
|
full_text=accumulated,
|
|
has_code_blocks="```" in accumulated,
|
|
ctx=block.ctx,
|
|
)
|
|
accumulated = ""
|
|
yield block
|
|
else:
|
|
yield block
|
|
# Edge case: stream ends without ResponseEndBlock.
|
|
if accumulated:
|
|
yield TextDone(
|
|
full_text=accumulated,
|
|
has_code_blocks="```" in accumulated,
|
|
)
|
|
|
|
return _inner()
|
|
|
|
return _transform
|
|
|
|
|
|
def only_agent(agent_name: str | None) -> object:
|
|
"""Filter to blocks from a specific agent.
|
|
|
|
Pass ``None`` to include all agents (no filtering).
|
|
|
|
Usage::
|
|
|
|
stream = only_agent("coder.researcher")(renderer.stream(...))
|
|
|
|
:param agent_name: Agent name to filter by, or ``None`` for all.
|
|
"""
|
|
|
|
def _transform(stream: AsyncIterator[AnyBlock]) -> AsyncIterator[AnyBlock]:
|
|
async def _inner() -> AsyncIterator[AnyBlock]:
|
|
async for block in stream:
|
|
if agent_name is None or block.ctx.agent == agent_name:
|
|
yield block
|
|
|
|
return _inner()
|
|
|
|
return _transform
|