feat: add multi-capability model registration and new providers

[Kamilbenkirane]

Summary

This PR adds multi-capability model registration support and introduces new providers and capabilities.

Multi-Capability Model Registration

• register_models() now accepts capability parameter to support registering the same model for multiple capabilities
• Models are merged when duplicate (id, provider) pairs are detected
• Capabilities and parameter constraints are accumulated across registrations
• Validates display_name consistency when merging

New Provider & Capability

• Added BYTEDANCE provider support
• Added BYTEDANCE credential management
• Added IMAGE_INTELLIGENCE capability

Error Handling

• Added _handle_error_response() method to Client base class for better API error handling

Code Quality

• Updated Model class field order (display_name before capabilities)
• Updated all tests for new registration pattern
• Consistent function signature style

Changes

• src/celeste/models.py: Multi-capability registration with merging logic
• src/celeste/core.py: Added BYTEDANCE provider and IMAGE_INTELLIGENCE capability
• src/celeste/credentials.py: Added BYTEDANCE credential support
• src/celeste/client.py: Added error handling method
• tests/unit_tests/test_models.py: Updated all tests for new pattern

Testing

All tests pass and linting/type checking passes.

[claude]

Pull Request Review: Multi-Capability Model Registration

Summary

This PR introduces multi-capability model registration support, new providers (ByteDance), and improved error handling. The implementation is well-designed with good test coverage.

Code Quality & Best Practices ✅

Strengths

1. Clean API Design: The register_models() function now requires an explicit capability parameter, making the registration intent clear and preventing ambiguity.

2. Proper Merging Logic: The model merging implementation (lines 46-66 in models.py) correctly handles duplicate registrations by:

   • Using setdefault() to initialize models
   • Validating display_name consistency
   • Accumulating capabilities and parameter constraints

3. Comprehensive Test Coverage: Added excellent test cases including:

   • Multi-capability registration and merging
   • Display name validation
   • Parameter constraint accumulation across registrations

4. Type Safety: Good use of type hints and Pydantic models throughout.

Suggestions for Improvement

1. Error Handling in _handle_error_response() (client.py:63-71)

Issue: The error handling has a potential flaw:

def _handle_error_response(self, response: httpx.Response) -> None:
    """Handle error responses from provider APIs"""
    if not response.is_success:
        try:
            error_data = response.json()
            error_msg = error_data.get("error", {}).get("message", response.text)
            raise ValueError(f"{self.provider.value} API error: {error_msg}")
        except (AttributeError, JSONDecodeError):
            response.raise_for_status()

Problems:

• Catches AttributeError but none of the operations should raise it (.json() raises JSONDecodeError, .get() doesn't raise AttributeError)
• The try/except block catches the ValueError that's being raised, which seems unintentional
• Generic ValueError doesn't convey enough context about HTTP errors

Recommendation:

def _handle_error_response(self, response: httpx.Response) -> None:
    """Handle error responses from provider APIs"""
    if not response.is_success:
        try:
            error_data = response.json()
            error_msg = error_data.get("error", {}).get("message", response.text)
        except JSONDecodeError:
            error_msg = response.text
        
        # Raise a more specific exception or use httpx.HTTPStatusError
        raise httpx.HTTPStatusError(
            f"{self.provider.value} API error: {error_msg}",
            request=response.request,
            response=response
        )

Alternatively, let response.raise_for_status() handle it and only add custom message formatting.

2. Model Field Ordering (models.py:12-16)

The PR changes field order from:

# Before
capabilities: set[Capability]
display_name: str

to:

# After  
display_name: str
capabilities: set[Capability]

While this works fine, it would be helpful to document why this ordering is preferred (e.g., for readability, importance, or API consistency).

3. Parameter Constraints Merging (models.py:66)

Observation:

registered.parameter_constraints.update(model.parameter_constraints)

This uses dict.update() which means later registrations will overwrite parameter constraints from earlier registrations for the same parameter name.

Question: Is this the intended behavior?

Scenarios:

• If Model A registers with temperature: Range(0, 1) for TEXT_GENERATION
• Then Model A registers with temperature: Range(0, 2) for IMAGE_GENERATION
• The constraint becomes Range(0, 2) for both capabilities

Recommendation: Either:

1. Document this behavior explicitly if intentional
2. Validate that constraints match when merging: if param in registered.parameter_constraints and registered.parameter_constraints[param] != new_constraint: raise ValueError(...)
3. Store per-capability constraints if different capabilities need different constraints

Potential Bugs 🐛

1. Missing Usage of _handle_error_response()

The method is added but never called in the codebase. Check if it should be integrated into existing client implementations.

Action: Search for response.raise_for_status() calls in client implementations and replace with self._handle_error_response(response).

2. Test Assertion Changes

In test_models.py:51, the test now uses:

assert get_model(single_model.id, single_model.provider) == single_model

But after the change, get_model() returns a different instance (the merged one from the registry), not the original single_model. The assertion happens to pass because Pydantic model equality is based on field values, but this could be fragile.

Recommendation: Use explicit field comparisons as done elsewhere in the tests:

retrieved = get_model(single_model.id, single_model.provider)
assert retrieved is not None
assert retrieved.id == single_model.id
assert retrieved.display_name == single_model.display_name

Performance Considerations ⚡

No significant concerns. The implementation is efficient:

• Uses dictionary lookups O(1)
• setdefault() is appropriate
• Set operations for capabilities are fast
• Dict updates for constraints are reasonable

Minor optimization: Consider using frozenset for capabilities if they never need mutation after creation, but this is optional.

Security Concerns 🔒

1. Credential Management

New BYTEDANCE credentials are added consistently across the codebase. ✅

2. Error Message Exposure

In _handle_error_response(), the error messages from provider APIs are included in exceptions:

error_msg = error_data.get("error", {}).get("message", response.text)

Consideration: Ensure provider error messages don't contain sensitive information (API keys, internal IDs, etc.). Consider sanitizing or truncating response.text if it could be large or sensitive.

Test Coverage ✅

Excellent test coverage! The PR includes:

• ✅ Multi-capability registration test
• ✅ Display name validation test
• ✅ Capability accumulation verification
• ✅ Updated all existing tests to use new signature
• ✅ Parameter constraint tests

Suggestion: Add a test for the parameter constraint merging behavior:

def test_parameter_constraints_merge_on_multi_capability_registration(self) -> None:
    """Verify parameter constraints are merged when registering same model for multiple capabilities."""
    model_v1 = Model(
        id="test-model",
        provider=Provider.OPENAI,
        display_name="Test",
        parameter_constraints={"temp": Str(), "tokens": Str()},
    )
    model_v2 = Model(
        id="test-model", 
        provider=Provider.OPENAI,
        display_name="Test",
        parameter_constraints={"style": Str()},
    )
    
    register_models(model_v1, Capability.TEXT_GENERATION)
    register_models(model_v2, Capability.IMAGE_GENERATION)
    
    retrieved = get_model("test-model", Provider.OPENAI)
    assert retrieved is not None
    assert len(retrieved.parameter_constraints) == 3
    assert "temp" in retrieved.supported_parameters
    assert "tokens" in retrieved.supported_parameters
    assert "style" in retrieved.supported_parameters

Recommendations Summary

High Priority

1. 🔴 Fix _handle_error_response() exception handling logic
2. 🔴 Integrate _handle_error_response() into actual client code or remove if unused
3. 🟡 Document or validate parameter constraint merging behavior

Medium Priority

4. 🟡 Add test for parameter constraint merging
5. 🟡 Document the field ordering rationale
6. 🟡 Consider sanitizing error messages in _handle_error_response()

Low Priority

7. 🟢 Update test assertions to use explicit field comparisons

Conclusion

This is a well-designed PR with good architecture and comprehensive tests. The multi-capability registration feature is implemented cleanly. The main concerns are around the error handling implementation which should be refined before merging.

Overall: Approve with minor changes requested ✨

[claude]

Code Review: Multi-Capability Model Registration

Summary

This PR introduces a valuable feature allowing models to be registered for multiple capabilities with proper merging logic. The implementation is generally solid with good test coverage.

---

✅ Strengths

1. Well-designed merging logic - The register_models() function elegantly handles duplicate registrations
2. Comprehensive test coverage - Tests thoroughly cover the new multi-capability registration pattern
3. Non-breaking for internal code - All tests updated consistently
4. Good validation - Display name consistency check prevents accidental misconfigurations

---

🔍 Code Quality & Best Practices

Breaking API Change (Location: src/celeste/models.py:28)

The register_models() function signature changed from accepting only models to requiring a capability parameter. This is a BREAKING CHANGE for external packages.

Recommendations:

• Document this breaking change in release notes
• Consider versioning strategy (bump to 0.1.0 instead of patch)
• Verify no published packages depend on the old signature

Subtle Logic Issue (Location: src/celeste/models.py:46-55)

The setdefault() implementation creates unused Model instances when the key already exists. Consider using explicit if-check instead.

Error Handling (Location: src/celeste/client.py:63-78)

The new _handle_error_response() method is not called anywhere. Either document its intended use for subclasses or remove if premature.

---

🐛 Potential Bugs

Parameter Constraint Merging (src/celeste/models.py:66)

What happens if the same parameter is registered with different constraints for different capabilities? The dict.update() will overwrite previous constraints. Please clarify if this is intentional and add test cases.

---

🚀 Performance

Minor: setdefault() creates unused instances, but overall performance impact is negligible.

---

🔒 Security

• No direct security issues
• New BYTEDANCE credentials follow existing secure patterns
• Error messages might expose sensitive API responses - consider sanitization

---

🧪 Test Coverage

Excellent coverage overall. Missing test cases:

1. Conflicting parameter constraints across capabilities
2. Edge cases with empty capabilities

---

📝 Documentation

Please add:

1. Migration guide for the breaking change
2. Multi-capability registration examples in docstrings
3. Explanation of why capability is a parameter vs using model.capabilities

---

🎯 Recommendations

High Priority

1. Document the breaking API change and versioning
2. Clarify parameter constraint merging behavior
3. Document _handle_error_response() usage or remove

Medium Priority

1. Optimize setdefault() usage
2. Add migration examples

---

✅ Conclusion

This is a well-implemented feature with solid test coverage. Main concerns are the breaking API change documentation and parameter constraint merging behavior clarification. Code quality is high and follows project conventions.

Recommendation: Request clarifications on constraint merging and ensure breaking changes are documented before merging.