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