Use aws_env in more places and lint.

Author: lkacenja

python_components/document_inference/lambda_function.py

@@ -27,9 +27,13 @@ def handler(event, context):
         all_models = helpers.get_models("models.json")
         helpers.validate_model(all_models, event["model_name"])
         helpers.logger.info("Model is valid")
-        api_key = helpers.get_secret(all_models[event["model_name"]]["key"], local_mode, aws_env)
+        api_key = helpers.get_secret(
+            all_models[event["model_name"]]["key"], local_mode, aws_env
+        )
         asap_creds_user = helpers.get_secret(API_USER_NAME_SECRET, local_mode, aws_env)
-        asap_creds_password = helpers.get_secret(API_PASSWORD_SECRET, local_mode, aws_env)
+        asap_creds_password = helpers.get_secret(
+            API_PASSWORD_SECRET, local_mode, aws_env
+        )
         page_limit_label = (
             "unlimited" if event["page_limit"] == 0 else event["page_limit"]
         )

python_components/evaluation/evaluation/exception/evaluation.py

@@ -46,6 +46,7 @@ async def evaluate(self, document: Document) -> List[Result]:
             self.inference_model_name,
             "exception",
             self.local_mode,
+            self.aws_env,
             self.page_limit,
         )
         logger.info("Exception check complete. Performing related evaluations.")

python_components/evaluation/evaluation/summary/evaluation.py

@@ -26,6 +26,7 @@ def evaluate(self, document: Document) -> List[Result]:
             self.inference_model_name,
             "summary",
             self.local_mode,
+            self.aws_env,
             self.page_limit,
         )
         logger.info("Summarization complete. Performing related evaluations.")

python_components/evaluation/evaluation/utility/asap_inference.py

@@ -24,6 +24,7 @@ def get_inference_for_document(
     inference_model_name: str,
     inference_type: str,
     local_mode: bool,
+    aws_env: str,
     page_number: int,
 ) -> None:
     logger.info(f"Performing inference type {inference_type} for {document.url}...")
@@ -40,7 +41,7 @@ def get_inference_for_document(
         session = boto3.Session()
         client = session.client("lambda")
         response = client.get_function_url_config(
-            FunctionName="asap-pdf-document-inference-evaluation-staging",
+            FunctionName=f"asap-pdf-document-inference-evaluation-{aws_env}",
         )
         if "FunctionArn" not in response.keys():
             raise RuntimeError(

python_components/evaluation/evaluation/utility/document.py

@@ -42,6 +42,7 @@ def __init__(
         self.commit_sha = commit_sha
         self.page_limit = kwargs.get("page_limit", 7)
         self.local_mode = kwargs.get("local_mode", False)
+        self.aws_env = kwargs.get("aws_env", "staging")
         now = datetime.datetime.now()
         metric_run_date = (
             now.strftime("%Y-%m-%d %H:%M:%S") + f".{now.microsecond // 1000:03d}"

python_components/evaluation/evaluation/utility/google_sheet.py

@@ -5,20 +5,18 @@
 from google.oauth2.service_account import Credentials
 from googleapiclient.discovery import build
 
-GOOGLE_EVAL_SERVICE_ACCOUNT_CREDS = (
-    "/asap-pdf/production/GOOGLE_SERVICE_ACCOUNT-20250605155250934400000001"
-)
-GOOGLE_EVAL_SHEET_ID = (
-    "/asap-pdf/production/GOOGLE_SHEET_ID_EVALUATION-20250605155250934400000003"
-)
+GOOGLE_EVAL_SERVICE_ACCOUNT_CREDS = "asap-pdf/{AWS_ENV}/GOOGLE_SERVICE_ACCOUNT"
+GOOGLE_EVAL_SHEET_ID = "asap-pdf/{AWS_ENV}/GOOGLE_SHEET_ID_EVALUATION"
 RANGE_NAME = "Results!A:L"
 SCOPES = ["https://www.googleapis.com/auth/spreadsheets"]
 
 
-def append_to_google_sheet(results: List[dict], local_mode: bool) -> None:
+def append_to_google_sheet(results: List[dict], local_mode: bool, aws_env: str) -> None:
 
-    creds_json = json.loads(get_secret(GOOGLE_EVAL_SERVICE_ACCOUNT_CREDS, local_mode))
-    sheet_id = get_secret(GOOGLE_EVAL_SHEET_ID, local_mode)
+    creds_json = json.loads(
+        get_secret(GOOGLE_EVAL_SERVICE_ACCOUNT_CREDS, local_mode, aws_env)
+    )
+    sheet_id = get_secret(GOOGLE_EVAL_SHEET_ID, local_mode, aws_env)
 
     logger.info(f"Appending results to {sheet_id} range {RANGE_NAME}...")
 

python_components/evaluation/lambda_function.py

@@ -48,6 +48,7 @@ def handler(event, context):
             event["commit_sha"],
             delta,
             local_mode=local_mode,
+            aws_env=aws_env,
         )
         exception_eval_wrapper = exception.EvaluationWrapper(
             eval_model,
@@ -56,6 +57,7 @@ def handler(event, context):
             event["commit_sha"],
             delta,
             local_mode=local_mode,
+            aws_env=aws_env,
         )
         output = []
         for document_dict in event["documents"]:
@@ -78,7 +80,7 @@ def handler(event, context):
                 output.extend(results)
         if "output_google_sheet" in event.keys():
             utility.helpers.logger.info("Writing eval results to Google Sheet")
-            utility.google_sheet.append_to_google_sheet(output, local_mode)
+            utility.google_sheet.append_to_google_sheet(output, local_mode, aws_env)
             return {
                 "statusCode": 200,
                 "body": "Wrote evaluation results to Google Sheet.",