Skip to content

feat(core): make capability optional in create_client#95

Merged
Kamilbenkirane merged 1 commit into
mainfrom
feat/optional-capability
Dec 22, 2025
Merged

feat(core): make capability optional in create_client#95
Kamilbenkirane merged 1 commit into
mainfrom
feat/optional-capability

Conversation

@Kamilbenkirane

Copy link
Copy Markdown
Member

Summary

  • Add capability inference from model when capability is not explicitly provided
  • Models with single capability auto-infer (e.g., SeeDreamIMAGE_GENERATION)
  • Models with multiple capabilities raise clear error asking for explicit capability
  • Add LIST_MODELS and GET_MODEL endpoints to provider configs
  • Remove deprecated Mixtral 8x7B and 8x22B models

Usage

# Before: required capability
celeste.create_client(capability=Capability.IMAGE_GENERATION, model="seedream-4-5-251128")

# After: inferred from model (SeeDream only has IMAGE_GENERATION)
celeste.create_client(model="seedream-4-5-251128")

# Multi-capability models still require explicit capability
celeste.create_client(model="gpt-4o")  # Raises ValueError with clear message

Test plan

  • Tested single-capability model inference (SeeDream)
  • Tested multi-capability model error (GPT-4o)
  • All tests pass
  • Lint passes

- Add capability inference from model when not explicitly provided
- Models with single capability auto-infer (e.g., SeeDream → IMAGE_GENERATION)
- Models with multiple capabilities require explicit capability parameter
- Add _infer_capability() helper for clean error messages
- Update _resolve_model() to accept optional capability
- Add LIST_MODELS and GET_MODEL endpoints to provider configs
- Remove deprecated Mixtral 8x7B and 8x22B models
- Update test to reflect new behavior (ModelNotFoundError vs ValueError)
@claude

claude Bot commented Dec 22, 2025

Copy link
Copy Markdown

PR Review: Make capability optional in create_client

Summary

This PR introduces a helpful UX improvement by making the capability parameter optional when it can be unambiguously inferred from the model. The implementation is generally sound, but there are several areas that could be improved for robustness, test coverage, and consistency.

Code Quality & Best Practices

✅ Strengths

  1. Clear error messages: The error messages in _infer_capability are descriptive and helpful (e.g., listing available capabilities when ambiguous)
  2. Backward compatibility: Existing code using explicit capability continues to work without changes
  3. Documentation: Docstrings updated appropriately to reflect the new optional behavior
  4. Code organization: New helper function _infer_capability is well-scoped and focused

⚠️ Areas for Improvement

1. Inconsistent behavior in get_model function (models.py:71-101)

The function now accepts provider: Provider | None = None and searches across all providers when None. However, this introduces non-deterministic behavior:

if len(matches) > 1:
    warnings.warn(
        f"Model '{model_id}' found in multiple providers: {providers}. "
        f"Using '{matches[0].provider.value}'.",
        UserWarning,
        stacklevel=2,
    )
    return matches[0]

Issue: _models.values() returns a dict_values object, and the order of matches depends on insertion order. This could lead to unpredictable behavior across different Python environments or when models are registered in different orders.

Suggestion: Sort matches consistently (e.g., by provider name) before returning:

matches = sorted(matches, key=lambda m: m.provider.value)

2. Logic flow concern in _resolve_model (init.py:52-77)

When model is a string, get_model(model, provider) is called with potentially None provider. This means:

  • If provider is None, it searches all providers and may issue a warning
  • If multiple providers have the same model ID, the first match is used

Issue: This diverges from the original behavior where provider was required for string model IDs. While the PR description mentions this is intentional, it's not clear if this is the desired UX.

Suggestion: Consider if this silent fallback is desirable, or if it would be better to raise an error when model ID is ambiguous and provider is not specified.

3. Edge case: Empty capabilities set (init.py:40-49)

if len(model.capabilities) > 1:
    # ... error for multiple capabilities
msg = f"Model '{model.id}' has no registered capabilities"
raise ValueError(msg)

Issue: The error for models with no capabilities is only raised if they don't have exactly 1 or >1 capabilities. While this is logically correct, a model with no capabilities reaching this point indicates a configuration error.

Suggestion: This is fine, but consider if this should be an assertion/programming error rather than a ValueError since it indicates a bug in model registration.

Potential Bugs

🐛 Critical: Capability mismatch not validated (init.py:113-116)

When capability is inferred from the model, there's no validation that the resolved model actually supports the inferred capability. While this should always be true, there's no defensive check.

resolved_capability = (
    capability if capability else _infer_capability(resolved_model)
)

Scenario: If a model's capabilities are modified after registration (though unlikely), or if there's a race condition, this could cause issues downstream.

Suggestion: Add a validation check:

if resolved_capability not in resolved_model.capabilities:
    raise UnsupportedCapabilityError(...)

🐛 Minor: Warning stacklevel in get_model (models.py:98)

The stacklevel=2 in the warning may not point to the correct caller location in all cases, especially when called through _resolve_modelcreate_client.

Suggestion: Consider stacklevel=3 or higher to ensure the warning points to the user's code, not internal celeste functions.

Performance Considerations

Good

  • The changes don't introduce any performance regressions
  • The list comprehension in get_model is efficient for typical registry sizes

💡 Minor optimization opportunity

In get_model, if provider is specified, the function still constructs the tuple key and performs a dictionary lookup. This is already O(1) and optimal.

Security Concerns

No security issues identified

  • No user input is used unsafely
  • No injection vulnerabilities
  • Appropriate type checking and validation

Test Coverage

⚠️ Significant gaps in test coverage

Missing tests for new functionality:

  1. _infer_capability function - No direct tests for:

    • Single capability model (should return that capability)
    • Multiple capability model (should raise ValueError with clear message)
    • Zero capability model (should raise ValueError)
  2. get_model with provider=None - No tests for:

    • Model found in single provider when provider=None
    • Model found in multiple providers when provider=None (warning behavior)
    • Warning message content and stacklevel correctness
  3. Integration test for capability inference - No test showing:

    # This should work now but isn't tested
    celeste.create_client(model="seedream-4-5-251128")
  4. Edge cases - Missing tests for:

    • What happens when both capability and model are None?
    • Behavior when inferred capability differs from explicitly provided capability (should use explicit)

Existing tests updated:

The test test_create_client_string_model_without_provider_raises_error was renamed to test_create_client_string_model_not_found_raises_error, which correctly reflects the new behavior. However:

Issue: The test no longer validates the original requirement that string model without provider should error. This behavior has changed, which is fine, but it's not explicitly tested that provider-less lookup now works.

📝 Recommended additional tests:

def test_infer_capability_single_capability():
    """Test capability inference for single-capability models."""
    model = Model(
        id="single-cap",
        provider=Provider.OPENAI,
        capabilities={Capability.TEXT_GENERATION},
        display_name="Single Cap",
    )
    assert _infer_capability(model) == Capability.TEXT_GENERATION

def test_infer_capability_multiple_raises_error():
    """Test that multi-capability models raise clear error."""
    model = Model(
        id="multi-cap",
        provider=Provider.OPENAI,
        capabilities={Capability.TEXT_GENERATION, Capability.EMBEDDINGS},
        display_name="Multi Cap",
    )
    with pytest.raises(ValueError, match="supports multiple capabilities"):
        _infer_capability(model)

def test_get_model_without_provider_warns_on_duplicate():
    """Test warning when model exists in multiple providers."""
    # Register same model ID for different providers
    # Call get_model with provider=None
    # Assert warning is issued with correct message
    pass

def test_create_client_infers_capability_from_model():
    """Test that capability is correctly inferred from single-capability model."""
    # Mock a single-capability model
    # Call create_client(model=...) without capability
    # Assert correct capability is used
    pass

Additional Notes

📋 Endpoint additions

The PR adds LIST_MODELS and GET_MODEL endpoints to provider configs. These additions are:

  • Consistent across providers
  • Properly formatted
  • However, not used anywhere in this PR

Question: Are these endpoint additions for future functionality? If so, consider splitting them into a separate PR to keep changes focused. If they're needed for this PR, there should be code using them.

🗑️ Model removal

Removal of deprecated Mixtral models (open-mixtral-8x7b and open-mixtral-8x22b) is:

  • Clean and straightforward
  • Properly documented in PR description
  • No migration path provided, but this is acceptable for deprecated models

Suggestion: If these models were publicly documented or commonly used, consider adding a deprecation notice in the changelog or migration guide.

Recommendations

High Priority

  1. ✅ Add comprehensive tests for _infer_capability function
  2. ✅ Add tests for get_model with provider=None including warning behavior
  3. ⚠️ Add integration test demonstrating the main use case (capability inference)
  4. ⚠️ Make get_model deterministic by sorting matches when multiple providers found
  5. ⚠️ Consider if ambiguous model IDs should error instead of warn when provider is None

Medium Priority

  1. Review warning stacklevel to ensure it points to user code
  2. Add defensive capability validation after inference
  3. Consider documenting the provider selection behavior when ambiguous

Low Priority

  1. Clarify if endpoint additions are needed for this PR or future work
  2. Consider if zero-capability models should use a different error type

Conclusion

This is a solid improvement to the developer experience. The core implementation is sound, but would benefit significantly from additional test coverage, particularly around the new capability inference logic and the modified get_model behavior. The non-deterministic ordering in multi-provider scenarios should be addressed to ensure consistent behavior.

Overall Assessment: ✅ Approve with suggestions for improvement

The changes are ready to merge after addressing the test coverage gaps and considering the determinism concern in get_model.

@Kamilbenkirane Kamilbenkirane merged commit 578dc27 into main Dec 22, 2025
11 checks passed
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