Add Bedrock runtime classification support

- Add SSM params for classification model ID and prompt template
- Add bedrock_runtime_invoke IAM policy
- Rename bedrock_invoke to bedrock_data_automation_invoke
- Remove document splitting override from BDA config

Author: laurencegoolsby

infra/app-docai/app-config/env-config/document_data_extraction.tf

@@ -52,13 +52,13 @@ locals {
       }
     }
 
-    override_configuration = {
-      document = {
-        splitter = {
-          state = "ENABLED"
-        }
-      }
-    }
+    # override_configuration = {
+    #   document = {
+    #     splitter = {
+    #       state = "ENABLED"
+    #     }
+    #   }
+    # }
 
   } : null
 }

infra/app-docai/service/document_data_extraction.tf

@@ -87,7 +87,6 @@ module "documentai" {
   source = "../../modules/document-data-extraction/resources"
 
   standard_output_configuration = local.document_data_extraction_config.standard_output_configuration
-  override_configuration        = local.document_data_extraction_config.override_configuration
   tags                          = local.tags
 
   blueprints = concat(
@@ -350,6 +349,44 @@ resource "aws_dynamodb_table" "document_batches" {
   tags = local.tags
 }
 
+#-------------------
+# Bedrock Classification Config (SSM)
+#-------------------
+resource "aws_ssm_parameter" "bedrock_classification_model_id" {
+  count = local.document_data_extraction_config != null ? 1 : 0
+
+  name  = "/service/${local.service_name}/bedrock/classification-model-id"
+  type  = "String"
+  value = "anthropic.claude-3-haiku-20240307-v1:0"
+
+  lifecycle {
+    ignore_changes = [value]
+  }
+}
+
+# <<DOCUMENT_TYPES>> in the classification prompt needs to be dynamically 
+# replaced with the document types that BDA is configured to extract. Store prompt 
+# in SSM Parameter Store; application reads and update it at runtime.
+resource "aws_ssm_parameter" "bedrock_classification_prompt" {
+  count = local.document_data_extraction_config != null ? 1 : 0
+
+  name  = "/service/${local.service_name}/bedrock/classification-prompt"
+  type  = "String"
+  value = <<-EOT
+Analyze this image. Respond in JSON only:
+{"document_type": "string", "confidence": float 0-1, "document_count": int}
+ONLY use one of these exact values for document_type: <<DOCUMENT_TYPES>>
+Do not create new categories. If unsure, use 'other_document'.
+If it's not a document, use 'not_a_document'.
+document_count: how many separate documents are visible in this image?
+EOT
+
+  lifecycle {
+    ignore_changes = [value]
+  }
+}
+
+
 #-------------------
 # IAM Policies
 #-------------------
@@ -385,7 +422,7 @@ resource "aws_iam_policy" "dynamodb_read_write" {
   })
 }
 
-resource "aws_iam_policy" "bedrock_invoke" {
+resource "aws_iam_policy" "bedrock_data_automation_invoke" {
   count = local.document_data_extraction_config != null ? 1 : 0
 
   name = "${local.prefix}bedrock-invoke"
@@ -403,4 +440,29 @@ resource "aws_iam_policy" "bedrock_invoke" {
       Effect = "Allow"
     }]
   })
+}
+
+
+resource "aws_iam_policy" "bedrock_runtime_invoke" {
+  count = local.document_data_extraction_config != null ? 1 : 0
+
+  name = "${local.prefix}bedrock-runtime-invoke"
+  policy = jsonencode({
+    Version = "2012-10-17"
+    Statement = [
+      {
+        Action   = "bedrock:InvokeModel"
+        Resource = "arn:aws:bedrock:${data.aws_region.current.name}::foundation-model/*"
+        Effect   = "Allow"
+      },
+      {
+        Action = "ssm:GetParameter"
+        Resource = [
+          aws_ssm_parameter.bedrock_classification_model_id[0].arn,
+          aws_ssm_parameter.bedrock_classification_prompt[0].arn,
+        ]
+        Effect = "Allow"
+      }
+    ]
+  })
 }

infra/app-docai/service/main.tf

@@ -134,12 +134,14 @@ module "service" {
       storage_access = module.storage.access_policy_arn
     },
     module.app_config.enable_document_data_extraction ? {
-      documentai_input_bucket_access   = module.documentai_input_bucket[0].access_policy_arn
-      documentai_output_bucket_access  = module.documentai_output_bucket[0].access_policy_arn,
-      documentai_metrics_bucket_access = module.documentai_metrics_bucket[0].access_policy_arn,
-      documentai_bedrock_access        = module.documentai[0].access_policy_arn,
-      documentai_dynamodb_access       = aws_iam_policy.dynamodb_read_write[0].arn
-      documentai_sqs_send_message      = aws_iam_policy.sqs_send_message[0].arn
+      documentai_input_bucket_access            = module.documentai_input_bucket[0].access_policy_arn
+      documentai_output_bucket_access           = module.documentai_output_bucket[0].access_policy_arn,
+      documentai_metrics_bucket_access          = module.documentai_metrics_bucket[0].access_policy_arn,
+      documentai_bedrock_access                 = module.documentai[0].access_policy_arn,
+      documentai_dynamodb_access                = aws_iam_policy.dynamodb_read_write[0].arn
+      documentai_sqs_send_message               = aws_iam_policy.sqs_send_message[0].arn
+      documentai_bedrock_data_automation_invoke = aws_iam_policy.bedrock_data_automation_invoke[0].arn
+      documentai_bedrock_runtime_invoke         = aws_iam_policy.bedrock_runtime_invoke[0].arn
     } : {},
     module.app_config.enable_identity_provider ? {
       identity_provider_access = module.identity_provider_client[0].access_policy_arn,