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
 10class ConfigProperty:
 11    def __init__(
 12        self,
 13        type_: type,
 14        default: Any = None,
 15        env_var: Optional[str] = None,
 16        default_lambda: Optional[Callable[[], Any]] = None,
 17        sensitive: bool = False,
 18        sensitive_keys: Optional[List[str]] = None,
 19    ):
 20        self.type = type_
 21        self.default = default
 22        self.env_var = env_var
 23        self.default_lambda = default_lambda
 24        self.sensitive = sensitive
 25        self.sensitive_keys = sensitive_keys
 26
 27
 28class Config:
 29    _shared_instance = None
 30
 31    def __init__(self, properties: Dict[str, ConfigProperty] | None = None):
 32        self._properties: Dict[str, ConfigProperty] = properties or {
 33            "user_id": ConfigProperty(
 34                str,
 35                env_var="KILN_USER_ID",
 36                default_lambda=_get_user_id,
 37            ),
 38            "autosave_runs": ConfigProperty(
 39                bool,
 40                env_var="KILN_AUTOSAVE_RUNS",
 41                default=True,
 42            ),
 43            "open_ai_api_key": ConfigProperty(
 44                str,
 45                env_var="OPENAI_API_KEY",
 46                sensitive=True,
 47            ),
 48            "groq_api_key": ConfigProperty(
 49                str,
 50                env_var="GROQ_API_KEY",
 51                sensitive=True,
 52            ),
 53            "ollama_base_url": ConfigProperty(
 54                str,
 55                env_var="OLLAMA_BASE_URL",
 56            ),
 57            "bedrock_access_key": ConfigProperty(
 58                str,
 59                env_var="AWS_ACCESS_KEY_ID",
 60                sensitive=True,
 61            ),
 62            "bedrock_secret_key": ConfigProperty(
 63                str,
 64                env_var="AWS_SECRET_ACCESS_KEY",
 65                sensitive=True,
 66            ),
 67            "open_router_api_key": ConfigProperty(
 68                str,
 69                env_var="OPENROUTER_API_KEY",
 70                sensitive=True,
 71            ),
 72            "fireworks_api_key": ConfigProperty(
 73                str,
 74                env_var="FIREWORKS_API_KEY",
 75                sensitive=True,
 76            ),
 77            "fireworks_account_id": ConfigProperty(
 78                str,
 79                env_var="FIREWORKS_ACCOUNT_ID",
 80            ),
 81            "projects": ConfigProperty(
 82                list,
 83                default_lambda=lambda: [],
 84            ),
 85            "custom_models": ConfigProperty(
 86                list,
 87                default_lambda=lambda: [],
 88            ),
 89            "openai_compatible_providers": ConfigProperty(
 90                list,
 91                default_lambda=lambda: [],
 92                sensitive_keys=["api_key"],
 93            ),
 94        }
 95        self._settings = self.load_settings()
 96
 97    @classmethod
 98    def shared(cls):
 99        if cls._shared_instance is None:
100            cls._shared_instance = cls()
101        return cls._shared_instance
102
103    # Get a value, mockable for testing
104    def get_value(self, name: str) -> Any:
105        try:
106            return self.__getattr__(name)
107        except AttributeError:
108            return None
109
110    def __getattr__(self, name: str) -> Any:
111        if name == "_properties":
112            return super().__getattribute__("_properties")
113        if name not in self._properties:
114            return super().__getattribute__(name)
115
116        property_config = self._properties[name]
117
118        # Check if the value is in settings
119        if name in self._settings:
120            value = self._settings[name]
121            return value if value is None else property_config.type(value)
122
123        # Check environment variable
124        if property_config.env_var and property_config.env_var in os.environ:
125            value = os.environ[property_config.env_var]
126            return property_config.type(value)
127
128        # Use default value or default_lambda
129        if property_config.default_lambda:
130            value = property_config.default_lambda()
131        else:
132            value = property_config.default
133
134        return None if value is None else property_config.type(value)
135
136    def __setattr__(self, name, value):
137        if name in ("_properties", "_settings"):
138            super().__setattr__(name, value)
139        elif name in self._properties:
140            self.update_settings({name: value})
141        else:
142            raise AttributeError(f"Config has no attribute '{name}'")
143
144    @classmethod
145    def settings_path(cls, create=True):
146        settings_dir = os.path.join(Path.home(), ".kiln_ai")
147        if create and not os.path.exists(settings_dir):
148            os.makedirs(settings_dir)
149        return os.path.join(settings_dir, "settings.yaml")
150
151    @classmethod
152    def load_settings(cls):
153        if not os.path.isfile(cls.settings_path(create=False)):
154            return {}
155        with open(cls.settings_path(), "r") as f:
156            settings = yaml.safe_load(f.read()) or {}
157        return settings
158
159    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
160        if not hide_sensitive:
161            return self._settings
162
163        settings = {
164            k: "[hidden]"
165            if k in self._properties and self._properties[k].sensitive
166            else v
167            for k, v in self._settings.items()
168        }
169        # 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
170        for key, value in settings.items():
171            if key in self._properties and self._properties[key].sensitive_keys:
172                sensitive_keys = self._properties[key].sensitive_keys or []
173                for sensitive_key in sensitive_keys:
174                    if isinstance(value, list):
175                        for item in value:
176                            if sensitive_key in item:
177                                item[sensitive_key] = "[hidden]"
178
179        return settings
180
181    def save_setting(self, name: str, value: Any):
182        self.update_settings({name: value})
183
184    def update_settings(self, new_settings: Dict[str, Any]):
185        # Lock to prevent race conditions in multi-threaded scenarios
186        with threading.Lock():
187            # Fresh load to avoid clobbering changes from other instances
188            current_settings = self.load_settings()
189            current_settings.update(new_settings)
190            # remove None values
191            current_settings = {
192                k: v for k, v in current_settings.items() if v is not None
193            }
194            with open(self.settings_path(), "w") as f:
195                yaml.dump(current_settings, f)
196            self._settings = current_settings
197
198
199def _get_user_id():
200    try:
201        return getpass.getuser() or "unknown_user"
202    except Exception:
203        return "unknown_user"