fix(capabilities): add explicit usage and finish_reason types to Output classes

[Kamilbenkirane]

Summary

• Override usage and finish_reason field types in capability Output classes
• Ensures proper Pydantic serialization with capability-specific subclass fields
• Without explicit type annotations, Pydantic uses base schema (empty {})

Changes

• text-generation/io.py - Add TextGenerationUsage and TextGenerationFinishReason to Output
• image-generation/io.py - Add ImageGenerationUsage and ImageGenerationFinishReason to Output
• video-generation/io.py - Add VideoGenerationUsage to Output
• speech-generation/io.py - Add SpeechGenerationUsage and SpeechGenerationFinishReason to Output
• Bump all capability versions to 0.3.7

Test plan

• Verify TextGenerationOutput.model_dump() includes usage fields
• Pre-commit hooks pass
• Type checks pass

[claude]

Code Review Summary

This PR addresses a real Pydantic serialization issue and the fix is correct. The changes properly ensure that capability-specific subclass fields are included in serialization. Overall, this is a solid fix with good scope and testing approach.

✅ Strengths

1. Correct Problem Identification: The PR correctly identifies that Pydantic uses the base class schema when subclass fields aren't explicitly annotated, resulting in empty {} during serialization.

2. Minimal, Focused Changes: The fix is surgical - only adding the necessary field overrides without refactoring or over-engineering.

3. Consistent Application: The fix is applied uniformly across all 4 capability packages (text, image, video, speech).

4. Proper Version Bumping: All packages correctly bumped to 0.3.7.

5. Test Coverage: The test plan mentions verification of TextGenerationOutput.model_dump() including usage fields.

---

🔍 Code Quality Observations

1. Consistency in Field Definitions

The changes follow a consistent pattern:

usage: SpecificUsage = Field(default_factory=SpecificUsage)
finish_reason: SpecificFinishReason | None = None

This matches the base class signature in src/celeste/io.py:32-33, which is good for maintainability.

2. Import Additions

All files correctly add from pydantic import Field - this is necessary since the base class imports aren't automatically available in subclasses.

3. Type Safety

The explicit type annotations improve type safety and IDE autocomplete, as tools can now properly infer the specific usage/finish_reason types rather than the base types.

---

💡 Potential Improvements

1. Video Generation Missing finish_reason

Issue: video-generation/io.py only adds usage but not finish_reason:

class VideoGenerationOutput(Output[VideoArtifact]):
    usage: VideoGenerationUsage = Field(default_factory=VideoGenerationUsage)
    # Missing: finish_reason override

Question: Is this intentional? If video generation doesn't have finish reasons, this is fine. Otherwise, you should define a VideoGenerationFinishReason class and add it here for consistency.

Impact: Low (video generation may not need finish reasons), but worth clarifying.

---

2. Missing FinishReason Class Definition

Observation: Looking at the diff, I don't see where VideoGenerationFinishReason is defined (or if it exists at all).

For reference, checking the other capabilities:

• ImageGenerationFinishReason - should be defined in image-generation/io.py
• SpeechGenerationFinishReason - should be defined in speech-generation/io.py
• TextGenerationFinishReason - should be defined in text-generation/io.py

Recommendation: Verify that all the *FinishReason classes used in the Output classes are actually defined. The diff doesn't show their definitions, so they may have been added in earlier commits - just confirm they exist.

---

3. Test Coverage Verification

The test plan mentions:

✅ Verify TextGenerationOutput.model_dump() includes usage fields

Recommendation: Consider adding a simple unit test to prevent regression:

def test_output_serialization_includes_usage_fields():
    """Ensure capability-specific usage fields are serialized."""
    output = TextGenerationOutput(
        content="test",
        usage=TextGenerationUsage(input_tokens=10, output_tokens=5),
    )
    
    dumped = output.model_dump()
    
    # Should include specific usage fields, not empty {}
    assert "input_tokens" in dumped["usage"]
    assert "output_tokens" in dumped["usage"]
    assert dumped["usage"]["input_tokens"] == 10

This would catch the exact issue this PR fixes if it ever regresses.

---

🔒 Security Considerations

No security concerns identified. This is purely a serialization fix with no external inputs or security-sensitive operations.

---

⚡ Performance Considerations

Negligible impact. The changes:

• Add minimal overhead (just field validation during instantiation)
• Don't change runtime behavior beyond serialization correctness
• Use default_factory efficiently (no extra allocations unless needed)

---

🧪 Testing Recommendations

1. Add serialization tests as mentioned above to prevent regression
2. Test JSON schema generation: Verify that Output.model_json_schema() now includes capability-specific fields
3. Integration testing: The existing integration tests should already validate this works end-to-end

---

📋 Best Practices Assessment

✅ Following Python 3.12+ conventions (union types with |)
✅ Proper Pydantic usage with Field(default_factory=...))
✅ Consistent code style across all capability packages
✅ Version management - all affected packages bumped together
✅ Clear PR description explaining the "why" behind the change

---

✨ Final Verdict

LGTM with minor questions - This is a solid fix for a legitimate Pydantic serialization issue. The only clarification needed is:

1. Confirm VideoGenerationFinishReason isn't needed (or add it if it is)
2. Verify all referenced *FinishReason classes are defined

Consider adding a unit test for serialization to prevent future regression, but it's not blocking.

Great job keeping the changes minimal and focused! 🎯