feat: improve throttling detection and retry handling in assessment service

Author: Bob Strahan

lib/idp_common_pkg/idp_common/assessment/granular_service.py

@@ -1734,11 +1734,16 @@ def process_document_section(self, document: Document, section_id: str) -> Docum
                     # Store the primary exception for easy access by caller
                     document.metadata["primary_exception"] = primary_exception
 
-                    # Cache successful task results (only when some tasks fail - for retry scenarios)
+                # Check for any failed tasks (both exceptions and unsuccessful results)
+                failed_results = [r for r in all_task_results if not r.success]
+                any_failures = bool(failed_task_exceptions or failed_results)
+
+                # Cache successful tasks only when there are failures (for retry optimization)
+                if any_failures:
                     successful_results = [r for r in all_task_results if r.success]
                     if successful_results:
                         logger.info(
-                            f"Caching {len(successful_results)} successful assessment task results for document {document.id} section {section_id} due to {len(failed_task_exceptions)} failed tasks (retry scenario)"
+                            f"Caching {len(successful_results)} successful assessment task results for document {document.id} section {section_id} due to {len(failed_results)} failed results + {len(failed_task_exceptions)} failed exceptions (retry scenario)"
                         )
                         self._cache_successful_assessment_tasks(
                             document.id,
@@ -1748,7 +1753,7 @@ def process_document_section(self, document: Document, section_id: str) -> Docum
                         )
                     else:
                         logger.warning(
-                            f"No successful assessment task results to cache for document {document.id} section {section_id} - all {len(failed_task_exceptions)} tasks failed"
+                            f"No successful assessment task results to cache for document {document.id} section {section_id} - all tasks failed"
                         )
                 else:
                     # All new tasks succeeded - no need to cache since there won't be retries
@@ -1868,6 +1873,69 @@ def process_document_section(self, document: Document, section_id: str) -> Docum
             logger.error(error_msg)
             document.status = Status.FAILED
             document.errors.append(error_msg)
+
+            # Check if this is a throttling exception and populate metadata for retry handling
+            if self._is_throttling_exception(e):
+                logger.info(
+                    f"Populating metadata for throttling exception: {type(e).__name__}"
+                )
+                document.metadata = document.metadata or {}
+                document.metadata["failed_assessment_tasks"] = {
+                    "granular_processing": {
+                        "exception_type": type(e).__name__,
+                        "exception_message": str(e),
+                        "exception_class": e.__class__.__module__
+                        + "."
+                        + e.__class__.__name__,
+                        "is_throttling": True,
+                    }
+                }
+                document.metadata["primary_exception"] = e
+
+        # Additional check: if document status is FAILED and contains throttling errors,
+        # populate metadata even if no exceptions were thrown
+        if (
+            document.status == Status.FAILED
+            and document.errors
+            and not hasattr(document, "metadata")
+            or not document.metadata
+            or "failed_assessment_tasks" not in document.metadata
+        ):
+            # Check if any errors contain throttling keywords
+            throttling_keywords = [
+                "throttlingexception",
+                "provisionedthroughputexceededexception",
+                "servicequotaexceededexception",
+                "toomanyrequestsexception",
+                "requestlimitexceeded",
+                "too many tokens",
+                "please wait before trying again",
+                "reached max retries",
+            ]
+
+            has_throttling_error = False
+            throttling_error_msg = None
+            for error_msg in document.errors:
+                error_lower = str(error_msg).lower()
+                if any(keyword in error_lower for keyword in throttling_keywords):
+                    has_throttling_error = True
+                    throttling_error_msg = error_msg
+                    break
+
+            if has_throttling_error:
+                logger.info(
+                    f"Populating metadata for throttling error found in document.errors: {throttling_error_msg}"
+                )
+                document.metadata = document.metadata or {}
+                document.metadata["failed_assessment_tasks"] = {
+                    "document_level_error": {
+                        "exception_type": "ThrottlingError",
+                        "exception_message": throttling_error_msg,
+                        "exception_class": "DocumentLevelThrottlingError",
+                        "is_throttling": True,
+                    }
+                }
+
         return document
 
     def assess_document(self, document: Document) -> Document:

patterns/pattern-2/src/assessment_function/index.py

@@ -12,12 +12,81 @@
 from idp_common import s3
 from assessment_validator import AssessmentValidator
 
+# Custom exception for throttling scenarios
+class ThrottlingException(Exception):
+    """Exception raised when throttling is detected in document processing results"""
+    pass
+
+# Throttling detection constants
+THROTTLING_KEYWORDS = [
+    "throttlingexception",
+    "provisionedthroughputexceededexception", 
+    "servicequotaexceededexception",
+    "toomanyrequestsexception",
+    "requestlimitexceeded",
+    "too many tokens",
+    "please wait before trying again",
+    "reached max retries"
+]
+
+THROTTLING_EXCEPTIONS = [
+    "ThrottlingException",
+    "ProvisionedThroughputExceededException",
+    "ServiceQuotaExceededException", 
+    "TooManyRequestsException",
+    "RequestLimitExceeded"
+]
+
 # Configuration will be loaded in handler function
 
 logger = logging.getLogger()
 logger.setLevel(os.environ.get("LOG_LEVEL", "INFO"))
 logging.getLogger('idp_common.bedrock.client').setLevel(os.environ.get("BEDROCK_LOG_LEVEL", "INFO"))
 
+def is_throttling_exception(exception):
+    """
+    Check if an exception is related to throttling.
+    
+    Args:
+        exception: The exception to check
+        
+    Returns:
+        bool: True if the exception is throttling-related, False otherwise
+    """
+    from botocore.exceptions import ClientError
+    
+    if isinstance(exception, ClientError):
+        error_code = exception.response.get('Error', {}).get('Code', '')
+        return error_code in THROTTLING_EXCEPTIONS
+    
+    exception_name = type(exception).__name__
+    exception_message = str(exception).lower()
+    
+    return (
+        exception_name in THROTTLING_EXCEPTIONS or
+        any(keyword in exception_message for keyword in THROTTLING_KEYWORDS)
+    )
+
+def check_document_for_throttling_errors(document):
+    """
+    Check if a document has throttling errors in its errors field.
+    
+    Args:
+        document: The document object to check
+        
+    Returns:
+        tuple: (has_throttling_errors: bool, first_throttling_error: str or None)
+    """
+    if document.status != Status.FAILED or not document.errors:
+        return False, None
+    
+    for error_msg in document.errors:
+        error_lower = str(error_msg).lower()
+        if any(keyword in error_lower for keyword in THROTTLING_KEYWORDS):
+            return True, error_msg
+    
+    return False, None
+
 def handler(event, context):
     """
     Lambda handler for document assessment.
@@ -85,55 +154,39 @@ def handler(event, context):
         t1 = time.time()
         logger.info(f"Total assessment time: {t1-t0:.2f} seconds")
 
-        # Check for failed assessment tasks that might require retry
-        if (hasattr(updated_document, 'metadata') and 
-            updated_document.metadata and 
-            'failed_assessment_tasks' in updated_document.metadata):
-            
-            failed_tasks = updated_document.metadata['failed_assessment_tasks']
-            throttling_tasks = {
-                task_id: task_info for task_id, task_info in failed_tasks.items()
-                if task_info.get('is_throttling', False)
-            }
-            
-            logger.warning(
-                f"Assessment completed with {len(failed_tasks)} failed tasks, "
-                f"{len(throttling_tasks)} due to throttling"
-            )
-            
-            if throttling_tasks:
-                logger.info(
-                    f"Throttling detected in {len(throttling_tasks)} tasks. "
-                    f"Successful tasks have been cached for retry."
+        # Check for failed assessment tasks that might require retry (granular assessment)
+        if hasattr(updated_document, 'metadata') and updated_document.metadata:
+            failed_tasks = updated_document.metadata.get('failed_assessment_tasks', {})
+            if failed_tasks:
+                throttling_tasks = {
+                    task_id: task_info for task_id, task_info in failed_tasks.items()
+                    if task_info.get('is_throttling', False)
+                }
+                
+                logger.warning(
+                    f"Assessment completed with {len(failed_tasks)} failed tasks, "
+                    f"{len(throttling_tasks)} due to throttling"
                 )
+                
+                if throttling_tasks:
+                    logger.info(
+                        f"Throttling detected in {len(throttling_tasks)} tasks. "
+                        f"Successful tasks have been cached for retry."
+                    )
+        
+        # Check for throttling errors in document status and errors field
+        has_throttling, throttling_error = check_document_for_throttling_errors(updated_document)
+        if has_throttling:
+            logger.error(f"Throttling error detected in document errors: {throttling_error}")
+            logger.error("Raising ThrottlingException to trigger Step Functions retry")
+            raise ThrottlingException(f"Throttling detected in document processing: {throttling_error}")
 
     except Exception as e:
         t1 = time.time()
         logger.error(f"Assessment failed after {t1-t0:.2f} seconds: {str(e)}")
 
         # Check if this is a throttling exception that should trigger retry
-        from botocore.exceptions import ClientError
-        throttling_exceptions = [
-            "ThrottlingException",
-            "ProvisionedThroughputExceededException",
-            "ServiceQuotaExceededException", 
-            "TooManyRequestsException",
-            "RequestLimitExceeded"
-        ]
-        
-        is_throttling = False
-        if isinstance(e, ClientError):
-            error_code = e.response.get('Error', {}).get('Code', '')
-            is_throttling = error_code in throttling_exceptions
-        else:
-            exception_name = type(e).__name__
-            exception_message = str(e).lower()
-            is_throttling = (
-                exception_name in throttling_exceptions or
-                any(throttle_term.lower() in exception_message for throttle_term in throttling_exceptions)
-            )
-        
-        if is_throttling:
+        if is_throttling_exception(e):
             logger.error(f"Throttling exception detected: {type(e).__name__}. This will trigger state machine retry.")
             # Update document status before re-raising
             document_service.update_document(docStatus)