From fd22fe50790f38608ce40da097bf9b8cd2c93b65 Mon Sep 17 00:00:00 2001
From: Bo Xu <shifucun@users.noreply.github.com>
Date: Thu, 23 May 2024 22:24:49 -0700
Subject: [PATCH] resolve comments

---
 nl_server/config.py   | 4 ++++
 nl_server/loader.py   | 5 +++++
 nl_server/registry.py | 6 +++---
 nl_server/routes.py   | 8 +++++---
 4 files changed, 17 insertions(+), 6 deletions(-)

diff --git a/nl_server/config.py b/nl_server/config.py
index 9c0a4a496b..c1e6a14084 100644
--- a/nl_server/config.py
+++ b/nl_server/config.py
@@ -22,8 +22,11 @@
 CUSTOM_DC_INDEX: str = 'custom_ft'
 
 # App Config constants.
+# Key for the registry that holds model objects/client and embedding store.
 REGISTRY_KEY: str = 'NL_REGISTRY'
+# Key for the catalog object that holds model and index configs.
 CATALOG_KEY: str = 'NL_CATALOG'
+# Key for the environment config object.
 ENV_KEY: str = 'NL_ENV'
 
 # Query to use to check index health if this is the default index.
@@ -92,6 +95,7 @@ class VertexAIIndexConfig(IndexConfig):
 
 
 # This represents the full catalog of models and indexes.
+# Only a subset of them are enabled/used by the Env config at runtime.
 @dataclass(kw_only=True)
 class Catalog:
   version: str = None
diff --git a/nl_server/loader.py b/nl_server/loader.py
index 90baba6ba3..56524cfd80 100644
--- a/nl_server/loader.py
+++ b/nl_server/loader.py
@@ -37,6 +37,11 @@
 
 @dataclass
 class Env:
+  """
+  A class to represent the NL server environment config.
+
+  This object is used with the Catalog object to configure the server.
+  """
   default_indexes: List[str]
   enabled_indexes: List[str]
   vertex_ai_model_info: Dict[str, any]
diff --git a/nl_server/registry.py b/nl_server/registry.py
index 73b3702656..51b8567cdc 100644
--- a/nl_server/registry.py
+++ b/nl_server/registry.py
@@ -33,10 +33,10 @@
 from shared.lib.custom_dc_util import is_custom_dc
 
 
-#
-# A class to hold embeddings stores and models.
-#
 class Registry:
+  """
+  A class to hold runtime model handle/client objects and embeddings stores.
+  """
 
   # Input is parsed embeddings config.
   def __init__(self, catalog: Catalog):
diff --git a/nl_server/routes.py b/nl_server/routes.py
index bf29839168..d430910a2e 100644
--- a/nl_server/routes.py
+++ b/nl_server/routes.py
@@ -48,10 +48,12 @@ def healthz():
   if default_embeddings_loaded:
     return 'OK', 200
 
-  default_index_type = current_app.config[config.ENV_KEY].default_indexes[0]
-  if not default_index_type:
-    logging.warning('Health Check Failed: Default index name empty!')
+  nl_env = current_app.config[config.ENV_KEY]
+  if not nl_env.default_indexes:
+    logging.error('Health Check Failed: Default index name empty!')
     return 'Service Unavailable', 500
+
+  default_index_type = nl_env.default_indexes[0]
   nl_embeddings: Embeddings = current_app.config[config.REGISTRY_KEY].get_index(
       default_index_type)
   if nl_embeddings: