fix(kb): resolve provider API key from env var over stored key

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 {}
 

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

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