Coverage for packages/registry/src/langgate/registry/config.py: 81%

1"""Configuration handling for the registry.""" 

2 

3import importlib.resources 

4import json 

5from pathlib import Path 

6from typing import Any 

7 

8from dotenv import load_dotenv 

9 

10from langgate.core.logging import get_logger 

11from langgate.core.schemas.config import ConfigSchema 

12from langgate.core.utils.config_utils import load_yaml_config, resolve_path 

13 

14logger = get_logger(__name__) 

15 

16 

17class RegistryConfig: 

18 """Configuration handler for the registry.""" 

19 

20 def __init__( 

21 self, 

22 models_data_path: Path | None = None, 

23 config_path: Path | None = None, 

24 env_file_path: Path | None = None, 

25 ): 

26 """ 

27 Args: 

28 models_data_path: Path to the models data JSON file 

29 config_path: Path to the main configuration YAML file 

30 env_file_path: Path to a `.env` file for environment variables 

31 """ 

32 # Set up default paths 

33 cwd = Path.cwd() 

34 # Get package resource paths 

35 registry_resources = importlib.resources.files("langgate.registry") 

36 core_resources = importlib.resources.files("langgate.core") 

37 default_models_path = Path( 

38 str(registry_resources.joinpath("data", "default_models.json")) 

39 ) 

40 default_config_path = Path( 

41 str(core_resources.joinpath("data", "default_config.yaml")) 

42 ) 

43 

44 # Define default paths with priorities 

45 # Models data: args > env > cwd > package_dir 

46 cwd_models_path = cwd / "langgate_models.json" 

47 

48 # Config: args > env > cwd > package_dir 

49 cwd_config_path = cwd / "langgate_config.yaml" 

50 

51 # Env file: args > env > cwd 

52 cwd_env_path = cwd / ".env" 

53 

54 # Resolve paths using priority order 

55 self.models_data_path = resolve_path( 

56 "LANGGATE_MODELS", 

57 models_data_path, 

58 cwd_models_path if cwd_models_path.exists() else default_models_path, 

59 "models_data_path", 

60 ) 

61 

62 self.config_path = resolve_path( 

63 "LANGGATE_CONFIG", 

64 config_path, 

65 cwd_config_path if cwd_config_path.exists() else default_config_path, 

66 "config_path", 

67 ) 

68 

69 self.env_file_path = resolve_path( 

70 "LANGGATE_ENV_FILE", env_file_path, cwd_env_path, "env_file_path" 

71 ) 

72 

73 # Load environment variables from .env file if it exists 

74 if self.env_file_path.exists(): 

75 load_dotenv(self.env_file_path) 

76 logger.debug("loaded_env_file", path=str(self.env_file_path)) 

77 

78 # Initialize data structures 

79 self.models_data: dict[str, dict[str, Any]] = {} 

80 self.global_config: dict[str, Any] = {} 

81 self.service_config: dict[str, dict[str, Any]] = {} 

82 self.model_mappings: dict[str, dict[str, Any]] = {} 

83 

84 # Load configuration 

85 self._load_config() 

86 

87 def _load_config(self) -> None: 

88 """Load configuration from files.""" 

89 try: 

90 # Load model data 

91 self._load_model_data() 

92 

93 # Load main configuration 

94 self._load_main_config() 

95 

96 except Exception: 

97 logger.exception( 

98 "failed_to_load_config", 

99 models_data_path=str(self.models_data_path), 

100 config_path=str(self.config_path), 

101 ) 

102 raise 

103 

104 def _load_model_data(self) -> None: 

105 """Load model data from JSON file.""" 

106 try: 

107 with open(self.models_data_path) as f: 

108 self.models_data = json.load(f) 

109 logger.info( 

110 "loaded_model_data", 

111 models_data_path=str(self.models_data_path), 

112 model_count=len(self.models_data), 

113 ) 

114 except FileNotFoundError: 

115 logger.warning( 

116 "model_data_file_not_found", 

117 models_data_path=str(self.models_data_path), 

118 ) 

119 self.models_data = {} 

120 

121 def _load_main_config(self) -> None: 

122 """Load main configuration from YAML file.""" 

123 config = load_yaml_config(self.config_path, ConfigSchema, logger) 

124 

125 if config: 

126 # Extract validated data 

127 self.global_config = { 

128 "default_params": config.default_params, 

129 } 

130 

131 # Extract service provider config 

132 self.service_config = { 

133 k: v.model_dump(exclude_none=True) for k, v in config.services.items() 

134 } 

135 

136 # Process model mappings 

137 self._process_model_mappings(config.models) 

138 else: 

139 self._set_empty_config() 

140 

141 def _set_empty_config(self) -> None: 

142 """Set empty/default config state, typically used in error scenarios.""" 

143 self.global_config = {"default_params": {}} 

144 self.service_config = {} 

145 self.model_mappings = {} 

146 

147 def _process_model_mappings(self, models_config) -> None: 

148 """Process model mappings from validated configuration. 

149 

150 Args: 

151 models_config: List of validated model configurations 

152 """ 

153 self.model_mappings = {} 

154 

155 for model_config in models_config: 

156 model_data = model_config.model_dump(exclude_none=True) 

157 model_id = model_data["id"] 

158 service = model_data["service"] 

159 

160 # Store mapping info with proper type handling 

161 self.model_mappings[model_id] = { 

162 "service_provider": service["provider"], 

163 "service_model_id": service["model_id"], 

164 "override_params": model_data.get("override_params", {}), 

165 "remove_params": model_data.get("remove_params", []), 

166 "rename_params": model_data.get("rename_params", {}), 

167 "name": model_data.get("name"), 

168 "description": model_data.get("description"), 

169 }