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        in_memory: bool = False,
 24    ):
 25        self.type = type_
 26        self.default = default
 27        self.env_var = env_var
 28        self.default_lambda = default_lambda
 29        self.sensitive = sensitive
 30        self.sensitive_keys = sensitive_keys
 31        self.in_memory = in_memory
 32
 33
 34class Config:
 35    _shared_instance = None
 36
 37    def __init__(self, properties: Dict[str, ConfigProperty] | None = None):
 38        self._properties: Dict[str, ConfigProperty] = properties or {
 39            "user_id": ConfigProperty(
 40                str,
 41                env_var="KILN_USER_ID",
 42                default_lambda=_get_user_id,
 43            ),
 44            "autosave_runs": ConfigProperty(
 45                bool,
 46                env_var="KILN_AUTOSAVE_RUNS",
 47                default=True,
 48                in_memory=True,
 49            ),
 50            "open_ai_api_key": ConfigProperty(
 51                str,
 52                env_var="OPENAI_API_KEY",
 53                sensitive=True,
 54            ),
 55            "groq_api_key": ConfigProperty(
 56                str,
 57                env_var="GROQ_API_KEY",
 58                sensitive=True,
 59            ),
 60            "ollama_base_url": ConfigProperty(
 61                str,
 62                env_var="OLLAMA_BASE_URL",
 63            ),
 64            "docker_model_runner_base_url": ConfigProperty(
 65                str,
 66                env_var="DOCKER_MODEL_RUNNER_BASE_URL",
 67            ),
 68            "bedrock_access_key": ConfigProperty(
 69                str,
 70                env_var="AWS_ACCESS_KEY_ID",
 71                sensitive=True,
 72            ),
 73            "bedrock_secret_key": ConfigProperty(
 74                str,
 75                env_var="AWS_SECRET_ACCESS_KEY",
 76                sensitive=True,
 77            ),
 78            "open_router_api_key": ConfigProperty(
 79                str,
 80                env_var="OPENROUTER_API_KEY",
 81                sensitive=True,
 82            ),
 83            "fireworks_api_key": ConfigProperty(
 84                str,
 85                env_var="FIREWORKS_API_KEY",
 86                sensitive=True,
 87            ),
 88            "fireworks_account_id": ConfigProperty(
 89                str,
 90                env_var="FIREWORKS_ACCOUNT_ID",
 91            ),
 92            "anthropic_api_key": ConfigProperty(
 93                str,
 94                env_var="ANTHROPIC_API_KEY",
 95                sensitive=True,
 96            ),
 97            "gemini_api_key": ConfigProperty(
 98                str,
 99                env_var="GEMINI_API_KEY",
100                sensitive=True,
101            ),
102            "projects": ConfigProperty(
103                list,
104                default_lambda=lambda: [],
105            ),
106            "azure_openai_api_key": ConfigProperty(
107                str,
108                env_var="AZURE_OPENAI_API_KEY",
109                sensitive=True,
110            ),
111            "azure_openai_endpoint": ConfigProperty(
112                str,
113                env_var="AZURE_OPENAI_ENDPOINT",
114            ),
115            "huggingface_api_key": ConfigProperty(
116                str,
117                env_var="HUGGINGFACE_API_KEY",
118                sensitive=True,
119            ),
120            "vertex_project_id": ConfigProperty(
121                str,
122                env_var="VERTEX_PROJECT_ID",
123            ),
124            "vertex_location": ConfigProperty(
125                str,
126                env_var="VERTEX_LOCATION",
127            ),
128            "together_api_key": ConfigProperty(
129                str,
130                env_var="TOGETHERAI_API_KEY",
131                sensitive=True,
132            ),
133            "wandb_api_key": ConfigProperty(
134                str,
135                env_var="WANDB_API_KEY",
136                sensitive=True,
137            ),
138            "wandb_entity": ConfigProperty(
139                str,
140                env_var="WANDB_ENTITY",
141            ),
142            "siliconflow_cn_api_key": ConfigProperty(
143                str,
144                env_var="SILICONFLOW_CN_API_KEY",
145                sensitive=True,
146            ),
147            "wandb_base_url": ConfigProperty(
148                str,
149                env_var="WANDB_BASE_URL",
150            ),
151            # Legacy custom models, replaced by user_model_registry below
152            "custom_models": ConfigProperty(
153                list,
154                default_lambda=lambda: [],
155            ),
156            "user_model_registry": ConfigProperty(
157                list,
158                default_lambda=lambda: [],
159            ),
160            "openai_compatible_providers": ConfigProperty(
161                list,
162                default_lambda=lambda: [],
163                sensitive_keys=["api_key"],
164            ),
165            "cerebras_api_key": ConfigProperty(
166                str,
167                env_var="CEREBRAS_API_KEY",
168                sensitive=True,
169            ),
170            "featherless_ai_api_key": ConfigProperty(
171                str,
172                env_var="FEATHERLESS_AI_API_KEY",
173                sensitive=True,
174            ),
175            "typesafe_api_key": ConfigProperty(
176                str,
177                env_var="TYPESAFE_API_KEY",
178                sensitive=True,
179            ),
180            "kiln_copilot_api_key": ConfigProperty(
181                str,
182                env_var="KILN_COPILOT_API_KEY",
183                sensitive=True,
184            ),
185            "enable_demo_tools": ConfigProperty(
186                bool,
187                env_var="ENABLE_DEMO_TOOLS",
188                default=False,
189            ),
190            "kiln_local_api_host": ConfigProperty(
191                str,
192                env_var="KILN_LOCAL_API_HOST",
193                default="127.0.0.1",
194                in_memory=True,
195            ),
196            "kiln_local_api_port": ConfigProperty(
197                int,
198                env_var="KILN_LOCAL_API_PORT",
199                default=8757,
200                in_memory=True,
201            ),
202            # Allow the user to set the path to lookup MCP server commands, like npx.
203            "custom_mcp_path": ConfigProperty(
204                str,
205                env_var="CUSTOM_MCP_PATH",
206            ),
207            # Allow the user to set secrets for MCP servers, the key is mcp_server_id::key_name
208            MCP_SECRETS_KEY: ConfigProperty(
209                dict[str, str],
210                sensitive=True,
211            ),
212            "git_sync_projects": ConfigProperty(
213                dict,
214                default_lambda=lambda: {},
215                sensitive_keys=["pat_token", "oauth_token"],
216            ),
217            # has the user indicated it's for personal or work use?
218            "user_type": ConfigProperty(
219                str,  # "personal" or "work"
220            ),
221            # if the user has provided their work contact
222            "work_use_contact": ConfigProperty(
223                str,
224            ),
225            # if the user has provided their personal contact
226            "personal_use_contact": ConfigProperty(
227                str,
228            ),
229        }
230        self._lock = threading.Lock()
231        self._in_memory_settings: Dict[str, Any] = {}
232        self._settings = self.load_settings()
233
234    @classmethod
235    def shared(cls):
236        if cls._shared_instance is None:
237            cls._shared_instance = cls()
238        return cls._shared_instance
239
240    # Get a value, mockable for testing
241    def get_value(self, name: str) -> Any:
242        try:
243            return self.__getattr__(name)
244        except AttributeError:
245            return None
246
247    def kiln_local_api_base_url(self) -> str:
248        return f"http://{self.kiln_local_api_host}:{self.kiln_local_api_port}"
249
250    def __getattr__(self, name: str) -> Any:
251        if name == "_properties":
252            return super().__getattribute__("_properties")
253        if name not in self._properties:
254            return super().__getattribute__(name)
255
256        property_config = self._properties[name]
257
258        if property_config.in_memory:
259            if name in self._in_memory_settings:
260                value = self._in_memory_settings[name]
261                return value if value is None else property_config.type(value)
262        else:
263            if name in self._settings:
264                value = self._settings[name]
265                return value if value is None else property_config.type(value)
266
267        # Check environment variable
268        if property_config.env_var and property_config.env_var in os.environ:
269            value = os.environ[property_config.env_var]
270            return property_config.type(value)
271
272        # Use default value or default_lambda
273        if property_config.default_lambda:
274            value = property_config.default_lambda()
275        else:
276            value = property_config.default
277
278        return None if value is None else property_config.type(value)
279
280    def __setattr__(self, name, value):
281        if name in ("_properties", "_settings", "_lock", "_in_memory_settings"):
282            super().__setattr__(name, value)
283        elif name in self._properties:
284            if self._properties[name].in_memory:
285                with self._lock:
286                    self._in_memory_settings[name] = value
287            else:
288                self.update_settings({name: value})
289        else:
290            raise AttributeError(f"Config has no attribute '{name}'")
291
292    @classmethod
293    def settings_dir(cls, create=True) -> str:
294        settings_dir = os.path.join(Path.home(), ".kiln_ai")
295        if create:
296            os.makedirs(settings_dir, exist_ok=True)
297        return settings_dir
298
299    @classmethod
300    def settings_path(cls, create=True) -> str:
301        settings_dir = cls.settings_dir(create)
302        return os.path.join(settings_dir, "settings.yaml")
303
304    @classmethod
305    def load_settings(cls):
306        if not os.path.isfile(cls.settings_path(create=False)):
307            return {}
308        with open(cls.settings_path(), "r") as f:
309            settings = yaml.safe_load(f.read()) or {}
310        return settings
311
312    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
313        with self._lock:
314            filtered_disk = {
315                k: v
316                for k, v in self._settings.items()
317                if k not in self._properties or not self._properties[k].in_memory
318            }
319            combined = {**filtered_disk, **self._in_memory_settings}
320
321        if not hide_sensitive:
322            return combined
323
324        settings = {
325            k: "[hidden]"
326            if k in self._properties and self._properties[k].sensitive
327            else copy.deepcopy(v)
328            for k, v in combined.items()
329        }
330        # Hide sensitive keys in nested structures (lists of dicts, or dicts of dicts)
331        for key, value in settings.items():
332            if key in self._properties and self._properties[key].sensitive_keys:
333                sensitive_keys = self._properties[key].sensitive_keys or []
334                for sensitive_key in sensitive_keys:
335                    if isinstance(value, list):
336                        for item in value:
337                            if isinstance(item, dict) and sensitive_key in item:
338                                item[sensitive_key] = "[hidden]"
339                    elif isinstance(value, dict):
340                        if sensitive_key in value:
341                            value[sensitive_key] = "[hidden]"
342                        for item in value.values():
343                            if isinstance(item, dict) and sensitive_key in item:
344                                item[sensitive_key] = "[hidden]"
345
346        return settings
347
348    def save_setting(self, name: str, value: Any):
349        self.update_settings({name: value})
350
351    def update_settings(self, new_settings: Dict[str, Any]):
352        with self._lock:
353            in_memory_updates = {
354                k: v
355                for k, v in new_settings.items()
356                if k in self._properties and self._properties[k].in_memory
357            }
358            disk_updates = {
359                k: v
360                for k, v in new_settings.items()
361                if k not in self._properties or not self._properties[k].in_memory
362            }
363
364            if in_memory_updates:
365                self._in_memory_settings.update(in_memory_updates)
366
367            if disk_updates:
368                # Fresh load to avoid clobbering changes from other instances
369                current_settings = self.load_settings()
370                current_settings.update(disk_updates)
371                # remove None values
372                current_settings = {
373                    k: v for k, v in current_settings.items() if v is not None
374                }
375                with open(self.settings_path(), "w") as f:
376                    yaml.dump(current_settings, f)
377                self._settings = current_settings
378
379
380def _get_user_id():
381    try:
382        return getpass.getuser() or "unknown_user"
383    except Exception:
384        return "unknown_user"