refactor(providers): standardize usage field mapping and add streaming mixins

packages/capabilities/image-generation/src/celeste_image_generation/providers/byteplus/streaming.py

@@ -1,18 +1,22 @@
 """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
 
 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,40 +25,53 @@ 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"),
                 error.get("message"),
             )
             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
 

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()
 
 

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.

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 }

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."""

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 }

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.
 

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"]