You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by "john-jac (via GitHub)" <gi...@apache.org> on 2026/04/28 17:29:49 UTC
[PR] Add `BedrockCreateGuardrailOperator` [airflow]
john-jac opened a new pull request, #66035:
URL: https://github.com/apache/airflow/pull/66035
## What
Add an operator to create Amazon Bedrock guardrails for implementing safety controls in generative AI applications.
## Use Case
Teams deploying LLM-powered applications need to enforce content safety policies (block hate speech, filter PII, deny off-topic discussions) as part of their deployment pipeline. A guardrail is created with content filters and topic policies, then associated with the Bedrock model. Having this in a DAG makes safety policy changes versioned, auditable, and promotable through dev/staging/prod:
```
create_guardrail >> configure_model_endpoint >> run_eval_tests >> notify_team
```
The `create_guardrail` API creates a DRAFT version that is immediately usable — no separate version step needed for dev/test workflows.
## How
- New `BedrockCreateGuardrailOperator` using `AwsBaseOperator[BedrockHook]` pattern
- Supports content filters, denied topics, word filters, PII masking, and contextual grounding
- Idempotent: `if_exists="skip"` handles `ConflictException`
- Uses `prune_dict` for clean kwargs building
- Returns the guardrail ID
## Files
- `operators/bedrock_guardrail.py` — operator
- `test_bedrock_guardrail.py` — unit tests (4 tests)
- `example_bedrock_guardrail.py` — system test with teardown
- `bedrock_guardrail.rst` — docs
- `provider.yaml` + `get_provider_info.py` — updated
## Testing
- Unit tests: 4 passed
- mypy: no issues
- Static checks: 43 passed, 0 PR-related failures
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on code in PR #66035:
URL: https://github.com/apache/airflow/pull/66035#discussion_r3161157213
##########
providers/amazon/src/airflow/providers/amazon/aws/hooks/bedrock.py:
##########
@@ -38,6 +38,15 @@ def __init__(self, *args, **kwargs) -> None:
kwargs["client_type"] = self.client_type
super().__init__(*args, **kwargs)
+ def get_guardrail_id_by_name(self, guardrail_name: str) -> str | None:
Review Comment:
This method needs to have unit tests associated to it
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "john-jac (via GitHub)" <gi...@apache.org>.
john-jac commented on PR #66035:
URL: https://github.com/apache/airflow/pull/66035#issuecomment-4340782114
Closing to resubmit with clean branch. Too many accumulated fixes.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "john-jac (via GitHub)" <gi...@apache.org>.
john-jac commented on PR #66035:
URL: https://github.com/apache/airflow/pull/66035#issuecomment-4345951343
> CI is failing
Should be good now. I'm still playing with my pre-submit CI process but it keeps missing things. It's a process.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "john-jac (via GitHub)" <gi...@apache.org>.
john-jac commented on PR #66035:
URL: https://github.com/apache/airflow/pull/66035#issuecomment-4344656215
Thanks @SameerMesiah97 and @vincbeck for the review! Addressed all feedback:
1. Moved guardrail tests into `test_bedrock.py` (no separate file)
2. Added hook unit tests for `get_guardrail_id_by_name` (`test_found` + `test_not_found`)
3. Unhappy path tests already added: `test_execute_fail_on_conflict` and `test_execute_skip_existing_not_found`
4. Operator uses `hook.get_guardrail_id_by_name()` instead of inline paginator
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on code in PR #66035:
URL: https://github.com/apache/airflow/pull/66035#discussion_r3156397266
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock_guardrail.py:
##########
Review Comment:
We try to have one file per service. Here Guardrails are still in the bedrock service, they are not a separate service (like S3 tables for instance). So I would rather have it in the existing `bedrock.py` file (and same for system test)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck merged PR #66035:
URL: https://github.com/apache/airflow/pull/66035
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "john-jac (via GitHub)" <gi...@apache.org>.
john-jac closed pull request #66035: Add `BedrockCreateGuardrailOperator`
URL: https://github.com/apache/airflow/pull/66035
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "SameerMesiah97 (via GitHub)" <gi...@apache.org>.
SameerMesiah97 commented on code in PR #66035:
URL: https://github.com/apache/airflow/pull/66035#discussion_r3156481359
##########
providers/amazon/tests/unit/amazon/aws/operators/test_bedrock_guardrail.py:
##########
@@ -0,0 +1,111 @@
+# 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 unittest import mock
+
+from botocore.exceptions import ClientError
+
+from airflow.providers.amazon.aws.hooks.bedrock import BedrockHook
+from airflow.providers.amazon.aws.operators.bedrock_guardrail import BedrockCreateGuardrailOperator
+
+from unit.amazon.aws.utils.test_template_fields import validate_template_fields
+
+GUARDRAIL_NAME = "test-guardrail"
+GUARDRAIL_ID = "abc123"
+
+
+class TestBedrockCreateGuardrailOperator:
+ def setup_method(self):
+ self.operator = BedrockCreateGuardrailOperator(
+ task_id="create_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",
+ 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"}]
+ },
+ topic_policy_config={
+ "topicsConfig": [
+ {"name": "competitors", "definition": "Discussion of competitors", "type": "DENY"}
+ ]
+ },
+ )
+ 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 = 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 "topicPolicyConfig" in call_kwargs
+ assert result == GUARDRAIL_ID
+
+ @mock.patch.object(BedrockHook, "conn", new_callable=mock.PropertyMock)
+ def test_execute_skip_existing(self, mock_conn):
+ mock_client = mock.MagicMock()
+ mock_client.create_guardrail.side_effect = ClientError(
+ {"Error": {"Code": "ConflictException", "Message": "Already exists"}},
+ "CreateGuardrail",
+ )
+ paginator = mock.MagicMock()
+ paginator.paginate.return_value = [{"guardrails": [{"name": GUARDRAIL_NAME, "id": GUARDRAIL_ID}]}]
+ mock_client.get_paginator.return_value = paginator
+ mock_conn.return_value = mock_client
+
+ result = self.operator.execute({})
+
+ assert result == GUARDRAIL_ID
+
+ def test_template_fields(self):
Review Comment:
I think you should additional tests for the following scenarios:
1) When `skip` is set to `fail` when there is a conflict and an error is raised.
2) When a conflict arises and no `guardrail_id` is found.
There may be other "unhappy" paths besides this so I would urge you to review your PR again and ensure that all important scenarios are covered.
##########
providers/amazon/src/airflow/providers/amazon/aws/operators/bedrock.py:
##########
@@ -1025,3 +1025,113 @@ def execute(self, context: Context) -> str:
)
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: dict[str, str] = {
+ "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)
+ paginator = self.hook.conn.get_paginator("list_guardrails")
+ for page in paginator.paginate():
+ for g in page.get("guardrails", []):
+ if g["name"] == self.guardrail_name:
+ guardrail_id = g["id"]
+ break
+ else:
+ continue
+ break
+ else:
+ raise
+ else:
+ raise
Review Comment:
I think it would be better to simplify this operator and keep it lightweight. I would create a method like this in the `BedrockHook`:
```
def get_guardrail_id_by_name(self, name: str) -> str | None:
paginator = self.conn.get_paginator("list_guardrails")
for page in paginator.paginate():
for g in page.get("guardrails", []):
if g.get("name") == name:
return g.get("id")
return None
```
And call it in the operator. And I believe the operator should be streamlined like this:
```
def execute(self, context: Context) -> str:
kwargs = self._build_payload()
try:
guardrail_id = self.hook.create_guardrail(**kwargs)
except ClientError as e:
error_code = e.response.get("Error", {}).get("Code")
if error_code == "ConflictException" and self.if_exists == "skip":
self.log.info("Guardrail %s already exists, skipping.", self.name)
guardrail_id = self.hook.get_guardrail_id_by_name(self.name)
if not guardrail_id:
raise
else:
raise
self.log.info("Guardrail %s: %s", self.name, guardrail_id)
return guardrail_id
```
You don't have to implement `_build_payload` for the `kwargs` but it would be cleaner if you did.
Right now, this is far too messy and hard to reason about.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org
Re: [PR] Add `BedrockCreateGuardrailOperator` [airflow]
Posted by "vincbeck (via GitHub)" <gi...@apache.org>.
vincbeck commented on PR #66035:
URL: https://github.com/apache/airflow/pull/66035#issuecomment-4345075763
CI is failing
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: commits-unsubscribe@airflow.apache.org
For queries about this service, please contact Infrastructure at:
users@infra.apache.org