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: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
