add custom model arn for classification and extraction

Author: Bob Strahan

VERSION

@@ -1 +1 @@
-0.3.2
+0.3.3-alpha

patterns/pattern-2/template.yaml

@@ -96,6 +96,8 @@ Parameters:
 Conditions:
   IsSummarizationEnabled: !Equals [!Ref IsSummarizationEnabled, "true"]
   HasGuardrailConfig: !And [!Not [!Equals [!Ref BedrockGuardrailId, ""]], !Not [!Equals [!Ref BedrockGuardrailVersion, ""]]]
+  HasCustomClassificationModelARN : !Not [!Equals [!Ref CustomClassificationModelARN , ""]]
+  HasCustomExtractionModelARN : !Not [!Equals [!Ref CustomExtractionModelARN , ""]]
 
 
 Resources:
@@ -248,7 +250,19 @@ Resources:
               model:
                 type: string
                 description: Model identifier
-                enum: ["us.amazon.nova-lite-v1:0", "us.amazon.nova-pro-v1:0", "us.amazon.nova-premier-v1:0", "us.anthropic.claude-3-haiku-20240307-v1:0", "us.anthropic.claude-3-5-sonnet-20241022-v2:0", "us.anthropic.claude-3-7-sonnet-20250219-v1:0", "us.anthropic.claude-sonnet-4-20250514-v1:0", "us.anthropic.claude-opus-4-20250514-v1:0"]
+                enum: 
+                  - !If 
+                    - HasCustomClassificationModelARN 
+                    - !Ref CustomClassificationModelARN
+                    - !Ref AWS::NoValue
+                  - "us.amazon.nova-lite-v1:0"
+                  - "us.amazon.nova-pro-v1:0"
+                  - "us.amazon.nova-premier-v1:0"
+                  - "us.anthropic.claude-3-haiku-20240307-v1:0"
+                  - "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
+                  - "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
+                  - "us.anthropic.claude-sonnet-4-20250514-v1:0"
+                  - "us.anthropic.claude-opus-4-20250514-v1:0"
                 order: 1
               classificationMethod:
                 type: string
@@ -299,7 +313,19 @@ Resources:
               model:
                 type: string
                 description: Model identifier
-                enum: ["us.amazon.nova-lite-v1:0", "us.amazon.nova-pro-v1:0", "us.amazon.nova-premier-v1:0", "us.anthropic.claude-3-haiku-20240307-v1:0", "us.anthropic.claude-3-5-sonnet-20241022-v2:0", "us.anthropic.claude-3-7-sonnet-20250219-v1:0", "us.anthropic.claude-sonnet-4-20250514-v1:0", "us.anthropic.claude-opus-4-20250514-v1:0"]
+                enum: 
+                  - !If 
+                    - HasCustomExtractionModelARN
+                    - !Ref CustomExtractionModelARN
+                    - !Ref AWS::NoValue
+                  - "us.amazon.nova-lite-v1:0"
+                  - "us.amazon.nova-pro-v1:0"
+                  - "us.amazon.nova-premier-v1:0"
+                  - "us.anthropic.claude-3-haiku-20240307-v1:0"
+                  - "us.anthropic.claude-3-5-sonnet-20241022-v2:0"
+                  - "us.anthropic.claude-3-7-sonnet-20250219-v1:0"
+                  - "us.anthropic.claude-sonnet-4-20250514-v1:0"
+                  - "us.anthropic.claude-opus-4-20250514-v1:0"
                 order: 1
               temperature:
                 type: number
@@ -446,6 +472,8 @@ Resources:
       ServiceToken: !Ref UpdateConfigurationFunctionArn
       Default: !Ref ConfigurationDefaultS3Uri
       ConfigLibraryHash: !Ref ConfigLibraryHash
+      CustomClassificationModelARN: !Ref CustomClassificationModelARN
+      CustomExtractionModelARN: !Ref CustomExtractionModelARN
 
   OCRFunction:
     Type: AWS::Serverless::Function

src/lambda/update_configuration/index.py

@@ -126,6 +126,21 @@ def handler(event: Dict[str, Any], context: Any) -> None:
             # Update Default configuration
             if 'Default' in properties:
                 resolved_default = resolve_content(properties['Default'])
+                
+                # Apply custom model ARNs if provided
+                if isinstance(resolved_default, dict):
+                    # Replace classification model if CustomClassificationModelARN is provided and not empty
+                    if 'CustomClassificationModelARN' in properties and properties['CustomClassificationModelARN'].strip():
+                        if 'classification' in resolved_default:
+                            resolved_default['classification']['model'] = properties['CustomClassificationModelARN']
+                            logger.info(f"Updated classification model to: {properties['CustomClassificationModelARN']}")
+                    
+                    # Replace extraction model if CustomExtractionModelARN is provided and not empty
+                    if 'CustomExtractionModelARN' in properties and properties['CustomExtractionModelARN'].strip():
+                        if 'extraction' in resolved_default:
+                            resolved_default['extraction']['model'] = properties['CustomExtractionModelARN']
+                            logger.info(f"Updated extraction model to: {properties['CustomExtractionModelARN']}")
+                
                 update_configuration('Default', resolved_default)
 
             # Update Custom configuration if provided and not empty
@@ -152,4 +167,4 @@ def handler(event: Dict[str, Any], context: Any) -> None:
         physical_id = generate_physical_id(event['StackId'], event['LogicalResourceId'])
         cfnresponse.send(event, context, cfnresponse.FAILED, {
             'Error': str(e)
-        }, physical_id)
+        }, physical_id)

template.yaml

@@ -82,12 +82,12 @@ Parameters:
   Pattern2CustomClassificationModelARN:
     Type: String
     Default: ""
-    Description: (NOT YET SUPPORTED) Specify Custom Fine Tuned Classification Model ARN. Leave blank to use pretrained model specified in the selected configuration.
+    Description: Specify Custom Fine Tuned Classification Model ARN. Leave blank to use pretrained model specified in the selected configuration.
 
   Pattern2CustomExtractionModelARN:
     Type: String
     Default: ""
-    Description: (NOT YET SUPPORTED) Specify Custom Fine Tuned Extraction Model ARN. Leave blank to use pretrained model specified in the selected configuration.
+    Description: Specify Custom Fine Tuned Extraction Model ARN. Leave blank to use pretrained model specified in the selected configuration.
 
   # Pattern 3 Parameters
 
@@ -292,8 +292,6 @@ Rules:
       - Assert: !Not [!Equals [!Ref Pattern3UDOPModelArtifactPath, ""]]
         AssertDescription: UDOP Model Artifact Path is required when IDPPattern is 'Pattern1'
 
-  # Pattern 2 rules - Custom model validation now handled in pattern templates
-
 Conditions:
   IsPattern1:
     !Equals [