Add support for stream responses

[gtopper]

ML-11875
ML-11877

[Copilot]

Pull request overview

This PR adds comprehensive streaming response support to the storey data processing library, enabling steps to yield multiple chunks of data incrementally rather than returning a single result. The implementation includes new primitive types for streaming, modifications to existing flow steps, a new Collector step for aggregating streams, and extensive test coverage.

Changes:

• Added streaming primitives (StreamChunk, StreamCompletion, StreamingError) and modified Map, MapClass, Complete, Reduce, and ParallelExecution steps to support generator functions
• Introduced Collector step to aggregate streaming chunks back into single events
• Updated AwaitableResult and AsyncAwaitableResult to return generators for streaming responses

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 10 comments.

Show a summary per file

File	Description
storey/dtypes.py	Adds StreamChunk, StreamCompletion, and StreamingError classes for streaming support
storey/flow.py	Adds _StreamingStepMixin and updates Map, MapClass, Complete, Reduce, ParallelExecution, and Choice to handle streaming
storey/sources.py	Updates AwaitableResult and AsyncAwaitableResult to support streaming generators
storey/steps/collector.py	Implements new Collector step for aggregating streaming chunks
storey/steps/init.py	Exports the new Collector step
storey/init.py	Exports StreamingError and Collector for public API
tests/test_streaming.py	Comprehensive test suite covering streaming primitives, Map/MapClass streaming, Collector, Complete, error handling, graph splits, and cyclic graphs
tests/test_flow.py	Refactors cycle creation to use cleaner .to() API instead of direct _outlets manipulation

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull request overview

Copilot reviewed 8 out of 8 changed files in this pull request and generated no new comments.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[alxtkr77]

tests/test_streaming.py - Test Duplication

The 56 tests have significant structural duplication - almost every test has both a sync and async version with ~90% identical code (27 pairs).

Consider using pytest parametrization to consolidate:

@pytest.fixture(params=["sync", "async"])
def flow_context(request):
    if request.param == "sync":
        return SyncEmitSource, lambda f: f()
    else:
        return AsyncEmitSource, lambda f: asyncio.run(f())

def test_collector_basic(self, flow_context):
    source_cls, run = flow_context
    # Single implementation handles both

This would reduce tests from 56 to ~30 while maintaining the same coverage, and make future maintenance easier.

[alxtkr77]

flow.py - Duplicate _is_generator method

The _is_generator() method is defined identically in two places:

• _StreamingStepMixin (line ~201)
• ParallelExecutionRunnable (line ~1448)

Consider reusing the mixin method or extracting to a shared utility to avoid duplication.

[gtopper]

tests/test_streaming.py - Test Duplication

The 56 tests have significant structural duplication - almost every test has both a sync and async version with ~90% identical code (27 pairs).

Consider using pytest parametrization to consolidate:

@pytest.fixture(params=["sync", "async"])
def flow_context(request):
    if request.param == "sync":
        return SyncEmitSource, lambda f: f()
    else:
        return AsyncEmitSource, lambda f: asyncio.run(f())

def test_collector_basic(self, flow_context):
    source_cls, run = flow_context
    # Single implementation handles both

This would reduce tests from 56 to ~30 while maintaining the same coverage, and make future maintenance easier.

Yeah, this is sort of true. My AI also wanted to do this. The problem is, the APIs are too divergent, and the result of parameterization isn't very good. I.e. the above code snippet doesn't cover all the conditionals needed.

[royischoss]

Hey LGTM two minor comments

[royischoss]

LGTM 👍