feat(bda): fix issue 1100

fix(bda): updated blueprint schema, added schema support

Author: dineshSajwan

lambda/aws-bedrock-data-automation/bda_blueprint/create_schema.py

@@ -10,9 +10,9 @@
 #  and limitations under the License.
 #
 from typing import Annotated, List, Dict, Type
-from custom_blueprint_schema import create_schema_fields,custom_blue_print
+from custom_blueprint_schema import create_schema_fields, custom_blue_print
 
-def create_schema(fields: List[Dict[str, str]]) -> Type[custom_blue_print]:
+def create_schema(fields: List[Dict[str, str]], document_class: str = "", document_description: str = "") -> Type[custom_blue_print]:
     """
     Create a schema class based on a list of field configurations.
     
@@ -25,6 +25,8 @@ def create_schema(fields: List[Dict[str, str]]) -> Type[custom_blue_print]:
                     "alias": "field alias"
                 }
             ]
+        document_class (str): The class of the document (default: "Notice of Assessment")
+        document_description (str): Description of the document (default: "Notice of Assessment Document")
     
     Returns:
         Type[BaseBlueprintSchema]: A new schema class
@@ -35,18 +37,24 @@ def create_schema(fields: List[Dict[str, str]]) -> Type[custom_blue_print]:
     for field in fields:
         field_name = field["name"]
         annotations[field_name] = Annotated[
-            List[str],
+            str,
             create_schema_fields(
                 description=field["description"],
                 alias=field["alias"]
             )
         ]
 
     # Create a new schema class dynamically
-    return type(
+    schema_class = type(
         "DynamicSchema",
         (custom_blue_print,),
         {
             "__annotations__": annotations
         }
-    )
+    )
+    
+    # Set the class and description in the model_config
+    schema_class.model_config["json_schema_extra"]["class"] = document_class
+    schema_class.model_config["json_schema_extra"]["description"] = document_description
+    
+    return schema_class

lambda/aws-bedrock-data-automation/bda_blueprint/custom_blueprint_schema.py

@@ -31,8 +31,8 @@ class custom_blue_print(BaseModel):
         "json_schema_extra": {
             # Specify JSON Schema version and default metadata
             "$schema": "http://json-schema.org/draft-07/schema#",
-            "description": "default",
-            "documentClass": "default",
+            "description": "",
+            "class": "",
         },
     }
 
@@ -53,6 +53,9 @@ def model_json_schema(cls, *args, **kwargs) -> dict:
 
         # Remove top-level title if present
         schema.pop("title", None)
+        
+        # Remove required field if present
+        schema.pop("required", None)
 
         # Process each field in the model
         properties = {}
@@ -67,21 +70,20 @@ def model_json_schema(cls, *args, **kwargs) -> dict:
 
 def create_schema_fields(description: str, alias: str | None = None):
     """
-    Creates a field configuration for extractive schema fields.
+    Creates a field configuration for explicit schema fields.
     
     Args:
-        description (str): Description of what should be extracted from the document
+        description (str): Instruction of what should be extracted from the document
         alias (str | None): Optional alternative name for the field in the output
         
     Returns:
-        Field: Configured Pydantic Field with extractive properties
+        Field: Configured Pydantic Field with explicit properties
     """
     return Field(
-        default_factory=list,
-        description=description,
         alias=alias,
         json_schema_extra={
-            "inferenceType": "extractive",  
-            "items": {"type": "string"},   
+            "type": "string",
+            "inferenceType": "explicit",
+            "instruction": description
         },
     )

lambda/aws-bedrock-data-automation/bda_blueprint/lambda.py

@@ -158,13 +158,19 @@ def handler(event, context: LambdaContext):
             case "create":
                 logger.info("create blueprint")
 
+                # Check if schema_file_name is provided
                 if 'schema_file_name' in blueprint_details:
                     input_key = blueprint_details['schema_file_name']
 
                     logger.info(f"Retrieving schema from S3: {input_bucket}/{input_key}")
                     schema_content = get_schema(input_bucket, input_key)
 
-                if 'schema_fields' in blueprint_details:
+                    # Convert schema_content to a JSON string if it's a dictionary
+                    if isinstance(schema_content, dict):
+                        schema_content = json.dumps(schema_content)
+                
+                # Only use schema_fields if schema_file_name is not provided
+                elif 'schema_fields' in blueprint_details:
                     schema_fields = blueprint_details['schema_fields']
 
                     # Validate schema_fields format
@@ -178,16 +184,32 @@ def handler(event, context: LambdaContext):
 
                     # Create schema using AWS BDA format
                     try:
-                        schema_content = create_schema(schema_fields)
+                        # Get document class and description from blueprint_details if provided
+                        document_class = blueprint_details.get('document_class', "")
+                        document_description = blueprint_details.get('document_description', "")
+                        
+                        # Create schema with the provided fields, class, and description
+                        schema_class = create_schema(
+                            schema_fields, 
+                            document_class=document_class, 
+                            document_description=document_description
+                        )
+                        
+                        # Convert the schema class to a JSON schema dictionary and then to a JSON string
+                        schema_json = schema_class.model_json_schema()
+                        schema_content = json.dumps(schema_json)
+
                     except Exception as e:
-                        print("Error creating schema")
+                        logger.error("Error creating schema", extra={"error": str(e)})
                         return {
                             'statusCode': 500,
                             'body': json.dumps({
                                 'message': 'Error creating schema',
                                 'error': str(e)
                             })
                         }
+                else:
+                    raise ValueError("Either schema_file_name or schema_fields must be provided")
 
 
                 response= create_blueprint(schema_content,blueprint_details)

lambda/aws-bedrock-data-automation/bda_project/lambda.py

@@ -136,4 +136,3 @@ def handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Any]:
             })
         }
 
-

lambda/aws-bedrock-data-automation/bda_project/project_config.py

@@ -107,20 +107,29 @@ def __init__(self, project_details: Dict[str, Any]):
 
     def _validate_required_fields(self) -> None:
         """Validate required fields are present in project_details"""
-        required_fields = ['projectName','modality']
+        required_fields = ['project_name', 'modality']
         missing_fields = [field for field in required_fields 
                          if not self.project_details.get(field)]
         if missing_fields:
             raise ValueError(f"Missing required fields: {', '.join(missing_fields)}")
 
         # Validate project stage
-        if self.project_details['projectStage'] not in [e.value for e in ProjectStage]:
-            raise ValueError(f"Invalid projectStage. Must be one of: {[e.value for e in ProjectStage]}")
+        project_stage = self.project_details.get('project_stage')
+        if project_stage and project_stage not in [e.value for e in ProjectStage]:
+            raise ValueError(f"Invalid project_stage. Must be one of: {[e.value for e in ProjectStage]}")
 
     def _get_standard_output_config(self) -> Dict[str, Any]:
         """Get standard output configuration if present"""
         config = self.project_details.get('standardOutputConfiguration', {})
         modality = self.project_details.get('modality')
+        
+        # If modality is not provided, return empty dict
+        if not modality:
+            return {}
+            
+        # Convert modality to lowercase for case-insensitive comparison
+        modality = modality.lower()
+            
         if modality == Modality.DOCUMENT.value:
             return {'document': self._get_document_config(config.get('document', {}))}
         elif modality == Modality.IMAGE.value:
@@ -234,9 +243,9 @@ def _get_audio_config(self, config: Dict[str, Any]) -> Dict[str, Any]:
     def project_config(self) -> Dict[str, Any]:
         """Get complete project configuration"""
         config = {
-            'projectName': self.project_details['projectName'],
-            'projectDescription': self.project_details.get('projectDescription','sample description'),
-            'projectStage': self.project_details.get('projectStage','LIVE')
+            'projectName': self.project_details.get('project_name'),
+            'projectDescription': self.project_details.get('project_description', 'sample description'),
+            'projectStage': self.project_details.get('project_stage', 'LIVE')
         }
 
         # Add standard output configuration if present
@@ -268,7 +277,7 @@ def project_config(self) -> Dict[str, Any]:
 
     def __str__(self) -> str:
         """String representation of project configuration"""
-        return f"ProjectConfig(name={self.project_details['projectName']}, stage={self.project_details['projectStage']})"
+        return f"ProjectConfig(name={self.project_details.get('project_name')}, stage={self.project_details.get('project_stage')})"
 
 class ListProjectsConfig:
     """Configuration class for listing Bedrock Data Automation projects"""

lambda/aws-bedrock-data-automation/data_processing/data_processing.py

@@ -100,6 +100,7 @@ def invoke_data_automation_async(
         notification_config: Optional[NotificationConfig] = None,
         data_automation_profile_arn: Optional[str] = None,
         tags: Optional[list] = None,
+        **kwargs  # To handle any additional parameters that might be passed but not used
     ) -> Dict[str, Any]:
         """
         Invoke data automation asynchronously
@@ -162,7 +163,11 @@ def invoke_data_automation_async(
 
             # Add data automation profile ARN if provided
             if data_automation_profile_arn:
+                # The AWS documentation shows this parameter as dataAutomationProfileArn
                 request_params['dataAutomationProfileArn'] = data_automation_profile_arn
+                logger.info("Added data automation profile ARN", extra={
+                    "data_automation_profile_arn": data_automation_profile_arn
+                })
 
             # Add encryption configuration if provided
             if encryption_config:
@@ -191,6 +196,7 @@ def invoke_data_automation_async(
             logger.info("Invoking data automation", extra={
                 "request_params": request_params
             })
+            
             response = self.client.invoke_data_automation_async(**request_params)
 
             logger.info("Successfully invoked data automation", extra={

lambda/aws-bedrock-data-automation/data_processing/lambda.py

@@ -75,7 +75,7 @@ def get_env_var(var_name: str, default: str = None) -> str:
         raise ValueError(f"Environment variable {var_name} is not set")
     return value
 
-#@logger.inject_lambda_context
+@logger.inject_lambda_context
 def handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Any]:
     """
     Lambda handler to process EventBridge events and invoke data automation
@@ -200,7 +200,7 @@ def handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Any]:
                 "notification_config": configs['notification_config']
             })
 
-        # Check and add data_automation_profile_arn
+        # Add data_automation_profile_arn to configs if provided
         if profile_arn := detail.get('data_automation_profile_arn'):
             configs['data_automation_profile_arn'] = profile_arn
             logger.info("Data automation profile ARN added", extra={
@@ -258,3 +258,5 @@ def handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Any]:
                 'error': str(e)
             }
         }
+
+

lambda/aws-bedrock-data-automation/data_result/data_automation_result.py

@@ -150,12 +150,17 @@ async def get_results(self, invocation_arn: str, wait: bool = False) -> Dict[str
                             "path": segment["custom_output_path"]
                         })
 
-                    # Add standard output
-                    standard_result = self._read_s3_json(segment["standard_output_path"])
-                    results.append(standard_result)
-                    logger.debug("Added standard output", extra={
-                        "path": segment["standard_output_path"]
-                    })
+                    # Add standard output if path exists
+                    if "standard_output_path" in segment:
+                        standard_result = self._read_s3_json(segment["standard_output_path"])
+                        results.append(standard_result)
+                        logger.debug("Added standard output", extra={
+                            "path": segment["standard_output_path"]
+                        })
+                    else:
+                        logger.warning("Missing standard_output_path in segment", extra={
+                            "segment_id": segment.get("segment_id", "unknown")
+                        })
 
             logger.info("Successfully retrieved results", extra={
                 "invoke_arn": invocation_arn,

lambda/aws-bedrock-data-automation/data_result/lambda.py

@@ -150,3 +150,5 @@ def handler(event: Dict[str, Any], context: LambdaContext) -> Dict[str, Any]:
     Main Lambda handler that wraps the async handler
     """
     return asyncio.run(data_automation_status(event, context))
+
+

layer/requirements.txt

@@ -1,3 +1,3 @@
-boto3>=1.37.5
-botocore>=1.35.5
+boto3>=1.38.1
+botocore>=1.38.2