Repository: lastmile-ai/aiconfig (language: python)

Issue:
Make `get_api_key_from_environment` return nullable
Make `get_api_key_from_environment` return nullable


Sometimes key is not defined or set, we need to do this to unblock https://github.com/lastmile-ai/aiconfig/pull/769

The functionality is still unchanged because the optional param `required` defaults to true

## Test Plan
Pass all automated tests (which it didn't before)

---
Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/lastmile-ai/aiconfig/pull/772).
* #769
* __->__ #772


Discussion / hints:

Fix patch:
diff --git a/extensions/Gemini/python/src/aiconfig_extension_gemini/Gemini.py b/extensions/Gemini/python/src/aiconfig_extension_gemini/Gemini.py
index 9766ee4d8..8927ee5f6 100644
--- a/extensions/Gemini/python/src/aiconfig_extension_gemini/Gemini.py
+++ b/extensions/Gemini/python/src/aiconfig_extension_gemini/Gemini.py
@@ -1,10 +1,8 @@
 # Define a Model Parser for LLama-Guard
 from typing import TYPE_CHECKING, Dict, List, Optional, Any
 import copy
-import json
 
 import google.generativeai as genai
-from google.generativeai.types import content_types
 from google.protobuf.json_format import MessageToDict
 
 from aiconfig import (
@@ -108,6 +106,9 @@ class GeminiModelParser(ParameterizedModelParser):
     def __init__(self, id: str = "gemini-pro"):
         super().__init__()
         self.model = genai.GenerativeModel(id)
+        # TODO: Define API key here so we only call get_api_key_from_environment once
+        # Pass it in explicit to genai.configure(api_key=self.api_key), similar to 
+        # https://github.com/zhayujie/chatgpt-on-wechat/blob/eb809055d49efcc2036f7d89088aea7b50097162/bot/gemini/google_gemini_bot.py#L37
 
     def id(self) -> str:
         """
diff --git a/python/src/aiconfig/util/config_utils.py b/python/src/aiconfig/util/config_utils.py
index 0b2eb6eb1..8bf9111dd 100644
--- a/python/src/aiconfig/util/config_utils.py
+++ b/python/src/aiconfig/util/config_utils.py
@@ -1,6 +1,7 @@
 import copy
+import dotenv
 import os
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
 
 if TYPE_CHECKING:
     pass
@@ -10,10 +11,28 @@
     from ..schema import AIConfig
 
 
-def get_api_key_from_environment(api_key_name: str):
-    if api_key_name not in os.environ:
-        raise Exception("Missing API key '{}' in environment".format(api_key_name))
+def get_api_key_from_environment(
+    api_key_name: str,
+    required: bool = True) -> Union[str, None]:
+    """Get the API key if it exists, return None or error if it doesn't
+
+    Args:
+        api_key_name (str): The keyname that we're trying to import from env variable
+        required (bool, optional): If this is true, we raise an error if the 
+            key is not found
+
+    Returns:
+        Union[str, None]: the value of the key. If `required` is false, this can be None
+    """
+    dotenv.load_dotenv()
+    if required:
+        _get_api_key_from_environment_required(api_key_name)
+    return os.getenv(api_key_name)
 
+
+def _get_api_key_from_environment_required(api_key_name: str) -> str:
+    if api_key_name not in os.environ:
+        raise KeyError(f"Missing API key '{api_key_name}' in environment")
     return os.environ[api_key_name]
 
 


Test patch:
diff --git a/python/tests/test_util/test_config_util.py b/python/tests/test_util/test_config_util.py
index a831d4fba..260dcc2b4 100644
--- a/python/tests/test_util/test_config_util.py
+++ b/python/tests/test_util/test_config_util.py
@@ -1,7 +1,7 @@
 from aiconfig.util.config_utils import get_api_key_from_environment
 
 
-def test_get_api_key_from_environment():
+def test_maybe_get_api_key_from_environment():
     key = "TEST_API_KEY"
     try:
         get_api_key_from_environment(key)