fixed pen test findings, fixed configuraiton loading mechanism, updated sanitized parmeters to ignore prints documents in content

Author: sathyaws

docs/discovery.md

@@ -1236,6 +1236,9 @@ def discovery_with_fallback(discovery_service, document_key, ground_truth_key=No
 **Discovery Output Format**
 - Output format is configuration via View/Edit configuration. JSON format should follow custom classes format.  
 - Output in any other format will result in failure.
+**Production Usage**
+- We recommend not to use the Discovery module in production. This is to reduce the risk of any hallunication during the document discovery. 
+- We recommend to use discovery module in your lower environment to discovery and construct the configurations. Export the tested configuration to production deployment.
 
 
 The Discovery module provides a powerful foundation for understanding and processing new document types. By following these guidelines and best practices, you can effectively leverage the module to bootstrap document processing workflows and continuously improve their accuracy and coverage.

lib/idp_common_pkg/idp_common/bedrock/client.py

@@ -858,6 +858,9 @@ def _sanitize_messages_for_logging(self, messages: List[Dict[str, Any]]) -> List
                     elif isinstance(content_item, dict) and 'bytes' in content_item:
                         # Handle raw binary format
                         content_item['bytes'] = '[binary_data]'
+                    elif isinstance(content_item, dict) and 'document' in content_item:
+                        # Handle different image format used by some models
+                        content_item['document'] = '[document_data]'
 
         return sanitized
 

lib/idp_common_pkg/idp_common/discovery/classes_discovery.py

@@ -1,6 +1,5 @@
 # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 # SPDX-License-Identifier: MIT-0
-import base64
 import json
 import logging
 import os
@@ -10,6 +9,7 @@
 from botocore.exceptions import ClientError
 
 from idp_common import bedrock, image
+from idp_common.config import ConfigurationReader
 from idp_common.utils.s3util import S3Util
 
 logger = logging.getLogger(__name__)
@@ -20,15 +20,28 @@ def __init__(
         self,
         input_bucket: str,
         input_prefix: str,
-        config: Optional[dict] = None,
-        region: Optional[str] = "us-west-2",
+        region: Optional[str] = None,
     ):
         self.input_bucket = input_bucket
         self.input_prefix = input_prefix
-        self.region = region or os.environ.get("AWS_REGION", "us-east-1")
+        self.region = region or os.environ.get("AWS_REGION")
 
         # Load configuration
-        self.config = config or self._load_default_config()
+        self.configuration_table_name = os.environ.get("CONFIGURATION_TABLE_NAME")
+
+        if not self.configuration_table_name:
+            raise ValueError(
+                "Configuration table name not provided. Set CONFIGURATION_TABLE_NAME environment variable."
+            )
+
+        try:
+            self.config_reader = ConfigurationReader(
+                table_name=self.configuration_table_name
+            )
+            self.config = self.config_reader.get_merged_configuration()
+        except Exception as e:
+            logger.error(f"Failed to load configuration from DynamoDB: {e}")
+            raise Exception(f"Failed to load configuration from DynamoDB: {str(e)}")
 
         # Get discovery configuration
         self.discovery_config = self.config.get("discovery", {})
@@ -40,87 +53,14 @@ def __init__(
         # Initialize Bedrock client using the common pattern
         self.bedrock_client = bedrock.BedrockClient(region=self.region)
 
-        self.configuration_table_name = os.environ.get("CONFIGURATION_TABLE_NAME", "")
-        dynamodb = boto3.resource("dynamodb")
-        self.configuration_table = dynamodb.Table(self.configuration_table_name)
+        if self.configuration_table_name:
+            dynamodb = boto3.resource("dynamodb")
+            self.configuration_table = dynamodb.Table(self.configuration_table_name)
+        else:
+            self.configuration_table = None
 
         return
 
-    def _load_default_config(self):
-        """Load default discovery configuration."""
-        return {
-            "discovery": {
-                "without_ground_truth": {
-                    "model_id": "anthropic.claude-3-sonnet-20240229-v1:0",
-                    "temperature": 1.0,
-                    "top_p": 0.1,
-                    "max_tokens": 10000,
-                    "system_prompt": "You are an expert in processing forms. Extracting data from images and documents. Analyze forms line by line to identify field names, data types, and organizational structure. Focus on creating comprehensive blueprints for document processing without extracting actual values.",
-                    "user_prompt": """This image contains forms data. Analyze the form line by line.
-Image may contains multiple pages, process all the pages. 
-Form may contain multiple name value pair in one line. 
-Extract all the names in the form including the name value pair which doesn't have value. 
-Organize them into groups, extract field_name, data_type and field description.
-Field_name should be less than 60 characters, should not have space use '-' instead of space.
-field_description is a brief description of the field and the location of the field like box number or line number in the form and section of the form.
-Field_name should be unique within the group.
-Add two fields document_class and document_description. 
-For document_class generate a short name based on the document content like W4, I-9, Paystub. 
-For document_description generate a description about the document in less than 50 words. 
-Group the fields based on the section they are grouped in the form. Group should have attributeType as "group".
-If the group repeats and follows table format, update the attributeType as "list".      
-Do not extract the values.
-Return the extracted data in JSON format.""",
-                },
-                "with_ground_truth": {
-                    "model_id": "anthropic.claude-3-sonnet-20240229-v1:0",
-                    "temperature": 1.0,
-                    "top_p": 0.1,
-                    "max_tokens": 10000,
-                    "system_prompt": "You are an expert in processing forms. Extracting data from images and documents. Use provided ground truth data as reference to optimize field extraction and ensure consistency with expected document structure and field definitions.",
-                    "user_prompt": """This image contains unstructured data. Analyze the data line by line using the provided ground truth as reference.                        
-<GROUND_TRUTH_REFERENCE>
-{ground_truth_json}
-</GROUND_TRUTH_REFERENCE>
-Ground truth reference JSON has the fields we are interested in extracting from the document/image. Use the ground truth to optimize field extraction. Match field names, data types, and groupings from the reference.
-Image may contain multiple pages, process all pages.
-Extract all field names including those without values.
-Do not change the group name and field name from ground truth in the extracted data json.
-Add field_description field for every field which will contain instruction to LLM to extract the field data from the image/document. Add data_type field for every field. 
-Add two fields document_class and document_description. 
-For document_class generate a short name based on the document content like W4, I-9, Paystub. 
-For document_description generate a description about the document in less than 50 words.
-If the group repeats and follows table format, update the attributeType as "list".      
-Do not extract the values.""",
-                },
-                "output_format": {
-                    "sample_json": """{
-    "document_class" : "Form-1040",
-    "document_description" : "Brief summary of the document",
-    "groups" : [
-        {
-            "name" : "PersonalInformation",
-            "description" : "Personal information of Tax payer",
-            "attributeType" : "group",
-            "groupAttributes" : [
-                {
-                    "name": "FirstName",
-                    "dataType" : "string",
-                    "description" : "First Name of Taxpayer"
-                },
-                {
-                    "name": "Age",
-                    "dataType" : "number",
-                    "description" : "Age of Taxpayer"
-                }
-            ]
-        }
-    ]
-}"""
-                },
-            }
-        }
-
     """
         Recursively convert all values to strings
         """
@@ -328,13 +268,6 @@ def _remove_duplicates(self, groups):
             group["groupAttributes"] = groupAttributesArray
         return groups
 
-    def _parse_s3_uri(self, s3_uri: str):
-        """Parse S3 URI to extract bucket and key."""
-        if not s3_uri.startswith("s3://"):
-            raise ValueError("Invalid S3 URI format")
-        parts = s3_uri[5:].split("/", 1)
-        return parts[0], parts[1] if len(parts) > 1 else ""
-
     def _load_ground_truth(self, bucket: str, key: str):
         """Load ground truth JSON data from S3."""
         try:
@@ -347,9 +280,7 @@ def _load_ground_truth(self, bucket: str, key: str):
     def _extract_data_from_document(self, document_content, file_extension):
         try:
             # Get configuration for without ground truth
-            model_id = self.without_gt_config.get(
-                "model_id", "anthropic.claude-3-sonnet-20240229-v1:0"
-            )
+            model_id = self.without_gt_config.get("model_id", "us.amazon.nova-pro-v1:0")
             system_prompt = self.without_gt_config.get(
                 "system_prompt",
                 "You are an expert in processing forms. Extracting data from images and documents",
@@ -424,9 +355,7 @@ def _extract_data_from_document_with_ground_truth(
         """Extract data from document using ground truth as reference."""
         try:
             # Get configuration for with ground truth
-            model_id = self.with_gt_config.get(
-                "model_id", "anthropic.claude-3-sonnet-20240229-v1:0"
-            )
+            model_id = self.with_gt_config.get("model_id", "us.amazon.nova-pro-v1:0")
             system_prompt = self.with_gt_config.get(
                 "system_prompt",
                 "You are an expert in processing forms. Extracting data from images and documents",
@@ -483,10 +412,6 @@ def _extract_data_from_document_with_ground_truth(
             logger.error(f"Error extracting data with Bedrock using ground truth: {e}")
             return None
 
-    def _get_base64_image(self, file_in_bytes):
-        """Get base64 encoded image data from sample image."""
-        return base64.b64encode(file_in_bytes).decode("utf-8")
-
     def _prompt_classes_discovery_with_ground_truth(self, ground_truth_data):
         ground_truth_json = json.dumps(ground_truth_data, indent=2)
         sample_output_format = self._sample_output_format()

lib/idp_common_pkg/tests/integration/test_discovery_config_integration.py

@@ -113,13 +113,22 @@ def test_end_to_end_config_flow_without_ground_truth(
         mock_extract_text.return_value = json.dumps(expected_result)
         mock_s3_get_bytes.return_value = b"mock document content"
 
-        # Initialize ClassesDiscovery with YAML config
-        discovery = ClassesDiscovery(
-            input_bucket=self.test_bucket,
-            input_prefix=self.test_prefix,
-            config=self.config_dict,
-            region=self.test_region,
-        )
+        # Mock ConfigurationReader to return config_dict
+        with patch(
+            "idp_common.discovery.classes_discovery.ConfigurationReader"
+        ) as mock_config_reader:
+            mock_reader_instance = mock_config_reader.return_value
+            mock_reader_instance.get_merged_configuration.return_value = (
+                self.config_dict
+            )
+
+            with patch.dict("os.environ", {"CONFIGURATION_TABLE_NAME": "test-table"}):
+                # Initialize ClassesDiscovery with YAML config
+                discovery = ClassesDiscovery(
+                    input_bucket=self.test_bucket,
+                    input_prefix=self.test_prefix,
+                    region=self.test_region,
+                )
 
         # Execute discovery without ground truth
         result = discovery.discovery_classes_with_document(
@@ -218,13 +227,22 @@ def mock_s3_side_effect(bucket, key):
 
         mock_s3_get_bytes.side_effect = mock_s3_side_effect
 
-        # Initialize ClassesDiscovery with YAML config
-        discovery = ClassesDiscovery(
-            input_bucket=self.test_bucket,
-            input_prefix=self.test_prefix,
-            config=self.config_dict,
-            region=self.test_region,
-        )
+        # Mock ConfigurationReader to return config_dict
+        with patch(
+            "idp_common.discovery.classes_discovery.ConfigurationReader"
+        ) as mock_config_reader:
+            mock_reader_instance = mock_config_reader.return_value
+            mock_reader_instance.get_merged_configuration.return_value = (
+                self.config_dict
+            )
+
+            with patch.dict("os.environ", {"CONFIGURATION_TABLE_NAME": "test-table"}):
+                # Initialize ClassesDiscovery with YAML config
+                discovery = ClassesDiscovery(
+                    input_bucket=self.test_bucket,
+                    input_prefix=self.test_prefix,
+                    region=self.test_region,
+                )
 
         # Execute discovery with ground truth
         result = discovery.discovery_classes_with_document_and_ground_truth(
@@ -290,13 +308,22 @@ def test_config_validation_and_defaults(
             }
         }
 
-        # Initialize with incomplete config
-        discovery = ClassesDiscovery(
-            input_bucket=self.test_bucket,
-            input_prefix=self.test_prefix,
-            config=incomplete_config,
-            region=self.test_region,
-        )
+        # Mock ConfigurationReader to return incomplete config
+        with patch(
+            "idp_common.discovery.classes_discovery.ConfigurationReader"
+        ) as mock_config_reader:
+            mock_reader_instance = mock_config_reader.return_value
+            mock_reader_instance.get_merged_configuration.return_value = (
+                incomplete_config
+            )
+
+            with patch.dict("os.environ", {"CONFIGURATION_TABLE_NAME": "test-table"}):
+                # Initialize with incomplete config
+                discovery = ClassesDiscovery(
+                    input_bucket=self.test_bucket,
+                    input_prefix=self.test_prefix,
+                    region=self.test_region,
+                )
 
         # Verify that missing fields get default values
         without_gt_config = discovery.without_gt_config