From 848bcf1b3d1ac0eefff6e6ec63b8dff5cb201bf0 Mon Sep 17 00:00:00 2001 From: Gabriel Bressan Date: Tue, 7 Jul 2026 15:47:35 -0400 Subject: [PATCH] fix(kb): resolve provider API key from env var over stored key Knowledge bases persisted the OpenAI/Gemini API key at creation time (in KB params, reranking params, and the embedding model's args). When keys are rotated, every KB kept using the stale stored key and failed with 401 Incorrect API key. Make the environment variable take precedence over any stored key on all KB key-resolution paths, so a rotated key can be updated in one place without recreating knowledge bases: - llm_client.py: embedding model prefers {PROVIDER}_API_KEY env var (gemini also honors GOOGLE_API_KEY) before the stored api_key. - controller.py: reranking model uses the same env-first resolution; add missing `import os`. - openai_handler.py: add `_resolve_api_key()` (env-first, keyed off api_key_name so provider subclasses stay correct) and route the predict/describe/validate/finetune sites through it. Follow-up (discussed, not in this PR): stop persisting the secret at creation and scrub existing rows so env becomes the sole source of truth and these overrides can be removed. Co-Authored-By: Claude Opus 4.8 --- .../handlers/openai_handler/openai_handler.py | 19 +++++++++++++++---- .../interfaces/knowledge_base/controller.py | 13 +++++++++++-- .../interfaces/knowledge_base/llm_client.py | 12 +++++++++++- 3 files changed, 37 insertions(+), 7 deletions(-) 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