File indexing completed on 2026-08-12 08:24:56
0001 """Registry for stack-specific configuration and stack implementation classes.
0002
0003 Project: AID2E v0.0.0 - AI assisted Detector Design for EIC
0004 Homepage: https://aid2e.github.io/aid2e-framework
0005 Repository: https://github.com/aid2e/AID2E-framework.git
0006 """
0007
0008 from typing import Dict, Type, Any
0009 from pydantic import BaseModel
0010
0011
0012 class StackRegistry:
0013 """
0014 Unified registry for experimental stack configuration models + interfaces
0015 """
0016 _env_configs: Dict[str, Type[BaseModel]] = {}
0017 _env_loaders: Dict[str, Type[Any]] = {}
0018 _design_configs: Dict[str, Type[BaseModel]] = {}
0019 _design_loaders: Dict[str, Type[Any]] = {}
0020 _workflow_configs: Dict[str, Type[BaseModel]] = {}
0021 _experimental_stacks: Dict[str, Type[Any]] = {}
0022
0023 @classmethod
0024 def register_stack(
0025 cls,
0026 name: str,
0027 env_config: Type[BaseModel],
0028 env_loader: Type[Any],
0029 design_config: Type[BaseModel],
0030 design_loader: Type[Any],
0031 workflow_config: Type[BaseModel],
0032 experimental_stack: Type[Any],
0033 ) -> None:
0034 """Register a stack type and its configuration/implementation pair."""
0035 cls._env_configs[name] = env_config
0036 cls._env_loaders[name] = env_loader
0037 cls._design_configs[name] = design_config
0038 cls._design_loaders[name] = design_loader
0039 cls._workflow_configs[name] = workflow_config
0040 cls._experimental_stacks[name] = experimental_stack
0041
0042 @classmethod
0043 def get_env_config(cls, name: str) -> Type[BaseModel]:
0044 """Get the environment config model for a stack."""
0045 if name not in cls._env_configs:
0046 raise KeyError(f"Stack config model not registered: {name}")
0047 return cls._env_configs[name]
0048
0049 @classmethod
0050 def get_env_loader(cls, name: str) -> Type[Any]:
0051 """Get the environment config loader for a stack."""
0052 if name not in cls._env_configs:
0053 raise KeyError(f"Stack config loader not registered: {name}")
0054 return cls._env_loaders[name]
0055
0056 @classmethod
0057 def get_design_config(cls, name: str) -> Type[BaseModel]:
0058 """Get the design config model for a stack."""
0059 if name not in cls._design_configs:
0060 raise KeyError(f"Stack config model not registered: {name}")
0061 return cls._design_configs[name]
0062
0063 @classmethod
0064 def get_design_loader(cls, name: str) -> Type[Any]:
0065 """Get the design config loader for a stack."""
0066 if name not in cls._design_configs:
0067 raise KeyError(f"Stack config loader not registered: {name}")
0068 return cls._design_loaders[name]
0069
0070 @classmethod
0071 def get_workflow_config(cls, name: str) -> Type[BaseModel]:
0072 """Get the workflow config model for a stack."""
0073 if name not in cls._workflow_configs:
0074 raise KeyError(f"Stack config model not registered: {name}")
0075 return cls._workflow_configs[name]
0076
0077 @classmethod
0078 def get_experimental_stack(cls, name: str) -> Type[Any]:
0079 """Get the stack implementation class for a stack name."""
0080 if name not in cls._experimental_stacks:
0081 raise KeyError(f"Experimental stack not registered: {name}")
0082 return cls._experimental_stacks[name]
0083
0084 @classmethod
0085 def list_registered_stacks(cls) -> Dict[str, Dict[str, Type[Any]]]:
0086 return {
0087 name: {
0088 "env_config": cls._env_configs[name],
0089 "env_loader": cls._env_loaders[name],
0090 "design_config" : cls._design_configs[name],
0091 "design_loader" : cls._design_loaders[name],
0092 "workflow_config" : cls._workflow_configs[name],
0093 "experimental_stack": cls._experimental_stacks[name],
0094 }
0095 for name in cls._env_configs
0096 }
