diff --git a/src/celeste/modalities/audio/providers/elevenlabs/client.py b/src/celeste/modalities/audio/providers/elevenlabs/client.py index 395d5051..4ae2c263 100644 --- a/src/celeste/modalities/audio/providers/elevenlabs/client.py +++ b/src/celeste/modalities/audio/providers/elevenlabs/client.py @@ -26,47 +26,14 @@ class ElevenLabsAudioStream(_ElevenLabsTextToSpeechStream, AudioStream): """ElevenLabs streaming for audio modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> AudioChunk | None: - """Parse binary audio chunk from stream event.""" - chunk_data = self._parse_chunk_content(event_data) - if not chunk_data: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - # Chunk with usage/finish_reason only (no audio) - return AudioChunk( - content=b"", - finish_reason=finish_reason, - usage=usage, - metadata={"event_data": event_data}, - ) - - return AudioChunk( - content=chunk_data, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - 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) - # Get mime_type from output_format parameter via client output_format = self._parameters.get("output_format") client: ElevenLabsAudioClient = self._client mime_type = client._map_output_format_to_mime_type(output_format) return AudioArtifact(data=audio_bytes, mime_type=mime_type) - def _aggregate_event_data(self, chunks: list[AudioChunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class ElevenLabsAudioClient(ElevenLabsTextToSpeechMixin, AudioClient): """ElevenLabs audio client (TTS).""" diff --git a/src/celeste/modalities/audio/providers/gradium/client.py b/src/celeste/modalities/audio/providers/gradium/client.py index d963c16c..2c5148d3 100644 --- a/src/celeste/modalities/audio/providers/gradium/client.py +++ b/src/celeste/modalities/audio/providers/gradium/client.py @@ -26,47 +26,14 @@ class GradiumAudioStream(_GradiumTextToSpeechStream, AudioStream): """Gradium streaming for audio modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> AudioChunk | None: - """Parse binary audio chunk from stream event.""" - chunk_data = self._parse_chunk_content(event_data) - if not chunk_data: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - # Chunk with usage/finish_reason only (no audio) - return AudioChunk( - content=b"", - finish_reason=finish_reason, - usage=usage, - metadata={"event_data": event_data}, - ) - - return AudioChunk( - content=chunk_data, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - 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) - # Get mime_type from output_format parameter via client output_format = self._parameters.get("output_format") client: GradiumAudioClient = self._client mime_type = client._map_output_format_to_mime_type(output_format) return AudioArtifact(data=audio_bytes, mime_type=mime_type) - def _aggregate_event_data(self, chunks: list[AudioChunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class GradiumAudioClient(GradiumTextToSpeechMixin, AudioClient): """Gradium audio client (TTS).""" diff --git a/src/celeste/modalities/audio/streaming.py b/src/celeste/modalities/audio/streaming.py index 32fe00ad..cada41fc 100644 --- a/src/celeste/modalities/audio/streaming.py +++ b/src/celeste/modalities/audio/streaming.py @@ -22,6 +22,8 @@ class AudioStream(Stream[AudioOutput, AudioParameters, AudioChunk]): _usage_class = AudioUsage _finish_reason_class = AudioFinishReason + _chunk_class = AudioChunk + _empty_content = b"" def __init__( self, @@ -39,28 +41,6 @@ def _aggregate_content(self, chunks: list[AudioChunk]) -> AudioArtifact: """Aggregate content from chunks into AudioArtifact.""" ... - def _aggregate_usage(self, chunks: list[AudioChunk]) -> AudioUsage: - """Aggregate usage across chunks.""" - for chunk in reversed(chunks): - if chunk.usage: - return chunk.usage - return AudioUsage() - - def _aggregate_finish_reason( - self, - chunks: list[AudioChunk], - ) -> AudioFinishReason | None: - """Aggregate finish reason across chunks.""" - for chunk in reversed(chunks): - if chunk.finish_reason: - return chunk.finish_reason - return None - - @abstractmethod - def _aggregate_event_data(self, chunks: list[AudioChunk]) -> list[dict[str, Any]]: - """Collect raw events.""" - ... - def _build_stream_metadata( self, raw_events: list[dict[str, Any]] ) -> dict[str, Any]: diff --git a/src/celeste/modalities/images/providers/byteplus/client.py b/src/celeste/modalities/images/providers/byteplus/client.py index df8ac543..7827381f 100644 --- a/src/celeste/modalities/images/providers/byteplus/client.py +++ b/src/celeste/modalities/images/providers/byteplus/client.py @@ -82,14 +82,11 @@ def _aggregate_usage(self, chunks: list[ImageChunk]) -> ImageUsage: return super()._aggregate_usage(chunks) def _aggregate_event_data(self, chunks: list[ImageChunk]) -> list[dict[str, Any]]: - """Collect metadata events (skip content-only events).""" + """Prepend completed_event_data, then delegate to base.""" events: list[dict[str, Any]] = [] if self._completed_event_data is not None: events.append(self._completed_event_data) - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) + events.extend(super()._aggregate_event_data(chunks)) return events diff --git a/src/celeste/modalities/images/providers/ollama/client.py b/src/celeste/modalities/images/providers/ollama/client.py index c707dff8..2d6a7ab5 100644 --- a/src/celeste/modalities/images/providers/ollama/client.py +++ b/src/celeste/modalities/images/providers/ollama/client.py @@ -32,14 +32,14 @@ def _parse_chunk(self, event_data: dict[str, Any]) -> ImageChunk | None: if not b64_image: return ImageChunk( content=ImageArtifact(data=b""), - metadata=self._parse_chunk_metadata(event_data), + metadata={"event_data": event_data}, ) return ImageChunk( content=ImageArtifact(data=b64_image), finish_reason=self._get_chunk_finish_reason(event_data), usage=self._get_chunk_usage(event_data), - metadata=self._parse_chunk_metadata(event_data), + metadata={"event_data": event_data}, ) def _aggregate_content(self, chunks: list[ImageChunk]) -> ImageArtifact: @@ -50,10 +50,6 @@ def _aggregate_content(self, chunks: list[ImageChunk]) -> ImageArtifact: msg = "No image in stream" raise ValueError(msg) - def _aggregate_event_data(self, chunks: list[ImageChunk]) -> list[dict[str, Any]]: - """Collect metadata from chunks.""" - return [chunk.metadata for chunk in chunks if chunk.metadata] - class OllamaImagesClient(OllamaGenerateClient, ImagesClient): """Ollama images client (generate only, no edit support yet).""" diff --git a/src/celeste/modalities/images/providers/openai/client.py b/src/celeste/modalities/images/providers/openai/client.py index ab1d805c..e440bf79 100644 --- a/src/celeste/modalities/images/providers/openai/client.py +++ b/src/celeste/modalities/images/providers/openai/client.py @@ -26,44 +26,10 @@ class OpenAIImagesStream(_OpenAIImagesStream, ImagesStream): """OpenAI streaming for images modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> ImageChunk | None: - """Parse one SSE event into a typed chunk.""" - b64_json = self._parse_chunk_content(event_data) - if not b64_json: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - # Chunk with usage/finish_reason only (no image) - return ImageChunk( - content=ImageArtifact(data=b""), - finish_reason=finish_reason, - usage=usage, - metadata={"event_data": event_data}, - ) - - artifact = ImageArtifact(data=b64_json) - - return ImageChunk( - content=artifact, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - def _aggregate_content(self, chunks: list[ImageChunk]) -> ImageArtifact: """Aggregate image content from chunks.""" return chunks[-1].content - def _aggregate_event_data(self, chunks: list[ImageChunk]) -> list[dict[str, Any]]: - """Collect metadata events (skip content-only events).""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class OpenAIImagesClient(OpenAIImagesMixin, ImagesClient): """OpenAI images client.""" diff --git a/src/celeste/modalities/images/streaming.py b/src/celeste/modalities/images/streaming.py index f2acebc1..5d1deccb 100644 --- a/src/celeste/modalities/images/streaming.py +++ b/src/celeste/modalities/images/streaming.py @@ -18,6 +18,8 @@ class ImagesStream(Stream[ImageOutput, ImageParameters, ImageChunk]): _usage_class = ImageUsage _finish_reason_class = ImageFinishReason + _chunk_class = ImageChunk + _empty_content = ImageArtifact(data=b"") def __init__( self, @@ -31,33 +33,15 @@ def __init__( 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) + @abstractmethod def _aggregate_content(self, chunks: list[ImageChunk]) -> ImageArtifact: """Aggregate content from chunks into raw content (modality-specific).""" ... - def _aggregate_usage(self, chunks: list[ImageChunk]) -> ImageUsage: - """Aggregate usage across chunks (universal).""" - for chunk in reversed(chunks): - if chunk.usage: - return chunk.usage - return ImageUsage() - - def _aggregate_finish_reason( - self, - chunks: list[ImageChunk], - ) -> ImageFinishReason | None: - """Aggregate finish reason across chunks (universal).""" - for chunk in reversed(chunks): - if chunk.finish_reason: - return chunk.finish_reason - return None - - @abstractmethod - def _aggregate_event_data(self, chunks: list[ImageChunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" - ... - def _build_stream_metadata( self, raw_events: list[dict[str, Any]] ) -> dict[str, Any]: diff --git a/src/celeste/modalities/text/providers/anthropic/client.py b/src/celeste/modalities/text/providers/anthropic/client.py index 13f24643..d9f46554 100644 --- a/src/celeste/modalities/text/providers/anthropic/client.py +++ b/src/celeste/modalities/text/providers/anthropic/client.py @@ -32,42 +32,21 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._message_start: dict[str, Any] | None = None def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk.""" + """Parse one SSE event into a typed chunk (captures message_start).""" event_type = event_data.get("type") if event_type == "message_start": message = event_data.get("message") if isinstance(message, dict): self._message_start = message return None - - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" - - return TextChunk( - content=content, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - - def _aggregate_content(self, chunks: list[TextChunk]) -> str: - """Aggregate streamed text content.""" - return "".join(chunk.content for chunk in chunks) + return super()._parse_chunk(event_data) def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" + """Prepend message_start, then delegate to base.""" events: list[dict[str, Any]] = [] if self._message_start is not None: events.append({"type": "message_start", "message": self._message_start}) - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) + events.extend(super()._aggregate_event_data(chunks)) return events diff --git a/src/celeste/modalities/text/providers/cohere/client.py b/src/celeste/modalities/text/providers/cohere/client.py index 6f9aba4f..08f0d459 100644 --- a/src/celeste/modalities/text/providers/cohere/client.py +++ b/src/celeste/modalities/text/providers/cohere/client.py @@ -12,7 +12,6 @@ from ...client import TextClient from ...io import ( - TextChunk, TextInput, TextOutput, ) @@ -24,36 +23,6 @@ class CohereTextStream(_CohereChatStream, TextStream): """Cohere streaming for text modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk.""" - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" - - return TextChunk( - content=content, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - - def _aggregate_content(self, chunks: list[TextChunk]) -> str: - """Aggregate streamed text content.""" - return "".join(chunk.content for chunk in chunks) - - def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class CohereTextClient(CohereChatClient, TextClient): """Cohere text client.""" diff --git a/src/celeste/modalities/text/providers/deepseek/client.py b/src/celeste/modalities/text/providers/deepseek/client.py index b308e159..38c25fac 100644 --- a/src/celeste/modalities/text/providers/deepseek/client.py +++ b/src/celeste/modalities/text/providers/deepseek/client.py @@ -11,7 +11,6 @@ from ...client import TextClient from ...io import ( - TextChunk, TextInput, TextOutput, ) @@ -23,36 +22,6 @@ class DeepSeekTextStream(_DeepSeekChatStream, TextStream): """DeepSeek streaming for text modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk.""" - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" - - return TextChunk( - content=content, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - - def _aggregate_content(self, chunks: list[TextChunk]) -> str: - """Aggregate streamed text content.""" - return "".join(chunk.content for chunk in chunks) - - def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect metadata events.""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class DeepSeekTextClient(DeepSeekChatClient, TextClient): """DeepSeek text client.""" diff --git a/src/celeste/modalities/text/providers/google/client.py b/src/celeste/modalities/text/providers/google/client.py index 3a5b917f..2ab31afa 100644 --- a/src/celeste/modalities/text/providers/google/client.py +++ b/src/celeste/modalities/text/providers/google/client.py @@ -15,7 +15,6 @@ from ...client import TextClient from ...io import ( - TextChunk, TextInput, TextOutput, ) @@ -27,36 +26,6 @@ class GoogleTextStream(_GoogleGenerateContentStream, TextStream): """Google streaming for text modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk.""" - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" - - return TextChunk( - content=content, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - - def _aggregate_content(self, chunks: list[TextChunk]) -> str: - """Aggregate streamed text content.""" - return "".join(chunk.content for chunk in chunks) - - def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect metadata events.""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class GoogleTextClient(GoogleGenerateContentClient, TextClient): """Google text client.""" diff --git a/src/celeste/modalities/text/providers/groq/client.py b/src/celeste/modalities/text/providers/groq/client.py index 1d2b5363..031ba819 100644 --- a/src/celeste/modalities/text/providers/groq/client.py +++ b/src/celeste/modalities/text/providers/groq/client.py @@ -10,7 +10,6 @@ from ...client import TextClient from ...io import ( - TextChunk, TextInput, TextOutput, ) @@ -22,36 +21,6 @@ class GroqTextStream(_GroqChatStream, TextStream): """Groq streaming for text modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk.""" - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" - - return TextChunk( - content=content, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - - def _aggregate_content(self, chunks: list[TextChunk]) -> str: - """Aggregate streamed text content.""" - return "".join(chunk.content for chunk in chunks) - - def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect metadata events.""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class GroqTextClient(GroqChatClient, TextClient): """Groq text client.""" diff --git a/src/celeste/modalities/text/providers/mistral/client.py b/src/celeste/modalities/text/providers/mistral/client.py index a502ca34..ea03a552 100644 --- a/src/celeste/modalities/text/providers/mistral/client.py +++ b/src/celeste/modalities/text/providers/mistral/client.py @@ -12,7 +12,6 @@ from ...client import TextClient from ...io import ( - TextChunk, TextInput, TextOutput, ) @@ -24,36 +23,6 @@ class MistralTextStream(_MistralChatStream, TextStream): """Mistral streaming for text modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk.""" - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" - - return TextChunk( - content=content, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - - def _aggregate_content(self, chunks: list[TextChunk]) -> str: - """Aggregate streamed text content.""" - return "".join(chunk.content for chunk in chunks) - - def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect metadata events.""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class MistralTextClient(MistralChatClient, TextClient): """Mistral text client.""" diff --git a/src/celeste/modalities/text/providers/moonshot/client.py b/src/celeste/modalities/text/providers/moonshot/client.py index 9bcde032..f9073c5f 100644 --- a/src/celeste/modalities/text/providers/moonshot/client.py +++ b/src/celeste/modalities/text/providers/moonshot/client.py @@ -12,7 +12,6 @@ from ...client import TextClient from ...io import ( - TextChunk, TextInput, TextOutput, ) @@ -24,36 +23,6 @@ class MoonshotTextStream(_MoonshotChatStream, TextStream): """Moonshot streaming for text modality.""" - def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk.""" - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" - - return TextChunk( - content=content, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - - def _aggregate_content(self, chunks: list[TextChunk]) -> str: - """Aggregate streamed text content.""" - return "".join(chunk.content for chunk in chunks) - - def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect metadata events.""" - events: list[dict[str, Any]] = [] - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class MoonshotTextClient(MoonshotChatClient, TextClient): """Moonshot text client.""" diff --git a/src/celeste/modalities/text/providers/openresponses/client.py b/src/celeste/modalities/text/providers/openresponses/client.py index 99ae5702..90b8714d 100644 --- a/src/celeste/modalities/text/providers/openresponses/client.py +++ b/src/celeste/modalities/text/providers/openresponses/client.py @@ -31,41 +31,20 @@ def __init__(self, *args: Any, **kwargs: Any) -> None: self._response_data: dict[str, Any] | None = None def _parse_chunk(self, event_data: dict[str, Any]) -> TextChunk | None: - """Parse one SSE event into a typed chunk.""" + """Parse one SSE event into a typed chunk (captures response.completed).""" event_type = event_data.get("type") if event_type == "response.completed": response = event_data.get("response") if isinstance(response, dict): self._response_data = response - - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._get_chunk_usage(event_data) - finish_reason = self._get_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" - - return TextChunk( - content=content, - finish_reason=self._get_chunk_finish_reason(event_data), - usage=self._get_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - - def _aggregate_content(self, chunks: list[TextChunk]) -> str: - """Aggregate streamed text content.""" - return "".join(chunk.content for chunk in chunks) + return super()._parse_chunk(event_data) def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" + """Prepend response_data, then delegate to base.""" events: list[dict[str, Any]] = [] if self._response_data is not None: events.append(self._response_data) - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) + events.extend(super()._aggregate_event_data(chunks)) return events diff --git a/src/celeste/modalities/text/streaming.py b/src/celeste/modalities/text/streaming.py index c0d41ae7..de24eed7 100644 --- a/src/celeste/modalities/text/streaming.py +++ b/src/celeste/modalities/text/streaming.py @@ -1,6 +1,5 @@ """Text streaming primitives.""" -from abc import abstractmethod from collections.abc import AsyncIterator, Callable from typing import Any, Unpack @@ -17,6 +16,8 @@ class TextStream(Stream[TextOutput, TextParameters, TextChunk]): _usage_class = TextUsage _finish_reason_class = TextFinishReason + _chunk_class = TextChunk + _empty_content = "" def __init__( self, @@ -30,32 +31,9 @@ def __init__( self._transform_output = transform_output self._client = client - @abstractmethod def _aggregate_content(self, chunks: list[TextChunk]) -> str: """Aggregate content from chunks into raw text.""" - ... - - def _aggregate_usage(self, chunks: list[TextChunk]) -> TextUsage: - """Aggregate usage across chunks (universal).""" - for chunk in reversed(chunks): - if chunk.usage: - return chunk.usage - return TextUsage() - - def _aggregate_finish_reason( - self, - chunks: list[TextChunk], - ) -> TextFinishReason | None: - """Aggregate finish reason across chunks (universal).""" - for chunk in reversed(chunks): - if chunk.finish_reason: - return chunk.finish_reason - return None - - @abstractmethod - def _aggregate_event_data(self, chunks: list[TextChunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" - ... + return "".join(chunk.content for chunk in chunks) def _build_stream_metadata( self, raw_events: list[dict[str, Any]] diff --git a/src/celeste/protocols/chatcompletions/streaming.py b/src/celeste/protocols/chatcompletions/streaming.py index c68bb69b..95656b74 100644 --- a/src/celeste/protocols/chatcompletions/streaming.py +++ b/src/celeste/protocols/chatcompletions/streaming.py @@ -40,7 +40,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: if not isinstance(delta, dict): return None - return delta.get("content") + return delta.get("content") or None def _parse_chunk_usage( self, event_data: dict[str, Any] diff --git a/src/celeste/protocols/openresponses/streaming.py b/src/celeste/protocols/openresponses/streaming.py index d838f58b..d44d7955 100644 --- a/src/celeste/protocols/openresponses/streaming.py +++ b/src/celeste/protocols/openresponses/streaming.py @@ -26,7 +26,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: """Extract content from SSE event.""" event_type = event_data.get("type") if event_type == "response.output_text.delta": - return event_data.get("delta") + return event_data.get("delta") or None return None def _parse_chunk_usage( diff --git a/src/celeste/providers/anthropic/messages/streaming.py b/src/celeste/providers/anthropic/messages/streaming.py index c9315386..07de3c90 100644 --- a/src/celeste/providers/anthropic/messages/streaming.py +++ b/src/celeste/providers/anthropic/messages/streaming.py @@ -28,7 +28,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: if event_type == "content_block_delta": delta = event_data.get("delta", {}) if delta.get("type") == "text_delta": - return delta.get("text") + return delta.get("text") or None return None diff --git a/src/celeste/providers/byteplus/images/streaming.py b/src/celeste/providers/byteplus/images/streaming.py index b9e2e284..0cf20b74 100644 --- a/src/celeste/providers/byteplus/images/streaming.py +++ b/src/celeste/providers/byteplus/images/streaming.py @@ -36,7 +36,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: url: str | None = event_data.get("url") if url: return url - return event_data.get("b64_json") + return event_data.get("b64_json") or None return None diff --git a/src/celeste/providers/cohere/chat/streaming.py b/src/celeste/providers/cohere/chat/streaming.py index b828eed1..a7941302 100644 --- a/src/celeste/providers/cohere/chat/streaming.py +++ b/src/celeste/providers/cohere/chat/streaming.py @@ -26,7 +26,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: delta = event_data.get("delta", {}) message = delta.get("message", {}) content = message.get("content", {}) - return content.get("text") + return content.get("text") or None return None diff --git a/src/celeste/providers/elevenlabs/text_to_speech/streaming.py b/src/celeste/providers/elevenlabs/text_to_speech/streaming.py index f94069a5..f4bbd92c 100644 --- a/src/celeste/providers/elevenlabs/text_to_speech/streaming.py +++ b/src/celeste/providers/elevenlabs/text_to_speech/streaming.py @@ -24,7 +24,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> bytes | None: Returns audio bytes if present, None otherwise. """ - return event_data.get("data") + return event_data.get("data") or None def _parse_chunk_usage( self, event_data: dict[str, Any] diff --git a/src/celeste/providers/google/generate_content/streaming.py b/src/celeste/providers/google/generate_content/streaming.py index 6e7cd096..220ce85a 100644 --- a/src/celeste/providers/google/generate_content/streaming.py +++ b/src/celeste/providers/google/generate_content/streaming.py @@ -29,7 +29,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: parts = content.get("parts", []) if parts: - return parts[0].get("text") + return parts[0].get("text") or None return None diff --git a/src/celeste/providers/google/interactions/streaming.py b/src/celeste/providers/google/interactions/streaming.py index 375665d5..582f9051 100644 --- a/src/celeste/providers/google/interactions/streaming.py +++ b/src/celeste/providers/google/interactions/streaming.py @@ -35,7 +35,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: delta_type = delta.get("type") if delta_type == "text": - return delta.get("text", "") + return delta.get("text") or None # Note: thought deltas are not returned as content (modality-specific handling) # Handle interaction complete events diff --git a/src/celeste/providers/gradium/text_to_speech/streaming.py b/src/celeste/providers/gradium/text_to_speech/streaming.py index baaf6365..f0fbdc94 100644 --- a/src/celeste/providers/gradium/text_to_speech/streaming.py +++ b/src/celeste/providers/gradium/text_to_speech/streaming.py @@ -23,7 +23,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> bytes | None: Returns audio bytes if present, None otherwise. """ - return event_data.get("data") + return event_data.get("data") or None def _parse_chunk_usage( self, event_data: dict[str, Any] diff --git a/src/celeste/providers/mistral/chat/streaming.py b/src/celeste/providers/mistral/chat/streaming.py index 6c8aac10..7b6bcdb9 100644 --- a/src/celeste/providers/mistral/chat/streaming.py +++ b/src/celeste/providers/mistral/chat/streaming.py @@ -38,7 +38,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: text_parts.append(block.get("text", "")) return "".join(text_parts) if text_parts else None - return content_delta + return content_delta or None def _parse_chunk_finish_reason( self, event_data: dict[str, Any] diff --git a/src/celeste/providers/ollama/generate/streaming.py b/src/celeste/providers/ollama/generate/streaming.py index 103e1aa9..d8d2204b 100644 --- a/src/celeste/providers/ollama/generate/streaming.py +++ b/src/celeste/providers/ollama/generate/streaming.py @@ -14,12 +14,11 @@ class OllamaGenerateStream: - _parse_chunk_content(event_data) - Extract image from NDJSON event - _parse_chunk_usage(event_data) - Extract usage from NDJSON event - _parse_chunk_finish_reason(event_data) - Extract finish reason from NDJSON event - - _parse_chunk_metadata(event_data) - Extract progress metadata from NDJSON event """ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: """Extract base64 image from NDJSON event.""" - return event_data.get("image") + return event_data.get("image") or None def _parse_chunk_usage( self, event_data: dict[str, Any] @@ -37,12 +36,5 @@ def _parse_chunk_finish_reason( return FinishReason(reason="completed") return None - def _parse_chunk_metadata(self, event_data: dict[str, Any]) -> dict[str, Any]: - """Extract progress metadata from NDJSON event.""" - return { - "completed": event_data.get("completed"), - "total": event_data.get("total"), - } - __all__ = ["OllamaGenerateStream"] diff --git a/src/celeste/providers/openai/images/streaming.py b/src/celeste/providers/openai/images/streaming.py index 334cb787..7ddb9bc1 100644 --- a/src/celeste/providers/openai/images/streaming.py +++ b/src/celeste/providers/openai/images/streaming.py @@ -38,7 +38,7 @@ def _parse_chunk_content(self, event_data: dict[str, Any]) -> str | None: "image_generation.completed", "image_edit.completed", ): - return event_data.get("b64_json") + return event_data.get("b64_json") or None return None diff --git a/src/celeste/streaming.py b/src/celeste/streaming.py index f438e80d..e258a877 100644 --- a/src/celeste/streaming.py +++ b/src/celeste/streaming.py @@ -28,6 +28,8 @@ 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]] + _empty_content: ClassVar[Any] def __init__( self, @@ -44,10 +46,31 @@ def __init__( self._portal: BlockingPortal | None = None self._portal_cm: AbstractContextManager[BlockingPortal] | None = None - @abstractmethod + def _parse_chunk_content(self, event_data: dict[str, Any]) -> Any | None: # noqa: ANN401 + """Parse content from chunk event. Override in provider mixin.""" + return None + + def _wrap_chunk_content(self, raw_content: Any) -> Any: # noqa: ANN401 + """Wrap raw content into chunk content type. Override for type transformation.""" + return raw_content + def _parse_chunk(self, event: dict[str, Any]) -> Chunk | None: """Parse SSE event into Chunk (returns None to filter lifecycle events).""" - ... + content = self._parse_chunk_content(event) + usage = self._get_chunk_usage(event) + finish_reason = self._get_chunk_finish_reason(event) + if content is None: + 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( # type: ignore[return-value] + content=content, + finish_reason=finish_reason, + usage=usage, + metadata={"event_data": event}, + ) @abstractmethod def _parse_output(self, chunks: list[Chunk], **parameters: Unpack[Params]) -> Out: # type: ignore[misc] @@ -88,6 +111,29 @@ def _build_stream_metadata( """Build metadata for streaming. Providers override to filter content.""" return {"raw_events": raw_events} + def _aggregate_usage(self, chunks: list[Chunk]) -> Usage: + """Aggregate usage across chunks (last chunk with usage wins).""" + for chunk in reversed(chunks): + if chunk.usage: + return chunk.usage + return self._usage_class() + + def _aggregate_finish_reason(self, chunks: list[Chunk]) -> FinishReason | None: + """Aggregate finish reason across chunks (last chunk with finish_reason wins).""" + for chunk in reversed(chunks): + if chunk.finish_reason: + return chunk.finish_reason + return None + + def _aggregate_event_data(self, chunks: list[Chunk]) -> list[dict[str, Any]]: + """Collect raw event_data from chunk metadata.""" + events: list[dict[str, Any]] = [] + for chunk in chunks: + event_data = chunk.metadata.get("event_data") + if isinstance(event_data, dict): + events.append(event_data) + return events + def __repr__(self) -> str: """Developer-friendly representation showing stream state.""" if self._output: diff --git a/templates/modalities/{modality_slug}/providers/{provider_slug}/client.py.template b/templates/modalities/{modality_slug}/providers/{provider_slug}/client.py.template index da143fdf..109701c2 100644 --- a/templates/modalities/{modality_slug}/providers/{provider_slug}/client.py.template +++ b/templates/modalities/{modality_slug}/providers/{provider_slug}/client.py.template @@ -25,50 +25,6 @@ from .parameters import {PROVIDER}_PARAMETER_MAPPERS class {Provider}{Modality}Stream(_{Provider}{Api}Stream, {Modality}Stream): """{Provider} streaming for {modality} modality.""" - def __init__(self, *args: Any, **kwargs: Any) -> None: # noqa: ANN401 - super().__init__(*args, **kwargs) - self._message_start: dict[str, Any] | None = None - - def _parse_chunk_usage( - self, event_data: dict[str, Any] - ) -> {Modality}Usage | None: - """Parse and wrap usage from SSE event.""" - usage = super()._parse_chunk_usage(event_data) - if usage: - return {Modality}Usage(**usage) - return None - - def _parse_chunk_finish_reason( - self, event_data: dict[str, Any] - ) -> {Modality}FinishReason | None: - """Parse and wrap finish reason from SSE event.""" - finish_reason = super()._parse_chunk_finish_reason(event_data) - if finish_reason: - return {Modality}FinishReason(reason=finish_reason.reason) - return None - - def _parse_chunk(self, event_data: dict[str, Any]) -> {Modality}Chunk | None: - """Parse one SSE event into a typed chunk.""" - # TODO: Handle provider-specific start events (e.g., message_start) - # if event_data.get("type") == "message_start": - # self._message_start = event_data.get("message") - # return None - - content = self._parse_chunk_content(event_data) - if content is None: - usage = self._parse_chunk_usage(event_data) - finish_reason = self._parse_chunk_finish_reason(event_data) - if usage is None and finish_reason is None: - return None - content = "" # Or appropriate default for modality - - return {Modality}Chunk( - content=content, - finish_reason=self._parse_chunk_finish_reason(event_data), - usage=self._parse_chunk_usage(event_data), - metadata={"event_data": event_data}, - ) - def _aggregate_content(self, chunks: list[{Modality}Chunk]) -> Any: """Aggregate streamed content (modality-specific). @@ -80,17 +36,6 @@ class {Provider}{Modality}Stream(_{Provider}{Api}Stream, {Modality}Stream): msg = "Implement content aggregation" raise NotImplementedError(msg) - def _aggregate_event_data(self, chunks: list[{Modality}Chunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" - events: list[dict[str, Any]] = [] - if self._message_start is not None: - events.append({"type": "message_start", "message": self._message_start}) - for chunk in chunks: - event_data = chunk.metadata.get("event_data") - if isinstance(event_data, dict): - events.append(event_data) - return events - class {Provider}{Modality}Client({Provider}{Api}Mixin, {Modality}Client): """{Provider} {modality} client.""" diff --git a/templates/modalities/{modality_slug}/streaming.py.template b/templates/modalities/{modality_slug}/streaming.py.template index 999cd740..f6806f4b 100644 --- a/templates/modalities/{modality_slug}/streaming.py.template +++ b/templates/modalities/{modality_slug}/streaming.py.template @@ -19,6 +19,9 @@ from .parameters import {Modality}Parameters class {Modality}Stream(Stream[{Modality}Output, {Modality}Parameters, {Modality}Chunk]): """Streaming for {modality} modality.""" + _chunk_class = {Modality}Chunk + _empty_content = "" # TODO: Set appropriate empty content for modality + def __init__( self, sse_iterator: AsyncIterator[dict[str, Any]], @@ -35,28 +38,6 @@ class {Modality}Stream(Stream[{Modality}Output, {Modality}Parameters, {Modality} """Aggregate content from chunks into raw content (modality-specific).""" ... - def _aggregate_usage(self, chunks: list[{Modality}Chunk]) -> {Modality}Usage: - """Aggregate usage across chunks (universal).""" - for chunk in reversed(chunks): - if chunk.usage: - return chunk.usage - return {Modality}Usage() - - def _aggregate_finish_reason( - self, - chunks: list[{Modality}Chunk], - ) -> {Modality}FinishReason | None: - """Aggregate finish reason across chunks (universal).""" - for chunk in reversed(chunks): - if chunk.finish_reason: - return chunk.finish_reason - return None - - @abstractmethod - def _aggregate_event_data(self, chunks: list[{Modality}Chunk]) -> list[dict[str, Any]]: - """Collect raw events (filtering happens in _build_stream_metadata).""" - ... - def _build_stream_metadata( self, raw_events: list[dict[str, Any]] ) -> dict[str, Any]: diff --git a/tests/unit_tests/test_streaming.py b/tests/unit_tests/test_streaming.py index d0ca23f7..679a085d 100644 --- a/tests/unit_tests/test_streaming.py +++ b/tests/unit_tests/test_streaming.py @@ -361,26 +361,7 @@ 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 any(method in error_msg for method in ["_parse_chunk", "_parse_output"]) - - async def test_subclass_without_parse_chunk_fails( - self, mock_sse_iterator: AsyncMock - ) -> None: - """Subclass without _parse_chunk implementation must fail instantiation.""" - - # Arrange - class IncompleteStream(Stream[ConcreteOutput, Parameters, Chunk]): - """Missing _parse_chunk implementation.""" - - def _parse_output( # type: ignore[override] - self, chunks: list[Chunk], **parameters: Unpack[Parameters] - ) -> ConcreteOutput: - """Implement _parse_output but not _parse_chunk.""" - return ConcreteOutput(content="test") - - # Act & Assert - with pytest.raises(TypeError, match=r".*_parse_chunk.*"): - IncompleteStream(mock_sse_iterator) # type: ignore[abstract] + assert "_parse_output" in error_msg async def test_subclass_without_parse_output_fails( self, mock_sse_iterator: AsyncMock