From 9c3fa492b099f6e53b9b7bf49542dece357469f7 Mon Sep 17 00:00:00 2001
From: Christie Ellks <calinc@google.com>
Date: Thu, 20 Feb 2025 11:36:51 -0800
Subject: [PATCH 1/6] Use api key helper during app creation.

---
 server/__init__.py | 46 +++++++++++-----------------------------------
 1 file changed, 11 insertions(+), 35 deletions(-)

diff --git a/server/__init__.py b/server/__init__.py
index cc0686e26f..5e3787306b 100644
--- a/server/__init__.py
+++ b/server/__init__.py
@@ -56,9 +56,7 @@ def _get_api_key(env_keys=[], gcp_project='', gcp_path=''):
 
   Returns:
       API key if it exists, otherwise an empty string.
-  
-  TODO: use this method everywhere else in this file
-  """
+    """
   # Try to get the key from the environment
   for k in env_keys:
     if os.environ.get(k):
@@ -73,6 +71,7 @@ def _get_api_key(env_keys=[], gcp_project='', gcp_path=''):
       secret_response = secret_client.access_secret_version(name=secret_name)
       return secret_response.payload.data.decode('UTF-8').replace('\n', '')
     except NotFound:
+      logging.warning(f'No secret found for {gcp_project}:{gcp_path}')
       return ''
 
   # If key is not found, return an empty string
@@ -377,34 +376,17 @@ def create_app(nl_root=DEFAULT_NL_ROOT):
     app.config['MAPS_API_KEY'] = ''
   else:
     # Get the API key from environment first.
-    if os.environ.get('MAPS_API_KEY'):
-      app.config['MAPS_API_KEY'] = os.environ.get('MAPS_API_KEY')
-    elif os.environ.get('maps_api_key'):
-      app.config['MAPS_API_KEY'] = os.environ.get('maps_api_key')
-    else:
-      secret_client = secretmanager.SecretManagerServiceClient()
-      secret_name = secret_client.secret_version_path(cfg.SECRET_PROJECT,
-                                                      'maps-api-key', 'latest')
-      secret_response = secret_client.access_secret_version(name=secret_name)
-      app.config['MAPS_API_KEY'] = secret_response.payload.data.decode('UTF-8')
+    app.config['MAPS_API_KEY'] = _get_api_key(['MAPS_API_KEY', 'maps_api_key'],
+                                              cfg.SECRET_PROJECT,
+                                              'maps-api-key')
 
   if cfg.LOCAL:
     app.config['LOCAL'] = True
 
   # Need to fetch the API key for non gcp environment.
   if cfg.LOCAL or cfg.WEBDRIVER or cfg.INTEGRATION:
-    # Get the API key from environment first.
-    if os.environ.get('DC_API_KEY'):
-      app.config['DC_API_KEY'] = os.environ.get('DC_API_KEY')
-    elif os.environ.get('dc_api_key'):
-      app.config['DC_API_KEY'] = os.environ.get('dc_api_key')
-    else:
-      secret_client = secretmanager.SecretManagerServiceClient()
-      secret_name = secret_client.secret_version_path(cfg.SECRET_PROJECT,
-                                                      'mixer-api-key', 'latest')
-      secret_response = secret_client.access_secret_version(name=secret_name)
-      app.config['DC_API_KEY'] = secret_response.payload.data.decode(
-          'UTF-8').replace('\n', '')
+    app.config['DC_API_KEY'] = _get_api_key(['DC_API_KEY', 'dc_api_key'],
+                                            cfg.SECRET_PROJECT, 'mixer-api-key')
 
   # Initialize translations
   babel = Babel(app, default_domain='all')
@@ -422,18 +404,12 @@ def create_app(nl_root=DEFAULT_NL_ROOT):
     else:
       app.config['NL_TABLE'] = None
 
-    # Get the API key from environment first.
     if cfg.USE_LLM:
       app.config['LLM_PROMPT_TEXT'] = llm_prompt.get_prompts()
-      if os.environ.get('LLM_API_KEY'):
-        app.config['LLM_API_KEY'] = os.environ.get('LLM_API_KEY')
-      else:
-        secret_client = secretmanager.SecretManagerServiceClient()
-        secret_name = secret_client.secret_version_path(cfg.SECRET_PROJECT,
-                                                        'palm-api-key',
-                                                        'latest')
-        secret_response = secret_client.access_secret_version(name=secret_name)
-        app.config['LLM_API_KEY'] = secret_response.payload.data.decode('UTF-8')
+      app.config['LLM_API_KEY'] = _get_api_key(['LLM_API_KEY'],
+                                               cfg.SECRET_PROJECT,
+                                               'palm-api-key')
+
     app.config[
         'NL_BAD_WORDS'] = EMPTY_BANNED_WORDS if cfg.CUSTOM else load_bad_words(
         )

From 7f6a63e99e6c6f3beea9d1e584cea485abd04262 Mon Sep 17 00:00:00 2001
From: Christie Ellks <calinc@google.com>
Date: Thu, 20 Feb 2025 12:36:52 -0800
Subject: [PATCH 2/6] Check env var DC_API_KEY is set when running in LITE
 mode.

---
 server/__init__.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/server/__init__.py b/server/__init__.py
index 5e3787306b..92103e3715 100644
--- a/server/__init__.py
+++ b/server/__init__.py
@@ -75,6 +75,8 @@ def _get_api_key(env_keys=[], gcp_project='', gcp_path=''):
       return ''
 
   # If key is not found, return an empty string
+  logging.warning(
+      f'No secret found for project: {gcp_project}; path:{gcp_path}')
   return ''
 
 
@@ -314,7 +316,8 @@ def create_app(nl_root=DEFAULT_NL_ROOT):
   app.config.from_object(cfg)
 
   # Check DC_API_KEY is set for local dev.
-  if cfg.CUSTOM and cfg.LOCAL and not os.environ.get('DC_API_KEY'):
+  if (cfg.LITE or
+      (cfg.CUSTOM and cfg.LOCAL)) and not os.environ.get('DC_API_KEY'):
     raise Exception(
         'Set environment variable DC_API_KEY for local custom DC development')
 

From 4b4ef0686c1e26755bf214fdfbbbd25a00f41b8b Mon Sep 17 00:00:00 2001
From: Christie Ellks <calinc@google.com>
Date: Thu, 20 Feb 2025 12:39:20 -0800
Subject: [PATCH 3/6] Update logging message.

---
 server/__init__.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/server/__init__.py b/server/__init__.py
index 92103e3715..50729059bf 100644
--- a/server/__init__.py
+++ b/server/__init__.py
@@ -71,7 +71,7 @@ def _get_api_key(env_keys=[], gcp_project='', gcp_path=''):
       secret_response = secret_client.access_secret_version(name=secret_name)
       return secret_response.payload.data.decode('UTF-8').replace('\n', '')
     except NotFound:
-      logging.warning(f'No secret found for {gcp_project}:{gcp_path}')
+      logging.warning(f'No key found at {gcp_project}:{gcp_path}')
       return ''
 
   # If key is not found, return an empty string

From 18c47c0a82717e91da6d50c23393c8fac1f4df80 Mon Sep 17 00:00:00 2001
From: Christie Ellks <calinc@google.com>
Date: Thu, 20 Feb 2025 12:39:20 -0800
Subject: [PATCH 4/6] Update logging message.

---
 server/__init__.py | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/server/__init__.py b/server/__init__.py
index 92103e3715..398e1f7b41 100644
--- a/server/__init__.py
+++ b/server/__init__.py
@@ -71,12 +71,11 @@ def _get_api_key(env_keys=[], gcp_project='', gcp_path=''):
       secret_response = secret_client.access_secret_version(name=secret_name)
       return secret_response.payload.data.decode('UTF-8').replace('\n', '')
     except NotFound:
-      logging.warning(f'No secret found for {gcp_project}:{gcp_path}')
+      logging.warning(f'No key found at {gcp_project}:{gcp_path}')
       return ''
 
   # If key is not found, return an empty string
-  logging.warning(
-      f'No secret found for project: {gcp_project}; path:{gcp_path}')
+  logging.warning(f'No key found for project: {gcp_project}; path:{gcp_path}')
   return ''
 
 

From afef3636ff6f0b633451c221123dddef858f7c3e Mon Sep 17 00:00:00 2001
From: Christie Ellks <calinc@google.com>
Date: Thu, 20 Feb 2025 12:55:17 -0800
Subject: [PATCH 5/6] Update logging message.

---
 server/__init__.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/server/__init__.py b/server/__init__.py
index 398e1f7b41..f6640e3cb5 100644
--- a/server/__init__.py
+++ b/server/__init__.py
@@ -75,7 +75,9 @@ def _get_api_key(env_keys=[], gcp_project='', gcp_path=''):
       return ''
 
   # If key is not found, return an empty string
-  logging.warning(f'No key found for project: {gcp_project}; path:{gcp_path}')
+  logging.warning(
+      f'No key found in environment variables: {env_keys} nor project-path: {gcp_project}-{gcp_path}'
+  )
   return ''
 
 

From e5bdf4522c6e21186fce078faf1420fce226f341 Mon Sep 17 00:00:00 2001
From: Christie Ellks <calinc@google.com>
Date: Thu, 20 Feb 2025 14:01:33 -0800
Subject: [PATCH 6/6] Updated logging messages.

---
 server/__init__.py | 5 +++--
 1 file changed, 3 insertions(+), 2 deletions(-)

diff --git a/server/__init__.py b/server/__init__.py
index f6640e3cb5..441b27b2d6 100644
--- a/server/__init__.py
+++ b/server/__init__.py
@@ -71,12 +71,13 @@ def _get_api_key(env_keys=[], gcp_project='', gcp_path=''):
       secret_response = secret_client.access_secret_version(name=secret_name)
       return secret_response.payload.data.decode('UTF-8').replace('\n', '')
     except NotFound:
-      logging.warning(f'No key found at {gcp_project}:{gcp_path}')
+      logging.warning(
+          f'No key found at {gcp_path} of the configured GCP project.')
       return ''
 
   # If key is not found, return an empty string
   logging.warning(
-      f'No key found in environment variables: {env_keys} nor project-path: {gcp_project}-{gcp_path}'
+      f'No key found in the [{",".join(env_keys)}] environment variable(s), nor at "{gcp_path}" of the configured GCP project.'
   )
   return ''