diff --git a/src/celeste/client.py b/src/celeste/client.py index f4b962d9..d4c220b8 100644 --- a/src/celeste/client.py +++ b/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, ) diff --git a/src/celeste/modalities/audio/providers/elevenlabs/client.py b/src/celeste/modalities/audio/providers/elevenlabs/client.py index 4ae2c263..50c76814 100644 --- a/src/celeste/modalities/audio/providers/elevenlabs/client.py +++ b/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) diff --git a/src/celeste/modalities/audio/providers/gradium/client.py b/src/celeste/modalities/audio/providers/gradium/client.py index 2c5148d3..029a03e4 100644 --- a/src/celeste/modalities/audio/providers/gradium/client.py +++ b/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) diff --git a/src/celeste/modalities/audio/streaming.py b/src/celeste/modalities/audio/streaming.py index cada41fc..33ec4f4d 100644 --- a/src/celeste/modalities/audio/streaming.py +++ b/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"] diff --git a/src/celeste/modalities/images/streaming.py b/src/celeste/modalities/images/streaming.py index 5d1deccb..17af9b22 100644 --- a/src/celeste/modalities/images/streaming.py +++ b/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"] diff --git a/src/celeste/modalities/text/streaming.py b/src/celeste/modalities/text/streaming.py index de24eed7..cbc59e2e 100644 --- a/src/celeste/modalities/text/streaming.py +++ b/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"] diff --git a/src/celeste/providers/elevenlabs/text_to_speech/client.py b/src/celeste/providers/elevenlabs/text_to_speech/client.py index 0322e329..294ad335 100644 --- a/src/celeste/providers/elevenlabs/text_to_speech/client.py +++ b/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. diff --git a/src/celeste/providers/gradium/text_to_speech/client.py b/src/celeste/providers/gradium/text_to_speech/client.py index 3481b4f5..a409054b 100644 --- a/src/celeste/providers/gradium/text_to_speech/client.py +++ b/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. diff --git a/src/celeste/streaming.py b/src/celeste/streaming.py index e258a877..05a79b58 100644 --- a/src/celeste/streaming.py +++ b/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).""" diff --git a/templates/modalities/{modality_slug}/streaming.py.template b/templates/modalities/{modality_slug}/streaming.py.template index f6806f4b..336d9b1f 100644 --- a/templates/modalities/{modality_slug}/streaming.py.template +++ b/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"] diff --git a/tests/unit_tests/test_stream_metadata_from_response_data.py b/tests/unit_tests/test_stream_metadata_from_response_data.py index 4e337542..8855c1bd 100644 --- a/tests/unit_tests/test_stream_metadata_from_response_data.py +++ b/tests/unit_tests/test_stream_metadata_from_response_data.py @@ -52,7 +52,11 @@ async def test_openai_stream_builds_metadata_from_inner_response_data() -> None: stream = OpenAITextStream( _async_iter(events), transform_output=lambda x, **_: x, - client=client, + stream_metadata={ + "model": client.model.id, + "provider": client.provider, + "modality": client.modality, + }, **{}, ) @@ -95,7 +99,11 @@ async def test_google_stream_builds_metadata_from_event_response_data() -> None: stream = GoogleTextStream( _async_iter([event]), transform_output=lambda x, **_: x, - client=client, + stream_metadata={ + "model": client.model.id, + "provider": client.provider, + "modality": client.modality, + }, **{}, ) @@ -138,7 +146,11 @@ async def test_openai_stream_filters_content_only_events() -> None: stream = OpenAITextStream( _async_iter(events), transform_output=lambda x, **_: x, - client=client, + stream_metadata={ + "model": client.model.id, + "provider": client.provider, + "modality": client.modality, + }, **{}, ) @@ -180,7 +192,11 @@ async def test_openai_stream_aggregates_usage_from_last_event() -> None: stream = OpenAITextStream( _async_iter(events), transform_output=lambda x, **_: x, - client=client, + stream_metadata={ + "model": client.model.id, + "provider": client.provider, + "modality": client.modality, + }, **{}, ) diff --git a/tests/unit_tests/test_streaming.py b/tests/unit_tests/test_streaming.py index 679a085d..367df7ae 100644 --- a/tests/unit_tests/test_streaming.py +++ b/tests/unit_tests/test_streaming.py @@ -40,6 +40,10 @@ def __init__( self._chunk_events = chunk_events or [] self._filter_none = filter_none + def _aggregate_content(self, chunks: list[Chunk]) -> str: + """Aggregate content from chunks.""" + return "".join(chunk.content for chunk in chunks) + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: """Parse event to Chunk or None (for lifecycle events).""" # If filter_none enabled, only return chunks for events in chunk_events @@ -361,23 +365,21 @@ async def test_cannot_instantiate_base_stream_directly( # Verify error mentions missing abstract methods error_msg = str(exc_info.value) assert "Stream" in error_msg - assert "_parse_output" in error_msg + assert "_aggregate_content" in error_msg - async def test_subclass_without_parse_output_fails( + async def test_subclass_without_aggregate_content_fails( self, mock_sse_iterator: AsyncMock ) -> None: - """Subclass without _parse_output implementation must fail instantiation.""" + """Subclass without _aggregate_content implementation must fail instantiation.""" # Arrange class IncompleteStream(Stream[ConcreteOutput, Parameters, Chunk]): - """Missing _parse_output implementation.""" + """Missing _aggregate_content implementation.""" - def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: - """Implement _parse_chunk but not _parse_output.""" - return Chunk(content="test") + pass # Act & Assert - with pytest.raises(TypeError, match=r".*_parse_output.*"): + with pytest.raises(TypeError, match=r".*_aggregate_content.*"): IncompleteStream(mock_sse_iterator) # type: ignore[abstract] @@ -626,6 +628,10 @@ class TypedOutput(Output[str]): class TypedStream(Stream[TypedOutput, Parameters, TypedChunk]): """Stream using typed classes.""" + def _aggregate_content(self, chunks: list[TypedChunk]) -> str: # type: ignore[override] + """Aggregate content from chunks.""" + return "".join(chunk.content for chunk in chunks) + def _parse_chunk(self, event: dict[str, Any]) -> TypedChunk | None: """Parse event to typed chunk.""" content = event.get("delta")