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