feat(pattern-1): enhance HITL wait function with document compression support

Author: Bob Strahan

patterns/pattern-1/src/hitl-wait-function/index.py

@@ -5,6 +5,7 @@
 import os
 import logging
 from datetime import datetime, timezone, timedelta
+from idp_common.models import Document
 
 logger = logging.getLogger()
 logger.setLevel(os.environ.get('LOG_LEVEL', 'INFO'))
@@ -25,12 +26,26 @@ def lambda_handler(event, context):
     """
     logger.info(f"Processing event: {json.dumps(event)}")
 
-    # Extract document information
-    document = event.get('Payload', {}).get('Result', {}).get('document', {})
-    if not document:
-        document = event.get('Payload', {}).get('ProcessingResult', {}).get('document', {})
-    if not document:
-        document = event.get('Payload', {}).get('document', {})
+    # Extract document information using load_document to handle compression
+    document_data = event.get('Payload', {}).get('Result', {}).get('document', {})
+    if not document_data:
+        document_data = event.get('Payload', {}).get('ProcessingResult', {}).get('document', {})
+    if not document_data:
+        document_data = event.get('Payload', {}).get('document', {})
+    
+    # Get working bucket for decompression
+    working_bucket = os.environ.get('WORKING_BUCKET')
+    if not working_bucket:
+        logger.warning("WORKING_BUCKET environment variable not set")
+    
+    # Load document using utility method to handle compression/decompression
+    try:
+        document_obj = Document.load_document(document_data, working_bucket, logger)
+        # Convert back to dict for the rest of the function
+        document = document_obj.to_dict()
+    except Exception as e:
+        logger.error(f"Error loading document: {str(e)}")
+        raise
 
     document_id = document.get('id')
     workflow_execution_arn = document.get('workflow_execution_arn')
@@ -185,4 +200,4 @@ def store_token_in_tracking_table(table, token_id, task_token, document_id, exec
         table.put_item(Item=item)
         logger.info(f"Stored token information: {token_id}")
     except Exception as e:
-        logger.error(f"Error storing token {token_id}: {str(e)}")
+        logger.error(f"Error storing token {token_id}: {str(e)}")

patterns/pattern-1/template.yaml

@@ -995,6 +995,7 @@ Resources:
           DYNAMODB_TABLE: !Ref BDAMetadataTable
           LOG_LEVEL: !Ref LogLevel
           SAGEMAKER_A2I_REVIEW_PORTAL_URL: !Ref SageMakerA2IReviewPortalURL
+          WORKING_BUCKET: !Ref WorkingBucket
       LoggingConfig:
         LogGroup: !Ref HITLWaitFunctionLogGroup
       Policies:
@@ -1003,6 +1004,8 @@ Resources:
             TableName: !Ref TrackingTable
         - DynamoDBCrudPolicy:
             TableName: !Ref BDAMetadataTable
+        - S3ReadPolicy:
+            BucketName: !Ref WorkingBucket
         - Statement:
           - Effect: Allow
             Action: