This is an automated email from the ASF dual-hosted git repository.
vincbeck pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0217a7b1048 Add `BedrockCreateGuardrailOperator` (#66035)
0217a7b1048 is described below
commit 0217a7b1048037895a41c30151314cf199ea8f14
Author: John Jackson <[email protected]>
AuthorDate: Wed Apr 29 14:07:08 2026 -0700
Add `BedrockCreateGuardrailOperator` (#66035)
---
providers/amazon/docs/operators/bedrock.rst | 15 +++
.../airflow/providers/amazon/aws/hooks/bedrock.py | 9 ++
.../providers/amazon/aws/operators/bedrock.py | 105 +++++++++++++++++++-
.../system/amazon/aws/example_bedrock_guardrail.py | 80 +++++++++++++++
.../tests/unit/amazon/aws/hooks/test_bedrock.py | 26 +++++
.../unit/amazon/aws/operators/test_bedrock.py | 108 +++++++++++++++++++++
6 files changed, 342 insertions(+), 1 deletion(-)
diff --git a/providers/amazon/docs/operators/bedrock.rst
b/providers/amazon/docs/operators/bedrock.rst
index 5494a34c82d..374cf6b4a0d 100644
--- a/providers/amazon/docs/operators/bedrock.rst
+++ b/providers/amazon/docs/operators/bedrock.rst
@@ -339,6 +339,21 @@ the service confirms that the job is in the queue, use
TargetState.SCHEDULED.
:start-after: [START howto_sensor_bedrock_batch_inference_scheduled]
:end-before: [END howto_sensor_bedrock_batch_inference_scheduled]
+.. _howto/operator:BedrockCreateGuardrailOperator:
+
+Create a Guardrail
+~~~~~~~~~~~~~~~~~~
+
+To create an Amazon Bedrock guardrail, use
+:class:`~airflow.providers.amazon.aws.operators.bedrock.BedrockCreateGuardrailOperator`.
+
+.. exampleinclude::
/../../amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
+ :language: python
+ :dedent: 4
+ :start-after: [START howto_operator_bedrock_create_guardrail]
+ :end-before: [END howto_operator_bedrock_create_guardrail]
+
+
Reference
---------
diff --git a/providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py
b/providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py
index 0c2fd1a11cf..c8fd53a9b7f 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py
@@ -38,6 +38,15 @@ class BedrockHook(AwsBaseHook):
kwargs["client_type"] = self.client_type
super().__init__(*args, **kwargs)
+ def get_guardrail_id_by_name(self, guardrail_name: str) -> str | None:
+ """Get the guardrail ID by name, or None if not found."""
+ paginator = self.conn.get_paginator("list_guardrails")
+ for page in paginator.paginate():
+ for g in page.get("guardrails", []):
+ if g.get("name") == guardrail_name:
+ return g["id"]
+ return None
+
class BedrockRuntimeHook(AwsBaseHook):
"""
diff --git
a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py
b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py
index 9f7378de03b..e561c4b1d6f 100644
--- a/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py
+++ b/providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py
@@ -19,7 +19,7 @@ from __future__ import annotations
import json
from collections.abc import Sequence
from time import sleep
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Literal
from botocore.exceptions import ClientError
@@ -1025,3 +1025,106 @@ class
BedrockBatchInferenceOperator(AwsBaseOperator[BedrockHook]):
)
return job_arn
+
+
+class BedrockCreateGuardrailOperator(AwsBaseOperator[BedrockHook]):
+ """
+ Create an Amazon Bedrock guardrail to implement safeguards for generative
AI applications.
+
+ A guardrail can filter harmful content, block denied topics, filter words,
+ and mask sensitive information. The guardrail is created with a DRAFT
version
+ that is immediately usable.
+
+ .. seealso::
+ For more information on how to use this operator, take a look at the
guide:
+ :ref:`howto/operator:BedrockCreateGuardrailOperator`
+
+ :param guardrail_name: The name of the guardrail (1-50 chars).
+ :param blocked_input_messaging: Message returned when the guardrail blocks
an input prompt.
+ :param blocked_outputs_messaging: Message returned when the guardrail
blocks a model response.
+ :param description: A description of the guardrail.
+ :param topic_policy_config: Topic policy configuration dict.
+ :param content_policy_config: Content filter policy configuration dict.
+ :param word_policy_config: Word filter policy configuration dict.
+ :param sensitive_information_policy_config: Sensitive information policy
configuration dict.
+ :param contextual_grounding_policy_config: Contextual grounding policy
configuration dict.
+ :param kms_key_id: ARN of the KMS key to encrypt the guardrail.
+ :param tags: Tags to attach to the guardrail.
+ :param if_exists: Behavior when a guardrail with the same name already
exists.
+ ``"fail"`` raises an error, ``"skip"`` returns the existing guardrail
ID.
+ """
+
+ aws_hook_class = BedrockHook
+ template_fields: Sequence[str] = aws_template_fields(
+ "guardrail_name",
+ "blocked_input_messaging",
+ "blocked_outputs_messaging",
+ "description",
+ )
+ template_fields_renderers = {
+ "topic_policy_config": "json",
+ "content_policy_config": "json",
+ "word_policy_config": "json",
+ "sensitive_information_policy_config": "json",
+ }
+
+ def __init__(
+ self,
+ *,
+ guardrail_name: str,
+ blocked_input_messaging: str,
+ blocked_outputs_messaging: str,
+ description: str | None = None,
+ topic_policy_config: dict[str, Any] | None = None,
+ content_policy_config: dict[str, Any] | None = None,
+ word_policy_config: dict[str, Any] | None = None,
+ sensitive_information_policy_config: dict[str, Any] | None = None,
+ contextual_grounding_policy_config: dict[str, Any] | None = None,
+ kms_key_id: str | None = None,
+ tags: list[dict[str, str]] | None = None,
+ if_exists: Literal["fail", "skip"] = "skip",
+ **kwargs,
+ ) -> None:
+ super().__init__(**kwargs)
+ self.guardrail_name = guardrail_name
+ self.blocked_input_messaging = blocked_input_messaging
+ self.blocked_outputs_messaging = blocked_outputs_messaging
+ self.description = description
+ self.topic_policy_config = topic_policy_config
+ self.content_policy_config = content_policy_config
+ self.word_policy_config = word_policy_config
+ self.sensitive_information_policy_config =
sensitive_information_policy_config
+ self.contextual_grounding_policy_config =
contextual_grounding_policy_config
+ self.kms_key_id = kms_key_id
+ self.tags = tags
+ self.if_exists = if_exists
+
+ def execute(self, context: Context) -> str:
+ kwargs: dict[str, Any] = prune_dict(
+ {
+ "name": self.guardrail_name,
+ "blockedInputMessaging": self.blocked_input_messaging,
+ "blockedOutputsMessaging": self.blocked_outputs_messaging,
+ "description": self.description,
+ "topicPolicyConfig": self.topic_policy_config,
+ "contentPolicyConfig": self.content_policy_config,
+ "wordPolicyConfig": self.word_policy_config,
+ "sensitiveInformationPolicyConfig":
self.sensitive_information_policy_config,
+ "contextualGroundingPolicyConfig":
self.contextual_grounding_policy_config,
+ "kmsKeyId": self.kms_key_id,
+ "tags": self.tags,
+ }
+ )
+ try:
+ response = self.hook.conn.create_guardrail(**kwargs)
+ guardrail_id = response["guardrailId"]
+ except ClientError as e:
+ if e.response["Error"]["Code"] == "ConflictException" and
self.if_exists == "skip":
+ self.log.info("Guardrail %s already exists, skipping.",
self.guardrail_name)
+ guardrail_id =
self.hook.get_guardrail_id_by_name(self.guardrail_name)
+ if guardrail_id is None:
+ raise
+ else:
+ raise
+ self.log.info("Guardrail %s: %s", self.guardrail_name, guardrail_id)
+ return guardrail_id
diff --git
a/providers/amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
b/providers/amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
new file mode 100644
index 00000000000..3473dce0f4f
--- /dev/null
+++ b/providers/amazon/tests/system/amazon/aws/example_bedrock_guardrail.py
@@ -0,0 +1,80 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+from __future__ import annotations
+
+from datetime import datetime
+
+from airflow.providers.amazon.aws.operators.bedrock import
BedrockCreateGuardrailOperator
+from airflow.providers.common.compat.sdk import DAG, chain
+
+from system.amazon.aws.utils import ENV_ID_KEY, SystemTestContextBuilder
+from tests_common.test_utils.version_compat import AIRFLOW_V_3_0_PLUS
+
+if AIRFLOW_V_3_0_PLUS:
+ from airflow.sdk import TriggerRule, task
+else:
+ from airflow.decorators import task # type: ignore[attr-defined,no-redef]
+ from airflow.utils.trigger_rule import TriggerRule # type:
ignore[no-redef,attr-defined]
+
+DAG_ID = "example_bedrock_guardrail"
+
+sys_test_context_task = SystemTestContextBuilder().build()
+
+with DAG(
+ dag_id=DAG_ID,
+ schedule=None,
+ start_date=datetime(2024, 1, 1),
+ catchup=False,
+) as dag:
+ test_context = sys_test_context_task()
+ env_id = test_context[ENV_ID_KEY]
+
+ # [START howto_operator_bedrock_create_guardrail]
+ create_guardrail = BedrockCreateGuardrailOperator(
+ task_id="create_guardrail",
+ guardrail_name=f"{env_id}-guardrail",
+ blocked_input_messaging="Sorry, this input is not allowed.",
+ blocked_outputs_messaging="Sorry, I cannot respond to that.",
+ content_policy_config={
+ "filtersConfig": [
+ {"type": "HATE", "inputStrength": "HIGH", "outputStrength":
"HIGH"},
+ {"type": "VIOLENCE", "inputStrength": "HIGH",
"outputStrength": "HIGH"},
+ ]
+ },
+ )
+ # [END howto_operator_bedrock_create_guardrail]
+
+ @task(trigger_rule=TriggerRule.ALL_DONE)
+ def delete_guardrail(guardrail_id: str):
+ """Delete the guardrail."""
+ import boto3
+
+
boto3.client("bedrock").delete_guardrail(guardrailIdentifier=guardrail_id)
+
+ chain(
+ test_context,
+ create_guardrail,
+ delete_guardrail(guardrail_id=create_guardrail.output),
+ )
+
+ from tests_common.test_utils.watcher import watcher
+
+ list(dag.tasks) >> watcher()
+
+from tests_common.test_utils.system_tests import get_test_run # noqa: E402
+
+test_run = get_test_run(dag)
diff --git a/providers/amazon/tests/unit/amazon/aws/hooks/test_bedrock.py
b/providers/amazon/tests/unit/amazon/aws/hooks/test_bedrock.py
index 12b027c5299..19fbdf41043 100644
--- a/providers/amazon/tests/unit/amazon/aws/hooks/test_bedrock.py
+++ b/providers/amazon/tests/unit/amazon/aws/hooks/test_bedrock.py
@@ -16,6 +16,8 @@
# under the License.
from __future__ import annotations
+from unittest import mock
+
import pytest
from airflow.providers.amazon.aws.hooks.bedrock import (
@@ -39,3 +41,27 @@ class TestBedrockHooks:
def test_bedrock_hooks(self, test_hook, service_name):
assert test_hook.conn is not None
assert test_hook.conn.meta.service_model.service_name == service_name
+
+
+class TestBedrockHookGetGuardrailIdByName:
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_found(self, mock_conn):
+ mock_client = mock.MagicMock()
+ mock_paginator = mock.MagicMock()
+ mock_paginator.paginate.return_value = [{"guardrails": [{"name":
"my-guardrail", "id": "abc123"}]}]
+ mock_client.get_paginator.return_value = mock_paginator
+ mock_conn.return_value = mock_client
+
+ hook = BedrockHook()
+ assert hook.get_guardrail_id_by_name("my-guardrail") == "abc123"
+
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_not_found(self, mock_conn):
+ mock_client = mock.MagicMock()
+ mock_paginator = mock.MagicMock()
+ mock_paginator.paginate.return_value = [{"guardrails": []}]
+ mock_client.get_paginator.return_value = mock_paginator
+ mock_conn.return_value = mock_client
+
+ hook = BedrockHook()
+ assert hook.get_guardrail_id_by_name("nonexistent") is None
diff --git a/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
b/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
index c51bacf93e3..517ccf70dcf 100644
--- a/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
+++ b/providers/amazon/tests/unit/amazon/aws/operators/test_bedrock.py
@@ -30,6 +30,7 @@ from airflow.providers.amazon.aws.hooks.bedrock import
BedrockAgentHook, Bedrock
from airflow.providers.amazon.aws.operators.bedrock import (
BedrockBatchInferenceOperator,
BedrockCreateDataSourceOperator,
+ BedrockCreateGuardrailOperator,
BedrockCreateKnowledgeBaseOperator,
BedrockCreateProvisionedModelThroughputOperator,
BedrockCustomizeModelOperator,
@@ -765,3 +766,110 @@ class TestBedrockBatchInferenceOperator:
def test_template_fields(self):
validate_template_fields(self.operator)
+
+
+GUARDRAIL_NAME = "test-guardrail"
+GUARDRAIL_ID = "abc123"
+
+
+class TestBedrockCreateGuardrailOperator:
+ def setup_method(self):
+ self.operator = BedrockCreateGuardrailOperator(
+ task_id="create_guardrail",
+ guardrail_name=GUARDRAIL_NAME,
+ blocked_input_messaging="Input blocked.",
+ blocked_outputs_messaging="Output blocked.",
+ )
+
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute(self, mock_conn):
+ mock_client = mock.MagicMock()
+ mock_client.create_guardrail.return_value = {
+ "guardrailId": GUARDRAIL_ID,
+ "guardrailArn":
"arn:aws:bedrock:us-east-1:123456789012:guardrail/abc123",
+ "version": "DRAFT",
+ }
+ mock_conn.return_value = mock_client
+
+ result = self.operator.execute({})
+
+ mock_client.create_guardrail.assert_called_once_with(
+ name=GUARDRAIL_NAME,
+ blockedInputMessaging="Input blocked.",
+ blockedOutputsMessaging="Output blocked.",
+ )
+ assert result == GUARDRAIL_ID
+
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute_with_policies(self, mock_conn):
+ op = BedrockCreateGuardrailOperator(
+ task_id="create_guardrail",
+ guardrail_name=GUARDRAIL_NAME,
+ blocked_input_messaging="Input blocked.",
+ blocked_outputs_messaging="Output blocked.",
+ description="Test guardrail",
+ content_policy_config={
+ "filtersConfig": [{"type": "HATE", "inputStrength": "HIGH",
"outputStrength": "HIGH"}]
+ },
+ )
+ mock_client = mock.MagicMock()
+ mock_client.create_guardrail.return_value = {"guardrailId":
GUARDRAIL_ID}
+ mock_conn.return_value = mock_client
+
+ result = op.execute({})
+
+ call_kwargs = mock_client.create_guardrail.call_args[1]
+ assert call_kwargs["name"] == GUARDRAIL_NAME
+ assert call_kwargs["description"] == "Test guardrail"
+ assert "contentPolicyConfig" in call_kwargs
+ assert result == GUARDRAIL_ID
+
+ @mock.patch.object(BedrockHook, "get_guardrail_id_by_name",
return_value=GUARDRAIL_ID)
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute_skip_existing(self, mock_conn, mock_get_id):
+ mock_client = mock.MagicMock()
+ mock_client.create_guardrail.side_effect = ClientError(
+ {"Error": {"Code": "ConflictException", "Message": "Already
exists"}},
+ "CreateGuardrail",
+ )
+ mock_conn.return_value = mock_client
+
+ result = self.operator.execute({})
+
+ assert result == GUARDRAIL_ID
+ mock_get_id.assert_called_once_with(GUARDRAIL_NAME)
+
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute_fail_on_conflict(self, mock_conn):
+ op = BedrockCreateGuardrailOperator(
+ task_id="create_guardrail",
+ guardrail_name=GUARDRAIL_NAME,
+ blocked_input_messaging="Input blocked.",
+ blocked_outputs_messaging="Output blocked.",
+ if_exists="fail",
+ )
+ mock_client = mock.MagicMock()
+ mock_client.create_guardrail.side_effect = ClientError(
+ {"Error": {"Code": "ConflictException", "Message": "Already
exists"}},
+ "CreateGuardrail",
+ )
+ mock_conn.return_value = mock_client
+
+ with pytest.raises(ClientError):
+ op.execute({})
+
+ @mock.patch.object(BedrockHook, "get_guardrail_id_by_name",
return_value=None)
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute_skip_existing_not_found(self, mock_conn, mock_get_id):
+ mock_client = mock.MagicMock()
+ mock_client.create_guardrail.side_effect = ClientError(
+ {"Error": {"Code": "ConflictException", "Message": "Already
exists"}},
+ "CreateGuardrail",
+ )
+ mock_conn.return_value = mock_client
+
+ with pytest.raises(ClientError):
+ self.operator.execute({})
+
+ def test_template_fields(self):
+ validate_template_fields(self.operator)
