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            "kiln_copilot_api_key": ConfigProperty(
171                str,
172                env_var="KILN_COPILOT_API_KEY",
173                sensitive=True,
174            ),
175            "enable_demo_tools": ConfigProperty(
176                bool,
177                env_var="ENABLE_DEMO_TOOLS",
178                default=False,
179            ),
180            "kiln_local_api_host": ConfigProperty(
181                str,
182                env_var="KILN_LOCAL_API_HOST",
183                default="127.0.0.1",
184                in_memory=True,
185            ),
186            "kiln_local_api_port": ConfigProperty(
187                int,
188                env_var="KILN_LOCAL_API_PORT",
189                default=8757,
190                in_memory=True,
191            ),
192            # Allow the user to set the path to lookup MCP server commands, like npx.
193            "custom_mcp_path": ConfigProperty(
194                str,
195                env_var="CUSTOM_MCP_PATH",
196            ),
197            # Allow the user to set secrets for MCP servers, the key is mcp_server_id::key_name
198            MCP_SECRETS_KEY: ConfigProperty(
199                dict[str, str],
200                sensitive=True,
201            ),
202            "git_sync_projects": ConfigProperty(
203                dict,
204                default_lambda=lambda: {},
205                sensitive_keys=["pat_token", "oauth_token"],
206            ),
207            # has the user indicated it's for personal or work use?
208            "user_type": ConfigProperty(
209                str,  # "personal" or "work"
210            ),
211            # if the user has provided their work contact
212            "work_use_contact": ConfigProperty(
213                str,
214            ),
215            # if the user has provided their personal contact
216            "personal_use_contact": ConfigProperty(
217                str,
218            ),
219        }
220        self._lock = threading.Lock()
221        self._in_memory_settings: Dict[str, Any] = {}
222        self._settings = self.load_settings()
223
224    @classmethod
225    def shared(cls):
226        if cls._shared_instance is None:
227            cls._shared_instance = cls()
228        return cls._shared_instance
229
230    # Get a value, mockable for testing
231    def get_value(self, name: str) -> Any:
232        try:
233            return self.__getattr__(name)
234        except AttributeError:
235            return None
236
237    def kiln_local_api_base_url(self) -> str:
238        return f"http://{self.kiln_local_api_host}:{self.kiln_local_api_port}"
239
240    def __getattr__(self, name: str) -> Any:
241        if name == "_properties":
242            return super().__getattribute__("_properties")
243        if name not in self._properties:
244            return super().__getattribute__(name)
245
246        property_config = self._properties[name]
247
248        if property_config.in_memory:
249            if name in self._in_memory_settings:
250                value = self._in_memory_settings[name]
251                return value if value is None else property_config.type(value)
252        else:
253            if name in self._settings:
254                value = self._settings[name]
255                return value if value is None else property_config.type(value)
256
257        # Check environment variable
258        if property_config.env_var and property_config.env_var in os.environ:
259            value = os.environ[property_config.env_var]
260            return property_config.type(value)
261
262        # Use default value or default_lambda
263        if property_config.default_lambda:
264            value = property_config.default_lambda()
265        else:
266            value = property_config.default
267
268        return None if value is None else property_config.type(value)
269
270    def __setattr__(self, name, value):
271        if name in ("_properties", "_settings", "_lock", "_in_memory_settings"):
272            super().__setattr__(name, value)
273        elif name in self._properties:
274            if self._properties[name].in_memory:
275                with self._lock:
276                    self._in_memory_settings[name] = value
277            else:
278                self.update_settings({name: value})
279        else:
280            raise AttributeError(f"Config has no attribute '{name}'")
281
282    @classmethod
283    def settings_dir(cls, create=True) -> str:
284        settings_dir = os.path.join(Path.home(), ".kiln_ai")
285        if create and not os.path.exists(settings_dir):
286            os.makedirs(settings_dir)
287        return settings_dir
288
289    @classmethod
290    def settings_path(cls, create=True) -> str:
291        settings_dir = cls.settings_dir(create)
292        return os.path.join(settings_dir, "settings.yaml")
293
294    @classmethod
295    def load_settings(cls):
296        if not os.path.isfile(cls.settings_path(create=False)):
297            return {}
298        with open(cls.settings_path(), "r") as f:
299            settings = yaml.safe_load(f.read()) or {}
300        return settings
301
302    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
303        with self._lock:
304            filtered_disk = {
305                k: v
306                for k, v in self._settings.items()
307                if k not in self._properties or not self._properties[k].in_memory
308            }
309            combined = {**filtered_disk, **self._in_memory_settings}
310
311        if not hide_sensitive:
312            return combined
313
314        settings = {
315            k: "[hidden]"
316            if k in self._properties and self._properties[k].sensitive
317            else copy.deepcopy(v)
318            for k, v in combined.items()
319        }
320        # Hide sensitive keys in nested structures (lists of dicts, or dicts of dicts)
321        for key, value in settings.items():
322            if key in self._properties and self._properties[key].sensitive_keys:
323                sensitive_keys = self._properties[key].sensitive_keys or []
324                for sensitive_key in sensitive_keys:
325                    if isinstance(value, list):
326                        for item in value:
327                            if isinstance(item, dict) and sensitive_key in item:
328                                item[sensitive_key] = "[hidden]"
329                    elif isinstance(value, dict):
330                        if sensitive_key in value:
331                            value[sensitive_key] = "[hidden]"
332                        for item in value.values():
333                            if isinstance(item, dict) and sensitive_key in item:
334                                item[sensitive_key] = "[hidden]"
335
336        return settings
337
338    def save_setting(self, name: str, value: Any):
339        self.update_settings({name: value})
340
341    def update_settings(self, new_settings: Dict[str, Any]):
342        with self._lock:
343            in_memory_updates = {
344                k: v
345                for k, v in new_settings.items()
346                if k in self._properties and self._properties[k].in_memory
347            }
348            disk_updates = {
349                k: v
350                for k, v in new_settings.items()
351                if k not in self._properties or not self._properties[k].in_memory
352            }
353
354            if in_memory_updates:
355                self._in_memory_settings.update(in_memory_updates)
356
357            if disk_updates:
358                # Fresh load to avoid clobbering changes from other instances
359                current_settings = self.load_settings()
360                current_settings.update(disk_updates)
361                # remove None values
362                current_settings = {
363                    k: v for k, v in current_settings.items() if v is not None
364                }
365                with open(self.settings_path(), "w") as f:
366                    yaml.dump(current_settings, f)
367                self._settings = current_settings
368
369
370def _get_user_id():
371    try:
372        return getpass.getuser() or "unknown_user"
373    except Exception:
374        return "unknown_user"