refactor: consolidate __init__, _build_stream_metadata, _parse_output into base Stream

src/celeste/client.py

@@ -225,7 +225,11 @@ def _stream(
         return stream_class(
             sse_iterator,
             transform_output=self._transform_output,
-            client=self,
+            stream_metadata={
+                "model": self.model.id,
+                "provider": self.provider,
+                "modality": self.modality,
+            },
             **parameters,
         )
 

src/celeste/modalities/audio/providers/elevenlabs/client.py

@@ -30,8 +30,9 @@ def _aggregate_content(self, chunks: list[AudioChunk]) -> AudioArtifact:
         """Aggregate audio content from chunks into AudioArtifact."""
         audio_bytes = b"".join(chunk.content for chunk in chunks if chunk.content)
         output_format = self._parameters.get("output_format")
-        client: ElevenLabsAudioClient = self._client
-        mime_type = client._map_output_format_to_mime_type(output_format)
+        mime_type = ElevenLabsTextToSpeechMixin._map_output_format_to_mime_type(
+            output_format
+        )
         return AudioArtifact(data=audio_bytes, mime_type=mime_type)
 
 

src/celeste/modalities/audio/providers/gradium/client.py

@@ -30,8 +30,9 @@ def _aggregate_content(self, chunks: list[AudioChunk]) -> AudioArtifact:
         """Aggregate audio content from chunks into AudioArtifact."""
         audio_bytes = b"".join(chunk.content for chunk in chunks if chunk.content)
         output_format = self._parameters.get("output_format")
-        client: GradiumAudioClient = self._client
-        mime_type = client._map_output_format_to_mime_type(output_format)
+        mime_type = GradiumTextToSpeechMixin._map_output_format_to_mime_type(
+            output_format
+        )
         return AudioArtifact(data=audio_bytes, mime_type=mime_type)
 
 

src/celeste/modalities/audio/streaming.py

@@ -1,11 +1,8 @@
 """Audio streaming primitives."""
 
 from abc import abstractmethod
-from collections.abc import AsyncIterator, Callable
-from typing import Any, Unpack
 
 from celeste.artifacts import AudioArtifact
-from celeste.client import ModalityClient
 from celeste.streaming import Stream
 
 from .io import (
@@ -23,50 +20,13 @@ class AudioStream(Stream[AudioOutput, AudioParameters, AudioChunk]):
     _usage_class = AudioUsage
     _finish_reason_class = AudioFinishReason
     _chunk_class = AudioChunk
+    _output_class = AudioOutput
     _empty_content = b""
 
-    def __init__(
-        self,
-        sse_iterator: AsyncIterator[dict[str, Any]],
-        transform_output: Callable[..., AudioArtifact],
-        client: ModalityClient,
-        **parameters: Unpack[AudioParameters],
-    ) -> None:
-        super().__init__(sse_iterator, **parameters)
-        self._transform_output = transform_output
-        self._client = client
-
     @abstractmethod
     def _aggregate_content(self, chunks: list[AudioChunk]) -> AudioArtifact:
         """Aggregate content from chunks into AudioArtifact."""
         ...
 
-    def _build_stream_metadata(
-        self, raw_events: list[dict[str, Any]]
-    ) -> dict[str, Any]:
-        """Build streaming metadata."""
-        return {
-            "model": self._client.model.id,
-            "provider": self._client.provider,
-            "modality": self._client.modality,
-            "raw_events": raw_events,
-        }
-
-    def _parse_output(
-        self,
-        chunks: list[AudioChunk],
-        **parameters: Unpack[AudioParameters],
-    ) -> AudioOutput:
-        """Assemble chunks into final output."""
-        raw_content = self._aggregate_content(chunks)
-        content: AudioArtifact = self._transform_output(raw_content, **parameters)
-        raw_events = self._aggregate_event_data(chunks)
-        return AudioOutput(
-            content=content,
-            usage=self._aggregate_usage(chunks),
-            finish_reason=self._aggregate_finish_reason(chunks),
-            metadata=self._build_stream_metadata(raw_events),
-        )
-
 
 __all__ = ["AudioStream"]

src/celeste/modalities/images/streaming.py

@@ -1,13 +1,10 @@
 """Images streaming primitives."""
 
 from abc import abstractmethod
-from collections.abc import AsyncIterator, Callable
-from typing import Any, Unpack
+from typing import Any
 
 from celeste.artifacts import ImageArtifact
-from celeste.client import ModalityClient
 from celeste.streaming import Stream
-from celeste.types import ImageContent
 
 from .io import ImageChunk, ImageFinishReason, ImageOutput, ImageUsage
 from .parameters import ImageParameters
@@ -19,20 +16,9 @@ class ImagesStream(Stream[ImageOutput, ImageParameters, ImageChunk]):
     _usage_class = ImageUsage
     _finish_reason_class = ImageFinishReason
     _chunk_class = ImageChunk
+    _output_class = ImageOutput
     _empty_content = ImageArtifact(data=b"")
 
-    def __init__(
-        self,
-        sse_iterator: AsyncIterator[dict[str, Any]],
-        transform_output: Callable[..., ImageContent],
-        client: ModalityClient,
-        **parameters: Unpack[ImageParameters],
-    ) -> None:
-        """Initialize stream with output transformation support."""
-        super().__init__(sse_iterator, **parameters)
-        self._transform_output = transform_output
-        self._client = client
-
     def _wrap_chunk_content(self, raw_content: Any) -> ImageArtifact:  # noqa: ANN401
         """Wrap raw content (base64 string) into ImageArtifact."""
         return ImageArtifact(data=raw_content)
@@ -42,36 +28,5 @@ def _aggregate_content(self, chunks: list[ImageChunk]) -> ImageArtifact:
         """Aggregate content from chunks into raw content (modality-specific)."""
         ...
 
-    def _build_stream_metadata(
-        self, raw_events: list[dict[str, Any]]
-    ) -> dict[str, Any]:
-        """Build streaming metadata. Provider API Stream overrides to filter content."""
-        return {
-            "model": self._client.model.id,
-            "provider": self._client.provider,
-            "modality": self._client.modality,
-            "raw_events": raw_events,
-        }
-
-    def _parse_output(
-        self,
-        chunks: list[ImageChunk],
-        **parameters: Unpack[ImageParameters],
-    ) -> ImageOutput:
-        """Assemble chunks into final output."""
-        if not chunks:
-            msg = "No chunks received from stream"
-            raise ValueError(msg)
-
-        raw_content = self._aggregate_content(chunks)
-        content: ImageContent = self._transform_output(raw_content, **parameters)
-        raw_events = self._aggregate_event_data(chunks)
-        return ImageOutput(
-            content=content,
-            usage=self._aggregate_usage(chunks),
-            finish_reason=self._aggregate_finish_reason(chunks),
-            metadata=self._build_stream_metadata(raw_events),
-        )
-
 
 __all__ = ["ImagesStream"]

src/celeste/modalities/text/streaming.py

@@ -1,11 +1,6 @@
 """Text streaming primitives."""
 
-from collections.abc import AsyncIterator, Callable
-from typing import Any, Unpack
-
-from celeste.client import ModalityClient
 from celeste.streaming import Stream
-from celeste.types import TextContent
 
 from .io import TextChunk, TextFinishReason, TextOutput, TextUsage
 from .parameters import TextParameters
@@ -17,50 +12,12 @@ class TextStream(Stream[TextOutput, TextParameters, TextChunk]):
     _usage_class = TextUsage
     _finish_reason_class = TextFinishReason
     _chunk_class = TextChunk
+    _output_class = TextOutput
     _empty_content = ""
 
-    def __init__(
-        self,
-        sse_iterator: AsyncIterator[dict[str, Any]],
-        transform_output: Callable[..., TextContent],
-        client: ModalityClient,
-        **parameters: Unpack[TextParameters],
-    ) -> None:
-        """Initialize stream with output transformation support."""
-        super().__init__(sse_iterator, **parameters)
-        self._transform_output = transform_output
-        self._client = client
-
     def _aggregate_content(self, chunks: list[TextChunk]) -> str:
         """Aggregate content from chunks into raw text."""
         return "".join(chunk.content for chunk in chunks)
 
-    def _build_stream_metadata(
-        self, raw_events: list[dict[str, Any]]
-    ) -> dict[str, Any]:
-        """Build streaming metadata. Provider API Stream overrides to filter content."""
-        return {
-            "model": self._client.model.id,
-            "provider": self._client.provider,
-            "modality": self._client.modality,
-            "raw_events": raw_events,
-        }
-
-    def _parse_output(
-        self,
-        chunks: list[TextChunk],
-        **parameters: Unpack[TextParameters],
-    ) -> TextOutput:
-        """Assemble chunks into final output."""
-        raw_content = self._aggregate_content(chunks)
-        content: TextContent = self._transform_output(raw_content, **parameters)
-        raw_events = self._aggregate_event_data(chunks)
-        return TextOutput(
-            content=content,
-            usage=self._aggregate_usage(chunks),
-            finish_reason=self._aggregate_finish_reason(chunks),
-            metadata=self._build_stream_metadata(raw_events),
-        )
-
 
 __all__ = ["TextStream"]

src/celeste/providers/elevenlabs/text_to_speech/client.py

@@ -159,8 +159,9 @@ def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason:
         """ElevenLabs TTS doesn't provide finish reasons."""
         return FinishReason(reason=None)
 
+    @staticmethod
     def _map_output_format_to_mime_type(
-        self, output_format: str | None
+        output_format: str | None,
     ) -> AudioMimeType:
         """Map ElevenLabs output_format to AudioMimeType.
 

src/celeste/providers/gradium/text_to_speech/client.py

@@ -163,8 +163,8 @@ async def _make_request(
             "output_format": output_format,
         }
 
+    @staticmethod
     def _map_output_format_to_mime_type(
-        self,
         output_format: str | None,
     ) -> AudioMimeType:
         """Map Gradium output_format to AudioMimeType.

src/celeste/streaming.py

@@ -1,7 +1,7 @@
 """Streaming support for Celeste."""
 
 from abc import ABC, abstractmethod
-from collections.abc import AsyncIterator
+from collections.abc import AsyncIterator, Callable
 from contextlib import AbstractContextManager, suppress
 from types import TracebackType
 from typing import Any, ClassVar, Self, Unpack
@@ -29,11 +29,14 @@ class Stream[Out: Output, Params: Parameters, Chunk: ChunkBase](ABC):
     _usage_class: ClassVar[type[Usage]] = Usage
     _finish_reason_class: ClassVar[type[FinishReason]] = FinishReason
     _chunk_class: ClassVar[type[ChunkBase]]
+    _output_class: ClassVar[type[Output]]
     _empty_content: ClassVar[Any]
 
     def __init__(
         self,
         sse_iterator: AsyncIterator[dict[str, Any]],
+        transform_output: Callable[..., Any] | None = None,
+        stream_metadata: dict[str, Any] | None = None,
         **parameters: Unpack[Params],  # type: ignore[misc]
     ) -> None:
         """Initialize stream with SSE iterator."""
@@ -42,6 +45,8 @@ def __init__(
         self._closed = False
         self._output: Out | None = None
         self._parameters = parameters
+        self._transform_output = transform_output
+        self._stream_metadata = stream_metadata or {}
         # Sync iteration state
         self._portal: BlockingPortal | None = None
         self._portal_cm: AbstractContextManager[BlockingPortal] | None = None
@@ -73,9 +78,25 @@ def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None:
         )
 
     @abstractmethod
+    def _aggregate_content(self, chunks: list[Chunk]) -> Any:  # noqa: ANN401
+        """Aggregate content from chunks into raw content (modality-specific)."""
+        ...
+
     def _parse_output(self, chunks: list[Chunk], **parameters: Unpack[Params]) -> Out:  # type: ignore[misc]
         """Parse final Output from accumulated chunks."""
-        ...
+        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(  # type: ignore[return-value]
+            content=content,
+            usage=self._aggregate_usage(chunks),
+            finish_reason=self._aggregate_finish_reason(chunks),
+            metadata=self._build_stream_metadata(raw_events),
+        )
 
     def _parse_chunk_usage(self, event_data: dict[str, Any]) -> RawUsage | None:
         """Parse usage from chunk event. Override in provider mixin."""
@@ -109,7 +130,7 @@ def _build_stream_metadata(
         self, raw_events: list[dict[str, Any]]
     ) -> dict[str, Any]:
         """Build metadata for streaming. Providers override to filter content."""
-        return {"raw_events": raw_events}
+        return {**self._stream_metadata, "raw_events": raw_events}
 
     def _aggregate_usage(self, chunks: list[Chunk]) -> Usage:
         """Aggregate usage across chunks (last chunk with usage wins)."""

templates/modalities/{modality_slug}/streaming.py.template

@@ -1,10 +1,8 @@
 """{Modality} streaming primitives."""
 
 from abc import abstractmethod
-from collections.abc import AsyncIterator, Callable
-from typing import Any, Unpack
+from typing import Any
 
-from celeste.client import ModalityClient
 from celeste.streaming import Stream
 
 from .io import (
@@ -19,51 +17,16 @@ from .parameters import {Modality}Parameters
 class {Modality}Stream(Stream[{Modality}Output, {Modality}Parameters, {Modality}Chunk]):
     """Streaming for {modality} modality."""
 
+    _usage_class = {Modality}Usage
+    _finish_reason_class = {Modality}FinishReason
     _chunk_class = {Modality}Chunk
+    _output_class = {Modality}Output
     _empty_content = ""  # TODO: Set appropriate empty content for modality
 
-    def __init__(
-        self,
-        sse_iterator: AsyncIterator[dict[str, Any]],
-        transform_output: Callable[..., {Content}],
-        client: ModalityClient,
-        **parameters: Unpack[{Modality}Parameters],
-    ) -> None:
-        super().__init__(sse_iterator, **parameters)
-        self._transform_output = transform_output
-        self._client = client
-
     @abstractmethod
     def _aggregate_content(self, chunks: list[{Modality}Chunk]) -> Any:
         """Aggregate content from chunks into raw content (modality-specific)."""
         ...
 
-    def _build_stream_metadata(
-        self, raw_events: list[dict[str, Any]]
-    ) -> dict[str, Any]:
-        """Build streaming metadata. Provider API Stream overrides to filter content."""
-        return {
-            "model": self._client.model.id,
-            "provider": self._client.provider,
-            "modality": self._client.modality,
-            "raw_events": raw_events,
-        }
-
-    def _parse_output(  # type: ignore[override]
-        self,
-        chunks: list[{Modality}Chunk],
-        **parameters: Unpack[{Modality}Parameters],
-    ) -> {Modality}Output:
-        """Assemble chunks into final output."""
-        raw_content = self._aggregate_content(chunks)
-        content: {Content} = self._transform_output(raw_content, **parameters)
-        raw_events = self._aggregate_event_data(chunks)
-        return {Modality}Output(
-            content=content,
-            usage=self._aggregate_usage(chunks),
-            finish_reason=self._aggregate_finish_reason(chunks),
-            metadata=self._build_stream_metadata(raw_events),
-        )
-
 
 __all__ = ["{Modality}Stream"]