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

Issue:
Set the Gemini api key in init instead of re-running every inference call
Set the Gemini api key in init instead of re-running every inference call


See comments from https://github.com/lastmile-ai/aiconfig/pull/772#discussion_r1443387445

## Test Plan
Run the Gemini cookbook to make sure it still works

---
Stack created with [Sapling](https://sapling-scm.com). Best reviewed with [ReviewStack](https://reviewstack.dev/lastmile-ai/aiconfig/pull/786).
* #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/extensions/Gemini/python/src/aiconfig_extension_gemini/Gemini.py b/extensions/Gemini/python/src/aiconfig_extension_gemini/Gemini.py
index 9766ee4d8..6852a3189 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)
+
     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)
+        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/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)