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        }
173        self._lock = threading.Lock()
174        self._settings = self.load_settings()
175
176    @classmethod
177    def shared(cls):
178        if cls._shared_instance is None:
179            cls._shared_instance = cls()
180        return cls._shared_instance
181
182    # Get a value, mockable for testing
183    def get_value(self, name: str) -> Any:
184        try:
185            return self.__getattr__(name)
186        except AttributeError:
187            return None
188
189    def __getattr__(self, name: str) -> Any:
190        if name == "_properties":
191            return super().__getattribute__("_properties")
192        if name not in self._properties:
193            return super().__getattribute__(name)
194
195        property_config = self._properties[name]
196
197        # Check if the value is in settings
198        if name in self._settings:
199            value = self._settings[name]
200            return value if value is None else property_config.type(value)
201
202        # Check environment variable
203        if property_config.env_var and property_config.env_var in os.environ:
204            value = os.environ[property_config.env_var]
205            return property_config.type(value)
206
207        # Use default value or default_lambda
208        if property_config.default_lambda:
209            value = property_config.default_lambda()
210        else:
211            value = property_config.default
212
213        return None if value is None else property_config.type(value)
214
215    def __setattr__(self, name, value):
216        if name in ("_properties", "_settings", "_lock"):
217            super().__setattr__(name, value)
218        elif name in self._properties:
219            self.update_settings({name: value})
220        else:
221            raise AttributeError(f"Config has no attribute '{name}'")
222
223    @classmethod
224    def settings_dir(cls, create=True):
225        settings_dir = os.path.join(Path.home(), ".kiln_ai")
226        if create and not os.path.exists(settings_dir):
227            os.makedirs(settings_dir)
228        return settings_dir
229
230    @classmethod
231    def settings_path(cls, create=True):
232        settings_dir = cls.settings_dir(create)
233        return os.path.join(settings_dir, "settings.yaml")
234
235    @classmethod
236    def load_settings(cls):
237        if not os.path.isfile(cls.settings_path(create=False)):
238            return {}
239        with open(cls.settings_path(), "r") as f:
240            settings = yaml.safe_load(f.read()) or {}
241        return settings
242
243    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
244        if not hide_sensitive:
245            return self._settings
246
247        settings = {
248            k: "[hidden]"
249            if k in self._properties and self._properties[k].sensitive
250            else v
251            for k, v in self._settings.items()
252        }
253        # 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
254        for key, value in settings.items():
255            if key in self._properties and self._properties[key].sensitive_keys:
256                sensitive_keys = self._properties[key].sensitive_keys or []
257                for sensitive_key in sensitive_keys:
258                    if isinstance(value, list):
259                        for item in value:
260                            if sensitive_key in item:
261                                item[sensitive_key] = "[hidden]"
262
263        return settings
264
265    def save_setting(self, name: str, value: Any):
266        self.update_settings({name: value})
267
268    def update_settings(self, new_settings: Dict[str, Any]):
269        # Lock to prevent race conditions in multi-threaded scenarios
270        with self._lock:
271            # Fresh load to avoid clobbering changes from other instances
272            current_settings = self.load_settings()
273            current_settings.update(new_settings)
274            # remove None values
275            current_settings = {
276                k: v for k, v in current_settings.items() if v is not None
277            }
278            with open(self.settings_path(), "w") as f:
279                yaml.dump(current_settings, f)
280            self._settings = current_settings
281
282
283def _get_user_id():
284    try:
285        return getpass.getuser() or "unknown_user"
286    except Exception:
287        return "unknown_user"