Set OpenAI API key on _OAITokenHolder init (mlflow#10095)

Signed-off-by: Daniel Lok <[email protected]>

Author: daniellok-db

mlflow/openai/utils.py

@@ -158,7 +158,13 @@ def __init__(self, api_type):
         self._api_token = None
         self._credential = None
         self._is_azure_ad = api_type in ("azure_ad", "azuread")
-        self._key_configured = bool(openai.api_key) or "OPENAI_API_KEY" in os.environ
+        self._key_configured = bool(openai.api_key)
+
+        # set the api key if it's not set. this is to deal with cases where the
+        # user sets the environment variable after importing the `openai` module
+        if not bool(openai.api_key) and "OPENAI_API_KEY" in os.environ:
+            openai.api_key = os.environ["OPENAI_API_KEY"]
+            self._key_configured = True
 
         if self._is_azure_ad and not self._key_configured:
             try:

tests/openai/test_openai_init.py

@@ -0,0 +1,26 @@
+import os
+from importlib import reload
+
+from mlflow.openai import _OAITokenHolder
+
+
+def test_set_api_key_on_tokenholder_init(monkeypatch):
+    # if the user sets the API key after the openai module,
+    # expect `openai.api_key` to not be set.
+    monkeypatch.delenv("OPENAI_API_KEY", False)
+    assert "OPENAI_API_KEY" not in os.environ
+
+    import openai
+
+    monkeypatch.setenv("OPENAI_API_KEY", "test-key")
+    assert openai.api_key is None
+
+    # when OAITokenHolder is initialized, expect it to set `openai.api_key`
+    token_holder = _OAITokenHolder("open_ai")
+    assert openai.api_key == "test-key"
+    assert token_holder._key_configured
+
+    # reload the module to simulate the env var being set before
+    # load. in this case we'd expect the API key to be present
+    reload(openai)
+    assert openai.api_key == "test-key"