adding maxPagesForClassification configuration in UpdateSchemaConfig

Author: Taniya Mathur

lib/idp_common_pkg/idp_common/classification/service.py

@@ -325,8 +325,8 @@ def _classify_pages_multimodal(self, document: Document) -> Document:
             )
             all_page_results = list(cached_page_classifications.values())
             combined_metering = {}
-            errors_lock = threading.Lock()
-            failed_page_exceptions = {}
+            errors_lock = threading.Lock()  # Thread safety for error collection
+            failed_page_exceptions = {}  # Store original exceptions for failed pages
 
             # Determine which pages need classification
             pages_to_classify = {}
@@ -343,12 +343,13 @@ def _classify_pages_multimodal(self, document: Document) -> Document:
                         page_id
                     ].confidence = cached_result.classification.confidence
 
-                    # Copy metadata to the page
+                    # Copy metadata (including boundary information) to the page
                     if hasattr(document.pages[page_id], "metadata"):
                         document.pages[
                             page_id
                         ].metadata = cached_result.classification.metadata
                     else:
+                        # If the page doesn't have a metadata attribute, add it
                         setattr(
                             document.pages[page_id],
                             "metadata",
@@ -403,12 +404,13 @@ def _classify_pages_multimodal(self, document: Document) -> Document:
                                 page_id
                             ].confidence = page_result.classification.confidence
 
-                            # Copy metadata to the page
+                            # Copy metadata (including boundary information) to the page
                             if hasattr(document.pages[page_id], "metadata"):
                                 document.pages[
                                     page_id
                                 ].metadata = page_result.classification.metadata
                             else:
+                                # If the page doesn't have a metadata attribute, add it
                                 setattr(
                                     document.pages[page_id],
                                     "metadata",
@@ -423,10 +425,12 @@ def _classify_pages_multimodal(self, document: Document) -> Document:
                                 combined_metering, page_metering
                             )
                         except Exception as e:
+                            # Capture exception details in the document object instead of raising
                             error_msg = f"Error classifying page {page_id}: {str(e)}"
                             logger.error(error_msg)
                             with errors_lock:
                                 document.errors.append(error_msg)
+                                # Store the original exception for later use
                                 failed_page_exceptions[page_id] = e
 
                             # Mark page as unclassified on error
@@ -436,11 +440,13 @@ def _classify_pages_multimodal(self, document: Document) -> Document:
                                 ].classification = "error (backoff/retry)"
                                 document.pages[page_id].confidence = 0.0
 
-                # Store failed page exceptions in document metadata
+                # Store failed page exceptions in document metadata for caller to access
                 if failed_page_exceptions:
                     logger.info(
                         f"Processing {len(failed_page_exceptions)} failed page exceptions for document {document.id}"
                     )
+
+                    # Store the first encountered exception as the primary failure cause
                     first_exception = next(iter(failed_page_exceptions.values()))
                     document.metadata = document.metadata or {}
                     document.metadata["failed_page_exceptions"] = {
@@ -453,13 +459,37 @@ def _classify_pages_multimodal(self, document: Document) -> Document:
                         }
                         for page_id, exc in failed_page_exceptions.items()
                     }
+                    # Store the primary exception for easy access by caller
                     document.metadata["primary_exception"] = first_exception
+
+                    # Cache successful page classifications (only when some pages fail - for retry scenarios)
+                    successful_results = [
+                        r
+                        for r in all_page_results
+                        if "error" not in r.classification.metadata
+                    ]
+                    if successful_results:
+                        logger.info(
+                            f"Caching {len(successful_results)} successful page classifications for document {document.id} due to {len(failed_page_exceptions)} failed pages (retry scenario)"
+                        )
+                        self._cache_successful_page_classifications(
+                            document, successful_results
+                        )
+                    else:
+                        logger.warning(
+                            f"No successful page classifications to cache for document {document.id} - all {len(failed_page_exceptions)} pages failed"
+                        )
+                else:
+                    # All pages succeeded - no need to cache since there won't be retries
+                    logger.info(
+                        f"All pages succeeded for document {document.id} - skipping cache (no retry needed)"
+                    )
             else:
                 logger.info(
                     f"All {len(cached_page_classifications)} page classifications found in cache"
                 )
 
-            # Group pages into sections
+            # Group pages into sections only if we have results
             document.sections = []
             sorted_results = self._sort_page_results(all_page_results)
 
@@ -506,13 +536,19 @@ def _classify_pages_multimodal(self, document: Document) -> Document:
             )
 
             t1 = time.time()
-            logger.info(f"Time taken for classification: {t1 - t0:.2f} seconds")
+            logger.info(
+                f"Document classified with {len(document.sections)} sections in {t1 - t0:.2f} seconds"
+            )
 
         except Exception as e:
-            logger.error(f"Error in multimodal classification: {str(e)}")
+            error_msg = f"Error classifying all document pages: {str(e)}"
             document = self._update_document_status(
-                document, success=False, error_message=str(e)
+                document, success=False, error_message=error_msg
             )
+            # Store the exception in metadata for caller to access
+            document.metadata = document.metadata or {}
+            document.metadata["primary_exception"] = e
+            # raise exception to enable client retries
             raise
 
         return document
@@ -1647,6 +1683,9 @@ def classify_document(self, document: Document) -> Document:
             if limited_document.id != document.id:  # Pages were actually limited
                 # Classify the limited document
                 if self.classification_method == self.TEXTBASED_HOLISTIC:
+                    logger.info(
+                        f"Classifying limited document with {len(limited_document.pages)} pages using holistic packet method"
+                    )
                     classified_limited = self.holistic_classify_document(
                         limited_document
                     )

lib/idp_common_pkg/tests/test_max_pages_classification.py renamed to lib/idp_common_pkg/tests/unit/test_max_pages_classification.py

@@ -91,8 +91,6 @@ Parameters:
       - "false"
     Description: "Enable Human In The Loop (A2I) for document review"
 
-
-
   SageMakerA2IReviewPortalURL:
     Type: String
     Default: ""
@@ -515,33 +513,38 @@ Resources:
                 description: "Classification methodology to use"
                 enum: ["multimodalPageLevelClassification", "textbasedHolisticClassification"]
                 order: 2
+              maxPagesForClassification:
+                type: string
+                description: "Number of pages to use for classification [When set to a number, forces the entire document to be assigned a single class]"
+                enum: ["ALL", "1", "2", "3", "5", "10"]
+                order: 3
               temperature:
                 type: number
                 minimum: 0
                 maximum: 1
                 description: Sampling temperature
-                order: 3
+                order: 4
               top_k:
                 type: integer
                 minimum: 1
                 description: Sampling Top K
-                order: 4
+                order: 5
               top_p:
                 type: number
                 description: Sampling Top P
-                order: 5
+                order: 6
               max_tokens:
                 type: number
                 description: Max tokens
-                order: 6
+                order: 7
               system_prompt:
                 type: string
                 description: System prompt
-                order: 7
+                order: 8
               task_prompt:
                 type: string
                 description: Task prompt - include placeholders {CLASS_NAMES_AND_DESCRIPTIONS} (replaced with the class names and descriptions for all specified classes), {FEW_SHOT_EXAMPLES} (replaced by classPrompt and image data from examples in class definitions), {DOCUMENT_TEXT} (replaced by the OCR output), and for multi-modal classification {DOCUMENT_IMAGE} (replaced by the page image attachment). Optionally use <<CACHEPOINT>> to separate static and dynamic elements of prompt for Bedrock prompt caching.
-                order: 8
+                order: 9
           extraction:
             order: 4
             type: object

patterns/pattern-2/template.yaml

@@ -105,8 +105,6 @@ Parameters:
       - "false"
     Description: "Enable Human In The Loop (A2I) for document review"
 
-
-
   SageMakerA2IReviewPortalURL:
     Type: String
     Default: ""

patterns/pattern-3/template.yaml

@@ -114,8 +114,6 @@ Parameters:
     Description: >-
      Select the configuration preset for Pattern 1. Each configuration contains pre-tuned settings for specific document processing scenarios - see https://github.com/aws-samples/sample-genai-idp/blob/main/config_library/README.md. Note: This selected configuration will be replaced by the Custom Configuration Path if specified.
 
-
-    
   # Pattern 2 Parameters
 
   Pattern2Configuration: