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._settings = self.load_settings()
142
143    @classmethod
144    def shared(cls):
145        if cls._shared_instance is None:
146            cls._shared_instance = cls()
147        return cls._shared_instance
148
149    # Get a value, mockable for testing
150    def get_value(self, name: str) -> Any:
151        try:
152            return self.__getattr__(name)
153        except AttributeError:
154            return None
155
156    def __getattr__(self, name: str) -> Any:
157        if name == "_properties":
158            return super().__getattribute__("_properties")
159        if name not in self._properties:
160            return super().__getattribute__(name)
161
162        property_config = self._properties[name]
163
164        # Check if the value is in settings
165        if name in self._settings:
166            value = self._settings[name]
167            return value if value is None else property_config.type(value)
168
169        # Check environment variable
170        if property_config.env_var and property_config.env_var in os.environ:
171            value = os.environ[property_config.env_var]
172            return property_config.type(value)
173
174        # Use default value or default_lambda
175        if property_config.default_lambda:
176            value = property_config.default_lambda()
177        else:
178            value = property_config.default
179
180        return None if value is None else property_config.type(value)
181
182    def __setattr__(self, name, value):
183        if name in ("_properties", "_settings"):
184            super().__setattr__(name, value)
185        elif name in self._properties:
186            self.update_settings({name: value})
187        else:
188            raise AttributeError(f"Config has no attribute '{name}'")
189
190    @classmethod
191    def settings_dir(cls, create=True):
192        settings_dir = os.path.join(Path.home(), ".kiln_ai")
193        if create and not os.path.exists(settings_dir):
194            os.makedirs(settings_dir)
195        return settings_dir
196
197    @classmethod
198    def settings_path(cls, create=True):
199        settings_dir = cls.settings_dir(create)
200        return os.path.join(settings_dir, "settings.yaml")
201
202    @classmethod
203    def load_settings(cls):
204        if not os.path.isfile(cls.settings_path(create=False)):
205            return {}
206        with open(cls.settings_path(), "r") as f:
207            settings = yaml.safe_load(f.read()) or {}
208        return settings
209
210    def settings(self, hide_sensitive=False) -> Dict[str, Any]:
211        if not hide_sensitive:
212            return self._settings
213
214        settings = {
215            k: "[hidden]"
216            if k in self._properties and self._properties[k].sensitive
217            else v
218            for k, v in self._settings.items()
219        }
220        # 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
221        for key, value in settings.items():
222            if key in self._properties and self._properties[key].sensitive_keys:
223                sensitive_keys = self._properties[key].sensitive_keys or []
224                for sensitive_key in sensitive_keys:
225                    if isinstance(value, list):
226                        for item in value:
227                            if sensitive_key in item:
228                                item[sensitive_key] = "[hidden]"
229
230        return settings
231
232    def save_setting(self, name: str, value: Any):
233        self.update_settings({name: value})
234
235    def update_settings(self, new_settings: Dict[str, Any]):
236        # Lock to prevent race conditions in multi-threaded scenarios
237        with threading.Lock():
238            # Fresh load to avoid clobbering changes from other instances
239            current_settings = self.load_settings()
240            current_settings.update(new_settings)
241            # remove None values
242            current_settings = {
243                k: v for k, v in current_settings.items() if v is not None
244            }
245            with open(self.settings_path(), "w") as f:
246                yaml.dump(current_settings, f)
247            self._settings = current_settings
248
249
250def _get_user_id():
251    try:
252        return getpass.getuser() or "unknown_user"
253    except Exception:
254        return "unknown_user"