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