@cdklabs/cdk-appmod-catalog-blueprints 1.5.0 → 1.6.0

package/lib/document-processing/resources/cleanup/test_handler.py

@@ -0,0 +1,436 @@
+"""
+Unit tests for cleanup Lambda handler
+
+Tests cover:
+- Batch delete with multiple chunks
+- Error handling for S3 delete failures
+- Non-blocking error handling
+
+Requirements: 8.4
+"""
+
+import json
+import unittest
+from unittest.mock import MagicMock, patch, call
+
+from handler import handler, delete_chunks_batch
+
+
+class TestCleanupHandler(unittest.TestCase):
+    """Test cases for cleanup Lambda handler"""
+    
+    def setUp(self):
+        """Set up test fixtures"""
+        self.document_id = "test-doc-123"
+        self.bucket = "test-bucket"
+        
+        # Sample chunks
+        self.chunks = [
+            {
+                'chunkId': f'{self.document_id}_chunk_0',
+                'bucket': self.bucket,
+                'key': f'chunks/{self.document_id}_chunk_0.pdf'
+            },
+            {
+                'chunkId': f'{self.document_id}_chunk_1',
+                'bucket': self.bucket,
+                'key': f'chunks/{self.document_id}_chunk_1.pdf'
+            },
+            {
+                'chunkId': f'{self.document_id}_chunk_2',
+                'bucket': self.bucket,
+                'key': f'chunks/{self.document_id}_chunk_2.pdf'
+            }
+        ]
+    
+    @patch('handler.s3_client')
+    def test_successful_cleanup(self, mock_s3):
+        """Test successful cleanup of all chunks"""
+        # Mock successful S3 delete
+        mock_s3.delete_objects.return_value = {
+            'Deleted': [
+                {'Key': chunk['key']} for chunk in self.chunks
+            ],
+            'Errors': []
+        }
+        
+        # Create event
+        event = {
+            'documentId': self.document_id,
+            'chunks': self.chunks
+        }
+        
+        # Call handler
+        result = handler(event, None)
+        
+        # Verify response
+        self.assertEqual(result['documentId'], self.document_id)
+        self.assertEqual(result['deletedChunks'], 3)
+        self.assertEqual(result['errors'], [])
+        
+        # Verify S3 delete was called
+        mock_s3.delete_objects.assert_called_once()
+        call_args = mock_s3.delete_objects.call_args
+        self.assertEqual(call_args[1]['Bucket'], self.bucket)
+        self.assertEqual(len(call_args[1]['Delete']['Objects']), 3)
+    
+    @patch('handler.s3_client')
+    def test_partial_delete_failure(self, mock_s3):
+        """Test cleanup with some S3 delete failures"""
+        # Mock partial S3 delete failure
+        mock_s3.delete_objects.return_value = {
+            'Deleted': [
+                {'Key': self.chunks[0]['key']},
+                {'Key': self.chunks[1]['key']}
+            ],
+            'Errors': [
+                {
+                    'Key': self.chunks[2]['key'],
+                    'Code': 'AccessDenied',
+                    'Message': 'Access Denied'
+                }
+            ]
+        }
+        
+        # Create event
+        event = {
+            'documentId': self.document_id,
+            'chunks': self.chunks
+        }
+        
+        # Call handler
+        result = handler(event, None)
+        
+        # Verify response
+        self.assertEqual(result['documentId'], self.document_id)
+        self.assertEqual(result['deletedChunks'], 2)
+        self.assertEqual(len(result['errors']), 1)
+        self.assertIn('AccessDenied', result['errors'][0])
+        self.assertIn(self.chunks[2]['key'], result['errors'][0])
+    
+    @patch('handler.s3_client')
+    def test_complete_delete_failure(self, mock_s3):
+        """Test cleanup when all S3 deletes fail"""
+        # Mock complete S3 delete failure
+        from botocore.exceptions import ClientError
+        mock_s3.delete_objects.side_effect = ClientError(
+            {
+                'Error': {
+                    'Code': 'NoSuchBucket',
+                    'Message': 'The specified bucket does not exist'
+                }
+            },
+            'DeleteObjects'
+        )
+        
+        # Create event
+        event = {
+            'documentId': self.document_id,
+            'chunks': self.chunks
+        }
+        
+        # Call handler
+        result = handler(event, None)
+        
+        # Verify response - should not fail workflow
+        self.assertEqual(result['documentId'], self.document_id)
+        self.assertEqual(result['deletedChunks'], 0)
+        self.assertEqual(len(result['errors']), 1)
+        self.assertIn('NoSuchBucket', result['errors'][0])
+    
+    @patch('handler.s3_client')
+    def test_batch_delete_with_multiple_chunks(self, mock_s3):
+        """Test batch delete with multiple chunks"""
+        # Create 5 chunks
+        chunks = [
+            {
+                'chunkId': f'{self.document_id}_chunk_{i}',
+                'bucket': self.bucket,
+                'key': f'chunks/{self.document_id}_chunk_{i}.pdf'
+            }
+            for i in range(5)
+        ]
+        
+        # Mock successful S3 delete
+        mock_s3.delete_objects.return_value = {
+            'Deleted': [{'Key': chunk['key']} for chunk in chunks],
+            'Errors': []
+        }
+        
+        # Create event
+        event = {
+            'documentId': self.document_id,
+            'chunks': chunks
+        }
+        
+        # Call handler
+        result = handler(event, None)
+        
+        # Verify response
+        self.assertEqual(result['documentId'], self.document_id)
+        self.assertEqual(result['deletedChunks'], 5)
+        self.assertEqual(result['errors'], [])
+    
+    @patch('handler.s3_client')
+    def test_batch_delete_respects_1000_limit(self, mock_s3):
+        """Test that batch delete respects S3's 1000 object limit"""
+        # Create 1500 chunks (should require 2 batches)
+        chunks = [
+            {
+                'chunkId': f'{self.document_id}_chunk_{i}',
+                'bucket': self.bucket,
+                'key': f'chunks/{self.document_id}_chunk_{i}.pdf'
+            }
+            for i in range(1500)
+        ]
+        
+        # Mock successful S3 delete
+        def mock_delete_objects(**kwargs):
+            objects = kwargs['Delete']['Objects']
+            return {
+                'Deleted': [{'Key': obj['Key']} for obj in objects],
+                'Errors': []
+            }
+        
+        mock_s3.delete_objects.side_effect = mock_delete_objects
+        
+        # Create event
+        event = {
+            'documentId': self.document_id,
+            'chunks': chunks
+        }
+        
+        # Call handler
+        result = handler(event, None)
+        
+        # Verify response
+        self.assertEqual(result['documentId'], self.document_id)
+        self.assertEqual(result['deletedChunks'], 1500)
+        self.assertEqual(result['errors'], [])
+        
+        # Verify S3 delete was called twice (2 batches)
+        self.assertEqual(mock_s3.delete_objects.call_count, 2)
+        
+        # Verify first batch has 1000 objects
+        first_call = mock_s3.delete_objects.call_args_list[0]
+        self.assertEqual(len(first_call[1]['Delete']['Objects']), 1000)
+        
+        # Verify second batch has 500 objects
+        second_call = mock_s3.delete_objects.call_args_list[1]
+        self.assertEqual(len(second_call[1]['Delete']['Objects']), 500)
+    
+    def test_missing_document_id(self):
+        """Test handler with missing documentId"""
+        event = {
+            'chunks': self.chunks
+        }
+        
+        result = handler(event, None)
+        
+        # Verify error response
+        self.assertIsNone(result['documentId'])
+        self.assertEqual(result['deletedChunks'], 0)
+        self.assertEqual(len(result['errors']), 1)
+        self.assertIn('documentId', result['errors'][0])
+    
+    def test_empty_chunks_array(self):
+        """Test handler with empty chunks array"""
+        event = {
+            'documentId': self.document_id,
+            'chunks': []
+        }
+        
+        result = handler(event, None)
+        
+        # Verify response
+        self.assertEqual(result['documentId'], self.document_id)
+        self.assertEqual(result['deletedChunks'], 0)
+        self.assertEqual(result['errors'], [])
+    
+    def test_chunks_missing_bucket_or_key(self):
+        """Test handler with chunks missing bucket or key"""
+        chunks = [
+            {
+                'chunkId': f'{self.document_id}_chunk_0',
+                'bucket': self.bucket,
+                # Missing 'key'
+            },
+            {
+                'chunkId': f'{self.document_id}_chunk_1',
+                # Missing 'bucket'
+                'key': f'chunks/{self.document_id}_chunk_1.pdf'
+            },
+            {
+                'chunkId': f'{self.document_id}_chunk_2',
+                'bucket': self.bucket,
+                'key': f'chunks/{self.document_id}_chunk_2.pdf'
+            }
+        ]
+        
+        event = {
+            'documentId': self.document_id,
+            'chunks': chunks
+        }
+        
+        # Mock S3 delete for the valid chunk
+        with patch('handler.s3_client') as mock_s3:
+            mock_s3.delete_objects.return_value = {
+                'Deleted': [{'Key': chunks[2]['key']}],
+                'Errors': []
+            }
+            
+            result = handler(event, None)
+            
+            # Verify only valid chunk was deleted
+            self.assertEqual(result['documentId'], self.document_id)
+            self.assertEqual(result['deletedChunks'], 1)
+    
+    @patch('handler.s3_client')
+    def test_multiple_buckets(self, mock_s3):
+        """Test cleanup with chunks in multiple buckets"""
+        chunks = [
+            {
+                'chunkId': f'{self.document_id}_chunk_0',
+                'bucket': 'bucket-1',
+                'key': f'chunks/{self.document_id}_chunk_0.pdf'
+            },
+            {
+                'chunkId': f'{self.document_id}_chunk_1',
+                'bucket': 'bucket-2',
+                'key': f'chunks/{self.document_id}_chunk_1.pdf'
+            },
+            {
+                'chunkId': f'{self.document_id}_chunk_2',
+                'bucket': 'bucket-1',
+                'key': f'chunks/{self.document_id}_chunk_2.pdf'
+            }
+        ]
+        
+        # Mock successful S3 delete
+        def mock_delete_objects(**kwargs):
+            objects = kwargs['Delete']['Objects']
+            return {
+                'Deleted': [{'Key': obj['Key']} for obj in objects],
+                'Errors': []
+            }
+        
+        mock_s3.delete_objects.side_effect = mock_delete_objects
+        
+        # Create event
+        event = {
+            'documentId': self.document_id,
+            'chunks': chunks
+        }
+        
+        # Call handler
+        result = handler(event, None)
+        
+        # Verify response
+        self.assertEqual(result['documentId'], self.document_id)
+        self.assertEqual(result['deletedChunks'], 3)
+        self.assertEqual(result['errors'], [])
+        
+        # Verify S3 delete was called twice (once per bucket)
+        self.assertEqual(mock_s3.delete_objects.call_count, 2)
+    
+    @patch('handler.s3_client')
+    def test_non_blocking_error_handling(self, mock_s3):
+        """Test that errors don't block workflow completion"""
+        # Mock S3 delete failure
+        from botocore.exceptions import ClientError
+        mock_s3.delete_objects.side_effect = ClientError(
+            {
+                'Error': {
+                    'Code': 'InternalError',
+                    'Message': 'Internal Server Error'
+                }
+            },
+            'DeleteObjects'
+        )
+        
+        # Create event
+        event = {
+            'documentId': self.document_id,
+            'chunks': self.chunks
+        }
+        
+        # Call handler - should not raise exception
+        result = handler(event, None)
+        
+        # Verify response contains error but doesn't fail
+        self.assertEqual(result['documentId'], self.document_id)
+        self.assertEqual(result['deletedChunks'], 0)
+        self.assertGreater(len(result['errors']), 0)
+        self.assertIn('InternalError', result['errors'][0])
+
+
+class TestDeleteChunksBatch(unittest.TestCase):
+    """Test cases for delete_chunks_batch function"""
+    
+    @patch('handler.s3_client')
+    def test_delete_chunks_batch_success(self, mock_s3):
+        """Test successful batch delete"""
+        chunk_keys = [
+            {
+                'bucket': 'test-bucket',
+                'key': 'chunks/chunk_0.pdf',
+                'chunkId': 'chunk_0'
+            },
+            {
+                'bucket': 'test-bucket',
+                'key': 'chunks/chunk_1.pdf',
+                'chunkId': 'chunk_1'
+            }
+        ]
+        
+        mock_s3.delete_objects.return_value = {
+            'Deleted': [
+                {'Key': 'chunks/chunk_0.pdf'},
+                {'Key': 'chunks/chunk_1.pdf'}
+            ],
+            'Errors': []
+        }
+        
+        deleted_count, errors = delete_chunks_batch(chunk_keys, 'test-doc')
+        
+        self.assertEqual(deleted_count, 2)
+        self.assertEqual(errors, [])
+    
+    @patch('handler.s3_client')
+    def test_delete_chunks_batch_with_errors(self, mock_s3):
+        """Test batch delete with some errors"""
+        chunk_keys = [
+            {
+                'bucket': 'test-bucket',
+                'key': 'chunks/chunk_0.pdf',
+                'chunkId': 'chunk_0'
+            },
+            {
+                'bucket': 'test-bucket',
+                'key': 'chunks/chunk_1.pdf',
+                'chunkId': 'chunk_1'
+            }
+        ]
+        
+        mock_s3.delete_objects.return_value = {
+            'Deleted': [
+                {'Key': 'chunks/chunk_0.pdf'}
+            ],
+            'Errors': [
+                {
+                    'Key': 'chunks/chunk_1.pdf',
+                    'Code': 'AccessDenied',
+                    'Message': 'Access Denied'
+                }
+            ]
+        }
+        
+        deleted_count, errors = delete_chunks_batch(chunk_keys, 'test-doc')
+        
+        self.assertEqual(deleted_count, 1)
+        self.assertEqual(len(errors), 1)
+        self.assertIn('AccessDenied', errors[0])
+
+
+if __name__ == '__main__':
+    unittest.main()

package/lib/document-processing/resources/default-bedrock-invoke/index.py

@@ -10,6 +10,67 @@ bedrock = boto3.client('bedrock-runtime')
 metrics = Metrics()
 tracer = Tracer()
 
+
+def parse_chunk_metadata(event):
+    """
+    Parse optional chunk metadata from the event payload.
+    
+    Returns a dictionary with chunk information if present, None otherwise.
+    Supports both direct chunk metadata and nested chunk object format.
+    """
+    # Check for direct chunkMetadata field
+    if 'chunkMetadata' in event:
+        return event['chunkMetadata']
+    
+    # Check for chunk object (from Map State iteration)
+    if 'chunk' in event:
+        chunk = event['chunk']
+        return {
+            'chunkIndex': chunk.get('chunkIndex', event.get('chunkIndex', 0)),
+            'totalChunks': event.get('totalChunks', 1),
+            'startPage': chunk.get('startPage', 0),
+            'endPage': chunk.get('endPage', 0),
+            'pageCount': chunk.get('pageCount', 0),
+            'estimatedTokens': chunk.get('estimatedTokens', 0),
+            'overlapPages': chunk.get('overlapPages', 0),
+        }
+    
+    return None
+
+
+def build_chunk_context_prompt(chunk_metadata):
+    """
+    Build a context prompt for chunk-aware processing.
+    
+    Args:
+        chunk_metadata: Dictionary containing chunk information
+        
+    Returns:
+        String with chunk context to prepend to the main prompt
+    """
+    if not chunk_metadata:
+        return ""
+    
+    chunk_index = chunk_metadata.get('chunkIndex', 0)
+    total_chunks = chunk_metadata.get('totalChunks', 1)
+    start_page = chunk_metadata.get('startPage', 0)
+    end_page = chunk_metadata.get('endPage', 0)
+    overlap_pages = chunk_metadata.get('overlapPages', 0)
+    
+    # Build context string
+    context_parts = [
+        f"You are analyzing chunk {chunk_index + 1} of {total_chunks} from pages {start_page + 1} to {end_page + 1}."
+    ]
+    
+    # Add overlap information if applicable
+    if overlap_pages > 0 and chunk_index > 0:
+        context_parts.append(
+            f"Note: This chunk includes {overlap_pages} overlapping pages from the previous chunk for context."
+        )
+    
+    return "\n".join(context_parts) + "\n\n"
+
+
 @metrics.log_metrics
 @tracer.capture_lambda_handler
 def handler(event, context):
@@ -18,20 +79,40 @@ def handler(event, context):
     tracer.put_annotation(key="documentId", value=event["documentId"])
     metrics.add_dimension(name="invoke_type", value=invoke_type)
     content_type = event["contentType"]
+    
+    # Parse optional chunk metadata
+    chunk_metadata = parse_chunk_metadata(event)
+    if chunk_metadata:
+        tracer.put_annotation(key="chunkIndex", value=str(chunk_metadata.get('chunkIndex', 0)))
+        tracer.put_annotation(key="totalChunks", value=str(chunk_metadata.get('totalChunks', 1)))
+        metrics.add_dimension(name="is_chunked", value="true")
+    else:
+        metrics.add_dimension(name="is_chunked", value="false")
+    
     # Format prompt if classification result exists
     prompt = os.environ['PROMPT']
     if 'classificationResult' in event:
         classification = event['classificationResult']['documentClassification']
         prompt = prompt.replace("[ACTUAL_CLASSIFICATION]", classification)
     
+    # Add chunk context to prompt if processing a chunk
+    chunk_context = build_chunk_context_prompt(chunk_metadata)
+    if chunk_context:
+        prompt = chunk_context + prompt
+    
     # Build content based on file type
     content = [{'type': 'text', 'text': prompt}]
     if content_type == 'file':
         content_location = event['content']['location']
         
         if content_location == 's3':
-            bucket = event['content']['bucket']
-            key = event['content']['key']
+            # Use chunk-specific S3 location if available, otherwise use original content
+            if chunk_metadata and 'bucket' in chunk_metadata and 'key' in chunk_metadata:
+                bucket = chunk_metadata['bucket']
+                key = chunk_metadata['key']
+            else:
+                bucket = event['content']['bucket']
+                key = event['content']['key']
             
             # Check file type
             ext = key.lower().split('.')[-1]
@@ -61,11 +142,12 @@ def handler(event, context):
         })
             
     # Invoke Bedrock
+    max_tokens = int(os.getenv('INVOKE_MAX_TOKENS', '1000'))
     response = bedrock.invoke_model(
         modelId=os.environ['MODEL_ID'],
         body=json.dumps({
             'anthropic_version': 'bedrock-2023-05-31',
-            'max_tokens': os.getenv('INVOKE_MAX_TOKENS', 1000),
+            'max_tokens': max_tokens,
             'messages': [{'role': 'user', 'content': content}]
         })
     )