fix whitespace

Author: danwaters

src/together/client.py

@@ -9,6 +9,7 @@
 from together.error import AuthenticationError
 from together.types import TogetherClient
 from together.utils import enforce_trailing_slash
+from together.utils.api_helpers import get_google_colab_secret
 
 
 class Together:
@@ -45,22 +46,8 @@ def __init__(
         if not api_key:
             api_key = os.environ.get("TOGETHER_API_KEY")
 
-        # If running in Google Colab, check notebook secrets
         if not api_key and "google.colab" in sys.modules:
-            if TYPE_CHECKING:
-                from google.colab import userdata  # type: ignore
-            else:
-                from google.colab import userdata
-            try:
-                api_key = userdata.get("TOGETHER_API_KEY")
-            except userdata.NotebookAccessError:
-                print(
-                    "The TOGETHER_API_KEY Colab secret was found, but notebook access is disabled. Please enable notebook "
-                    "access for the secret."
-                )
-            except userdata.SecretNotFoundError:
-                # warn and carry on
-                print("Colab: No Google Colab secret named TOGETHER_API_KEY was found.")
+            api_key = get_google_colab_secret("TOGETHER_API_KEY")
 
         if not api_key:
             raise AuthenticationError(
@@ -135,6 +122,9 @@ def __init__(
         if not api_key:
             api_key = os.environ.get("TOGETHER_API_KEY")
 
+        if not api_key and "google.colab" in sys.modules:
+            api_key = get_google_colab_secret("TOGETHER_API_KEY")
+
         if not api_key:
             raise AuthenticationError(
                 "The api_key client option must be set either by passing api_key to the client or by setting the "

src/together/utils/api_helpers.py

@@ -2,6 +2,7 @@
 
 import json
 import os
+import sys
 import platform
 from typing import TYPE_CHECKING, Any, Dict
 
@@ -12,6 +13,7 @@
 import together
 from together import error
 from together.utils._log import _console_log_level
+from together.utils import log_info
 
 
 def get_headers(
@@ -82,3 +84,41 @@ def default_api_key(api_key: str | None = None) -> str | None:
         return os.environ.get("TOGETHER_API_KEY")
 
     raise error.AuthenticationError(together.constants.MISSING_API_KEY_MESSAGE)
+
+
+def get_google_colab_secret(secret_name: str = "TOGETHER_API_KEY") -> str | None:
+    """
+    Checks to see if the user is running in Google Colab, and looks for the Together API Key secret.
+
+    Args:
+        secret_name (str, optional). Defaults to TOGETHER_API_KEY
+
+    Returns:
+        str: if the API key is found; None if an error occurred or the secret was not found.
+    """
+    # If running in Google Colab, check for Together in notebook secrets
+    if "google.colab" in sys.modules:
+        if TYPE_CHECKING:
+            from google.colab import userdata  # type: ignore
+        else:
+            from google.colab import userdata
+
+        try:
+            api_key = userdata.get(secret_name)
+            if not isinstance(api_key, str):
+                return None
+            else:
+                return str(api_key)
+        except userdata.NotebookAccessError:
+            log_info(
+                "The TOGETHER_API_KEY Colab secret was found, but notebook access is disabled. Please enable notebook "
+                "access for the secret."
+            )
+        except userdata.SecretNotFoundError:
+            # warn and carry on
+            log_info("Colab: No Google Colab secret named TOGETHER_API_KEY was found.")
+
+        return None
+
+    else:
+        return None