34 lines
959 B
Python
34 lines
959 B
Python
from __future__ import annotations
|
|
|
|
from typing import Dict, Type
|
|
|
|
from ..model_configuration import ModelConfig
|
|
from .base import BaseModelAdapter
|
|
|
|
|
|
class ModelRegistry:
|
|
"""
|
|
Maps model keys to adapter classes so configs can reference them declaratively.
|
|
"""
|
|
|
|
_registry: Dict[str, Type[BaseModelAdapter]] = {}
|
|
|
|
@classmethod
|
|
def register(cls, name: str):
|
|
def decorator(adapter_cls: Type[BaseModelAdapter]) -> Type[BaseModelAdapter]:
|
|
cls._registry[name] = adapter_cls
|
|
return adapter_cls
|
|
|
|
return decorator
|
|
|
|
@classmethod
|
|
def create(cls, name: str, config: ModelConfig) -> BaseModelAdapter:
|
|
if name not in cls._registry:
|
|
raise KeyError(f"ModelAdapter '{name}' is not registered.")
|
|
adapter_cls = cls._registry[name]
|
|
return adapter_cls(config)
|
|
|
|
@classmethod
|
|
def available(cls) -> Dict[str, Type[BaseModelAdapter]]:
|
|
return dict(cls._registry)
|