Skip to content

feat: input class registry and unregistered model fallback#99

Merged
Kamilbenkirane merged 3 commits into
mainfrom
feat/input-registry-and-model-fallback
Dec 24, 2025
Merged

feat: input class registry and unregistered model fallback#99
Kamilbenkirane merged 3 commits into
mainfrom
feat/input-registry-and-model-fallback

Conversation

@Kamilbenkirane

Copy link
Copy Markdown
Member

Summary

This PR introduces two core enhancements to Celeste's registration system:

Input Class Registry

Adds a capability-based Input class registry (register_input/get_input_class) that allows runtime discovery of the correct Input type for any capability. Each capability package now registers its Input class during initialization, enabling dynamic input type resolution.

Unregistered Model Fallback

Previously, using an unknown model ID would raise ModelNotFoundError. Now, when both provider and capability are explicitly specified, Celeste:

  1. Emits a UserWarning that parameter validation is disabled
  2. Creates a dynamic Model on-the-fly
  3. Proceeds with the request

This enables using new or unreleased provider models before they're added to Celeste's registry.

Additional Changes

  • BytePlus: Added Seedance 1.5 Pro video model (4-12s, 480p/720p, first/last frame support)
  • Cleanup: Fixed gradium py.typed marker, reordered google pyproject.toml

Test plan

  • Updated test to verify warning behavior for unregistered models
  • All 316 tests pass with 83% coverage
  • CI passes (ruff, mypy, bandit)

- Add register_input/get_input_class to celeste.io for capability-based
  Input class discovery
- Register Input classes in all capability packages during initialization
- Allow unregistered models when provider and capability are explicit,
  with warning about disabled parameter validation
- Export new functions from celeste package
- Update test to verify new warning behavior
- Add seedance-1-5-pro-251215 with 4-12s duration, 480p/720p resolution
- Fix gradium py.typed marker file (should be empty)
- Reorder google pyproject.toml sections
- Update .gitignore for temporary files
@claude

claude Bot commented Dec 24, 2025

Copy link
Copy Markdown

Code Review - PR #99: Input Class Registry and Unregistered Model Fallback

Overview

This PR introduces two significant enhancements to Celeste's registration system. The implementation is well-structured and addresses real user needs. Overall, this is a solid contribution with good test coverage (83%).


✅ Strengths

1. Input Class Registry Design

The new registry pattern in src/celeste/io.py:46-59 is clean and follows the existing models.py registry pattern:

  • Simple module-level dict with clear registration/retrieval functions
  • Proper error handling with descriptive KeyError messages
  • Consistent with existing codebase patterns

2. Unregistered Model Fallback

The fallback mechanism in src/celeste/__init__.py:72-95 is well thought out:

  • Emits appropriate warnings to users about disabled parameter validation
  • Requires explicit provider and capability to prevent misuse
  • Creates dynamic Model objects with sensible defaults (streaming=True)

3. Test Coverage

The updated test in tests/unit_tests/test_init.py:53-77 properly validates:

  • Warning emission behavior
  • The fallback code path
  • Prevention of ModelNotFoundError in the right scenarios

4. Consistent Package Updates

All four capability packages uniformly register their Input classes during initialization, maintaining symmetry across the codebase.


🔍 Issues & Concerns

CRITICAL: Missing Test Coverage for Input Registry

Location: src/celeste/io.py:46-59

The new register_input and get_input_class functions have zero test coverage. This is a significant gap.

Recommended tests:

# tests/unit_tests/test_io.py (new file)
def test_register_and_get_input_class():
    """Test input class registration and retrieval."""
    # Test successful registration and retrieval
    # Test KeyError when capability not registered
    # Test overwriting existing registration (if allowed)

Impact: Medium-High (core functionality untested)


MODERATE: Incomplete Test for Unregistered Model Fallback

Location: tests/unit_tests/test_init.py:53-77

The current test mocks get_client_class to raise NotImplementedError, so it never validates:

  1. The dynamically created Model object's structure
  2. That the Model is correctly passed to the client
  3. Edge cases like missing capability parameter

Recommended improvements:

def test_create_client_unregistered_model_creates_dynamic_model():
    """Test that dynamic model has correct attributes."""
    with (
        patch("celeste.get_model") as mock_get_model,
        patch("celeste.get_client_class") as mock_client_class,
        pytest.warns(UserWarning, match="not registered"),
    ):
        mock_get_model.return_value = None
        mock_client = Mock()
        mock_client_class.return_value = mock_client
        
        create_client(
            capability=Capability.TEXT_GENERATION,
            provider=Provider.OPENAI,
            model="new-model-2025",
        )
        
        # Verify the dynamic model passed to client
        call_args = mock_client_class.call_args
        dynamic_model = call_args.kwargs['model']
        assert dynamic_model.id == "new-model-2025"
        assert dynamic_model.provider == Provider.OPENAI
        assert Capability.TEXT_GENERATION in dynamic_model.capabilities
        assert dynamic_model.streaming is True

def test_unregistered_model_without_capability_raises_error():
    """Test that missing capability prevents fallback."""
    with patch("celeste.get_model") as mock_get_model:
        mock_get_model.return_value = None
        
        with pytest.raises(ValueError, match="Specify 'capability' explicitly"):
            create_client(provider=Provider.OPENAI, model="new-model")

Impact: Medium (partial validation of critical feature)


MINOR: Potential Registry Race Condition

Location: src/celeste/io.py:46

The module-level _inputs dict is not thread-safe. While Python's GIL provides some protection, concurrent registration from multiple threads could theoretically cause issues.

Consideration:

  • If Celeste is intended for multi-threaded environments, consider using threading.Lock
  • Document thread-safety expectations
  • The existing _models registry has the same pattern, so this is consistent but worth noting

Impact: Low (unlikely in practice, but worth documenting)


MINOR: Gitignore Changes

Location: .gitignore:162-165

Adding scripts/ to gitignore is overly broad:

  • Prevents committing any legitimate scripts directory
  • Consider more specific patterns: scripts/temp/, scripts/audit/, etc.
  • Or document why all scripts should be ignored

Impact: Low (may affect future contributors)


MINOR: Empty py.typed File

Location: packages/providers/gradium/src/celeste_gradium/py.typed

The PR removes content from this file, leaving it empty. This is actually correct per PEP 561 (empty marker file), but the change appears incidental.

Question: Was the previous "" content causing issues? If so, consider documenting this fix in the PR description.

Impact: Negligible (correctly implements PEP 561)


🛡️ Security Assessment

No Critical Security Issues Found

No injection vulnerabilities: Model IDs are validated through registry or explicitly provided
No credential exposure: No changes to authentication handling
Input validation: The fallback correctly warns about disabled parameter validation
Warning stacklevel: Properly set to 3 to show caller's location

Recommendation

The warning message at src/celeste/__init__.py:82-86 should emphasize security implications:

warnings.warn(
    f"Model '{model}' not registered in Celeste for provider {provider.value}. "
    "Parameter validation and safety checks are DISABLED. "
    "Use only with trusted model IDs from official provider documentation.",
    UserWarning,
    stacklevel=3,
)

🚀 Performance Considerations

Minimal overhead: Registry lookups are O(1) dict operations
Lazy loading: Input registration happens during package initialization
No blocking operations: All synchronous, no I/O in registration

No performance concerns identified.


📋 Additional Observations

1. BytePlus Model Addition

The new Seedance 1.5 Pro model (packages/capabilities/video-generation/src/celeste_video_generation/providers/byteplus/models.py:100-114) is well-specified:

  • Proper constraints (duration 4-12s vs 2-12s for v1.0)
  • Reduced resolution options (480p/720p only, no 1080p)
  • Follows existing naming conventions

2. Pyproject.toml Reordering

The change in packages/providers/google/pyproject.toml is purely cosmetic (moving entry-points before build-system). While this improves consistency, it's unrelated to the PR's stated purpose.

3. Type Safety

The code maintains strong typing throughout:

  • Generic type[Input] for registry values
  • Proper Pydantic model usage
  • No Any types introduced

📝 Recommendations Summary

Priority Issue Action
HIGH Missing input registry tests Add comprehensive tests for register_input/get_input_class
MEDIUM Incomplete fallback tests Expand tests to validate dynamic Model structure and error cases
LOW Security warning clarity Enhance warning message to emphasize disabled safety checks
LOW Gitignore specificity Use more specific patterns instead of blanket scripts/

✅ Approval Recommendation

Approve with minor revisions

This PR delivers valuable functionality and is well-implemented. The input registry pattern is clean, and the unregistered model fallback addresses a real user need.

Before merging:

  1. Add tests for the input registry functions (high priority)
  2. Expand unregistered model fallback test coverage (recommended)
  3. Consider enhancing the security warning message

The existing 316 passing tests and 83% coverage demonstrate overall quality, but the new registry functions deserve explicit test coverage.


Great work on this feature! The capability-based input discovery will simplify client implementations significantly. 🎉

@Kamilbenkirane Kamilbenkirane merged commit dfbc5fc into main Dec 24, 2025
11 checks passed
@claude

claude Bot commented Dec 24, 2025

Copy link
Copy Markdown

Pull Request Review: Input Class Registry and Unregistered Model Fallback

Summary

This PR introduces two important enhancements to Celeste's registration system: a capability-based Input class registry and a fallback mechanism for unregistered models. Overall, the implementation is well-designed and production-ready, with good test coverage (83%) and proper warning mechanisms.


Strengths

1. Clean Architecture

  • The input registry pattern mirrors the existing client and model registry patterns, maintaining architectural consistency
  • Registry implementation in src/celeste/io.py:46-59 is simple and effective
  • Proper separation of concerns with capability packages handling their own registration

2. Thoughtful Unregistered Model Handling

  • Graceful degradation with clear warnings when using unregistered models (src/celeste/init.py:82-94)
  • Requires explicit provider AND capability parameters, preventing accidental misuse
  • stacklevel=3 in warning is correctly calibrated for the call stack
  • Creates sensible defaults (streaming=True, empty parameter_constraints)

3. Consistent Implementation

  • All four capability packages updated identically with input registration
  • Proper version bumping across all packages (0.3.5 to 0.3.6)
  • Exports added to all for proper public API

4. Good Test Coverage

  • Test in tests/unit_tests/test_init.py:53-77 properly validates the warning behavior
  • Uses pytest.warns() correctly to assert the warning is emitted
  • Verifies the exception flow when client isn't registered

Code Quality Observations

Minor Issues

  1. Empty py.typed File - The gradium py.typed file was changed from containing quotes to being completely empty. This is correct for PEP 561 type marker files. Good cleanup.

  2. gitignore Additions - Added CHANGELOG_SINCE_RELEASE.md, MIGRATION_AUDIT_REPORT.md, and scripts/. Are these temporary development artifacts? Consider documenting if they're project-specific conventions.

  3. Google pyproject.toml Reordering - Moving project.entry-points after build-system is cosmetic but improves consistency with other provider packages.


Security and Safety

Excellent Security Practices

  • No credential exposure risks
  • Warning mechanism prevents silent failures
  • Validation is explicitly disabled (and warned about) for unregistered models
  • Requires explicit user intent (provider + capability both required)

Potential Concerns

Bypassing Parameter Validation: When using unregistered models, parameter validation is completely disabled. Users could pass invalid parameters that might cause runtime errors downstream. The warning clearly states "Parameter validation disabled" which is good mitigation. Consider documenting common pitfalls in docstrings or user guide.


Potential Bugs and Edge Cases

Critical: KeyError Risk in get_input_class

Location: src/celeste/io.py:54-59

When using the new unregistered model fallback, if someone calls get_input_class() with a capability that hasn't been loaded yet, this will raise a KeyError.

Recommendation: Call _load_from_entry_points() in get_input_class() before checking the registry (similar to list_models in src/celeste/models.py:118-120)

Minor: No Validation on register_input Overwrite

Location: src/celeste/io.py:49-51

Silently overwrites if called twice with different classes for the same capability. Consider adding a warning similar to the model registry's display name validation (src/celeste/models.py:60-64).


Performance Considerations

  • Registry lookups are O(1) dictionary operations
  • No unnecessary iterations or redundant operations
  • Lazy loading via entry points maintained
  • Package registration happens once at initialization

No performance concerns identified.


Test Coverage

What's Tested

  • Warning emission for unregistered models
  • Model not found error when provider not specified
  • Explicit model selection precedence

What Could Be Tested

  1. Input Registry Tests: No tests for register_input()/get_input_class()

    • Should test registration and retrieval
    • Should test KeyError when capability not registered
    • Should test duplicate registration behavior
  2. Unregistered Model Scenarios:

    • Missing capability parameter (should raise ValueError)
    • Missing provider parameter (should raise ModelNotFoundError)
    • Successful dynamic model creation and usage
  3. New BytePlus Model: seedance-1-5-pro-251215 has no specific tests


BytePlus Seedance 1.5 Pro Addition

Location: packages/capabilities/video-generation/src/celeste_video_generation/providers/byteplus/models.py:99-115

  • Follows existing model definition patterns exactly
  • Constraints are properly defined (duration 4-12s, resolution 480p/720p)
  • Supports first/last frame with correct MIME types
  • Note: Duration minimum increased from 2s to 4s (vs other models)
  • Note: Resolution options reduced (no 1080p support vs other models)

These constraints look intentional based on model capabilities.


Final Verdict

Approval Status: APPROVE WITH MINOR SUGGESTIONS

This PR is well-implemented and ready to merge with only minor, non-blocking suggestions.

Must Fix Before Merge: None

Should Fix (Recommended):

  1. Add lazy loading to get_input_class() to prevent KeyError with unregistered models
  2. Add warning on duplicate register_input() calls
  3. Add unit tests for input registry functions

Nice to Have:

  1. Document gitignore additions if they're project conventions
  2. Enhanced docstrings for new public APIs
  3. Integration test for BytePlus Seedance 1.5 Pro

Summary

Excellent work on this feature! The implementation is clean, well-tested, and follows established patterns in the codebase. The unregistered model fallback is particularly thoughtful in its safeguards. The input registry is a logical extension of the existing architecture.

The suggested fixes are minor and primarily defensive programming improvements. The PR demonstrates good software engineering practices with consistent version bumping, proper test coverage, and attention to backward compatibility.

Recommendation: Merge after addressing the get_input_class() lazy loading issue to prevent potential KeyError in edge cases.

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