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

Issue:
Make `get_api_key_from_environment` use Result-Ok
Make `get_api_key_from_environment` use Result-Ok


Nice, cool tips from Jonathan

Also deleted `test_config_util.py` because it wasn't being referenced anywhere

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


Discussion / hints:

Fix patch:
diff --git a/cookbooks/Gemini/Gemini.ipynb b/cookbooks/Gemini/Gemini.ipynb
index e5ed0a85c..ebdbc9e16 100644
--- a/cookbooks/Gemini/Gemini.ipynb
+++ b/cookbooks/Gemini/Gemini.ipynb
@@ -30,7 +30,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 1,
       "metadata": {
         "id": "Ppc9RiUGg_7c"
       },
@@ -54,19 +54,35 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 2,
       "metadata": {
         "id": "wLS1XJvyosT7"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "data": {
+            "text/plain": [
+              "True"
+            ]
+          },
+          "execution_count": 2,
+          "metadata": {},
+          "output_type": "execute_result"
+        }
+      ],
       "source": [
         "# Environment Variable (with dotenv)\n",
         "import dotenv\n",
         "import os\n",
         "dotenv.load_dotenv()\n",
         "\n",
-        "# Retrieve and set your Google API Key - https://ai.google.dev/tutorials/setup \n",
-        "# os.environ[\"GOOGLE_API_KEY\"] = \"your key\""
+        "# Create .env file at aiconfig/.env containing the following line:\n",
+        "# GOOGLE_API_KEY=<your key here>\n",
+        "# You can get your key from https://ai.google.dev/tutorials/setup \n",
+        "import openai\n",
+        "import dotenv\n",
+        "import os\n",
+        "dotenv.load_dotenv()"
       ]
     },
     {
@@ -85,11 +101,24 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 3,
       "metadata": {
         "id": "W_UGROg7jlmP"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stderr",
+          "output_type": "stream",
+          "text": [
+            "/Users/rossdancraig/.pyenv/versions/3.11.6/lib/python3.11/site-packages/pydantic/_internal/_fields.py:128: UserWarning: Field \"model_parsers\" has conflict with protected namespace \"model_\".\n",
+            "\n",
+            "You may be able to resolve this warning by setting `model_config['protected_namespaces'] = ()`.\n",
+            "  warnings.warn(\n",
+            "/Users/rossdancraig/.pyenv/versions/3.11.6/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
+            "  from .autonotebook import tqdm as notebook_tqdm\n"
+          ]
+        }
+      ],
       "source": [
         "from aiconfig import AIConfigRuntime\n",
         "import aiconfig_extension_gemini\n",
@@ -152,7 +181,7 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 4,
       "metadata": {
         "id": "YFBE95DUjBvU"
       },
@@ -168,11 +197,58 @@
     },
     {
       "cell_type": "code",
-      "execution_count": null,
+      "execution_count": 5,
       "metadata": {
         "id": "7-ci5weaIlnY"
       },
-      "outputs": [],
+      "outputs": [
+        {
+          "name": "stdout",
+          "output_type": "stream",
+          "text": [
+            "Gemini output: - Spanish: pato\n",
+            "- French: canard\n",
+            "- German: Ente\n",
+            "- Italian: anatra\n",
+            "- Portuguese: pato\n",
+            "- Russian: yтка\n",
+            "- Korean: 오리\n",
+            "- Japanese: 鴨 (かも)\n",
+            "- Chinese: 鸭子 (yāzi)\n",
+            "- Hindi: बत्तख (baṭṭakh)\n",
+            "- Arabic: بطة (baṭṭah) \n",
+            "\n",
+            "Gemini output: The Mandarin word for duck is 鸭子 (yāzi). It is pronounced in two syllables:\n",
+            "\n",
+            "1. **yā (鸭)**:\n",
+            "\n",
+            "   - Start by saying \"ya\" as in \"yard.\"\n",
+            "   - Then, raise the tone of your voice slightly and say \"ah\" as in \"father.\"\n",
+            "   - The resulting sound is a rising tone.\n",
+            "\n",
+            "2. **zi (子)**:\n",
+            "\n",
+            "   - Start by saying \"dz\" as in \"adze.\"\n",
+            "   - Then, immediately follow it with an \"ih\" sound as in \"sit.\"\n",
+            "   - The resulting sound is a neutral tone.\n",
+            "\n",
+            "To pronounce 鸭子 (yāzi) correctly, connect the two syllables smoothly and say them with the correct tones. Here's a step-by-step guide:\n",
+            "\n",
+            "1. Inhale deeply and open your mouth slightly.\n",
+            "2. Start with a rising tone on the first syllable \"yā.\"\n",
+            "3. Transition smoothly to a neutral tone on the second syllable \"zi.\"\n",
+            "4. Keep your tongue relaxed and let the air flow out naturally.\n",
+            "5. The overall pronunciation should sound like \"yah-dzih.\"\n",
+            "\n",
+            "Here are some tips for practicing the pronunciation:\n",
+            "\n",
+            "- Listen to native speakers pronouncing 鸭子 (yāzi) online or in person.\n",
+            "- Record yourself saying 鸭子 (yāzi) and compare it to the pronunciation of native speakers.\n",
+            "- Practice saying 鸭子 (yāzi) in different contexts, such as \"我要吃鸭子\" (wǒ yào chī yāzi, I want to eat duck) or \"鸭子在水里游泳\" (yāzi zài shuǐ lǐ yóuyǒng, the duck is swimming in the water).\n",
+            "- Be patient and keep practicing. With time and effort, you will be able to pronounce 鸭子 (yāzi) like a native speaker.\n"
+          ]
+        }
+      ],
       "source": [
         "# Config.run() handles the complex things for you\n",
         "# Run multiple different prompts and pass in parameters to specify variables to pass into the prompt\n",
@@ -264,4 +340,4 @@
   },
   "nbformat": 4,
   "nbformat_minor": 0
-}
\ No newline at end of file
+}
diff --git a/cookbooks/HuggingFace/hf.py b/cookbooks/HuggingFace/hf.py
index 8257b326e..b1b365753 100644
--- a/cookbooks/HuggingFace/hf.py
+++ b/cookbooks/HuggingFace/hf.py
@@ -153,7 +153,9 @@ def __init__(self, model_id: str = None, use_api_token=False):
         token = None
 
         if use_api_token:
-            token = get_api_key_from_environment("HUGGING_FACE_API_TOKEN")
+            # You are allowed to use Hugging Face for a bit before you get
+            # rate limited, in which case you will receive a clear error
+            token = get_api_key_from_environment("HUGGING_FACE_API_TOKEN", required=False).unwrap()
 
         self.client = InferenceClient(model_id, token=token)
 
diff --git a/cookbooks/HuggingFace/python/hf.py b/cookbooks/HuggingFace/python/hf.py
index 8257b326e..3958e5c4e 100644
--- a/cookbooks/HuggingFace/python/hf.py
+++ b/cookbooks/HuggingFace/python/hf.py
@@ -128,7 +128,7 @@ class HuggingFaceTextParser(ParameterizedModelParser):
     A model parser for HuggingFace text generation models.
     """
 
-    def __init__(self, model_id: str = None, use_api_token=False):
+    def __init__(self, model_id: str = None, use_api_token=True):
         """
         Args:
             model_id (str): The model ID of the model to use.
@@ -153,7 +153,9 @@ def __init__(self, model_id: str = None, use_api_token=False):
         token = None
 
         if use_api_token:
-            token = get_api_key_from_environment("HUGGING_FACE_API_TOKEN")
+            # You are allowed to use Hugging Face for a bit before you get
+            # rate limited, in which case you will receive a clear error
+            token = get_api_key_from_environment("HUGGING_FACE_API_TOKEN", required=False).unwrap()
 
         self.client = InferenceClient(model_id, token=token)
 
diff --git a/extensions/Gemini/python/src/aiconfig_extension_gemini/Gemini.py b/extensions/Gemini/python/src/aiconfig_extension_gemini/Gemini.py
index 9766ee4d8..922f9d0dd 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 (
@@ -109,6 +107,12 @@ def __init__(self, id: str = "gemini-pro"):
         super().__init__()
         self.model = genai.GenerativeModel(id)
 
+        # Note: we don't need to explicitly set the key as long as this is set
+        # as an env var genai.configure() will pick up the env var
+        # `GOOGLE_API_KEY`, it's just that we prefer not to call
+        # `get_api_key_from_environment` multiple times if we don't need to
+        self.api_key = get_api_key_from_environment("GOOGLE_API_KEY", required=False).unwrap()
+
     def id(self) -> str:
         """
         Returns an identifier for the model (e.g. llama-2, gpt-4, etc.).
@@ -325,9 +329,9 @@ async def run_inference(
             )
         )
 
-        # Auth check. don't need to explicitly set the key as long as this is set as an env var. genai.configure() will pick it up
-        get_api_key_from_environment("GOOGLE_API_KEY")
-        genai.configure()
+        if not self.api_key:
+            self.api_key = get_api_key_from_environment("GOOGLE_API_KEY", required=True).unwrap()
+        genai.configure(api_key=self.api_key)
 
         # TODO: check and handle api key here
         completion_data = await self.deserialize(prompt, aiconfig, parameters)
diff --git a/python/src/aiconfig/default_parsers/dalle.py b/python/src/aiconfig/default_parsers/dalle.py
index 47acccac3..1c1958982 100644
--- a/python/src/aiconfig/default_parsers/dalle.py
+++ b/python/src/aiconfig/default_parsers/dalle.py
@@ -163,7 +163,7 @@ async def run_inference(self, prompt: Prompt, aiconfig, _options, parameters) ->
         """
         # If needed, certify the API key and initialize the OpenAI client
         if not openai.api_key:
-            openai.api_key = get_api_key_from_environment("OPENAI_API_KEY")
+            openai.api_key = get_api_key_from_environment("OPENAI_API_KEY").unwrap()
         if not self.client:
             self.client = OpenAI(api_key=openai.api_key)
 
diff --git a/python/src/aiconfig/default_parsers/hf.py b/python/src/aiconfig/default_parsers/hf.py
index caedf246d..bb5ce57a3 100644
--- a/python/src/aiconfig/default_parsers/hf.py
+++ b/python/src/aiconfig/default_parsers/hf.py
@@ -145,7 +145,9 @@ def __init__(self, model_id: str = None, use_api_token=False):
         token = None
 
         if use_api_token:
-            token = get_api_key_from_environment("HUGGING_FACE_API_TOKEN")
+            # You are allowed to use Hugging Face for a bit before you get
+            # rate limited, in which case you will receive a clear error
+            token = get_api_key_from_environment("HUGGING_FACE_API_TOKEN", required=False).unwrap()
 
         self.client = InferenceClient(model_id, token=token)
 
diff --git a/python/src/aiconfig/default_parsers/openai.py b/python/src/aiconfig/default_parsers/openai.py
index e3e4d95ac..a7e6b03a1 100644
--- a/python/src/aiconfig/default_parsers/openai.py
+++ b/python/src/aiconfig/default_parsers/openai.py
@@ -235,7 +235,7 @@ async def run_inference(
         )
 
         if not openai.api_key:
-            openai.api_key = get_api_key_from_environment("OPENAI_API_KEY")
+            openai.api_key = get_api_key_from_environment("OPENAI_API_KEY").unwrap()
 
         completion_data = await self.deserialize(prompt, aiconfig, parameters)
         # if stream enabled in runtime options and config, then stream. Otherwise don't stream.
diff --git a/python/src/aiconfig/util/config_utils.py b/python/src/aiconfig/util/config_utils.py
index 0b2eb6eb1..2f3513a49 100644
--- a/python/src/aiconfig/util/config_utils.py
+++ b/python/src/aiconfig/util/config_utils.py
@@ -1,6 +1,8 @@
 import copy
+import dotenv
 import os
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Union
+from result import Result, Ok, Err
 
 if TYPE_CHECKING:
     pass
@@ -10,11 +12,30 @@
     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) -> Result[str | None, str]:
+    """Get the API key if it exists, return None or error if it doesn't
 
-    return os.environ[api_key_name]
+    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:
+        return _get_api_key_from_environment_required(api_key_name)
+    return Ok(os.getenv(api_key_name))
+
+
+def _get_api_key_from_environment_required(api_key_name: str) -> Result[str, str]:
+    try:
+        return Ok(os.environ[api_key_name])
+    except KeyError:
+        return Err(f"Missing API key '{api_key_name}' in environment")
 
 
 def extract_override_settings(config_runtime: "AIConfig", inference_settings: "InferenceSettings", model_id: str):


Test patch:
diff --git a/python/tests/test_util/test_config_util.py b/python/tests/test_util/test_config_util.py
deleted file mode 100644
index a831d4fba..000000000
--- a/python/tests/test_util/test_config_util.py
+++ /dev/null
@@ -1,11 +0,0 @@
-from aiconfig.util.config_utils import get_api_key_from_environment
-
-
-def test_get_api_key_from_environment():
-    key = "TEST_API_KEY"
-    try:
-        get_api_key_from_environment(key)
-    except Exception:
-        pass  # The expected exception was raised, so do nothing
-    else:
-        raise AssertionError("ExpectedException was not raised")