fix missing docstrings

apache · Apr 25, 2024 · be1caac · be1caac
1 parent 293744a
commit be1caac
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 11 deletions.
diff --git a/airflow/providers/amazon/aws/operators/bedrock.py b/airflow/providers/amazon/aws/operators/bedrock.py
@@ -374,6 +374,7 @@ class BedrockCreateKnowledgeBaseOperator(AwsBaseOperator[BedrockAgentHook]):
         being available, the operator will wait and retry.  (default: True) (templated)
     :param indexing_error_retry_delay: Seconds between retries if an index error is encountered. (default 5) (templated)
     :param indexing_error_max_attempts: Maximum number of times to retry when encountering an index error. (default 20) (templated)
+    :param create_knowledge_base_kwargs: Any additional optional parameters to pass to the API call. (templated)
 
     :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
     :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
@@ -402,6 +403,7 @@ class BedrockCreateKnowledgeBaseOperator(AwsBaseOperator[BedrockAgentHook]):
         "wait_for_indexing",
         "indexing_error_retry_delay",
         "indexing_error_max_attempts",
+        "create_knowledge_base_kwargs",
     )
 
     def __init__(
@@ -518,6 +520,7 @@ class BedrockCreateDataSourceOperator(AwsBaseOperator[BedrockAgentHook]):
     :param name: name for the Amazon Bedrock Data Source being created. (templated).
     :param bucket_name: The name of the Amazon S3 bucket to use for data source storage. (templated)
     :param knowledge_base_id: The unique identifier of the knowledge base to which to add the data source. (templated)
+    :param create_data_source_kwargs: Any additional optional parameters to pass to the API call. (templated)
 
     :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
@@ -532,6 +535,12 @@ class BedrockCreateDataSourceOperator(AwsBaseOperator[BedrockAgentHook]):
     """
 
     aws_hook_class = BedrockAgentHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "name",
+        "bucket_name",
+        "knowledge_base_id",
+        "create_data_source_kwargs",
+    )
 
     def __init__(
         self,
@@ -547,12 +556,6 @@ def __init__(
         self.bucket_name = bucket_name
         self.create_data_source_kwargs = create_data_source_kwargs or {}
 
-    template_fields: Sequence[str] = aws_template_fields(
-        "name",
-        "bucket_name",
-        "knowledge_base_id",
-    )
-
     def execute(self, context: Context) -> str:
         create_ds_response = self.hook.conn.create_data_source(
             name=self.name,
@@ -577,6 +580,7 @@ class BedrockIngestDataOperator(AwsBaseOperator[BedrockAgentHook]):
 
     :param knowledge_base_id: The unique identifier of the knowledge base to which to add the data source. (templated)
     :param data_source_id: The unique identifier of the data source to ingest. (templated)
+    :param ingest_data_kwargs: Any additional optional parameters to pass to the API call. (templated)
 
     :param wait_for_completion: Whether to wait for cluster to stop. (default: True)
     :param waiter_delay: Time in seconds to wait between status checks. (default: 60)
@@ -597,6 +601,11 @@ class BedrockIngestDataOperator(AwsBaseOperator[BedrockAgentHook]):
     """
 
     aws_hook_class = BedrockAgentHook
+    template_fields: Sequence[str] = aws_template_fields(
+        "knowledge_base_id",
+        "data_source_id",
+        "ingest_data_kwargs",
+    )
 
     def __init__(
         self,
@@ -619,11 +628,6 @@ def __init__(
         self.waiter_max_attempts = waiter_max_attempts
         self.deferrable = deferrable
 
-    template_fields: Sequence[str] = aws_template_fields(
-        "knowledge_base_id",
-        "data_source_id",
-    )
-
     def execute(self, context: Context) -> str:
         ingestion_job_id = self.hook.conn.start_ingestion_job(
             knowledgeBaseId=self.knowledge_base_id, dataSourceId=self.data_source_id

diff --git a/airflow/providers/amazon/aws/triggers/bedrock.py b/airflow/providers/amazon/aws/triggers/bedrock.py
@@ -66,6 +66,7 @@ class BedrockKnowledgeBaseActiveTrigger(AwsBaseWaiterTrigger):
     Trigger when a Bedrock Knowledge Base reaches the ACTIVE state.
 
     :param knowledge_base_id: The unique identifier of the knowledge base for which to get information.
+
     :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 60)
     :param waiter_max_attempts: The maximum number of attempts to be made. (default: 20)
     :param aws_conn_id: The Airflow connection used for AWS credentials.
@@ -102,6 +103,7 @@ class BedrockProvisionModelThroughputCompletedTrigger(AwsBaseWaiterTrigger):
     Trigger when a provisioned throughput job is complete.
 
     :param provisioned_model_id: The ARN or name of the provisioned throughput.
+
     :param waiter_delay: The amount of time in seconds to wait between attempts. (default: 120)
     :param waiter_max_attempts: The maximum number of attempts to be made. (default: 75)
     :param aws_conn_id: The Airflow connection used for AWS credentials.