diff --git a/mindsdb/integrations/handlers/openai_handler/openai_handler.py b/mindsdb/integrations/handlers/openai_handler/openai_handler.py index 54aa4d7b34b..33c38658658 100644 --- a/mindsdb/integrations/handlers/openai_handler/openai_handler.py +++ b/mindsdb/integrations/handlers/openai_handler/openai_handler.py @@ -74,6 +74,17 @@ def __init__(self, *args, **kwargs): self.api_key_name = getattr(self, "api_key_name", self.name) self.api_base = getattr(self, "api_base", OPENAI_API_BASE) + def _resolve_api_key(self, args: dict) -> Text: + """Resolve the API key with the environment variable taking precedence over + any key stored in the model args. This lets a rotated key be updated in one + place (the env) without recreating models/knowledge bases that stored the old + key. Falls back to the standard resolution chain when no env var is set. + """ + env_api_key = os.getenv(f"{self.api_key_name.upper()}_API_KEY") or os.getenv(f"{self.api_key_name.lower()}_api_key") + if env_api_key: + return env_api_key + return get_api_key(self.api_key_name, args, self.engine_storage) + def create_engine(self, connection_args: Dict) -> None: """ Validate the OpenAI API credentials on engine creation. @@ -231,7 +242,7 @@ def create(self, target, args: Dict = None, **kwargs: Any) -> None: args = args["using"] args["target"] = target try: - api_key = get_api_key(self.api_key_name, args, self.engine_storage) + api_key = self._resolve_api_key(args) connection_args = self.engine_storage.get_connection_args() api_base = ( args.get("api_base") @@ -434,7 +445,7 @@ def predict(self, df: pd.DataFrame, args: Optional[Dict] = None) -> pd.DataFrame # remove prompts without signal from completion queue prompts = [j for i, j in enumerate(prompts) if i not in empty_prompt_ids] - api_key = get_api_key(self.api_key_name, args, self.engine_storage) + api_key = self._resolve_api_key(args) api_args = {k: v for k, v in api_args.items() if v is not None} # filter out non-specified api args completion = self._completion(model_name, prompts, api_key, api_args, args, df) @@ -851,7 +862,7 @@ def describe(self, attribute: Optional[Text] = None) -> pd.DataFrame: # TODO: Update to use update() artifacts args = self.model_storage.json_get("args") - api_key = get_api_key(self.api_key_name, args, self.engine_storage) + api_key = self._resolve_api_key(args) if attribute == "args": return pd.DataFrame(args.items(), columns=["key", "value"]) elif attribute == "metadata": @@ -899,7 +910,7 @@ def finetune(self, df: Optional[pd.DataFrame] = None, args: Optional[Dict] = Non """ # noqa args = args if args else {} - api_key = get_api_key(self.api_key_name, args, self.engine_storage) + api_key = self._resolve_api_key(args) using_args = args.pop("using") if "using" in args else {} diff --git a/mindsdb/interfaces/knowledge_base/controller.py b/mindsdb/interfaces/knowledge_base/controller.py index d158917f72c..d134ef6cbd1 100644 --- a/mindsdb/interfaces/knowledge_base/controller.py +++ b/mindsdb/interfaces/knowledge_base/controller.py @@ -1,4 +1,5 @@ import copy +import os from typing import Dict, List, Optional, Any, Text, Tuple, Union import json import decimal @@ -91,9 +92,17 @@ def get_reranking_model_from_params(reranking_model_params: dict): # Work on a copy; do not mutate caller's dict params_copy = copy.deepcopy(reranking_model_params) - # Handle API key if not provided + # Handle API key. The environment variable takes precedence over any api_key + # stored in the reranking_model params, so a rotated key can be updated in one + # place (the env) without recreating knowledge bases that stored the old key. provider = params_copy.get("provider", "openai").lower() - if "api_key" not in params_copy: + env_var_names = [f"{provider.upper()}_API_KEY", f"{provider}_api_key"] + if provider in ("gemini", "google"): + env_var_names += ["GEMINI_API_KEY", "GOOGLE_API_KEY"] + env_api_key = next((os.getenv(name) for name in env_var_names if os.getenv(name)), None) + if env_api_key: + params_copy["api_key"] = env_api_key + elif "api_key" not in params_copy: params_copy["api_key"] = get_api_key(provider, params_copy, strict=False) # Handle model_name -> model alias for backward compatibility diff --git a/mindsdb/interfaces/knowledge_base/llm_client.py b/mindsdb/interfaces/knowledge_base/llm_client.py index 61591b62c17..6cd8688fe31 100644 --- a/mindsdb/interfaces/knowledge_base/llm_client.py +++ b/mindsdb/interfaces/knowledge_base/llm_client.py @@ -76,7 +76,17 @@ def __init__(self, params: dict = None, session=None): if self.provider == "google": self.provider = "gemini" - if "api_key" not in params: + # The environment variable takes precedence over any api_key stored in the + # knowledge base params. This lets a rotated key be updated in one place (the + # env) without recreating every knowledge base that stored the old key. + env_var_names = [f"{self.provider.upper()}_API_KEY", f"{self.provider.lower()}_api_key"] + if self.provider == "gemini": + # provider is normalized from "google" above; the conventional env var is GOOGLE_API_KEY + env_var_names += ["GOOGLE_API_KEY", "google_api_key"] + env_api_key = next((os.getenv(name) for name in env_var_names if os.getenv(name)), None) + if env_api_key: + params["api_key"] = env_api_key + elif "api_key" not in params: api_key = get_api_key(self.provider, params, strict=False) if api_key is not None: params["api_key"] = api_key