Skip to content

Commit

Permalink
Make get_api_key_from_environment return nullable
Browse files Browse the repository at this point in the history
Sometimes key is not defined or set, we need to do this to unblock #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)
  • Loading branch information
Rossdan Craig [email protected] committed Jan 5, 2024
1 parent aa1eff3 commit f24588e
Show file tree
Hide file tree
Showing 3 changed files with 27 additions and 7 deletions.
Original file line number Diff line number Diff line change
@@ -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 (
Expand Down Expand Up @@ -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:
"""
Expand Down
27 changes: 23 additions & 4 deletions python/src/aiconfig/util/config_utils.py
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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]


Expand Down
2 changes: 1 addition & 1 deletion python/tests/test_util/test_config_util.py
Original file line number Diff line number Diff line change
@@ -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)
Expand Down

0 comments on commit f24588e

Please sign in to comment.