From 05eee231c91b76ab811c662197fe22b1c99bd77a Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Fri, 19 Dec 2025 12:19:46 +0100 Subject: [PATCH 1/5] refactor(providers): standardize map_usage_fields as static methods Refactor all provider clients to expose map_usage_fields() as static methods, enabling streaming mixins to share usage parsing logic with clients for consistent usage field mapping across sync and streaming. Affected providers: - Anthropic Messages - BFL Images - BytePlus Images/Videos - Cohere Chat - ElevenLabs Text-to-Speech - Google GenerateContent/Imagen/Veo - Mistral Chat - OpenAI Audio/Images/Responses/Videos - xAI Responses --- .../src/celeste_anthropic/messages/client.py | 2 +- .../bfl/src/celeste_bfl/images/client.py | 2 +- .../src/celeste_byteplus/images/client.py | 17 ++++++++++++----- .../src/celeste_byteplus/videos/client.py | 13 +++++++++---- .../cohere/src/celeste_cohere/chat/client.py | 2 +- .../text_to_speech/client.py | 13 +++++++++++-- .../celeste_google/generate_content/client.py | 2 +- .../google/src/celeste_google/imagen/client.py | 13 ++++++++++--- .../google/src/celeste_google/veo/client.py | 11 ++++++++--- .../mistral/src/celeste_mistral/chat/client.py | 2 +- .../openai/src/celeste_openai/audio/client.py | 11 +++++++++-- .../openai/src/celeste_openai/images/client.py | 18 +++++++++++++----- .../src/celeste_openai/responses/client.py | 2 +- .../openai/src/celeste_openai/videos/client.py | 14 ++++++++++---- .../xai/src/celeste_xai/responses/client.py | 2 +- 15 files changed, 89 insertions(+), 35 deletions(-) diff --git a/packages/providers/anthropic/src/celeste_anthropic/messages/client.py b/packages/providers/anthropic/src/celeste_anthropic/messages/client.py index 28199c32..8d442aba 100644 --- a/packages/providers/anthropic/src/celeste_anthropic/messages/client.py +++ b/packages/providers/anthropic/src/celeste_anthropic/messages/client.py @@ -120,7 +120,7 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: """Extract usage data from Messages API response.""" usage_data = response_data.get("usage", {}) - return self.map_usage_fields(usage_data) + return AnthropicMessagesClient.map_usage_fields(usage_data) def _parse_content(self, response_data: dict[str, Any]) -> Any: """Parse content array from Messages API. diff --git a/packages/providers/bfl/src/celeste_bfl/images/client.py b/packages/providers/bfl/src/celeste_bfl/images/client.py index d24442f7..b417d7e1 100644 --- a/packages/providers/bfl/src/celeste_bfl/images/client.py +++ b/packages/providers/bfl/src/celeste_bfl/images/client.py @@ -136,7 +136,7 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, float | None]: def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, float | None]: """Extract usage data from BFL response.""" submit_metadata = response_data.get("_submit_metadata", {}) - return self.map_usage_fields(submit_metadata) + return BFLImagesClient.map_usage_fields(submit_metadata) def _parse_content(self, response_data: dict[str, Any]) -> Any: """Parse result from response.""" diff --git a/packages/providers/byteplus/src/celeste_byteplus/images/client.py b/packages/providers/byteplus/src/celeste_byteplus/images/client.py index 18875f85..411bfea5 100644 --- a/packages/providers/byteplus/src/celeste_byteplus/images/client.py +++ b/packages/providers/byteplus/src/celeste_byteplus/images/client.py @@ -68,19 +68,26 @@ def _make_stream_request( json_body=request_body, ) - def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: - """Extract usage data from Images API response. + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map BytePlus Images usage fields to unified names. - Returns dict that capability clients wrap in their specific Usage type. + Shared by client and streaming across all capabilities. """ - usage_data = response_data.get("usage", {}) - return { UsageField.TOTAL_TOKENS: usage_data.get("total_tokens"), UsageField.OUTPUT_TOKENS: usage_data.get("output_tokens"), UsageField.NUM_IMAGES: usage_data.get("generated_images"), } + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from Images API response. + + Returns dict that capability clients wrap in their specific Usage type. + """ + usage_data = response_data.get("usage", {}) + return BytePlusImagesClient.map_usage_fields(usage_data) + def _parse_content(self, response_data: dict[str, Any]) -> Any: """Parse images/data array from Images API response. diff --git a/packages/providers/byteplus/src/celeste_byteplus/videos/client.py b/packages/providers/byteplus/src/celeste_byteplus/videos/client.py index f04fcef8..8103db5c 100644 --- a/packages/providers/byteplus/src/celeste_byteplus/videos/client.py +++ b/packages/providers/byteplus/src/celeste_byteplus/videos/client.py @@ -114,16 +114,21 @@ async def _make_request( await asyncio.sleep(config.POLLING_INTERVAL) - def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, Any]: - """Extract usage data from BytePlus API response. + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, Any]: + """Map BytePlus Videos usage fields to unified names. - Returns dict with usage fields for capability clients to wrap in Usage type. + Shared by client and streaming across all capabilities. """ - usage_data = response_data.get("usage", {}) return { UsageField.TOTAL_TOKENS: usage_data.get("total_tokens"), } + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Extract usage data from BytePlus API response.""" + usage_data = response_data.get("usage", {}) + return BytePlusVideosClient.map_usage_fields(usage_data) + def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason: """BytePlus provides status but not structured finish reasons.""" return FinishReason(reason=None) diff --git a/packages/providers/cohere/src/celeste_cohere/chat/client.py b/packages/providers/cohere/src/celeste_cohere/chat/client.py index 4387ae03..5ee9b422 100644 --- a/packages/providers/cohere/src/celeste_cohere/chat/client.py +++ b/packages/providers/cohere/src/celeste_cohere/chat/client.py @@ -88,7 +88,7 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: """Extract usage data from Chat API response.""" usage_data = response_data.get("usage", {}) - return self.map_usage_fields(usage_data) + return CohereChatClient.map_usage_fields(usage_data) def _parse_content(self, response_data: dict[str, Any]) -> Any: """Parse content array from Chat API response message. diff --git a/packages/providers/elevenlabs/src/celeste_elevenlabs/text_to_speech/client.py b/packages/providers/elevenlabs/src/celeste_elevenlabs/text_to_speech/client.py index 1aab1bc0..d3f67445 100644 --- a/packages/providers/elevenlabs/src/celeste_elevenlabs/text_to_speech/client.py +++ b/packages/providers/elevenlabs/src/celeste_elevenlabs/text_to_speech/client.py @@ -127,10 +127,19 @@ async def _stream_binary_audio( if chunk: yield {"data": chunk} - def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: - """ElevenLabs TTS doesn't return usage in response body.""" + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map ElevenLabs usage fields to unified names. + + Shared by client and streaming across all capabilities. + ElevenLabs TTS doesn't return usage in response body. + """ return {} + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from ElevenLabs TTS response.""" + return ElevenLabsTextToSpeechClient.map_usage_fields(response_data) + def _map_output_format_to_mime_type( self, output_format: str | None ) -> AudioMimeType: diff --git a/packages/providers/google/src/celeste_google/generate_content/client.py b/packages/providers/google/src/celeste_google/generate_content/client.py index eaa078b3..a5a50418 100644 --- a/packages/providers/google/src/celeste_google/generate_content/client.py +++ b/packages/providers/google/src/celeste_google/generate_content/client.py @@ -87,7 +87,7 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: """Extract usage data from Gemini usageMetadata.""" usage_metadata = response_data.get("usageMetadata", {}) - return self.map_usage_fields(usage_metadata) + return GoogleGenerateContentClient.map_usage_fields(usage_metadata) def _parse_content(self, response_data: dict[str, Any]) -> Any: """Return all candidates from response. diff --git a/packages/providers/google/src/celeste_google/imagen/client.py b/packages/providers/google/src/celeste_google/imagen/client.py index d52f4695..5bbd9474 100644 --- a/packages/providers/google/src/celeste_google/imagen/client.py +++ b/packages/providers/google/src/celeste_google/imagen/client.py @@ -49,16 +49,23 @@ async def _make_request( json_body=request_body, ) - def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: - """Extract usage data from Imagen API response. + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map Google Imagen usage fields to unified names. + Shared by client and streaming across all capabilities. Imagen API doesn't provide token usage, but we can extract num_images. """ - predictions = response_data.get("predictions", []) + predictions = usage_data.get("predictions", []) return { UsageField.NUM_IMAGES: len(predictions), } + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from Imagen API response.""" + predictions = response_data.get("predictions", []) + return GoogleImagenClient.map_usage_fields({"predictions": predictions}) + def _parse_content(self, response_data: dict[str, Any]) -> Any: """Parse predictions from response. diff --git a/packages/providers/google/src/celeste_google/veo/client.py b/packages/providers/google/src/celeste_google/veo/client.py index 68a01f84..52c6c215 100644 --- a/packages/providers/google/src/celeste_google/veo/client.py +++ b/packages/providers/google/src/celeste_google/veo/client.py @@ -110,14 +110,19 @@ def _parse_content(self, response_data: dict[str, Any]) -> Any: raise ValueError(msg) return generated_samples[0].get("video", {}) - def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, Any]: - """Extract usage data from Veo API response. + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, Any]: + """Map Google Veo usage fields to unified names. + Shared by client and streaming across all capabilities. Veo API doesn't provide usage metadata. - Returns empty dict for capability clients to wrap in Usage type. """ return {} + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Extract usage data from Veo API response.""" + return GoogleVeoClient.map_usage_fields(response_data) + async def download_content(self, url: str) -> bytes: """Download video content from GCS URL. diff --git a/packages/providers/mistral/src/celeste_mistral/chat/client.py b/packages/providers/mistral/src/celeste_mistral/chat/client.py index 3ab98329..d20eabdd 100644 --- a/packages/providers/mistral/src/celeste_mistral/chat/client.py +++ b/packages/providers/mistral/src/celeste_mistral/chat/client.py @@ -83,7 +83,7 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: """Extract usage data from response.""" usage_data = response_data.get("usage", {}) - return self.map_usage_fields(usage_data) + return MistralChatClient.map_usage_fields(usage_data) def _parse_content(self, response_data: dict[str, Any]) -> Any: """Return choices from response.""" diff --git a/packages/providers/openai/src/celeste_openai/audio/client.py b/packages/providers/openai/src/celeste_openai/audio/client.py index a7bd8afd..42ff371b 100644 --- a/packages/providers/openai/src/celeste_openai/audio/client.py +++ b/packages/providers/openai/src/celeste_openai/audio/client.py @@ -48,13 +48,20 @@ async def _make_request( json_body=request_body, ) - def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: - """Audio API speech endpoint doesn't return usage in response body. + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map OpenAI Audio usage fields to unified names. + Shared by client and streaming across all capabilities. + Audio API speech endpoint doesn't return usage in response body. Usage may be available in response headers or streaming events. """ return {} + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from Audio API response.""" + return OpenAIAudioClient.map_usage_fields(response_data) + def _map_response_format_to_mime_type( self, response_format: str | None ) -> AudioMimeType: diff --git a/packages/providers/openai/src/celeste_openai/images/client.py b/packages/providers/openai/src/celeste_openai/images/client.py index 6c176663..0660fe81 100644 --- a/packages/providers/openai/src/celeste_openai/images/client.py +++ b/packages/providers/openai/src/celeste_openai/images/client.py @@ -83,19 +83,27 @@ def _make_stream_request( json_body=request_body, ) - def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: - """Extract usage data from Images API response. + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: + """Map OpenAI Images usage fields to unified names. - Returns dict that capability clients wrap in their specific Usage type. - gpt-image-1 returns usage, DALL-E models don't. + Shared by client and streaming across all capabilities. """ - usage_data = response_data.get("usage", {}) return { UsageField.INPUT_TOKENS: usage_data.get("input_tokens"), UsageField.OUTPUT_TOKENS: usage_data.get("output_tokens"), UsageField.TOTAL_TOKENS: usage_data.get("total_tokens"), } + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: + """Extract usage data from Images API response. + + Returns dict that capability clients wrap in their specific Usage type. + gpt-image-1 returns usage, DALL-E models don't. + """ + usage_data = response_data.get("usage", {}) + return OpenAIImagesClient.map_usage_fields(usage_data) + def _parse_content(self, response_data: dict[str, Any]) -> Any: """Parse data array from Images API response. diff --git a/packages/providers/openai/src/celeste_openai/responses/client.py b/packages/providers/openai/src/celeste_openai/responses/client.py index 7a25634c..d512d6f5 100644 --- a/packages/providers/openai/src/celeste_openai/responses/client.py +++ b/packages/providers/openai/src/celeste_openai/responses/client.py @@ -93,7 +93,7 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: """Extract usage data from Responses API response.""" usage_data = response_data.get("usage", {}) - return self.map_usage_fields(usage_data) + return OpenAIResponsesClient.map_usage_fields(usage_data) def _parse_content(self, response_data: dict[str, Any]) -> Any: """Parse output array from Responses API. diff --git a/packages/providers/openai/src/celeste_openai/videos/client.py b/packages/providers/openai/src/celeste_openai/videos/client.py index 9818066a..b9a8e159 100644 --- a/packages/providers/openai/src/celeste_openai/videos/client.py +++ b/packages/providers/openai/src/celeste_openai/videos/client.py @@ -151,15 +151,21 @@ async def _prepare_multipart_request( """ return {}, {} - def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, Any]: - """Extract usage data from Videos API response. + @staticmethod + def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, Any]: + """Map OpenAI Videos usage fields to unified names. - Returns dict with seconds for capability clients to wrap in Usage type. + Shared by client and streaming across all capabilities. + Videos API uses billing units (seconds), not tokens. """ return { - UsageField.BILLED_UNITS: response_data.get("seconds"), + UsageField.BILLED_UNITS: usage_data.get("seconds"), } + def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, Any]: + """Extract usage data from Videos API response.""" + return OpenAIVideosClient.map_usage_fields(response_data) + def _parse_finish_reason(self, response_data: dict[str, Any]) -> FinishReason: """Videos API doesn't provide finish reasons.""" return FinishReason(reason=None) diff --git a/packages/providers/xai/src/celeste_xai/responses/client.py b/packages/providers/xai/src/celeste_xai/responses/client.py index 9c62913e..c3083aac 100644 --- a/packages/providers/xai/src/celeste_xai/responses/client.py +++ b/packages/providers/xai/src/celeste_xai/responses/client.py @@ -95,7 +95,7 @@ def map_usage_fields(usage_data: dict[str, Any]) -> dict[str, int | None]: def _parse_usage(self, response_data: dict[str, Any]) -> dict[str, int | None]: """Extract usage data from Responses API response.""" usage_data = response_data.get("usage", {}) - return self.map_usage_fields(usage_data) + return XAIResponsesClient.map_usage_fields(usage_data) def _parse_content(self, response_data: dict[str, Any]) -> Any: """Return output array from response.""" From 8ff8e0e375ca5d95104616828584fc1b09e59a8c Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Fri, 19 Dec 2025 12:20:10 +0100 Subject: [PATCH 2/5] feat(providers): add OpenAI Images streaming mixin Add OpenAIImagesStream mixin for SSE parsing in OpenAI Images API: - Handle image_generation.partial_image and image_generation.completed - Handle image_edit.partial_image and image_edit.completed - Extract content, usage, and metadata from SSE events - Use static map_usage_fields for consistent usage parsing --- .../src/celeste_openai/images/streaming.py | 89 +++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 packages/providers/openai/src/celeste_openai/images/streaming.py diff --git a/packages/providers/openai/src/celeste_openai/images/streaming.py b/packages/providers/openai/src/celeste_openai/images/streaming.py new file mode 100644 index 00000000..c13b46e2 --- /dev/null +++ b/packages/providers/openai/src/celeste_openai/images/streaming.py @@ -0,0 +1,89 @@ +"""OpenAI Images SSE parsing for streaming.""" + +from typing import Any + +from .client import OpenAIImagesClient + + +class OpenAIImagesStream: + """Mixin for Images API SSE parsing. + + Provides shared implementation for capabilities using OpenAI Images API streaming: + - _parse_chunk() - Parse SSE event into raw chunk dict + + Handles all image streaming event types: + - image_generation.partial_image / image_generation.completed + - image_edit.partial_image / image_edit.completed + + Capability streams extend via super() to wrap results in typed Chunks. + + Usage: + class OpenAIImageGenerationStream(OpenAIImagesStream, ImageGenerationStream): + def _parse_chunk(self, event): + raw = super()._parse_chunk(event) + if not raw: + return None + return ImageGenerationChunk(...) + """ + + def _parse_chunk(self, event: dict[str, Any]) -> dict[str, Any] | None: + """Parse SSE event into raw chunk data. + + Returns dict with: + - content: b64_json string (raw, not decoded) + - is_partial: True for partial_image events, False for completed + - usage: usage dict from completed events (None for partials) + - metadata: size, quality, output_format, background, created_at, partial_image_index + - raw_event: original event dict + """ + event_type = event.get("type") + if not event_type: + return None + + # Handle partial image events (generation or edit) + if event_type in ("image_generation.partial_image", "image_edit.partial_image"): + b64_json = event.get("b64_json") + if not b64_json: + return None + return { + "content": b64_json, + "is_partial": True, + "usage": None, + "metadata": { + "size": event.get("size"), + "quality": event.get("quality"), + "output_format": event.get("output_format"), + "background": event.get("background"), + "created_at": event.get("created_at"), + "partial_image_index": event.get("partial_image_index"), + }, + "raw_event": event, + } + + # Handle completed events (generation or edit) + if event_type in ("image_generation.completed", "image_edit.completed"): + b64_json = event.get("b64_json") + if not b64_json: + return None + usage_data = event.get("usage") + usage = None + if usage_data: + usage = OpenAIImagesClient.map_usage_fields(usage_data) + return { + "content": b64_json, + "is_partial": False, + "usage": usage, + "metadata": { + "size": event.get("size"), + "quality": event.get("quality"), + "output_format": event.get("output_format"), + "background": event.get("background"), + "created_at": event.get("created_at"), + }, + "raw_event": event, + } + + return None + + +__all__ = ["OpenAIImagesStream"] From da7f9045a4f5a677f15349ce21d6f9ef9eaacb14 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Fri, 19 Dec 2025 12:21:21 +0100 Subject: [PATCH 3/5] feat(providers): add BytePlus Images streaming mixin Add BytePlusImagesStream mixin for SSE parsing in BytePlus Images API: - Handle image_generation.partial_succeeded with url or b64_json - Handle image_generation.partial_failed error events - Handle image_generation.completed with usage data - Use static map_usage_fields for consistent usage parsing Also fix missing dependencies in BytePlus pyproject.toml. --- packages/providers/byteplus/pyproject.toml | 1 + .../src/celeste_byteplus/images/streaming.py | 116 ++++++++++++++++++ 2 files changed, 117 insertions(+) create mode 100644 packages/providers/byteplus/src/celeste_byteplus/images/streaming.py diff --git a/packages/providers/byteplus/pyproject.toml b/packages/providers/byteplus/pyproject.toml index d23ac2fd..41c49c2b 100644 --- a/packages/providers/byteplus/pyproject.toml +++ b/packages/providers/byteplus/pyproject.toml @@ -5,6 +5,7 @@ description = "BytePlus provider package for Celeste AI" authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] license = {text = "Apache-2.0"} requires-python = ">=3.12" + dependencies = ["celeste-ai", "httpx"] [tool.uv.sources] celeste-ai = { workspace = true } diff --git a/packages/providers/byteplus/src/celeste_byteplus/images/streaming.py b/packages/providers/byteplus/src/celeste_byteplus/images/streaming.py new file mode 100644 index 00000000..799b3b3b --- /dev/null +++ b/packages/providers/byteplus/src/celeste_byteplus/images/streaming.py @@ -0,0 +1,116 @@ +"""BytePlus Images SSE parsing for streaming.""" + +from typing import Any + +from .client import BytePlusImagesClient + + +class BytePlusImagesStream: + """Mixin for BytePlus Images API SSE parsing. + + Provides shared implementation for capabilities using BytePlus Images API streaming: + - _parse_chunk() - Parse SSE event into raw chunk dict + + Handles all image streaming event types: + - image_generation.partial_succeeded - Partial image with url or b64_json + - image_generation.partial_failed - Error event + - image_generation.completed - Final event with usage only + + Capability streams extend via super() to wrap results in typed Chunks. + + Usage: + class BytePlusImageGenerationStream(BytePlusImagesStream, ImageGenerationStream): + def _parse_chunk(self, event): + raw = super()._parse_chunk(event) + if not raw: + return None + return ImageGenerationChunk(...) + """ + + def _parse_chunk(self, event: dict[str, Any]) -> dict[str, Any] | None: + """Parse SSE event into raw chunk data. + + Returns dict with: + - content_type: "url" or "b64_json" or None + - content: url string or b64_json string + - is_error: True for partial_failed events + - error: error dict for failed events + - usage: usage dict from completed event (None otherwise) + - metadata: model, created, image_index, size + - raw_event: original event dict + """ + event_type = event.get("type") + if not event_type: + return None + + # Handle successful partial image + if event_type == "image_generation.partial_succeeded": + url = event.get("url") + b64_json = event.get("b64_json") + + content_type = None + content = None + if url: + content_type = "url" + content = url + elif b64_json: + content_type = "b64_json" + content = b64_json + + if not content: + return None + + return { + "content_type": content_type, + "content": content, + "is_error": False, + "error": None, + "usage": None, + "metadata": { + "model": event.get("model"), + "created": event.get("created"), + "image_index": event.get("image_index"), + "size": event.get("size"), + }, + "raw_event": event, + } + + # Handle failed partial image + if event_type == "image_generation.partial_failed": + return { + "content_type": None, + "content": None, + "is_error": True, + "error": event.get("error"), + "usage": None, + "metadata": { + "model": event.get("model"), + "created": event.get("created"), + "image_index": event.get("image_index"), + }, + "raw_event": event, + } + + # Handle completed event (usage only, no image) + if event_type == "image_generation.completed": + usage_data = event.get("usage") + usage = None + if usage_data: + usage = BytePlusImagesClient.map_usage_fields(usage_data) + return { + "content_type": None, + "content": None, + "is_error": False, + "error": None, + "usage": usage, + "metadata": { + "model": event.get("model"), + "created": event.get("created"), + }, + "raw_event": event, + } + + return None + + +__all__ = ["BytePlusImagesStream"] From be82663b856c5a5a2e0b2ab0085fbb2c2b6f1504 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Fri, 19 Dec 2025 12:22:02 +0100 Subject: [PATCH 4/5] refactor(image-generation): migrate streaming to use provider mixins Update capability streaming implementations to use new provider mixins: - OpenAIImageGenerationStream now inherits from OpenAIImagesStream - BytePlusImageGenerationStream now inherits from BytePlusImagesStream Both now use super()._parse_chunk() to get raw parsed data, then wrap in typed ImageGenerationChunk with proper ImageArtifact handling. --- .../providers/byteplus/streaming.py | 65 ++++++++++++------- .../providers/openai/streaming.py | 61 +++++++---------- 2 files changed, 64 insertions(+), 62 deletions(-) diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/byteplus/streaming.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/byteplus/streaming.py index 3c06cea3..96575547 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/byteplus/streaming.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/byteplus/streaming.py @@ -1,10 +1,14 @@ """BytePlus streaming for image generation.""" +import base64 import logging from collections.abc import AsyncIterator from typing import Any +from celeste_byteplus.images.streaming import BytePlusImagesStream + from celeste.artifacts import ImageArtifact +from celeste.core import UsageField from celeste.mime_types import ImageMimeType from celeste_image_generation.io import ImageGenerationChunk, ImageGenerationUsage from celeste_image_generation.streaming import ImageGenerationStream @@ -12,7 +16,7 @@ logger = logging.getLogger(__name__) -class BytePlusImageGenerationStream(ImageGenerationStream): +class BytePlusImageGenerationStream(BytePlusImagesStream, ImageGenerationStream): """BytePlus streaming for image generation.""" def __init__(self, sse_iterator: AsyncIterator[dict[str, Any]]) -> None: @@ -21,28 +25,17 @@ def __init__(self, sse_iterator: AsyncIterator[dict[str, Any]]) -> None: self._completed_usage: ImageGenerationUsage | None = None def _parse_chunk(self, chunk_data: dict[str, Any]) -> ImageGenerationChunk | None: - """Parse chunk from SSE event.""" - event_type = chunk_data.get("type") - - if event_type == "image_generation.partial_succeeded": - url = chunk_data.get("url") - if not url: - logger.warning("partial_succeeded event missing URL") - return None - - artifact = ImageArtifact(url=url, mime_type=ImageMimeType.PNG) - return ImageGenerationChunk(content=artifact) - - if event_type == "image_generation.completed": - usage_data = chunk_data.get("usage") - if usage_data: - self._completed_usage = ImageGenerationUsage( - total_tokens=usage_data.get("total_tokens"), - ) + """Parse chunk from SSE event. + + Uses provider mixin to parse raw SSE event, then wraps in typed chunk. + """ + raw = super()._parse_chunk(chunk_data) + if not raw: return None - if event_type == "image_generation.partial_failed": - error = chunk_data.get("error", {}) + # Handle error events + if raw.get("is_error"): + error = raw.get("error", {}) logger.error( "Image generation failed: %s - %s", error.get("code"), @@ -50,11 +43,35 @@ def _parse_chunk(self, chunk_data: dict[str, Any]) -> ImageGenerationChunk | Non ) return None - logger.warning("Unknown event type: %s", event_type) - return None + # Handle completed event (usage only) + usage_data = raw.get("usage") + if usage_data: + self._completed_usage = ImageGenerationUsage( + total_tokens=usage_data.get(UsageField.TOTAL_TOKENS), + output_tokens=usage_data.get(UsageField.OUTPUT_TOKENS), + num_images=usage_data.get(UsageField.NUM_IMAGES), + ) + return None + + # Handle partial succeeded (image content) + content = raw.get("content") + content_type = raw.get("content_type") + if not content: + return None + + if content_type == "url": + artifact = ImageArtifact(url=content, mime_type=ImageMimeType.PNG) + else: # b64_json + image_data = base64.b64decode(content) + artifact = ImageArtifact(data=image_data) + + return ImageGenerationChunk(content=artifact) def _parse_usage(self, chunks: list[ImageGenerationChunk]) -> ImageGenerationUsage: - """Parse usage from chunks.""" + """Parse usage from chunks. + + Usage is stored from the completed event. + """ if self._completed_usage is not None: return self._completed_usage diff --git a/packages/capabilities/image-generation/src/celeste_image_generation/providers/openai/streaming.py b/packages/capabilities/image-generation/src/celeste_image_generation/providers/openai/streaming.py index b0a5794c..cbf01fa9 100644 --- a/packages/capabilities/image-generation/src/celeste_image_generation/providers/openai/streaming.py +++ b/packages/capabilities/image-generation/src/celeste_image_generation/providers/openai/streaming.py @@ -4,71 +4,56 @@ import logging from typing import Any +from celeste_openai.images.streaming import OpenAIImagesStream + from celeste.artifacts import ImageArtifact +from celeste.core import UsageField from celeste_image_generation.io import ImageGenerationChunk, ImageGenerationUsage from celeste_image_generation.streaming import ImageGenerationStream logger = logging.getLogger(__name__) -class OpenAIImageGenerationStream(ImageGenerationStream): +class OpenAIImageGenerationStream(OpenAIImagesStream, ImageGenerationStream): """OpenAI streaming for image generation.""" def _parse_chunk(self, chunk_data: dict[str, Any]) -> ImageGenerationChunk | None: """Parse chunk from SSE event. - OpenAI returns two event types: - - image_generation.partial_image: Progressive image chunks - - image_generation.completed: Final image with usage data + Uses provider mixin to parse raw SSE event, then wraps in typed chunk. """ - event_type = chunk_data.get("type") - - if event_type == "image_generation.partial_image": - # Partial image chunk - b64_json = chunk_data.get("b64_json") - if not b64_json: - return None - - image_data = base64.b64decode(b64_json) - artifact = ImageArtifact(data=image_data) - - return ImageGenerationChunk(content=artifact) - - if event_type == "image_generation.completed": - # Final image with usage - b64_json = chunk_data.get("b64_json") - if not b64_json: - return None + raw = super()._parse_chunk(chunk_data) + if not raw: + return None - image_data = base64.b64decode(b64_json) - artifact = ImageArtifact(data=image_data) + b64_json = raw.get("content") + if not b64_json: + return None - # Parse usage from completed event - usage_data = chunk_data.get("usage") - usage = None - if usage_data: - usage = ImageGenerationUsage( - total_tokens=usage_data.get("total_tokens"), - input_tokens=usage_data.get("input_tokens"), - output_tokens=usage_data.get("output_tokens"), - ) + image_data = base64.b64decode(b64_json) + artifact = ImageArtifact(data=image_data) - return ImageGenerationChunk(content=artifact, usage=usage) + # Parse usage from raw dict (already mapped to UsageField keys) + usage = None + usage_data = raw.get("usage") + if usage_data: + usage = ImageGenerationUsage( + total_tokens=usage_data.get(UsageField.TOTAL_TOKENS), + input_tokens=usage_data.get(UsageField.INPUT_TOKENS), + output_tokens=usage_data.get(UsageField.OUTPUT_TOKENS), + ) - logger.warning("Unknown event type: %s", event_type) - return None + return ImageGenerationChunk(content=artifact, usage=usage) def _parse_usage(self, chunks: list[ImageGenerationChunk]) -> ImageGenerationUsage: """Parse usage from chunks. Usage is only available in the final completed event. """ - # Look for usage in final chunk (completed event) for chunk in reversed(chunks): if chunk.usage is not None: return chunk.usage - # No usage found return ImageGenerationUsage() From ee8028e924d15a9579b192f0b21606c45033c443 Mon Sep 17 00:00:00 2001 From: kamilbenkirane Date: Fri, 19 Dec 2025 12:27:06 +0100 Subject: [PATCH 5/5] fix(bfl): add missing dependencies to pyproject.toml --- packages/providers/bfl/pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/packages/providers/bfl/pyproject.toml b/packages/providers/bfl/pyproject.toml index dbfe2280..2ec55f0f 100644 --- a/packages/providers/bfl/pyproject.toml +++ b/packages/providers/bfl/pyproject.toml @@ -5,6 +5,7 @@ description = "BFL (Black Forest Labs) provider package for Celeste AI" authors = [{name = "Kamilbenkirane", email = "kamil@withceleste.ai"}] license = {text = "Apache-2.0"} requires-python = ">=3.12" +dependencies = ["celeste-ai", "httpx"] [tool.uv.sources] celeste-ai = { workspace = true }