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            "custom_models": ConfigProperty(
149                list,
150                default_lambda=lambda: [],
151            ),
152            "openai_compatible_providers": ConfigProperty(
153                list,
154                default_lambda=lambda: [],
155                sensitive_keys=["api_key"],
156            ),
157            "cerebras_api_key": ConfigProperty(
158                str,
159                env_var="CEREBRAS_API_KEY",
160                sensitive=True,
161            ),
162            "enable_demo_tools": ConfigProperty(
163                bool,
164                env_var="ENABLE_DEMO_TOOLS",
165                default=False,
166            ),
167            # Allow the user to set the path to lookup MCP server commands, like npx.
168            "custom_mcp_path": ConfigProperty(
169                str,
170                env_var="CUSTOM_MCP_PATH",
171            ),
172            # Allow the user to set secrets for MCP servers, the key is mcp_server_id::key_name
173            MCP_SECRETS_KEY: ConfigProperty(
174                dict[str, str],
175                sensitive=True,
176            ),
177            # has the user indicated it's for personal or work use?
178            "user_type": ConfigProperty(
179                str,  # "personal" or "work"
180            ),
181            # if the user has provided their work contact
182            "work_use_contact": ConfigProperty(
183                str,
184            ),
185        }
186        self._lock = threading.Lock()
187        self._settings = self.load_settings()
188
189    @classmethod
190    def shared(cls):
191        if cls._shared_instance is None:
192            cls._shared_instance = cls()
193        return cls._shared_instance
194
195    # Get a value, mockable for testing
196    def get_value(self, name: str) -> Any:
197        try:
198            return self.__getattr__(name)
199        except AttributeError:
200            return None
201
202    def __getattr__(self, name: str) -> Any:
203        if name == "_properties":
204            return super().__getattribute__("_properties")
205        if name not in self._properties:
206            return super().__getattribute__(name)
207
208        property_config = self._properties[name]
209
210        # Check if the value is in settings
211        if name in self._settings:
212            value = self._settings[name]
213            return value if value is None else property_config.type(value)
214
215        # Check environment variable
216        if property_config.env_var and property_config.env_var in os.environ:
217            value = os.environ[property_config.env_var]
218            return property_config.type(value)
219
220        # Use default value or default_lambda
221        if property_config.default_lambda:
222            value = property_config.default_lambda()
223        else:
224            value = property_config.default
225
226        return None if value is None else property_config.type(value)
227
228    def __setattr__(self, name, value):
229        if name in ("_properties", "_settings", "_lock"):
230            super().__setattr__(name, value)
231        elif name in self._properties:
232            self.update_settings({name: value})
233        else:
234            raise AttributeError(f"Config has no attribute '{name}'")
235
236    @classmethod
237    def settings_dir(cls, create=True) -> str:
238        settings_dir = os.path.join(Path.home(), ".kiln_ai")
239        if create and not os.path.exists(settings_dir):
240            os.makedirs(settings_dir)
241        return settings_dir
242
243    @classmethod
244    def settings_path(cls, create=True) -> str:
245        settings_dir = cls.settings_dir(create)
246        return os.path.join(settings_dir, "settings.yaml")
247
248    @classmethod
249    def load_settings(cls):
250        if not os.path.isfile(cls.settings_path(create=False)):
251            return {}
252        with open(cls.settings_path(), "r") as f:
253            settings = yaml.safe_load(f.read()) or {}
254        return settings
255
256    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
257        if not hide_sensitive:
258            return self._settings
259
260        settings = {
261            k: "[hidden]"
262            if k in self._properties and self._properties[k].sensitive
263            else copy.deepcopy(v)
264            for k, v in self._settings.items()
265        }
266        # 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
267        for key, value in settings.items():
268            if key in self._properties and self._properties[key].sensitive_keys:
269                sensitive_keys = self._properties[key].sensitive_keys or []
270                for sensitive_key in sensitive_keys:
271                    if isinstance(value, list):
272                        for item in value:
273                            if sensitive_key in item:
274                                item[sensitive_key] = "[hidden]"
275
276        return settings
277
278    def save_setting(self, name: str, value: Any):
279        self.update_settings({name: value})
280
281    def update_settings(self, new_settings: Dict[str, Any]):
282        # Lock to prevent race conditions in multi-threaded scenarios
283        with self._lock:
284            # Fresh load to avoid clobbering changes from other instances
285            current_settings = self.load_settings()
286            current_settings.update(new_settings)
287            # remove None values
288            current_settings = {
289                k: v for k, v in current_settings.items() if v is not None
290            }
291            with open(self.settings_path(), "w") as f:
292                yaml.dump(current_settings, f)
293            self._settings = current_settings
294
295
296def _get_user_id():
297    try:
298        return getpass.getuser() or "unknown_user"
299    except Exception:
300        return "unknown_user"