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

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

@@ -0,0 +1,622 @@
+"""
+Unit tests for the default-bedrock-invoke Lambda handler.
+
+Tests chunk-aware classification and processing functionality including:
+- Parsing chunk metadata from events
+- Building chunk context prompts
+- Backward compatibility with non-chunked documents
+- Processing with chunk metadata
+- Processing without chunk metadata (backward compatibility)
+"""
+
+import pytest
+from unittest.mock import patch, MagicMock
+import json
+import os
+
+# Set required environment variables before importing the module
+os.environ['INVOKE_TYPE'] = 'classification'
+os.environ['PROMPT'] = 'Test prompt'
+os.environ['MODEL_ID'] = 'test-model'
+
+# Import after setting environment variables
+from index import parse_chunk_metadata, build_chunk_context_prompt, handler
+
+
+class TestParseChunkMetadata:
+    """Tests for parse_chunk_metadata function."""
+    
+    def test_returns_none_when_no_chunk_metadata(self):
+        """Test that None is returned when no chunk metadata is present."""
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'file',
+            'content': {'bucket': 'test', 'key': 'test.pdf'}
+        }
+        result = parse_chunk_metadata(event)
+        assert result is None
+    
+    def test_parses_direct_chunk_metadata(self):
+        """Test parsing direct chunkMetadata field."""
+        event = {
+            'documentId': 'doc-123',
+            'chunkMetadata': {
+                'chunkIndex': 2,
+                'totalChunks': 5,
+                'startPage': 100,
+                'endPage': 149,
+                'pageCount': 50,
+                'estimatedTokens': 75000,
+                'overlapPages': 5
+            }
+        }
+        result = parse_chunk_metadata(event)
+        
+        assert result is not None
+        assert result['chunkIndex'] == 2
+        assert result['totalChunks'] == 5
+        assert result['startPage'] == 100
+        assert result['endPage'] == 149
+        assert result['pageCount'] == 50
+        assert result['estimatedTokens'] == 75000
+        assert result['overlapPages'] == 5
+    
+    def test_parses_chunk_object_from_map_state(self):
+        """Test parsing chunk object from Step Functions Map State iteration."""
+        event = {
+            'documentId': 'doc-123',
+            'chunk': {
+                'chunkId': 'doc-123_chunk_1',
+                'chunkIndex': 1,
+                'startPage': 50,
+                'endPage': 99,
+                'pageCount': 50,
+                'estimatedTokens': 80000,
+                'bucket': 'test-bucket',
+                'key': 'chunks/doc-123_chunk_1.pdf'
+            },
+            'chunkIndex': 1,
+            'totalChunks': 3
+        }
+        result = parse_chunk_metadata(event)
+        
+        assert result is not None
+        assert result['chunkIndex'] == 1
+        assert result['totalChunks'] == 3
+        assert result['startPage'] == 50
+        assert result['endPage'] == 99
+    
+    def test_handles_missing_optional_fields(self):
+        """Test that missing optional fields default to 0."""
+        event = {
+            'documentId': 'doc-123',
+            'chunk': {
+                'chunkId': 'doc-123_chunk_0',
+                'startPage': 0,
+                'endPage': 49
+            },
+            'totalChunks': 2
+        }
+        result = parse_chunk_metadata(event)
+        
+        assert result is not None
+        assert result['chunkIndex'] == 0
+        assert result['overlapPages'] == 0
+        assert result['estimatedTokens'] == 0
+
+
+class TestBuildChunkContextPrompt:
+    """Tests for build_chunk_context_prompt function."""
+    
+    def test_returns_empty_string_when_no_metadata(self):
+        """Test that empty string is returned when no metadata is provided."""
+        result = build_chunk_context_prompt(None)
+        assert result == ""
+    
+    def test_builds_basic_chunk_context(self):
+        """Test building basic chunk context without overlap."""
+        chunk_metadata = {
+            'chunkIndex': 0,
+            'totalChunks': 3,
+            'startPage': 0,
+            'endPage': 49,
+            'overlapPages': 0
+        }
+        result = build_chunk_context_prompt(chunk_metadata)
+        
+        assert "You are analyzing chunk 1 of 3" in result
+        assert "from pages 1 to 50" in result
+        assert "overlapping" not in result
+    
+    def test_builds_chunk_context_with_overlap(self):
+        """Test building chunk context with overlap information."""
+        chunk_metadata = {
+            'chunkIndex': 1,
+            'totalChunks': 3,
+            'startPage': 45,
+            'endPage': 99,
+            'overlapPages': 5
+        }
+        result = build_chunk_context_prompt(chunk_metadata)
+        
+        assert "You are analyzing chunk 2 of 3" in result
+        assert "from pages 46 to 100" in result
+        assert "5 overlapping pages" in result
+    
+    def test_no_overlap_message_for_first_chunk(self):
+        """Test that first chunk doesn't mention overlap even if overlapPages > 0."""
+        chunk_metadata = {
+            'chunkIndex': 0,  # First chunk
+            'totalChunks': 3,
+            'startPage': 0,
+            'endPage': 49,
+            'overlapPages': 5  # Has overlap config but is first chunk
+        }
+        result = build_chunk_context_prompt(chunk_metadata)
+        
+        assert "You are analyzing chunk 1 of 3" in result
+        assert "overlapping" not in result  # First chunk shouldn't mention overlap
+    
+    def test_format_matches_design_spec(self):
+        """Test that format matches the design specification."""
+        chunk_metadata = {
+            'chunkIndex': 2,
+            'totalChunks': 5,
+            'startPage': 100,
+            'endPage': 149,
+            'overlapPages': 0
+        }
+        result = build_chunk_context_prompt(chunk_metadata)
+        
+        # Format should be: "You are analyzing chunk {N} of {total} from pages {start} to {end}"
+        expected_format = "You are analyzing chunk 3 of 5 from pages 101 to 150"
+        assert expected_format in result
+
+
+class TestHandlerChunkAwareness:
+    """Tests for handler function chunk awareness."""
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_handler_without_chunk_metadata(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test handler processes non-chunked documents correctly (backward compatibility)."""
+        # Setup mocks
+        mock_bedrock.invoke_model.return_value = {
+            'body': MagicMock(read=lambda: json.dumps({
+                'content': [{'text': '{"documentClassification": "INVOICE"}'}]
+            }).encode())
+        }
+        mock_s3.download_file = MagicMock()
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'data',
+            'content': {'data': 'Test document content'}
+        }
+        
+        # Call handler
+        with patch.dict(os.environ, {'INVOKE_TYPE': 'classification', 'PROMPT': 'Classify this:', 'MODEL_ID': 'test-model'}):
+            result = handler(event, None)
+        
+        # Verify result
+        assert result == {'documentClassification': 'INVOICE'}
+        
+        # Verify is_chunked dimension is set to false
+        mock_metrics.add_dimension.assert_any_call(name="is_chunked", value="false")
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_handler_with_chunk_metadata(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test handler processes chunked documents with context."""
+        # Setup mocks
+        mock_bedrock.invoke_model.return_value = {
+            'body': MagicMock(read=lambda: json.dumps({
+                'content': [{'text': '{"documentClassification": "CONTRACT"}'}]
+            }).encode())
+        }
+        mock_s3.download_file = MagicMock()
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'data',
+            'content': {'data': 'Test document content'},
+            'chunkMetadata': {
+                'chunkIndex': 1,
+                'totalChunks': 3,
+                'startPage': 50,
+                'endPage': 99,
+                'overlapPages': 5
+            }
+        }
+        
+        # Call handler
+        with patch.dict(os.environ, {'INVOKE_TYPE': 'classification', 'PROMPT': 'Classify this:', 'MODEL_ID': 'test-model'}):
+            result = handler(event, None)
+        
+        # Verify result
+        assert result == {'documentClassification': 'CONTRACT'}
+        
+        # Verify is_chunked dimension is set to true
+        mock_metrics.add_dimension.assert_any_call(name="is_chunked", value="true")
+        
+        # Verify chunk annotations are set
+        mock_tracer.put_annotation.assert_any_call(key="chunkIndex", value="1")
+        mock_tracer.put_annotation.assert_any_call(key="totalChunks", value="3")
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_prompt_includes_chunk_context(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test that the prompt sent to Bedrock includes chunk context."""
+        captured_body = None
+        
+        def capture_invoke(modelId, body):
+            nonlocal captured_body
+            captured_body = json.loads(body)
+            return {
+                'body': MagicMock(read=lambda: json.dumps({
+                    'content': [{'text': '{"documentClassification": "INVOICE"}'}]
+                }).encode())
+            }
+        
+        mock_bedrock.invoke_model = capture_invoke
+        mock_s3.download_file = MagicMock()
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'data',
+            'content': {'data': 'Test document content'},
+            'chunkMetadata': {
+                'chunkIndex': 2,
+                'totalChunks': 5,
+                'startPage': 100,
+                'endPage': 149,
+                'overlapPages': 0
+            }
+        }
+        
+        # Call handler
+        with patch.dict(os.environ, {'INVOKE_TYPE': 'classification', 'PROMPT': 'Classify this document:', 'MODEL_ID': 'test-model'}):
+            handler(event, None)
+        
+        # Verify the prompt includes chunk context
+        assert captured_body is not None
+        prompt_text = captured_body['messages'][0]['content'][0]['text']
+        assert "You are analyzing chunk 3 of 5" in prompt_text
+        assert "from pages 101 to 150" in prompt_text
+        assert "Classify this document:" in prompt_text
+
+
+if __name__ == '__main__':
+    pytest.main([__file__, '-v'])
+
+
+class TestProcessingChunkAwareness:
+    """Tests for handler function chunk awareness in processing mode."""
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_processing_without_chunk_metadata(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test processing handler processes non-chunked documents correctly (backward compatibility)."""
+        # Setup mocks
+        mock_bedrock.invoke_model.return_value = {
+            'body': MagicMock(read=lambda: json.dumps({
+                'content': [{'text': '{"documentClassification": "INVOICE", "result": {"entities": [{"type": "AMOUNT", "value": "$100.00"}]}}'}]
+            }).encode())
+        }
+        mock_s3.download_file = MagicMock()
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'data',
+            'content': {'data': 'Test document content'},
+            'classificationResult': {'documentClassification': 'INVOICE'}
+        }
+        
+        # Call handler with processing mode
+        with patch.dict(os.environ, {
+            'INVOKE_TYPE': 'processing',
+            'PROMPT': 'Extract entities from this [ACTUAL_CLASSIFICATION]:',
+            'MODEL_ID': 'test-model'
+        }):
+            result = handler(event, None)
+        
+        # Verify result
+        assert result['documentClassification'] == 'INVOICE'
+        assert 'result' in result
+        assert 'entities' in result['result']
+        
+        # Verify is_chunked dimension is set to false
+        mock_metrics.add_dimension.assert_any_call(name="is_chunked", value="false")
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_processing_with_chunk_metadata(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test processing handler processes chunked documents with context."""
+        # Setup mocks
+        mock_bedrock.invoke_model.return_value = {
+            'body': MagicMock(read=lambda: json.dumps({
+                'content': [{'text': '{"documentClassification": "CONTRACT", "result": {"entities": [{"type": "PARTY", "value": "Acme Corp", "page": 51}]}}'}]
+            }).encode())
+        }
+        mock_s3.download_file = MagicMock()
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'data',
+            'content': {'data': 'Test document content'},
+            'classificationResult': {'documentClassification': 'CONTRACT'},
+            'chunkMetadata': {
+                'chunkIndex': 1,
+                'totalChunks': 3,
+                'startPage': 50,
+                'endPage': 99,
+                'overlapPages': 5
+            }
+        }
+        
+        # Call handler with processing mode
+        with patch.dict(os.environ, {
+            'INVOKE_TYPE': 'processing',
+            'PROMPT': 'Extract entities from this [ACTUAL_CLASSIFICATION]:',
+            'MODEL_ID': 'test-model'
+        }):
+            result = handler(event, None)
+        
+        # Verify result
+        assert result['documentClassification'] == 'CONTRACT'
+        assert 'result' in result
+        
+        # Verify is_chunked dimension is set to true
+        mock_metrics.add_dimension.assert_any_call(name="is_chunked", value="true")
+        
+        # Verify chunk annotations are set
+        mock_tracer.put_annotation.assert_any_call(key="chunkIndex", value="1")
+        mock_tracer.put_annotation.assert_any_call(key="totalChunks", value="3")
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_processing_prompt_includes_chunk_context(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test that the processing prompt sent to Bedrock includes chunk context."""
+        captured_body = None
+        
+        def capture_invoke(modelId, body):
+            nonlocal captured_body
+            captured_body = json.loads(body)
+            return {
+                'body': MagicMock(read=lambda: json.dumps({
+                    'content': [{'text': '{"documentClassification": "INVOICE", "result": {"entities": []}}'}]
+                }).encode())
+            }
+        
+        mock_bedrock.invoke_model = capture_invoke
+        mock_s3.download_file = MagicMock()
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'data',
+            'content': {'data': 'Test document content'},
+            'classificationResult': {'documentClassification': 'INVOICE'},
+            'chunkMetadata': {
+                'chunkIndex': 2,
+                'totalChunks': 5,
+                'startPage': 100,
+                'endPage': 149,
+                'overlapPages': 0
+            }
+        }
+        
+        # Call handler with processing mode
+        with patch.dict(os.environ, {
+            'INVOKE_TYPE': 'processing',
+            'PROMPT': 'Extract entities from this [ACTUAL_CLASSIFICATION]:',
+            'MODEL_ID': 'test-model'
+        }):
+            handler(event, None)
+        
+        # Verify the prompt includes chunk context
+        assert captured_body is not None
+        prompt_text = captured_body['messages'][0]['content'][0]['text']
+        assert "You are analyzing chunk 3 of 5" in prompt_text
+        assert "from pages 101 to 150" in prompt_text
+        assert "Extract entities from this INVOICE:" in prompt_text
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_processing_prompt_includes_overlap_info(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test that the processing prompt includes overlap information when applicable."""
+        captured_body = None
+        
+        def capture_invoke(modelId, body):
+            nonlocal captured_body
+            captured_body = json.loads(body)
+            return {
+                'body': MagicMock(read=lambda: json.dumps({
+                    'content': [{'text': '{"documentClassification": "CONTRACT", "result": {"entities": []}}'}]
+                }).encode())
+            }
+        
+        mock_bedrock.invoke_model = capture_invoke
+        mock_s3.download_file = MagicMock()
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'data',
+            'content': {'data': 'Test document content'},
+            'classificationResult': {'documentClassification': 'CONTRACT'},
+            'chunkMetadata': {
+                'chunkIndex': 2,  # Not first chunk
+                'totalChunks': 4,
+                'startPage': 95,
+                'endPage': 144,
+                'overlapPages': 5
+            }
+        }
+        
+        # Call handler with processing mode
+        with patch.dict(os.environ, {
+            'INVOKE_TYPE': 'processing',
+            'PROMPT': 'Extract entities from this [ACTUAL_CLASSIFICATION]:',
+            'MODEL_ID': 'test-model'
+        }):
+            handler(event, None)
+        
+        # Verify the prompt includes overlap information
+        assert captured_body is not None
+        prompt_text = captured_body['messages'][0]['content'][0]['text']
+        assert "You are analyzing chunk 3 of 4" in prompt_text
+        assert "5 overlapping pages" in prompt_text
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_processing_uses_chunk_s3_location(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test that processing uses chunk-specific S3 location when available."""
+        captured_bucket = None
+        captured_key = None
+        
+        def capture_download(bucket, key, local_path):
+            nonlocal captured_bucket, captured_key
+            captured_bucket = bucket
+            captured_key = key
+            # Create a dummy file
+            with open(local_path, 'wb') as f:
+                f.write(b'%PDF-1.4 dummy content')
+        
+        mock_bedrock.invoke_model.return_value = {
+            'body': MagicMock(read=lambda: json.dumps({
+                'content': [{'text': '{"documentClassification": "INVOICE", "result": {"entities": []}}'}]
+            }).encode())
+        }
+        mock_s3.download_file = capture_download
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'file',
+            'content': {
+                'location': 's3',
+                'bucket': 'original-bucket',
+                'key': 'raw/document.pdf'
+            },
+            'classificationResult': {'documentClassification': 'INVOICE'},
+            'chunkMetadata': {
+                'chunkIndex': 1,
+                'totalChunks': 3,
+                'startPage': 50,
+                'endPage': 99,
+                'bucket': 'chunk-bucket',
+                'key': 'chunks/doc-123_chunk_1.pdf'
+            }
+        }
+        
+        # Call handler with processing mode
+        with patch.dict(os.environ, {
+            'INVOKE_TYPE': 'processing',
+            'PROMPT': 'Extract entities from this [ACTUAL_CLASSIFICATION]:',
+            'MODEL_ID': 'test-model'
+        }):
+            handler(event, None)
+        
+        # Verify chunk-specific S3 location was used
+        assert captured_bucket == 'chunk-bucket'
+        assert captured_key == 'chunks/doc-123_chunk_1.pdf'
+    
+    @patch('index.bedrock')
+    @patch('index.s3')
+    @patch('index.tracer')
+    @patch('index.metrics')
+    def test_processing_classification_replacement(self, mock_metrics, mock_tracer, mock_s3, mock_bedrock):
+        """Test that [ACTUAL_CLASSIFICATION] placeholder is replaced in processing prompt."""
+        captured_body = None
+        
+        def capture_invoke(modelId, body):
+            nonlocal captured_body
+            captured_body = json.loads(body)
+            return {
+                'body': MagicMock(read=lambda: json.dumps({
+                    'content': [{'text': '{"documentClassification": "RECEIPT", "result": {"entities": []}}'}]
+                }).encode())
+            }
+        
+        mock_bedrock.invoke_model = capture_invoke
+        mock_s3.download_file = MagicMock()
+        mock_tracer.put_annotation = MagicMock()
+        mock_tracer.capture_lambda_handler = lambda f: f
+        mock_metrics.log_metrics = lambda f: f
+        mock_metrics.add_dimension = MagicMock()
+        mock_metrics.add_metric = MagicMock()
+        
+        event = {
+            'documentId': 'doc-123',
+            'contentType': 'data',
+            'content': {'data': 'Test document content'},
+            'classificationResult': {'documentClassification': 'RECEIPT'}
+        }
+        
+        # Call handler with processing mode
+        with patch.dict(os.environ, {
+            'INVOKE_TYPE': 'processing',
+            'PROMPT': 'The document is classified as [ACTUAL_CLASSIFICATION]. Extract entities:',
+            'MODEL_ID': 'test-model'
+        }):
+            handler(event, None)
+        
+        # Verify the classification was replaced in the prompt
+        assert captured_body is not None
+        prompt_text = captured_body['messages'][0]['content'][0]['text']
+        assert "[ACTUAL_CLASSIFICATION]" not in prompt_text
+        assert "RECEIPT" in prompt_text