kiln_ai.utils.config

  1import getpass
  2import os
  3import threading
  4from pathlib import Path
  5from typing import Any, Callable, Dict, List, Optional
  6
  7import yaml
  8
  9# Configuration keys
 10MCP_SECRETS_KEY = "mcp_secrets"
 11
 12
 13class ConfigProperty:
 14    def __init__(
 15        self,
 16        type_: type,
 17        default: Any = None,
 18        env_var: Optional[str] = None,
 19        default_lambda: Optional[Callable[[], Any]] = None,
 20        sensitive: bool = False,
 21        sensitive_keys: Optional[List[str]] = None,
 22    ):
 23        self.type = type_
 24        self.default = default
 25        self.env_var = env_var
 26        self.default_lambda = default_lambda
 27        self.sensitive = sensitive
 28        self.sensitive_keys = sensitive_keys
 29
 30
 31class Config:
 32    _shared_instance = None
 33
 34    def __init__(self, properties: Dict[str, ConfigProperty] | None = None):
 35        self._properties: Dict[str, ConfigProperty] = properties or {
 36            "user_id": ConfigProperty(
 37                str,
 38                env_var="KILN_USER_ID",
 39                default_lambda=_get_user_id,
 40            ),
 41            "autosave_runs": ConfigProperty(
 42                bool,
 43                env_var="KILN_AUTOSAVE_RUNS",
 44                default=True,
 45            ),
 46            "open_ai_api_key": ConfigProperty(
 47                str,
 48                env_var="OPENAI_API_KEY",
 49                sensitive=True,
 50            ),
 51            "groq_api_key": ConfigProperty(
 52                str,
 53                env_var="GROQ_API_KEY",
 54                sensitive=True,
 55            ),
 56            "ollama_base_url": ConfigProperty(
 57                str,
 58                env_var="OLLAMA_BASE_URL",
 59            ),
 60            "docker_model_runner_base_url": ConfigProperty(
 61                str,
 62                env_var="DOCKER_MODEL_RUNNER_BASE_URL",
 63            ),
 64            "bedrock_access_key": ConfigProperty(
 65                str,
 66                env_var="AWS_ACCESS_KEY_ID",
 67                sensitive=True,
 68            ),
 69            "bedrock_secret_key": ConfigProperty(
 70                str,
 71                env_var="AWS_SECRET_ACCESS_KEY",
 72                sensitive=True,
 73            ),
 74            "open_router_api_key": ConfigProperty(
 75                str,
 76                env_var="OPENROUTER_API_KEY",
 77                sensitive=True,
 78            ),
 79            "fireworks_api_key": ConfigProperty(
 80                str,
 81                env_var="FIREWORKS_API_KEY",
 82                sensitive=True,
 83            ),
 84            "fireworks_account_id": ConfigProperty(
 85                str,
 86                env_var="FIREWORKS_ACCOUNT_ID",
 87            ),
 88            "anthropic_api_key": ConfigProperty(
 89                str,
 90                env_var="ANTHROPIC_API_KEY",
 91                sensitive=True,
 92            ),
 93            "gemini_api_key": ConfigProperty(
 94                str,
 95                env_var="GEMINI_API_KEY",
 96                sensitive=True,
 97            ),
 98            "projects": ConfigProperty(
 99                list,
100                default_lambda=lambda: [],
101            ),
102            "azure_openai_api_key": ConfigProperty(
103                str,
104                env_var="AZURE_OPENAI_API_KEY",
105                sensitive=True,
106            ),
107            "azure_openai_endpoint": ConfigProperty(
108                str,
109                env_var="AZURE_OPENAI_ENDPOINT",
110            ),
111            "huggingface_api_key": ConfigProperty(
112                str,
113                env_var="HUGGINGFACE_API_KEY",
114                sensitive=True,
115            ),
116            "vertex_project_id": ConfigProperty(
117                str,
118                env_var="VERTEX_PROJECT_ID",
119            ),
120            "vertex_location": ConfigProperty(
121                str,
122                env_var="VERTEX_LOCATION",
123            ),
124            "together_api_key": ConfigProperty(
125                str,
126                env_var="TOGETHERAI_API_KEY",
127                sensitive=True,
128            ),
129            "wandb_api_key": ConfigProperty(
130                str,
131                env_var="WANDB_API_KEY",
132                sensitive=True,
133            ),
134            "siliconflow_cn_api_key": ConfigProperty(
135                str,
136                env_var="SILICONFLOW_CN_API_KEY",
137                sensitive=True,
138            ),
139            "wandb_base_url": ConfigProperty(
140                str,
141                env_var="WANDB_BASE_URL",
142            ),
143            "custom_models": ConfigProperty(
144                list,
145                default_lambda=lambda: [],
146            ),
147            "openai_compatible_providers": ConfigProperty(
148                list,
149                default_lambda=lambda: [],
150                sensitive_keys=["api_key"],
151            ),
152            "cerebras_api_key": ConfigProperty(
153                str,
154                env_var="CEREBRAS_API_KEY",
155                sensitive=True,
156            ),
157            "enable_demo_tools": ConfigProperty(
158                bool,
159                env_var="ENABLE_DEMO_TOOLS",
160                default=False,
161            ),
162            # Allow the user to set the path to lookup MCP server commands, like npx.
163            "custom_mcp_path": ConfigProperty(
164                str,
165                env_var="CUSTOM_MCP_PATH",
166            ),
167            # Allow the user to set secrets for MCP servers, the key is mcp_server_id::key_name
168            MCP_SECRETS_KEY: ConfigProperty(
169                dict[str, str],
170                sensitive=True,
171            ),
172            # has the user indicated it's for personal or work use?
173            "user_type": ConfigProperty(
174                str,  # "personal" or "work"
175            ),
176            # if the user has provided their work contact
177            "work_use_contact": ConfigProperty(
178                str,
179            ),
180        }
181        self._lock = threading.Lock()
182        self._settings = self.load_settings()
183
184    @classmethod
185    def shared(cls):
186        if cls._shared_instance is None:
187            cls._shared_instance = cls()
188        return cls._shared_instance
189
190    # Get a value, mockable for testing
191    def get_value(self, name: str) -> Any:
192        try:
193            return self.__getattr__(name)
194        except AttributeError:
195            return None
196
197    def __getattr__(self, name: str) -> Any:
198        if name == "_properties":
199            return super().__getattribute__("_properties")
200        if name not in self._properties:
201            return super().__getattribute__(name)
202
203        property_config = self._properties[name]
204
205        # Check if the value is in settings
206        if name in self._settings:
207            value = self._settings[name]
208            return value if value is None else property_config.type(value)
209
210        # Check environment variable
211        if property_config.env_var and property_config.env_var in os.environ:
212            value = os.environ[property_config.env_var]
213            return property_config.type(value)
214
215        # Use default value or default_lambda
216        if property_config.default_lambda:
217            value = property_config.default_lambda()
218        else:
219            value = property_config.default
220
221        return None if value is None else property_config.type(value)
222
223    def __setattr__(self, name, value):
224        if name in ("_properties", "_settings", "_lock"):
225            super().__setattr__(name, value)
226        elif name in self._properties:
227            self.update_settings({name: value})
228        else:
229            raise AttributeError(f"Config has no attribute '{name}'")
230
231    @classmethod
232    def settings_dir(cls, create=True) -> str:
233        settings_dir = os.path.join(Path.home(), ".kiln_ai")
234        if create and not os.path.exists(settings_dir):
235            os.makedirs(settings_dir)
236        return settings_dir
237
238    @classmethod
239    def settings_path(cls, create=True) -> str:
240        settings_dir = cls.settings_dir(create)
241        return os.path.join(settings_dir, "settings.yaml")
242
243    @classmethod
244    def load_settings(cls):
245        if not os.path.isfile(cls.settings_path(create=False)):
246            return {}
247        with open(cls.settings_path(), "r") as f:
248            settings = yaml.safe_load(f.read()) or {}
249        return settings
250
251    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
252        if not hide_sensitive:
253            return self._settings
254
255        settings = {
256            k: "[hidden]"
257            if k in self._properties and self._properties[k].sensitive
258            else v
259            for k, v in self._settings.items()
260        }
261        # 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
262        for key, value in settings.items():
263            if key in self._properties and self._properties[key].sensitive_keys:
264                sensitive_keys = self._properties[key].sensitive_keys or []
265                for sensitive_key in sensitive_keys:
266                    if isinstance(value, list):
267                        for item in value:
268                            if sensitive_key in item:
269                                item[sensitive_key] = "[hidden]"
270
271        return settings
272
273    def save_setting(self, name: str, value: Any):
274        self.update_settings({name: value})
275
276    def update_settings(self, new_settings: Dict[str, Any]):
277        # Lock to prevent race conditions in multi-threaded scenarios
278        with self._lock:
279            # Fresh load to avoid clobbering changes from other instances
280            current_settings = self.load_settings()
281            current_settings.update(new_settings)
282            # remove None values
283            current_settings = {
284                k: v for k, v in current_settings.items() if v is not None
285            }
286            with open(self.settings_path(), "w") as f:
287                yaml.dump(current_settings, f)
288            self._settings = current_settings
289
290
291def _get_user_id():
292    try:
293        return getpass.getuser() or "unknown_user"
294    except Exception:
295        return "unknown_user"