Enhancing A2I template and adding logic for failed review tasks

Author: gudivt

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

@@ -673,7 +673,7 @@ def start_human_loop(
         FlowDefinitionArn = ssm.get_parameter(Name=f"/{os.environ.get('METRIC_NAMESPACE', 'IDP')}/FlowDefinitionArn")['Parameter']['Value']
         human_review_id = generate_random_string(2)
         response = a2i_runtime_client.start_human_loop(
-            HumanLoopName=f"review-bda-{execution_id}-{human_review_id}",
+            HumanLoopName=f"review-bda-{execution_id}-{record_number}-{page_id_num}",
             FlowDefinitionArn=FlowDefinitionArn,
             HumanLoopInput={"InputContent": json.dumps(human_loop_input)}
         )
@@ -929,6 +929,7 @@ def process_segments(
         else:
             if enable_hitl == 'true':
                 std_hitl = 'true'
+                # std_hitl = None # HITL for standard output blueprint match is disabled until we have option to choose Blueprint in A2I
             else:
                 std_hitl = None 
             # Process standard output if no custom output match
@@ -940,7 +941,8 @@ def process_segments(
             page_array = list(range(start_page, end_page + 1))
             item.update({
                 "page_array": page_array,
-                "hitl_triggered": std_hitl,
+                # "hitl_triggered": std_hitl,
+                "hitl_triggered": None,
                 "extraction_bp_name": "None",
                 "extracted_result": std_output
             })
@@ -950,31 +952,34 @@ def process_segments(
                 record_number=record_number,
                 bp_match=segment.get('custom_output_status'),
                 extraction_bp_name="None",
-                hitl_triggered=std_hitl,
+                # hitl_triggered=std_hitl,
+                hitl_triggered=None,
                 page_array=page_array,
                 review_portal_url=SAGEMAKER_A2I_REVIEW_PORTAL_URL
             )
 
-            hitl_triggered = std_hitl
-            if enable_hitl == 'true':
-                for page_number in range(start_page, end_page + 1):
-                    ImageUri = f"s3://{output_bucket}/{object_key}/pages/{page_number}/image.jpg"
-                    try:
-                        human_loop_response = start_human_loop(
-                            execution_id=execution_id,
-                            kv_pairs=[],
-                            source_image_uri=ImageUri,
-                            bounding_boxes=[],
-                            blueprintName="",
-                            bp_confidence=0.00,
-                            confidenceThreshold=confidence_threshold,
-                            page_id=page_number,
-                            page_indices=page_array,
-                            record_number=record_number
-                        )
-                        logger.info(f"Triggered human loop for page {page_number}: {human_loop_response}")
-                    except Exception as e:
-                        logger.error(f"Failed to start human loop for page {page_number}: {str(e)}")
+            # hitl_triggered = std_hitl
+            hitl_triggered = None
+            # if enable_hitl == 'true':
+            # # if std_hitl: # HITL for standard output blueprint match is disabled until we have option to choose Blueprint in A2I
+            #     for page_number in range(start_page, end_page + 1):
+            #         ImageUri = f"s3://{output_bucket}/{object_key}/pages/{page_number}/image.jpg"
+            #         try:
+            #             human_loop_response = start_human_loop(
+            #                 execution_id=execution_id,
+            #                 kv_pairs=[],
+            #                 source_image_uri=ImageUri,
+            #                 bounding_boxes=[],
+            #                 blueprintName="",
+            #                 bp_confidence=0.00,
+            #                 confidenceThreshold=confidence_threshold,
+            #                 page_id=page_number,
+            #                 page_indices=page_array,
+            #                 record_number=record_number
+            #             )
+            #             logger.info(f"Triggered human loop for page {page_number}: {human_loop_response}")
+            #         except Exception as e:
+            #             logger.error(f"Failed to start human loop for page {page_number}: {str(e)}")
 
         document.hitl_metadata.append(hitl_metadata)
 
@@ -1136,7 +1141,7 @@ def handler(event, context):
                             execution_id,
                             document
                         )
-                        if hitl_result:
+                        if hitl_result or hitl_triggered== "true":
                             hitl_triggered = "true"
                 elif isinstance(output_metadata, dict):
                     for asset_id, asset in output_metadata.items():
@@ -1149,7 +1154,7 @@ def handler(event, context):
                             execution_id,
                             document
                         )
-                        if hitl_result:
+                        if hitl_result or hitl_triggered== "true":
                             hitl_triggered = "true"
                 else:
                     logger.error("Unexpected output_metadata format in job_metadata.json")

patterns/pattern-1/src/processresults_function/index.py

@@ -108,6 +108,12 @@
                     "BackoffRate": 2
                 }
             ],
+            "Catch": [
+                {
+                    "ErrorEquals": ["HITLFailedException"],
+                    "Next": "HITLFailState"
+                }
+            ],
             "Next": "HITLStatusUpdate"
         },
         "HITLStatusUpdate": {
@@ -171,6 +177,11 @@
             "Type": "Fail",
             "Cause": "Data Automation Job Failed",
             "Error": "JobFailedException"
+        },
+        "HITLFailState": {
+            "Type": "Fail",
+            "Cause": "Human In The Loop Review Failed",
+            "Error": "HITLFailedException"
         }
     }
 }

patterns/pattern-1/statemachine/workflow.asl.json

@@ -878,6 +878,8 @@ Resources:
         detail:
           humanLoopStatus:
             - 'Completed'
+            - 'Failed'
+            - 'Stopped'
       State: 'ENABLED'
       Targets:
         - Arn: !GetAtt HITLProcessLambdaFunction.Arn