refactor(providers): standardize usage field mapping and add streaming mixins#86
Conversation
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.
PR Review: Standardize usage field mapping and add streaming mixinsSummaryThis 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. ✅ Strengths1. Excellent Architecture PatternThe new mixin approach is well-designed:
2. Consistent Usage MappingConverting
3. Documentation QualityBoth new mixin classes have excellent docstrings:
4. Dependency FixThe 🔍 Observations & Potential Issues1. BytePlus
|
| 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:
- Add unit tests for new streaming mixin classes
- 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! 🎉
PR Review: Refactor Usage Field Mapping and Add Streaming MixinsOverviewThis PR refactors ✅ StrengthsArchitecture & Design
Implementation Quality
🔍 Issues FoundHigh Priority1. BytePlus TOML Formatting 📝File: requires-python = ">=3.12"
dependencies = ["celeste-ai", "httpx"] # ← Leading space\!Issue: Leading space before Fix: Remove the leading space. 2. BytePlus Streaming mime_type Inconsistency 🖼️File: 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 Fix: Add Medium Priority3. Missing Error Handling for base64 Decoding
|
| 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:
- ✅ Fix TOML formatting (remove leading space in byteplus/pyproject.toml:8)
- ✅ Add mime_type to BytePlus base64 images for consistency
Follow-up PRs (Optional):
- Add error handling for base64 decoding
- Add logging for unknown event types
- Add unit tests for streaming mixins
Great work on this refactoring! The architecture is sound and the implementation is clean. 🚀
Summary
map_usage_fields()as static methods across all provider clientsOpenAIImagesStream)BytePlusImagesStream)Changes
Provider Refactoring
All provider clients now expose
map_usage_fields()as a static method, enabling:Affected providers:
New Streaming Mixins
OpenAIImagesStream- Handlesimage_generation.partial_image,image_generation.completed,image_edit.*eventsBytePlusImagesStream- Handlesimage_generation.partial_succeeded,image_generation.partial_failed,image_generation.completedeventsCapability Migrations
OpenAIImageGenerationStreamnow inherits fromOpenAIImagesStreamBytePlusImageGenerationStreamnow inherits fromBytePlusImagesStreamTest plan
make ci