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            "anthropic_api_key": ConfigProperty(
 82                str,
 83                env_var="ANTHROPIC_API_KEY",
 84                sensitive=True,
 85            ),
 86            "gemini_api_key": ConfigProperty(
 87                str,
 88                env_var="GEMINI_API_KEY",
 89                sensitive=True,
 90            ),
 91            "projects": ConfigProperty(
 92                list,
 93                default_lambda=lambda: [],
 94            ),
 95            "azure_openai_api_key": ConfigProperty(
 96                str,
 97                env_var="AZURE_OPENAI_API_KEY",
 98                sensitive=True,
 99            ),
100            "azure_openai_endpoint": ConfigProperty(
101                str,
102                env_var="AZURE_OPENAI_ENDPOINT",
103            ),
104            "huggingface_api_key": ConfigProperty(
105                str,
106                env_var="HUGGINGFACE_API_KEY",
107                sensitive=True,
108            ),
109            "vertex_project_id": ConfigProperty(
110                str,
111                env_var="VERTEX_PROJECT_ID",
112            ),
113            "vertex_location": ConfigProperty(
114                str,
115                env_var="VERTEX_LOCATION",
116            ),
117            "together_api_key": ConfigProperty(
118                str,
119                env_var="TOGETHERAI_API_KEY",
120                sensitive=True,
121            ),
122            "wandb_api_key": ConfigProperty(
123                str,
124                env_var="WANDB_API_KEY",
125                sensitive=True,
126            ),
127            "wandb_base_url": ConfigProperty(
128                str,
129                env_var="WANDB_BASE_URL",
130            ),
131            "custom_models": ConfigProperty(
132                list,
133                default_lambda=lambda: [],
134            ),
135            "openai_compatible_providers": ConfigProperty(
136                list,
137                default_lambda=lambda: [],
138                sensitive_keys=["api_key"],
139            ),
140        }
141        self._lock = threading.Lock()
142        self._settings = self.load_settings()
143
144    @classmethod
145    def shared(cls):
146        if cls._shared_instance is None:
147            cls._shared_instance = cls()
148        return cls._shared_instance
149
150    # Get a value, mockable for testing
151    def get_value(self, name: str) -> Any:
152        try:
153            return self.__getattr__(name)
154        except AttributeError:
155            return None
156
157    def __getattr__(self, name: str) -> Any:
158        if name == "_properties":
159            return super().__getattribute__("_properties")
160        if name not in self._properties:
161            return super().__getattribute__(name)
162
163        property_config = self._properties[name]
164
165        # Check if the value is in settings
166        if name in self._settings:
167            value = self._settings[name]
168            return value if value is None else property_config.type(value)
169
170        # Check environment variable
171        if property_config.env_var and property_config.env_var in os.environ:
172            value = os.environ[property_config.env_var]
173            return property_config.type(value)
174
175        # Use default value or default_lambda
176        if property_config.default_lambda:
177            value = property_config.default_lambda()
178        else:
179            value = property_config.default
180
181        return None if value is None else property_config.type(value)
182
183    def __setattr__(self, name, value):
184        if name in ("_properties", "_settings", "_lock"):
185            super().__setattr__(name, value)
186        elif name in self._properties:
187            self.update_settings({name: value})
188        else:
189            raise AttributeError(f"Config has no attribute '{name}'")
190
191    @classmethod
192    def settings_dir(cls, create=True):
193        settings_dir = os.path.join(Path.home(), ".kiln_ai")
194        if create and not os.path.exists(settings_dir):
195            os.makedirs(settings_dir)
196        return settings_dir
197
198    @classmethod
199    def settings_path(cls, create=True):
200        settings_dir = cls.settings_dir(create)
201        return os.path.join(settings_dir, "settings.yaml")
202
203    @classmethod
204    def load_settings(cls):
205        if not os.path.isfile(cls.settings_path(create=False)):
206            return {}
207        with open(cls.settings_path(), "r") as f:
208            settings = yaml.safe_load(f.read()) or {}
209        return settings
210
211    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
212        if not hide_sensitive:
213            return self._settings
214
215        settings = {
216            k: "[hidden]"
217            if k in self._properties and self._properties[k].sensitive
218            else v
219            for k, v in self._settings.items()
220        }
221        # 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
222        for key, value in settings.items():
223            if key in self._properties and self._properties[key].sensitive_keys:
224                sensitive_keys = self._properties[key].sensitive_keys or []
225                for sensitive_key in sensitive_keys:
226                    if isinstance(value, list):
227                        for item in value:
228                            if sensitive_key in item:
229                                item[sensitive_key] = "[hidden]"
230
231        return settings
232
233    def save_setting(self, name: str, value: Any):
234        self.update_settings({name: value})
235
236    def update_settings(self, new_settings: Dict[str, Any]):
237        # Lock to prevent race conditions in multi-threaded scenarios
238        with self._lock:
239            # Fresh load to avoid clobbering changes from other instances
240            current_settings = self.load_settings()
241            current_settings.update(new_settings)
242            # remove None values
243            current_settings = {
244                k: v for k, v in current_settings.items() if v is not None
245            }
246            with open(self.settings_path(), "w") as f:
247                yaml.dump(current_settings, f)
248            self._settings = current_settings
249
250
251def _get_user_id():
252    try:
253        return getpass.getuser() or "unknown_user"
254    except Exception:
255        return "unknown_user"