kiln_ai.utils.config

  1import copy
  2import getpass
  3import os
  4import threading
  5from pathlib import Path
  6from typing import Any, Callable, Dict, List, Optional
  7
  8import yaml
  9
 10# Configuration keys
 11MCP_SECRETS_KEY = "mcp_secrets"
 12
 13
 14class ConfigProperty:
 15    def __init__(
 16        self,
 17        type_: type,
 18        default: Any = None,
 19        env_var: Optional[str] = None,
 20        default_lambda: Optional[Callable[[], Any]] = None,
 21        sensitive: bool = False,
 22        sensitive_keys: Optional[List[str]] = None,
 23    ):
 24        self.type = type_
 25        self.default = default
 26        self.env_var = env_var
 27        self.default_lambda = default_lambda
 28        self.sensitive = sensitive
 29        self.sensitive_keys = sensitive_keys
 30
 31
 32class Config:
 33    _shared_instance = None
 34
 35    def __init__(self, properties: Dict[str, ConfigProperty] | None = None):
 36        self._properties: Dict[str, ConfigProperty] = properties or {
 37            "user_id": ConfigProperty(
 38                str,
 39                env_var="KILN_USER_ID",
 40                default_lambda=_get_user_id,
 41            ),
 42            "autosave_runs": ConfigProperty(
 43                bool,
 44                env_var="KILN_AUTOSAVE_RUNS",
 45                default=True,
 46            ),
 47            "open_ai_api_key": ConfigProperty(
 48                str,
 49                env_var="OPENAI_API_KEY",
 50                sensitive=True,
 51            ),
 52            "groq_api_key": ConfigProperty(
 53                str,
 54                env_var="GROQ_API_KEY",
 55                sensitive=True,
 56            ),
 57            "ollama_base_url": ConfigProperty(
 58                str,
 59                env_var="OLLAMA_BASE_URL",
 60            ),
 61            "docker_model_runner_base_url": ConfigProperty(
 62                str,
 63                env_var="DOCKER_MODEL_RUNNER_BASE_URL",
 64            ),
 65            "bedrock_access_key": ConfigProperty(
 66                str,
 67                env_var="AWS_ACCESS_KEY_ID",
 68                sensitive=True,
 69            ),
 70            "bedrock_secret_key": ConfigProperty(
 71                str,
 72                env_var="AWS_SECRET_ACCESS_KEY",
 73                sensitive=True,
 74            ),
 75            "open_router_api_key": ConfigProperty(
 76                str,
 77                env_var="OPENROUTER_API_KEY",
 78                sensitive=True,
 79            ),
 80            "fireworks_api_key": ConfigProperty(
 81                str,
 82                env_var="FIREWORKS_API_KEY",
 83                sensitive=True,
 84            ),
 85            "fireworks_account_id": ConfigProperty(
 86                str,
 87                env_var="FIREWORKS_ACCOUNT_ID",
 88            ),
 89            "anthropic_api_key": ConfigProperty(
 90                str,
 91                env_var="ANTHROPIC_API_KEY",
 92                sensitive=True,
 93            ),
 94            "gemini_api_key": ConfigProperty(
 95                str,
 96                env_var="GEMINI_API_KEY",
 97                sensitive=True,
 98            ),
 99            "projects": ConfigProperty(
100                list,
101                default_lambda=lambda: [],
102            ),
103            "azure_openai_api_key": ConfigProperty(
104                str,
105                env_var="AZURE_OPENAI_API_KEY",
106                sensitive=True,
107            ),
108            "azure_openai_endpoint": ConfigProperty(
109                str,
110                env_var="AZURE_OPENAI_ENDPOINT",
111            ),
112            "huggingface_api_key": ConfigProperty(
113                str,
114                env_var="HUGGINGFACE_API_KEY",
115                sensitive=True,
116            ),
117            "vertex_project_id": ConfigProperty(
118                str,
119                env_var="VERTEX_PROJECT_ID",
120            ),
121            "vertex_location": ConfigProperty(
122                str,
123                env_var="VERTEX_LOCATION",
124            ),
125            "together_api_key": ConfigProperty(
126                str,
127                env_var="TOGETHERAI_API_KEY",
128                sensitive=True,
129            ),
130            "wandb_api_key": ConfigProperty(
131                str,
132                env_var="WANDB_API_KEY",
133                sensitive=True,
134            ),
135            "wandb_entity": ConfigProperty(
136                str,
137                env_var="WANDB_ENTITY",
138            ),
139            "siliconflow_cn_api_key": ConfigProperty(
140                str,
141                env_var="SILICONFLOW_CN_API_KEY",
142                sensitive=True,
143            ),
144            "wandb_base_url": ConfigProperty(
145                str,
146                env_var="WANDB_BASE_URL",
147            ),
148            # Legacy custom models, replaced by user_model_registry below
149            "custom_models": ConfigProperty(
150                list,
151                default_lambda=lambda: [],
152            ),
153            "user_model_registry": ConfigProperty(
154                list,
155                default_lambda=lambda: [],
156            ),
157            "openai_compatible_providers": ConfigProperty(
158                list,
159                default_lambda=lambda: [],
160                sensitive_keys=["api_key"],
161            ),
162            "cerebras_api_key": ConfigProperty(
163                str,
164                env_var="CEREBRAS_API_KEY",
165                sensitive=True,
166            ),
167            "kiln_copilot_api_key": ConfigProperty(
168                str,
169                env_var="KILN_COPILOT_API_KEY",
170                sensitive=True,
171            ),
172            "enable_demo_tools": ConfigProperty(
173                bool,
174                env_var="ENABLE_DEMO_TOOLS",
175                default=False,
176            ),
177            # Allow the user to set the path to lookup MCP server commands, like npx.
178            "custom_mcp_path": ConfigProperty(
179                str,
180                env_var="CUSTOM_MCP_PATH",
181            ),
182            # Allow the user to set secrets for MCP servers, the key is mcp_server_id::key_name
183            MCP_SECRETS_KEY: ConfigProperty(
184                dict[str, str],
185                sensitive=True,
186            ),
187            # has the user indicated it's for personal or work use?
188            "user_type": ConfigProperty(
189                str,  # "personal" or "work"
190            ),
191            # if the user has provided their work contact
192            "work_use_contact": ConfigProperty(
193                str,
194            ),
195            # if the user has provided their personal contact
196            "personal_use_contact": ConfigProperty(
197                str,
198            ),
199        }
200        self._lock = threading.Lock()
201        self._settings = self.load_settings()
202
203    @classmethod
204    def shared(cls):
205        if cls._shared_instance is None:
206            cls._shared_instance = cls()
207        return cls._shared_instance
208
209    # Get a value, mockable for testing
210    def get_value(self, name: str) -> Any:
211        try:
212            return self.__getattr__(name)
213        except AttributeError:
214            return None
215
216    def __getattr__(self, name: str) -> Any:
217        if name == "_properties":
218            return super().__getattribute__("_properties")
219        if name not in self._properties:
220            return super().__getattribute__(name)
221
222        property_config = self._properties[name]
223
224        # Check if the value is in settings
225        if name in self._settings:
226            value = self._settings[name]
227            return value if value is None else property_config.type(value)
228
229        # Check environment variable
230        if property_config.env_var and property_config.env_var in os.environ:
231            value = os.environ[property_config.env_var]
232            return property_config.type(value)
233
234        # Use default value or default_lambda
235        if property_config.default_lambda:
236            value = property_config.default_lambda()
237        else:
238            value = property_config.default
239
240        return None if value is None else property_config.type(value)
241
242    def __setattr__(self, name, value):
243        if name in ("_properties", "_settings", "_lock"):
244            super().__setattr__(name, value)
245        elif name in self._properties:
246            self.update_settings({name: value})
247        else:
248            raise AttributeError(f"Config has no attribute '{name}'")
249
250    @classmethod
251    def settings_dir(cls, create=True) -> str:
252        settings_dir = os.path.join(Path.home(), ".kiln_ai")
253        if create and not os.path.exists(settings_dir):
254            os.makedirs(settings_dir)
255        return settings_dir
256
257    @classmethod
258    def settings_path(cls, create=True) -> str:
259        settings_dir = cls.settings_dir(create)
260        return os.path.join(settings_dir, "settings.yaml")
261
262    @classmethod
263    def load_settings(cls):
264        if not os.path.isfile(cls.settings_path(create=False)):
265            return {}
266        with open(cls.settings_path(), "r") as f:
267            settings = yaml.safe_load(f.read()) or {}
268        return settings
269
270    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
271        if not hide_sensitive:
272            return self._settings
273
274        settings = {
275            k: "[hidden]"
276            if k in self._properties and self._properties[k].sensitive
277            else copy.deepcopy(v)
278            for k, v in self._settings.items()
279        }
280        # Hide sensitive keys in lists. Could generalize this if we every have more types, but right not it's only needed for root elements of lists
281        for key, value in settings.items():
282            if key in self._properties and self._properties[key].sensitive_keys:
283                sensitive_keys = self._properties[key].sensitive_keys or []
284                for sensitive_key in sensitive_keys:
285                    if isinstance(value, list):
286                        for item in value:
287                            if sensitive_key in item:
288                                item[sensitive_key] = "[hidden]"
289
290        return settings
291
292    def save_setting(self, name: str, value: Any):
293        self.update_settings({name: value})
294
295    def update_settings(self, new_settings: Dict[str, Any]):
296        # Lock to prevent race conditions in multi-threaded scenarios
297        with self._lock:
298            # Fresh load to avoid clobbering changes from other instances
299            current_settings = self.load_settings()
300            current_settings.update(new_settings)
301            # remove None values
302            current_settings = {
303                k: v for k, v in current_settings.items() if v is not None
304            }
305            with open(self.settings_path(), "w") as f:
306                yaml.dump(current_settings, f)
307            self._settings = current_settings
308
309
310def _get_user_id():
311    try:
312        return getpass.getuser() or "unknown_user"
313    except Exception:
314        return "unknown_user"