Skip to content

Deduplicate streaming methods across all modalities #164

Description

@Kamilbenkirane

Problem

The streaming layer has significant code duplication across all three modalities (text, images, audio). Three methods — _aggregate_usage, _aggregate_finish_reason, _aggregate_event_data — are copy-pasted identically across TextStream, ImagesStream, and AudioStream. The _parse_chunk method follows the same flow in 10 of 14 provider streams. _aggregate_content is identical across all 8 text providers.

Estimated savings: ~290 lines across 14 files with zero behavioral changes.

Proposed Changes

1. Base Stream class — add hooks + concrete defaults

File: src/celeste/streaming.py

Add _chunk_class and _empty_content ClassVars (alongside existing _usage_class, _finish_reason_class).

Add _parse_chunk_content and _wrap_chunk_content hooks (mirroring existing _parse_chunk_usage pattern).

Change _parse_chunk from @abstractmethod to a concrete default implementing the universal pattern:

def _parse_chunk(self, event: dict[str, Any]) -> ChunkBase | None:
    content = self._parse_chunk_content(event)
    if content is None:
        usage = self._get_chunk_usage(event)
        finish_reason = self._get_chunk_finish_reason(event)
        if usage is None and finish_reason is None:
            return None
        content = self._empty_content
    else:
        content = self._wrap_chunk_content(content)
    return self._chunk_class(
        content=content,
        finish_reason=self._get_chunk_finish_reason(event),
        usage=self._get_chunk_usage(event),
        metadata={"event_data": event},
    )

Add 3 concrete aggregation methods:

  • _aggregate_usage — reverse scan, return last chunk.usage or _usage_class()
  • _aggregate_finish_reason — reverse scan, return last chunk.finish_reason
  • _aggregate_event_data — collect chunk.metadata["event_data"]

2. Update modality streams

TextStream (src/celeste/modalities/text/streaming.py):

  • Set _chunk_class = TextChunk, _empty_content = ""
  • Remove _aggregate_usage, _aggregate_finish_reason, @abstractmethod _aggregate_event_data
  • Change _aggregate_content from @abstractmethod to default: "".join(chunk.content for chunk in chunks)

ImagesStream (src/celeste/modalities/images/streaming.py):

  • Set _chunk_class = ImageChunk, _empty_content = ImageArtifact(data=b"")
  • Override _wrap_chunk_content to return ImageArtifact(data=raw_content)
  • Remove _aggregate_usage, _aggregate_finish_reason, @abstractmethod _aggregate_event_data
  • Keep @abstractmethod _aggregate_content (implementations differ across providers)

AudioStream (src/celeste/modalities/audio/streaming.py):

  • Set _chunk_class = AudioChunk, _empty_content = b""
  • Remove _aggregate_usage, _aggregate_finish_reason, @abstractmethod _aggregate_event_data
  • Keep @abstractmethod _aggregate_content (needs client-specific mime_type)

3. Clean up provider streams

Text — become empty classes (remove _parse_chunk, _aggregate_content, _aggregate_event_data):

  • groq/client.py, deepseek/client.py, mistral/client.py, cohere/client.py, moonshot/client.py, google/client.py

Text — simplify (keep custom _parse_chunk, remove _aggregate_content, simplify _aggregate_event_data to super() delegation):

  • anthropic/client.py — keeps _parse_chunk override (captures message_start)
  • openresponses/client.py — keeps _parse_chunk override (captures response.completed)

Images:

  • openai/client.py — remove _parse_chunk + _aggregate_event_data, keep _aggregate_content
  • byteplus/client.py — keep _parse_chunk override (custom error handling), simplify _aggregate_event_data to super()
  • ollama/client.py — keep all (different metadata structure)

Audio:

  • elevenlabs/client.py — remove _parse_chunk + _aggregate_event_data, keep _aggregate_content
  • gradium/client.py — remove _parse_chunk + _aggregate_event_data, keep _aggregate_content

4. Update tests

  • Add _chunk_class and _empty_content to ConcreteStream in test_streaming.py
  • Remove test_subclass_without_parse_chunk_fails (_parse_chunk no longer abstract)
  • Add tests for new default behaviors

Design Decisions

  • is None check (not not content) in _parse_chunk_parse_chunk_content returns None for "no content", providers never return empty strings/bytes as content
  • _wrap_chunk_content hook — only ImagesStream overrides it (to wrap in ImageArtifact), keeps base _parse_chunk universal
  • _aggregate_content stays at modality level — text joins strings, images take last chunk, audio concatenates bytes — fundamentally different operations

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Fields

    No fields configured for issues without a type.

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions