"""Base registry for model management."""
from typing import Any, Callable, Generic, TypeVar
T = TypeVar("T")
class ModelRegistry(Generic[T]):
"""Base registry for managing model adapters.
Usage:
@MyRegistry.register("model_name")
class MyModelAdapter:
...
registry = MyRegistry()
registry.load("model_name")
model = registry.active
"""
_models: dict[str, type[T]] = {}
def __init__(self):
self._active: T | None = None
self._active_name: str | None = None
@classmethod
def register(cls, name: str) -> Callable[[type[T]], type[T]]:
"""Decorator to register a model adapter.
Args:
name: Unique name for this model
Example:
@MyRegistry.register("mymodel")
class MyModelAdapter:
...
"""
def decorator(model_cls: type[T]) -> type[T]:
cls._models[name] = model_cls
return model_cls
return decorator
@classmethod
def available(cls) -> list[str]:
"""List available model names."""
return list(cls._models.keys())
@property
def active(self) -> T | None:
"""Get the currently active model."""
return self._active
@property
def active_name(self) -> str | None:
"""Get the name of the currently active model."""
return self._active_name
def load(self, name: str, **kwargs: Any) -> T:
"""Load a model by name.
If another model is currently loaded, it will be unloaded first.
Args:
name: Model name (must be registered)
**kwargs: Additional arguments passed to model constructor
Returns:
The loaded model instance
"""
if name not in self._models:
raise ValueError(
f"Unknown model: {name}. Available: {self.available()}"
)
# Unload current model
if self._active is not None:
self.unload()
# Create and load new model
model_cls = self._models[name]
model = model_cls(**kwargs)
model.load()
self._active = model
self._active_name = name
return model
def unload(self) -> None:
"""Unload the current model."""
if self._active is not None:
self._active.unload()
self._active = None
self._active_name = None
def is_loaded(self, name: str) -> bool:
"""Check if a specific model is loaded."""
return self._active_name == name
