-
Notifications
You must be signed in to change notification settings - Fork 22
Description
Problem
After #164 consolidated _aggregate_usage, _aggregate_finish_reason, _aggregate_event_data, and _parse_chunk into the base Stream class (~370 lines removed), three methods remain identically duplicated across all modality streams (TextStream, ImagesStream, AudioStream):
__init__— stores_transform_outputand_client(identical in all 3)_build_stream_metadata— addsmodel/provider/modalityfrom_client(identical in all 3)_parse_output— aggregates content → transforms → builds Output object (near-identical in all 3, only Output class differs)
Estimated savings: ~80 lines across 6 files with zero behavioral changes.
Proposed Changes
1. Move _transform_output and _stream_metadata to base Stream.__init__()
File: src/celeste/streaming.py
Add transform_output: Callable[..., Any] | None = None and stream_metadata: dict[str, Any] | None = None params to Stream.__init__(). Make both | None = None so test helpers (ConcreteStream, TypedStream) continue to work unchanged.
Instead of passing the entire client object, pass only the metadata the stream needs:
# ModalityClient._stream() call site
return stream_class(
sse_iterator,
transform_output=self._transform_output,
stream_metadata={"model": self.model.id, "provider": self.provider, "modality": self.modality},
**parameters,
)This eliminates the _client reference entirely — the stream only receives the data it actually uses. No circular import, no Any, no TYPE_CHECKING.
Remove the identical __init__ overrides from:
src/celeste/modalities/text/streaming.pysrc/celeste/modalities/images/streaming.pysrc/celeste/modalities/audio/streaming.py
2. Consolidate _build_stream_metadata() into base Stream
File: src/celeste/streaming.py
Replace the minimal base implementation:
def _build_stream_metadata(self, raw_events: list[dict[str, Any]]) -> dict[str, Any]:
return {**self._stream_metadata, "raw_events": raw_events}Remove the identical overrides from all 3 modality streams.
MRO safety: 8 provider stream mixins override _build_stream_metadata to filter events, then call super()._build_stream_metadata(filtered). Currently resolves: ProviderMixin → ModalityStream → base Stream. After removing the modality override, it resolves: ProviderMixin → base Stream directly. Behavior is identical because the modality override and the new base implementation produce the same output.
3. Pull _parse_output() into base Stream
File: src/celeste/streaming.py
Add _output_class: ClassVar[type[Output]] and abstract _aggregate_content() to base Stream. Change _parse_output from @abstractmethod to a concrete default:
def _parse_output(self, chunks, **parameters):
raw_content = self._aggregate_content(chunks)
content = self._transform_output(raw_content, **parameters) if self._transform_output else raw_content
raw_events = self._aggregate_event_data(chunks)
return self._output_class(
content=content,
usage=self._aggregate_usage(chunks),
finish_reason=self._aggregate_finish_reason(chunks),
metadata=self._build_stream_metadata(raw_events),
)Remove _parse_output from all 3 modality streams. Add _output_class ClassVar to each:
TextStream._output_class = TextOutputImagesStream._output_class = ImageOutputAudioStream._output_class = AudioOutput
4. Update ModalityClient._stream() call site
File: src/celeste/client.py
Change from passing client=self to passing stream_metadata=dict:
return stream_class(
sse_iterator,
transform_output=self._transform_output,
stream_metadata={"model": self.model.id, "provider": self.provider, "modality": self.modality},
**parameters,
)5. Update template and tests
templates/modalities/{modality_slug}/streaming.py.template— slim to ClassVars +_aggregate_contentonlytests/unit_tests/test_streaming.py— update abstract method assertions (_aggregate_contentreplaces_parse_outputas the required abstract method)
Result
Each modality stream reduces to ~10-15 lines (ClassVars + _aggregate_content):
class TextStream(Stream[TextOutput, TextParameters, TextChunk]):
_usage_class = TextUsage
_finish_reason_class = TextFinishReason
_chunk_class = TextChunk
_output_class = TextOutput
_empty_content = ""
def _aggregate_content(self, chunks: list[TextChunk]) -> str:
return "".join(chunk.content for chunk in chunks)Design Decisions
stream_metadatadict instead of client reference — the stream only uses_clientfor 3 metadata values (.model.id,.provider,.modality). Passing a dict eliminates the circular import betweenceleste.streamingandceleste.clientwithout resorting toAnyorTYPE_CHECKING| None = Nonefortransform_outputandstream_metadata— keeps baseStreamusable in tests without requiring a full client setup_output_classas ClassVar — consistent with existing pattern (_usage_class,_finish_reason_class,_chunk_class,_empty_content)- Defensive empty-chunks guard in base
_parse_output—ImagesStreamhad this; base__anext__already prevents it, but included for safety