Skip to content

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

Merged
Kamilbenkirane merged 5 commits into
mainfrom
refactor/streaming-mixins-and-usage-mapping
Dec 19, 2025
Merged

refactor(providers): standardize usage field mapping and add streaming mixins#86
Kamilbenkirane merged 5 commits into
mainfrom
refactor/streaming-mixins-and-usage-mapping

Conversation

@Kamilbenkirane

Copy link
Copy Markdown
Member

Summary

  • Standardize map_usage_fields() as static methods across all provider clients
  • Add OpenAI Images streaming mixin (OpenAIImagesStream)
  • Add BytePlus Images streaming mixin (BytePlusImagesStream)
  • Migrate image-generation capability streaming to use provider mixins
  • Fix BytePlus pyproject.toml missing dependencies

Changes

Provider Refactoring

All provider clients now expose map_usage_fields() as a static method, enabling:

  • Streaming mixins to share usage parsing logic with clients
  • Consistent usage field mapping across sync and streaming paths

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

New Streaming Mixins

  • OpenAIImagesStream - Handles image_generation.partial_image, image_generation.completed, image_edit.* events
  • BytePlusImagesStream - Handles image_generation.partial_succeeded, image_generation.partial_failed, image_generation.completed events

Capability Migrations

  • OpenAIImageGenerationStream now inherits from OpenAIImagesStream
  • BytePlusImageGenerationStream now inherits from BytePlusImagesStream

Test plan

  • Run make ci

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
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
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.
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.
@claude

claude Bot commented Dec 19, 2025

Copy link
Copy Markdown

PR Review: Standardize usage field mapping and add streaming mixins

Summary

This PR introduces a solid architectural improvement by extracting provider-specific SSE parsing logic into reusable mixins and standardizing usage field mapping across all provider clients. The refactoring follows good separation of concerns and DRY principles.

✅ Strengths

1. Excellent Architecture Pattern

The new mixin approach is well-designed:

  • Provider mixins (OpenAIImagesStream, BytePlusImagesStream) handle raw SSE parsing
  • Capability streams extend mixins via super() and wrap results in typed chunks
  • Clear separation between provider-specific and capability-specific logic
  • Good use of multiple inheritance that doesn't create diamond problems

2. Consistent Usage Mapping

Converting map_usage_fields() to static methods across all providers:

  • Enables code sharing between sync clients and streaming implementations
  • Consistent pattern applied across 14+ provider files
  • Proper use of UsageField constants for unified field names

3. Documentation Quality

Both new mixin classes have excellent docstrings:

  • Clear usage examples
  • Event type listings
  • Return value structure documentation
  • Inheritance patterns explained

4. Dependency Fix

The pyproject.toml fix for BytePlus package properly adds missing dependencies.

🔍 Observations & Potential Issues

1. BytePlus mime_type Assumption (packages/capabilities/image-generation/src/celeste_image_generation/providers/byteplus/streaming.py:63)

if content_type == "url":
    artifact = ImageArtifact(url=content, mime_type=ImageMimeType.PNG)

Concern: Hard-coding ImageMimeType.PNG for URL-based images may not always be correct. BytePlus could return JPEGs or other formats.

Recommendation:

  • Check if BytePlus API provides format information in the SSE event
  • Consider adding a _infer_mime_type_from_url() helper method
  • Or make the assumption explicit in a comment explaining why PNG is safe

2. Missing Mime Type for b64_json (packages/capabilities/image-generation/src/celeste_image_generation/providers/byteplus/streaming.py:65)

else:  # b64_json
    image_data = base64.b64decode(content)
    artifact = ImageArtifact(data=image_data)

Concern: When creating ImageArtifact from base64 data, no mime_type is specified. This could cause issues downstream if the mime type is needed for rendering or file saving.

Recommendation: Either add mime_type=ImageMimeType.PNG or document why it's optional/inferred elsewhere.

3. Error Handling Inconsistency

  • BytePlus streaming logs errors but returns None (line 39-44)
  • OpenAI streaming doesn't appear to handle error events explicitly

Question: Does OpenAI Images API emit error events in SSE streams? If so, should OpenAIImagesStream handle them?

4. Test Coverage Concern

The PR description mentions make ci was run, but:

  • No new unit tests were added for the new streaming mixins
  • Integration tests exist but may not cover the new mixin layer specifically
  • Static methods like map_usage_fields() should have unit tests to ensure correct field mapping

Recommendation: Consider adding unit tests for:

  • OpenAIImagesStream._parse_chunk() with various event types
  • BytePlusImagesStream._parse_chunk() with success/error/completed events
  • map_usage_fields() static methods to verify field transformations

5. Dependency Declaration Format (packages/providers/byteplus/pyproject.toml:8)

 dependencies = ["celeste-ai", "httpx"]

Minor Issue: Leading whitespace before dependencies. While this may work, it's inconsistent with standard TOML formatting.

Recommendation: Remove the leading space for consistency.

🚀 Performance Considerations

Good:

  • Static methods avoid instance overhead
  • Streaming implementation uses generators efficiently
  • No unnecessary object creation in hot paths

🔒 Security Considerations

Good:

  • Proper use of base64.b64decode() without unsafe evaluation
  • No SQL injection, XSS, or injection vulnerabilities introduced
  • Usage data extraction is safe

⚠️ Minor:

  • Consider validating that decoded base64 data is actually valid image data to prevent malicious payloads
  • URL-based artifacts could potentially point to malicious content (though this is likely handled at a higher layer)

📊 Code Quality

Aspect Rating Notes
Architecture ⭐⭐⭐⭐⭐ Excellent mixin pattern
Consistency ⭐⭐⭐⭐⭐ Uniform changes across providers
Documentation ⭐⭐⭐⭐⭐ Clear docstrings and comments
Error Handling ⭐⭐⭐⭐ Good but could be more consistent
Test Coverage ⭐⭐⭐ Integration tests exist, unit tests for mixins needed
Type Safety ⭐⭐⭐⭐⭐ Good use of type hints

📝 Recommendations Summary

High Priority:

  1. Add unit tests for new streaming mixin classes
  2. Clarify mime_type handling for BytePlus images (both URL and b64_json)

Medium Priority:
3. Fix TOML formatting (remove leading space in dependencies)
4. Document error handling strategy across providers
5. Consider adding validation for base64-decoded image data

Low Priority:
6. Add inline comment explaining PNG assumption for BytePlus URLs

✅ Approval

This is a well-executed refactoring that improves code organization and maintainability. The architectural pattern is sound and the implementation is clean. The issues identified are mostly minor and don't block the PR, but addressing them would make the code even more robust.

Recommendation: Approve with minor suggestions

Great work on this refactoring! 🎉

@Kamilbenkirane Kamilbenkirane merged commit 78292c2 into main Dec 19, 2025
11 checks passed
@claude

claude Bot commented Dec 19, 2025

Copy link
Copy Markdown

PR Review: Refactor Usage Field Mapping and Add Streaming Mixins

Overview

This PR refactors map_usage_fields() to static methods across all providers and introduces new streaming mixins for the image generation capability. Overall, this is a well-executed refactoring that improves code organization and reusability.


✅ Strengths

Architecture & Design

  • Excellent consistency: The static method pattern is applied uniformly across all 15 provider clients
  • Clean separation of concerns: The mixin pattern effectively shares SSE parsing logic between clients and streaming implementations
  • No breaking changes: The refactored methods were internal-only, maintaining backward compatibility
  • Type safety: Excellent use of type hints throughout with proper return type annotations

Implementation Quality

  • Correct static method usage: All providers properly call ClassName.map_usage_fields() instead of self.map_usage_fields()
  • Good documentation: Clear docstrings explaining purpose and usage patterns
  • Proper inheritance: MRO (Method Resolution Order) in multiple inheritance is correct

🔍 Issues Found

High Priority

1. BytePlus TOML Formatting 📝

File: packages/providers/byteplus/pyproject.toml:8

requires-python = ">=3.12"
 dependencies = ["celeste-ai", "httpx"]  # ← Leading space\!

Issue: Leading space before dependencies is non-standard and inconsistent with BFL's formatting.

Fix: Remove the leading space.

2. BytePlus Streaming mime_type Inconsistency 🖼️

File: packages/capabilities/image-generation/src/celeste_image_generation/providers/byteplus/streaming.py:63-66

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)  # ← Missing mime_type\!

Issue: URL-based images specify mime_type=ImageMimeType.PNG but base64 images don't, creating inconsistency.

Fix: Add mime_type=ImageMimeType.PNG to line 66 for consistency, or add a comment explaining why it's omitted.

Medium Priority

3. Missing Error Handling for base64 Decoding ⚠️

Files:

  • openai/streaming.py:31
  • byteplus/streaming.py:65
image_data = base64.b64decode(b64_json)  # Could raise binascii.Error

Issue: Malformed base64 data would raise an exception, crashing the stream.

Recommendation:

try:
    image_data = base64.b64decode(content)
except (binascii.Error, ValueError) as e:
    logger.error("Failed to decode base64 image: %s", e)
    return None

4. Silent Unknown Event Handling 🔇

Files:

  • openai/images/streaming.py:74
  • byteplus/images/streaming.py:112

Issue: Both mixins return None for unknown event types without logging.

Recommendation: Add logging before returning:

logger.warning("Unknown event type: %s", event_type)
return None

This helps detect API changes and debugging issues.

5. Missing Unit Tests 🧪

Files: New streaming mixin files

Issue: No dedicated unit tests for:

  • OpenAIImagesStream._parse_chunk() event type handling
  • BytePlusImagesStream._parse_chunk() event type handling
  • Static map_usage_fields() methods

Recommendation: Add tests covering:

  • All event types (partial_image, completed, edit events, errors)
  • Missing/malformed fields
  • Edge cases (empty content, invalid base64, unknown events)

6. Hard-coded PNG Assumption 🎨

File: byteplus/streaming.py:63

artifact = ImageArtifact(url=content, mime_type=ImageMimeType.PNG)

Issue: Assumes all BytePlus images are PNG. If the API returns JPEG/WebP, this would be incorrect.

Recommendation: Verify with BytePlus API docs whether format information is available in the event, or add a comment explaining the assumption.

Low Priority

7. OpenAI Streaming mime_type

Similar to issue #2, OpenAI streaming (openai/streaming.py:31) doesn't specify mime_type for decoded images. However, this is consistent with OpenAI's non-streaming client, so it's lower priority.


🧪 Testing

✅ PR mentions make ci passed
⚠️ Consider adding:

  • Unit tests for new streaming mixins
  • Tests for error scenarios
  • Tests for static method field mapping

🔒 Security & Performance

Security: ✅ Good

  • Safe use of base64.b64decode() (no eval/exec)
  • No injection risks (SQL, XSS, code)
  • Safe dictionary access with .get()

Performance: ✅ Good

  • Static methods avoid instance overhead
  • Efficient generator usage in streaming
  • No unnecessary object creation

📊 Overall Assessment

Aspect Rating Notes
Code Quality ⭐⭐⭐⭐ (4/5) Clean architecture, minor error handling gaps
Consistency ⭐⭐⭐⭐⭐ (5/5) Uniform changes across all providers
Documentation ⭐⭐⭐⭐⭐ (5/5) Excellent docstrings and examples
Testing ⭐⭐⭐ (3/5) Integration tests exist, missing unit tests
Type Safety ⭐⭐⭐⭐⭐ (5/5) Excellent use of type hints

🎯 Recommendation: Approve with Minor Fixes

This is a solid refactoring with a clean architecture. The mixin pattern is well-implemented and will make future provider additions cleaner.

Before Merge:

  1. ✅ Fix TOML formatting (remove leading space in byteplus/pyproject.toml:8)
  2. ✅ Add mime_type to BytePlus base64 images for consistency

Follow-up PRs (Optional):

  1. Add error handling for base64 decoding
  2. Add logging for unknown event types
  3. Add unit tests for streaming mixins

Great work on this refactoring! The architecture is sound and the implementation is clean. 🚀

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant