Source code for airflow.providers.amazon.aws.sensors.bedrock
## 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__importannotationsimportabcfromtypingimportTYPE_CHECKING,Any,Sequence,TypeVarfromairflow.configurationimportconffromairflow.exceptionsimportAirflowException,AirflowSkipExceptionfromairflow.providers.amazon.aws.hooks.bedrockimportBedrockAgentHook,BedrockHookfromairflow.providers.amazon.aws.sensors.base_awsimportAwsBaseSensorfromairflow.providers.amazon.aws.triggers.bedrockimport(BedrockCustomizeModelCompletedTrigger,BedrockIngestionJobTrigger,BedrockKnowledgeBaseActiveTrigger,BedrockProvisionModelThroughputCompletedTrigger,)fromairflow.providers.amazon.aws.utils.mixinsimportaws_template_fieldsifTYPE_CHECKING:fromairflow.utils.contextimportContext_GenericBedrockHook=TypeVar("_GenericBedrockHook",BedrockAgentHook,BedrockHook)
[docs]classBedrockBaseSensor(AwsBaseSensor[_GenericBedrockHook]):""" General sensor behavior for Amazon Bedrock. Subclasses must implement following methods: - ``get_state()`` Subclasses must set the following fields: - ``INTERMEDIATE_STATES`` - ``FAILURE_STATES`` - ``SUCCESS_STATES`` - ``FAILURE_MESSAGE`` :param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore module to be installed. (default: False, but can be overridden in config file by setting default_deferrable to True) """
[docs]defpoke(self,context:Context,**kwargs)->bool:state=self.get_state()ifstateinself.FAILURE_STATES:# TODO: remove this if block when min_airflow_version is set to higher than 2.7.1ifself.soft_fail:raiseAirflowSkipException(self.FAILURE_MESSAGE)raiseAirflowException(self.FAILURE_MESSAGE)returnstatenotinself.INTERMEDIATE_STATES
@abc.abstractmethod
[docs]defget_state(self)->str:"""Implement in subclasses."""
[docs]classBedrockCustomizeModelCompletedSensor(BedrockBaseSensor[BedrockHook]):""" Poll the state of the model customization job until it reaches a terminal state; fails if the job fails. .. seealso:: For more information on how to use this sensor, take a look at the guide: :ref:`howto/sensor:BedrockCustomizeModelCompletedSensor` :param job_name: The name of the Bedrock model customization job. :param deferrable: If True, the sensor will operate in deferrable mode. This mode requires aiobotocore module to be installed. (default: False, but can be overridden in config file by setting default_deferrable to True) :param poke_interval: Polling period in seconds to check for the status of the job. (default: 120) :param max_retries: Number of times before returning the current state. (default: 75) :param aws_conn_id: The Airflow connection used for AWS credentials. If this is ``None`` or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. :param verify: Whether or not to verify SSL certificates. See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """
[docs]classBedrockProvisionModelThroughputCompletedSensor(BedrockBaseSensor[BedrockHook]):""" Poll the provisioned model throughput job until it reaches a terminal state; fails if the job fails. .. seealso:: For more information on how to use this sensor, take a look at the guide: :ref:`howto/sensor:BedrockProvisionModelThroughputCompletedSensor` :param model_id: The ARN or name of the provisioned throughput. :param deferrable: If True, the sensor will operate in deferrable more. This mode requires aiobotocore module to be installed. (default: False, but can be overridden in config file by setting default_deferrable to True) :param poke_interval: Polling period in seconds to check for the status of the job. (default: 60) :param max_retries: Number of times before returning the current state (default: 20) :param aws_conn_id: The Airflow connection used for AWS credentials. If this is ``None`` or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. :param verify: Whether or not to verify SSL certificates. See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """
[docs]classBedrockKnowledgeBaseActiveSensor(BedrockBaseSensor[BedrockAgentHook]):""" Poll the Knowledge Base status until it reaches a terminal state; fails if creation fails. .. seealso:: For more information on how to use this sensor, take a look at the guide: :ref:`howto/sensor:BedrockKnowledgeBaseActiveSensor` :param knowledge_base_id: The unique identifier of the knowledge base for which to get information. (templated) :param deferrable: If True, the sensor will operate in deferrable more. This mode requires aiobotocore module to be installed. (default: False, but can be overridden in config file by setting default_deferrable to True) :param poke_interval: Polling period in seconds to check for the status of the job. (default: 5) :param max_retries: Number of times before returning the current state (default: 24) :param aws_conn_id: The Airflow connection used for AWS credentials. If this is ``None`` or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. :param verify: Whether or not to verify SSL certificates. See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """
[docs]classBedrockIngestionJobSensor(BedrockBaseSensor[BedrockAgentHook]):""" Poll the ingestion job status until it reaches a terminal state; fails if creation fails. .. seealso:: For more information on how to use this sensor, take a look at the guide: :ref:`howto/sensor:BedrockIngestionJobSensor` :param knowledge_base_id: The unique identifier of the knowledge base for which to get information. (templated) :param data_source_id: The unique identifier of the data source in the ingestion job. (templated) :param ingestion_job_id: The unique identifier of the ingestion job. (templated) :param deferrable: If True, the sensor will operate in deferrable more. This mode requires aiobotocore module to be installed. (default: False, but can be overridden in config file by setting default_deferrable to True) :param poke_interval: Polling period in seconds to check for the status of the job. (default: 60) :param max_retries: Number of times before returning the current state (default: 10) :param aws_conn_id: The Airflow connection used for AWS credentials. If this is ``None`` or empty then the default boto3 behaviour is used. If running Airflow in a distributed manner and aws_conn_id is None or empty, then default boto3 configuration would be used (and must be maintained on each worker node). :param region_name: AWS region_name. If not specified then the default boto3 behaviour is used. :param verify: Whether or not to verify SSL certificates. See: https://boto3.amazonaws.com/v1/documentation/api/latest/reference/core/session.html :param botocore_config: Configuration dictionary (key-values) for botocore client. See: https://botocore.amazonaws.com/v1/documentation/api/latest/reference/config.html """