apache-airflow-providers-google 12.0.0rc2__py3-none-any.whl → 14.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (163) hide show
  1. airflow/providers/google/LICENSE +0 -52
  2. airflow/providers/google/__init__.py +1 -1
  3. airflow/providers/google/ads/hooks/ads.py +27 -13
  4. airflow/providers/google/ads/transfers/ads_to_gcs.py +18 -4
  5. airflow/providers/google/assets/bigquery.py +17 -0
  6. airflow/providers/google/cloud/_internal_client/secret_manager_client.py +2 -3
  7. airflow/providers/google/cloud/hooks/alloy_db.py +736 -8
  8. airflow/providers/google/cloud/hooks/automl.py +10 -4
  9. airflow/providers/google/cloud/hooks/bigquery.py +125 -22
  10. airflow/providers/google/cloud/hooks/bigquery_dts.py +8 -8
  11. airflow/providers/google/cloud/hooks/bigtable.py +2 -3
  12. airflow/providers/google/cloud/hooks/cloud_batch.py +3 -4
  13. airflow/providers/google/cloud/hooks/cloud_build.py +4 -5
  14. airflow/providers/google/cloud/hooks/cloud_composer.py +3 -4
  15. airflow/providers/google/cloud/hooks/cloud_memorystore.py +3 -4
  16. airflow/providers/google/cloud/hooks/cloud_run.py +3 -4
  17. airflow/providers/google/cloud/hooks/cloud_sql.py +7 -3
  18. airflow/providers/google/cloud/hooks/cloud_storage_transfer_service.py +119 -7
  19. airflow/providers/google/cloud/hooks/compute.py +3 -3
  20. airflow/providers/google/cloud/hooks/datacatalog.py +3 -4
  21. airflow/providers/google/cloud/hooks/dataflow.py +12 -12
  22. airflow/providers/google/cloud/hooks/dataform.py +2 -3
  23. airflow/providers/google/cloud/hooks/datafusion.py +2 -2
  24. airflow/providers/google/cloud/hooks/dataplex.py +1032 -11
  25. airflow/providers/google/cloud/hooks/dataproc.py +4 -5
  26. airflow/providers/google/cloud/hooks/dataproc_metastore.py +3 -4
  27. airflow/providers/google/cloud/hooks/dlp.py +3 -4
  28. airflow/providers/google/cloud/hooks/gcs.py +7 -6
  29. airflow/providers/google/cloud/hooks/kms.py +2 -3
  30. airflow/providers/google/cloud/hooks/kubernetes_engine.py +8 -8
  31. airflow/providers/google/cloud/hooks/life_sciences.py +1 -1
  32. airflow/providers/google/cloud/hooks/managed_kafka.py +482 -0
  33. airflow/providers/google/cloud/hooks/natural_language.py +2 -3
  34. airflow/providers/google/cloud/hooks/os_login.py +2 -3
  35. airflow/providers/google/cloud/hooks/pubsub.py +6 -6
  36. airflow/providers/google/cloud/hooks/secret_manager.py +2 -3
  37. airflow/providers/google/cloud/hooks/spanner.py +2 -2
  38. airflow/providers/google/cloud/hooks/speech_to_text.py +2 -3
  39. airflow/providers/google/cloud/hooks/stackdriver.py +4 -4
  40. airflow/providers/google/cloud/hooks/tasks.py +3 -4
  41. airflow/providers/google/cloud/hooks/text_to_speech.py +2 -3
  42. airflow/providers/google/cloud/hooks/translate.py +236 -5
  43. airflow/providers/google/cloud/hooks/vertex_ai/auto_ml.py +9 -4
  44. airflow/providers/google/cloud/hooks/vertex_ai/batch_prediction_job.py +3 -4
  45. airflow/providers/google/cloud/hooks/vertex_ai/custom_job.py +4 -5
  46. airflow/providers/google/cloud/hooks/vertex_ai/dataset.py +3 -4
  47. airflow/providers/google/cloud/hooks/vertex_ai/endpoint_service.py +2 -3
  48. airflow/providers/google/cloud/hooks/vertex_ai/feature_store.py +3 -4
  49. airflow/providers/google/cloud/hooks/vertex_ai/generative_model.py +1 -181
  50. airflow/providers/google/cloud/hooks/vertex_ai/hyperparameter_tuning_job.py +3 -4
  51. airflow/providers/google/cloud/hooks/vertex_ai/model_service.py +2 -3
  52. airflow/providers/google/cloud/hooks/vertex_ai/pipeline_job.py +3 -4
  53. airflow/providers/google/cloud/hooks/vertex_ai/prediction_service.py +2 -3
  54. airflow/providers/google/cloud/hooks/video_intelligence.py +2 -3
  55. airflow/providers/google/cloud/hooks/vision.py +3 -4
  56. airflow/providers/google/cloud/hooks/workflows.py +2 -3
  57. airflow/providers/google/cloud/links/alloy_db.py +46 -0
  58. airflow/providers/google/cloud/links/bigquery.py +25 -0
  59. airflow/providers/google/cloud/links/dataplex.py +172 -2
  60. airflow/providers/google/cloud/links/kubernetes_engine.py +1 -2
  61. airflow/providers/google/cloud/links/managed_kafka.py +104 -0
  62. airflow/providers/google/cloud/links/translate.py +28 -0
  63. airflow/providers/google/cloud/log/gcs_task_handler.py +3 -3
  64. airflow/providers/google/cloud/log/stackdriver_task_handler.py +11 -10
  65. airflow/providers/google/cloud/openlineage/facets.py +67 -0
  66. airflow/providers/google/cloud/openlineage/mixins.py +438 -173
  67. airflow/providers/google/cloud/openlineage/utils.py +394 -61
  68. airflow/providers/google/cloud/operators/alloy_db.py +980 -69
  69. airflow/providers/google/cloud/operators/automl.py +83 -245
  70. airflow/providers/google/cloud/operators/bigquery.py +377 -74
  71. airflow/providers/google/cloud/operators/bigquery_dts.py +126 -13
  72. airflow/providers/google/cloud/operators/bigtable.py +1 -3
  73. airflow/providers/google/cloud/operators/cloud_base.py +1 -2
  74. airflow/providers/google/cloud/operators/cloud_batch.py +2 -4
  75. airflow/providers/google/cloud/operators/cloud_build.py +3 -5
  76. airflow/providers/google/cloud/operators/cloud_composer.py +5 -7
  77. airflow/providers/google/cloud/operators/cloud_memorystore.py +4 -6
  78. airflow/providers/google/cloud/operators/cloud_run.py +6 -5
  79. airflow/providers/google/cloud/operators/cloud_sql.py +20 -8
  80. airflow/providers/google/cloud/operators/cloud_storage_transfer_service.py +62 -8
  81. airflow/providers/google/cloud/operators/compute.py +3 -4
  82. airflow/providers/google/cloud/operators/datacatalog.py +9 -11
  83. airflow/providers/google/cloud/operators/dataflow.py +1 -112
  84. airflow/providers/google/cloud/operators/dataform.py +3 -5
  85. airflow/providers/google/cloud/operators/datafusion.py +1 -1
  86. airflow/providers/google/cloud/operators/dataplex.py +2046 -7
  87. airflow/providers/google/cloud/operators/dataproc.py +102 -17
  88. airflow/providers/google/cloud/operators/dataproc_metastore.py +7 -9
  89. airflow/providers/google/cloud/operators/dlp.py +17 -19
  90. airflow/providers/google/cloud/operators/gcs.py +14 -17
  91. airflow/providers/google/cloud/operators/kubernetes_engine.py +2 -2
  92. airflow/providers/google/cloud/operators/managed_kafka.py +788 -0
  93. airflow/providers/google/cloud/operators/natural_language.py +3 -5
  94. airflow/providers/google/cloud/operators/pubsub.py +39 -7
  95. airflow/providers/google/cloud/operators/speech_to_text.py +3 -5
  96. airflow/providers/google/cloud/operators/stackdriver.py +3 -5
  97. airflow/providers/google/cloud/operators/tasks.py +4 -6
  98. airflow/providers/google/cloud/operators/text_to_speech.py +2 -4
  99. airflow/providers/google/cloud/operators/translate.py +414 -5
  100. airflow/providers/google/cloud/operators/translate_speech.py +2 -4
  101. airflow/providers/google/cloud/operators/vertex_ai/auto_ml.py +9 -8
  102. airflow/providers/google/cloud/operators/vertex_ai/batch_prediction_job.py +4 -6
  103. airflow/providers/google/cloud/operators/vertex_ai/custom_job.py +6 -8
  104. airflow/providers/google/cloud/operators/vertex_ai/dataset.py +4 -6
  105. airflow/providers/google/cloud/operators/vertex_ai/endpoint_service.py +4 -6
  106. airflow/providers/google/cloud/operators/vertex_ai/generative_model.py +0 -322
  107. airflow/providers/google/cloud/operators/vertex_ai/hyperparameter_tuning_job.py +4 -6
  108. airflow/providers/google/cloud/operators/vertex_ai/model_service.py +4 -6
  109. airflow/providers/google/cloud/operators/vertex_ai/pipeline_job.py +4 -6
  110. airflow/providers/google/cloud/operators/video_intelligence.py +3 -5
  111. airflow/providers/google/cloud/operators/vision.py +4 -6
  112. airflow/providers/google/cloud/operators/workflows.py +5 -7
  113. airflow/providers/google/cloud/secrets/secret_manager.py +1 -2
  114. airflow/providers/google/cloud/sensors/bigquery_dts.py +3 -5
  115. airflow/providers/google/cloud/sensors/bigtable.py +2 -3
  116. airflow/providers/google/cloud/sensors/cloud_composer.py +32 -8
  117. airflow/providers/google/cloud/sensors/cloud_storage_transfer_service.py +39 -1
  118. airflow/providers/google/cloud/sensors/dataplex.py +4 -6
  119. airflow/providers/google/cloud/sensors/dataproc.py +2 -3
  120. airflow/providers/google/cloud/sensors/dataproc_metastore.py +1 -2
  121. airflow/providers/google/cloud/sensors/gcs.py +2 -4
  122. airflow/providers/google/cloud/sensors/pubsub.py +2 -3
  123. airflow/providers/google/cloud/sensors/workflows.py +3 -5
  124. airflow/providers/google/cloud/transfers/bigquery_to_gcs.py +5 -5
  125. airflow/providers/google/cloud/transfers/gcs_to_bigquery.py +10 -12
  126. airflow/providers/google/cloud/transfers/gcs_to_gcs.py +1 -1
  127. airflow/providers/google/cloud/transfers/gcs_to_sftp.py +36 -4
  128. airflow/providers/google/cloud/transfers/mssql_to_gcs.py +27 -2
  129. airflow/providers/google/cloud/transfers/mysql_to_gcs.py +27 -2
  130. airflow/providers/google/cloud/transfers/postgres_to_gcs.py +27 -2
  131. airflow/providers/google/cloud/transfers/sftp_to_gcs.py +34 -5
  132. airflow/providers/google/cloud/transfers/sql_to_gcs.py +15 -0
  133. airflow/providers/google/cloud/transfers/trino_to_gcs.py +25 -2
  134. airflow/providers/google/cloud/triggers/bigquery_dts.py +1 -2
  135. airflow/providers/google/cloud/triggers/cloud_batch.py +1 -2
  136. airflow/providers/google/cloud/triggers/cloud_build.py +1 -2
  137. airflow/providers/google/cloud/triggers/cloud_composer.py +13 -3
  138. airflow/providers/google/cloud/triggers/cloud_storage_transfer_service.py +102 -4
  139. airflow/providers/google/cloud/triggers/dataflow.py +2 -3
  140. airflow/providers/google/cloud/triggers/dataplex.py +1 -2
  141. airflow/providers/google/cloud/triggers/dataproc.py +2 -3
  142. airflow/providers/google/cloud/triggers/kubernetes_engine.py +1 -1
  143. airflow/providers/google/cloud/triggers/pubsub.py +1 -2
  144. airflow/providers/google/cloud/triggers/vertex_ai.py +7 -8
  145. airflow/providers/google/cloud/utils/credentials_provider.py +15 -8
  146. airflow/providers/google/cloud/utils/external_token_supplier.py +1 -0
  147. airflow/providers/google/common/auth_backend/google_openid.py +4 -4
  148. airflow/providers/google/common/consts.py +1 -2
  149. airflow/providers/google/common/hooks/base_google.py +8 -7
  150. airflow/providers/google/get_provider_info.py +186 -134
  151. airflow/providers/google/marketing_platform/hooks/analytics_admin.py +2 -3
  152. airflow/providers/google/marketing_platform/hooks/search_ads.py +1 -1
  153. airflow/providers/google/marketing_platform/operators/analytics_admin.py +5 -7
  154. {apache_airflow_providers_google-12.0.0rc2.dist-info → apache_airflow_providers_google-14.0.0.dist-info}/METADATA +41 -58
  155. {apache_airflow_providers_google-12.0.0rc2.dist-info → apache_airflow_providers_google-14.0.0.dist-info}/RECORD +157 -159
  156. airflow/providers/google/cloud/example_dags/example_facebook_ads_to_gcs.py +0 -141
  157. airflow/providers/google/cloud/example_dags/example_looker.py +0 -64
  158. airflow/providers/google/cloud/example_dags/example_presto_to_gcs.py +0 -194
  159. airflow/providers/google/cloud/example_dags/example_salesforce_to_gcs.py +0 -129
  160. airflow/providers/google/marketing_platform/example_dags/__init__.py +0 -16
  161. airflow/providers/google/marketing_platform/example_dags/example_display_video.py +0 -213
  162. {apache_airflow_providers_google-12.0.0rc2.dist-info → apache_airflow_providers_google-14.0.0.dist-info}/WHEEL +0 -0
  163. {apache_airflow_providers_google-12.0.0rc2.dist-info → apache_airflow_providers_google-14.0.0.dist-info}/entry_points.txt +0 -0
@@ -22,19 +22,17 @@ from __future__ import annotations
22
22
  from collections.abc import Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
- from google.protobuf.json_format import MessageToDict
26
-
27
25
  from airflow.exceptions import AirflowException
28
26
  from airflow.providers.google.cloud.hooks.speech_to_text import CloudSpeechToTextHook
29
27
  from airflow.providers.google.cloud.hooks.translate import CloudTranslateHook
30
28
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
31
29
  from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
32
30
  from airflow.providers.google.common.links.storage import FileDetailsLink
31
+ from google.protobuf.json_format import MessageToDict
33
32
 
34
33
  if TYPE_CHECKING:
35
- from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig
36
-
37
34
  from airflow.utils.context import Context
35
+ from google.cloud.speech_v1.types import RecognitionAudio, RecognitionConfig
38
36
 
39
37
 
40
38
  class CloudTranslateSpeechOperator(GoogleCloudBaseOperator):
@@ -23,12 +23,6 @@ from __future__ import annotations
23
23
  from collections.abc import Sequence
24
24
  from typing import TYPE_CHECKING
25
25
 
26
- from google.api_core.exceptions import NotFound
27
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
28
- from google.cloud.aiplatform import datasets
29
- from google.cloud.aiplatform.models import Model
30
- from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
31
-
32
26
  from airflow.exceptions import AirflowProviderDeprecationWarning
33
27
  from airflow.providers.google.cloud.hooks.vertex_ai.auto_ml import AutoMLHook
34
28
  from airflow.providers.google.cloud.links.vertex_ai import (
@@ -38,11 +32,15 @@ from airflow.providers.google.cloud.links.vertex_ai import (
38
32
  )
39
33
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
40
34
  from airflow.providers.google.common.deprecated import deprecated
35
+ from google.api_core.exceptions import NotFound
36
+ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
37
+ from google.cloud.aiplatform import datasets
38
+ from google.cloud.aiplatform.models import Model
39
+ from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
41
40
 
42
41
  if TYPE_CHECKING:
43
- from google.api_core.retry import Retry
44
-
45
42
  from airflow.utils.context import Context
43
+ from google.api_core.retry import Retry
46
44
 
47
45
 
48
46
  class AutoMLTrainingJobBaseOperator(GoogleCloudBaseOperator):
@@ -145,6 +143,7 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
145
143
  parent_model: str | None = None,
146
144
  window_stride_length: int | None = None,
147
145
  window_max_count: int | None = None,
146
+ holiday_regions: list[str] | None = None,
148
147
  **kwargs,
149
148
  ) -> None:
150
149
  super().__init__(
@@ -184,6 +183,7 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
184
183
  self.budget_milli_node_hours = budget_milli_node_hours
185
184
  self.window_stride_length = window_stride_length
186
185
  self.window_max_count = window_max_count
186
+ self.holiday_regions = holiday_regions
187
187
 
188
188
  def execute(self, context: Context):
189
189
  self.hook = AutoMLHook(
@@ -236,6 +236,7 @@ class CreateAutoMLForecastingTrainingJobOperator(AutoMLTrainingJobBaseOperator):
236
236
  sync=self.sync,
237
237
  window_stride_length=self.window_stride_length,
238
238
  window_max_count=self.window_max_count,
239
+ holiday_regions=self.holiday_regions,
239
240
  )
240
241
 
241
242
  if model:
@@ -24,10 +24,6 @@ from collections.abc import Sequence
24
24
  from functools import cached_property
25
25
  from typing import TYPE_CHECKING, Any
26
26
 
27
- from google.api_core.exceptions import NotFound
28
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
29
- from google.cloud.aiplatform_v1.types import BatchPredictionJob
30
-
31
27
  from airflow.configuration import conf
32
28
  from airflow.exceptions import AirflowException
33
29
  from airflow.providers.google.cloud.hooks.vertex_ai.batch_prediction_job import BatchPredictionJobHook
@@ -37,13 +33,15 @@ from airflow.providers.google.cloud.links.vertex_ai import (
37
33
  )
38
34
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
39
35
  from airflow.providers.google.cloud.triggers.vertex_ai import CreateBatchPredictionJobTrigger
36
+ from google.api_core.exceptions import NotFound
37
+ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
38
+ from google.cloud.aiplatform_v1.types import BatchPredictionJob
40
39
 
41
40
  if TYPE_CHECKING:
41
+ from airflow.utils.context import Context
42
42
  from google.api_core.retry import Retry
43
43
  from google.cloud.aiplatform import BatchPredictionJob as BatchPredictionJobObject, Model, explain
44
44
 
45
- from airflow.utils.context import Context
46
-
47
45
 
48
46
  class CreateBatchPredictionJobOperator(GoogleCloudBaseOperator):
49
47
  """
@@ -23,12 +23,6 @@ from collections.abc import Sequence
23
23
  from functools import cached_property
24
24
  from typing import TYPE_CHECKING, Any
25
25
 
26
- from google.api_core.exceptions import NotFound
27
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
28
- from google.cloud.aiplatform.models import Model
29
- from google.cloud.aiplatform_v1.types.dataset import Dataset
30
- from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
31
-
32
26
  from airflow.configuration import conf
33
27
  from airflow.exceptions import AirflowException, AirflowProviderDeprecationWarning
34
28
  from airflow.providers.google.cloud.hooks.vertex_ai.custom_job import CustomJobHook
@@ -44,8 +38,14 @@ from airflow.providers.google.cloud.triggers.vertex_ai import (
44
38
  CustomTrainingJobTrigger,
45
39
  )
46
40
  from airflow.providers.google.common.deprecated import deprecated
41
+ from google.api_core.exceptions import NotFound
42
+ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
43
+ from google.cloud.aiplatform.models import Model
44
+ from google.cloud.aiplatform_v1.types.dataset import Dataset
45
+ from google.cloud.aiplatform_v1.types.training_pipeline import TrainingPipeline
47
46
 
48
47
  if TYPE_CHECKING:
48
+ from airflow.utils.context import Context
49
49
  from google.api_core.retry import Retry
50
50
  from google.cloud.aiplatform import (
51
51
  CustomContainerTrainingJob,
@@ -53,8 +53,6 @@ if TYPE_CHECKING:
53
53
  CustomTrainingJob,
54
54
  )
55
55
 
56
- from airflow.utils.context import Context
57
-
58
56
 
59
57
  class CustomTrainingJobBaseOperator(GoogleCloudBaseOperator):
60
58
  """The base class for operators that launch Custom jobs on VertexAI."""
@@ -22,20 +22,18 @@ from __future__ import annotations
22
22
  from collections.abc import Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
- from google.api_core.exceptions import NotFound
26
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
27
- from google.cloud.aiplatform_v1.types import Dataset, ExportDataConfig, ImportDataConfig
28
-
29
25
  from airflow.providers.google.cloud.hooks.vertex_ai.dataset import DatasetHook
30
26
  from airflow.providers.google.cloud.links.vertex_ai import VertexAIDatasetLink, VertexAIDatasetListLink
31
27
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
28
+ from google.api_core.exceptions import NotFound
29
+ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
30
+ from google.cloud.aiplatform_v1.types import Dataset, ExportDataConfig, ImportDataConfig
32
31
 
33
32
  if TYPE_CHECKING:
33
+ from airflow.utils.context import Context
34
34
  from google.api_core.retry import Retry
35
35
  from google.protobuf.field_mask_pb2 import FieldMask
36
36
 
37
- from airflow.utils.context import Context
38
-
39
37
 
40
38
  class CreateDatasetOperator(GoogleCloudBaseOperator):
41
39
  """
@@ -23,10 +23,6 @@ from __future__ import annotations
23
23
  from collections.abc import Sequence
24
24
  from typing import TYPE_CHECKING
25
25
 
26
- from google.api_core.exceptions import NotFound
27
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
28
- from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint, endpoint_service
29
-
30
26
  from airflow.providers.google.cloud.hooks.vertex_ai.endpoint_service import EndpointServiceHook
31
27
  from airflow.providers.google.cloud.links.vertex_ai import (
32
28
  VertexAIEndpointLink,
@@ -34,13 +30,15 @@ from airflow.providers.google.cloud.links.vertex_ai import (
34
30
  VertexAIModelLink,
35
31
  )
36
32
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
33
+ from google.api_core.exceptions import NotFound
34
+ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
35
+ from google.cloud.aiplatform_v1.types import DeployedModel, Endpoint, endpoint_service
37
36
 
38
37
  if TYPE_CHECKING:
38
+ from airflow.utils.context import Context
39
39
  from google.api_core.retry import Retry
40
40
  from google.protobuf.field_mask_pb2 import FieldMask
41
41
 
42
- from airflow.utils.context import Context
43
-
44
42
 
45
43
  class CreateEndpointOperator(GoogleCloudBaseOperator):
46
44
  """
@@ -31,328 +31,6 @@ if TYPE_CHECKING:
31
31
  from airflow.utils.context import Context
32
32
 
33
33
 
34
- @deprecated(
35
- planned_removal_date="January 01, 2025",
36
- use_instead="TextGenerationModelPredictOperator",
37
- category=AirflowProviderDeprecationWarning,
38
- )
39
- class PromptLanguageModelOperator(GoogleCloudBaseOperator):
40
- """
41
- Uses the Vertex AI PaLM API to generate natural language text.
42
-
43
- :param project_id: Required. The ID of the Google Cloud project that the
44
- service belongs to (templated).
45
- :param location: Required. The ID of the Google Cloud location that the
46
- service belongs to (templated).
47
- :param prompt: Required. Inputs or queries that a user or a program gives
48
- to the Vertex AI PaLM API, in order to elicit a specific response (templated).
49
- :param pretrained_model: By default uses the pre-trained model `text-bison`,
50
- optimized for performing natural language tasks such as classification,
51
- summarization, extraction, content creation, and ideation.
52
- :param temperature: Temperature controls the degree of randomness in token
53
- selection. Defaults to 0.0.
54
- :param max_output_tokens: Token limit determines the maximum amount of text
55
- output. Defaults to 256.
56
- :param top_p: Tokens are selected from most probable to least until the sum
57
- of their probabilities equals the top_p value. Defaults to 0.8.
58
- :param top_k: A top_k of 1 means the selected token is the most probable
59
- among all tokens. Defaults to 0.4.
60
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
61
- :param impersonation_chain: Optional service account to impersonate using short-term
62
- credentials, or chained list of accounts required to get the access_token
63
- of the last account in the list, which will be impersonated in the request.
64
- If set as a string, the account must grant the originating account
65
- the Service Account Token Creator IAM role.
66
- If set as a sequence, the identities from the list must grant
67
- Service Account Token Creator IAM role to the directly preceding identity, with first
68
- account from the list granting this role to the originating account (templated).
69
- """
70
-
71
- template_fields = ("location", "project_id", "impersonation_chain", "prompt")
72
-
73
- def __init__(
74
- self,
75
- *,
76
- project_id: str,
77
- location: str,
78
- prompt: str,
79
- pretrained_model: str = "text-bison",
80
- temperature: float = 0.0,
81
- max_output_tokens: int = 256,
82
- top_p: float = 0.8,
83
- top_k: int = 40,
84
- gcp_conn_id: str = "google_cloud_default",
85
- impersonation_chain: str | Sequence[str] | None = None,
86
- **kwargs,
87
- ) -> None:
88
- super().__init__(**kwargs)
89
- self.project_id = project_id
90
- self.location = location
91
- self.prompt = prompt
92
- self.pretrained_model = pretrained_model
93
- self.temperature = temperature
94
- self.max_output_tokens = max_output_tokens
95
- self.top_p = top_p
96
- self.top_k = top_k
97
- self.gcp_conn_id = gcp_conn_id
98
- self.impersonation_chain = impersonation_chain
99
-
100
- def execute(self, context: Context):
101
- self.hook = GenerativeModelHook(
102
- gcp_conn_id=self.gcp_conn_id,
103
- impersonation_chain=self.impersonation_chain,
104
- )
105
-
106
- self.log.info("Submitting prompt")
107
- response = self.hook.prompt_language_model(
108
- project_id=self.project_id,
109
- location=self.location,
110
- prompt=self.prompt,
111
- pretrained_model=self.pretrained_model,
112
- temperature=self.temperature,
113
- max_output_tokens=self.max_output_tokens,
114
- top_p=self.top_p,
115
- top_k=self.top_k,
116
- )
117
-
118
- self.log.info("Model response: %s", response)
119
- self.xcom_push(context, key="prompt_response", value=response)
120
-
121
- return response
122
-
123
-
124
- @deprecated(
125
- planned_removal_date="January 01, 2025",
126
- use_instead="TextEmbeddingModelGetEmbeddingsOperator",
127
- category=AirflowProviderDeprecationWarning,
128
- )
129
- class GenerateTextEmbeddingsOperator(GoogleCloudBaseOperator):
130
- """
131
- Uses the Vertex AI PaLM API to generate natural language text.
132
-
133
- :param project_id: Required. The ID of the Google Cloud project that the
134
- service belongs to (templated).
135
- :param location: Required. The ID of the Google Cloud location that the
136
- service belongs to (templated).
137
- :param prompt: Required. Inputs or queries that a user or a program gives
138
- to the Vertex AI PaLM API, in order to elicit a specific response (templated).
139
- :param pretrained_model: By default uses the pre-trained model `textembedding-gecko`,
140
- optimized for performing text embeddings.
141
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
142
- :param impersonation_chain: Optional service account to impersonate using short-term
143
- credentials, or chained list of accounts required to get the access_token
144
- of the last account in the list, which will be impersonated in the request.
145
- If set as a string, the account must grant the originating account
146
- the Service Account Token Creator IAM role.
147
- If set as a sequence, the identities from the list must grant
148
- Service Account Token Creator IAM role to the directly preceding identity, with first
149
- account from the list granting this role to the originating account (templated).
150
- """
151
-
152
- template_fields = ("location", "project_id", "impersonation_chain", "prompt")
153
-
154
- def __init__(
155
- self,
156
- *,
157
- project_id: str,
158
- location: str,
159
- prompt: str,
160
- pretrained_model: str = "textembedding-gecko",
161
- gcp_conn_id: str = "google_cloud_default",
162
- impersonation_chain: str | Sequence[str] | None = None,
163
- **kwargs,
164
- ) -> None:
165
- super().__init__(**kwargs)
166
- self.project_id = project_id
167
- self.location = location
168
- self.prompt = prompt
169
- self.pretrained_model = pretrained_model
170
- self.gcp_conn_id = gcp_conn_id
171
- self.impersonation_chain = impersonation_chain
172
-
173
- def execute(self, context: Context):
174
- self.hook = GenerativeModelHook(
175
- gcp_conn_id=self.gcp_conn_id,
176
- impersonation_chain=self.impersonation_chain,
177
- )
178
-
179
- self.log.info("Generating text embeddings")
180
- response = self.hook.generate_text_embeddings(
181
- project_id=self.project_id,
182
- location=self.location,
183
- prompt=self.prompt,
184
- pretrained_model=self.pretrained_model,
185
- )
186
-
187
- self.log.info("Model response: %s", response)
188
- self.xcom_push(context, key="prompt_response", value=response)
189
-
190
- return response
191
-
192
-
193
- @deprecated(
194
- planned_removal_date="January 01, 2025",
195
- use_instead="GenerativeModelGenerateContentOperator",
196
- category=AirflowProviderDeprecationWarning,
197
- )
198
- class PromptMultimodalModelOperator(GoogleCloudBaseOperator):
199
- """
200
- Use the Vertex AI Gemini Pro foundation model to generate natural language text.
201
-
202
- :param project_id: Required. The ID of the Google Cloud project that the
203
- service belongs to (templated).
204
- :param location: Required. The ID of the Google Cloud location that the
205
- service belongs to (templated).
206
- :param prompt: Required. Inputs or queries that a user or a program gives
207
- to the Multi-modal model, in order to elicit a specific response (templated).
208
- :param generation_config: Optional. Generation configuration settings.
209
- :param safety_settings: Optional. Per request settings for blocking unsafe content.
210
- :param pretrained_model: By default uses the pre-trained model `gemini-pro`,
211
- supporting prompts with text-only input, including natural language
212
- tasks, multi-turn text and code chat, and code generation. It can
213
- output text and code.
214
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
215
- :param impersonation_chain: Optional service account to impersonate using short-term
216
- credentials, or chained list of accounts required to get the access_token
217
- of the last account in the list, which will be impersonated in the request.
218
- If set as a string, the account must grant the originating account
219
- the Service Account Token Creator IAM role.
220
- If set as a sequence, the identities from the list must grant
221
- Service Account Token Creator IAM role to the directly preceding identity, with first
222
- account from the list granting this role to the originating account (templated).
223
- """
224
-
225
- template_fields = ("location", "project_id", "impersonation_chain", "prompt")
226
-
227
- def __init__(
228
- self,
229
- *,
230
- project_id: str,
231
- location: str,
232
- prompt: str,
233
- generation_config: dict | None = None,
234
- safety_settings: dict | None = None,
235
- pretrained_model: str = "gemini-pro",
236
- gcp_conn_id: str = "google_cloud_default",
237
- impersonation_chain: str | Sequence[str] | None = None,
238
- **kwargs,
239
- ) -> None:
240
- super().__init__(**kwargs)
241
- self.project_id = project_id
242
- self.location = location
243
- self.prompt = prompt
244
- self.generation_config = generation_config
245
- self.safety_settings = safety_settings
246
- self.pretrained_model = pretrained_model
247
- self.gcp_conn_id = gcp_conn_id
248
- self.impersonation_chain = impersonation_chain
249
-
250
- def execute(self, context: Context):
251
- self.hook = GenerativeModelHook(
252
- gcp_conn_id=self.gcp_conn_id,
253
- impersonation_chain=self.impersonation_chain,
254
- )
255
- response = self.hook.prompt_multimodal_model(
256
- project_id=self.project_id,
257
- location=self.location,
258
- prompt=self.prompt,
259
- generation_config=self.generation_config,
260
- safety_settings=self.safety_settings,
261
- pretrained_model=self.pretrained_model,
262
- )
263
-
264
- self.log.info("Model response: %s", response)
265
- self.xcom_push(context, key="prompt_response", value=response)
266
-
267
- return response
268
-
269
-
270
- @deprecated(
271
- planned_removal_date="January 01, 2025",
272
- use_instead="GenerativeModelGenerateContentOperator",
273
- category=AirflowProviderDeprecationWarning,
274
- )
275
- class PromptMultimodalModelWithMediaOperator(GoogleCloudBaseOperator):
276
- """
277
- Use the Vertex AI Gemini Pro foundation model to generate natural language text.
278
-
279
- :param project_id: Required. The ID of the Google Cloud project that the
280
- service belongs to (templated).
281
- :param location: Required. The ID of the Google Cloud location that the
282
- service belongs to (templated).
283
- :param prompt: Required. Inputs or queries that a user or a program gives
284
- to the Multi-modal model, in order to elicit a specific response (templated).
285
- :param generation_config: Optional. Generation configuration settings.
286
- :param safety_settings: Optional. Per request settings for blocking unsafe content.
287
- :param pretrained_model: By default uses the pre-trained model `gemini-pro-vision`,
288
- supporting prompts with text-only input, including natural language
289
- tasks, multi-turn text and code chat, and code generation. It can
290
- output text and code.
291
- :param media_gcs_path: A GCS path to a media file such as an image or a video.
292
- Can be passed to the multi-modal model as part of the prompt. Used with vision models.
293
- :param mime_type: Validates the media type presented by the file in the media_gcs_path.
294
- :param gcp_conn_id: The connection ID to use connecting to Google Cloud.
295
- :param impersonation_chain: Optional service account to impersonate using short-term
296
- credentials, or chained list of accounts required to get the access_token
297
- of the last account in the list, which will be impersonated in the request.
298
- If set as a string, the account must grant the originating account
299
- the Service Account Token Creator IAM role.
300
- If set as a sequence, the identities from the list must grant
301
- Service Account Token Creator IAM role to the directly preceding identity, with first
302
- account from the list granting this role to the originating account (templated).
303
- """
304
-
305
- template_fields = ("location", "project_id", "impersonation_chain", "prompt")
306
-
307
- def __init__(
308
- self,
309
- *,
310
- project_id: str,
311
- location: str,
312
- prompt: str,
313
- media_gcs_path: str,
314
- mime_type: str,
315
- generation_config: dict | None = None,
316
- safety_settings: dict | None = None,
317
- pretrained_model: str = "gemini-pro-vision",
318
- gcp_conn_id: str = "google_cloud_default",
319
- impersonation_chain: str | Sequence[str] | None = None,
320
- **kwargs,
321
- ) -> None:
322
- super().__init__(**kwargs)
323
- self.project_id = project_id
324
- self.location = location
325
- self.prompt = prompt
326
- self.generation_config = generation_config
327
- self.safety_settings = safety_settings
328
- self.pretrained_model = pretrained_model
329
- self.media_gcs_path = media_gcs_path
330
- self.mime_type = mime_type
331
- self.gcp_conn_id = gcp_conn_id
332
- self.impersonation_chain = impersonation_chain
333
-
334
- def execute(self, context: Context):
335
- self.hook = GenerativeModelHook(
336
- gcp_conn_id=self.gcp_conn_id,
337
- impersonation_chain=self.impersonation_chain,
338
- )
339
- response = self.hook.prompt_multimodal_model_with_media(
340
- project_id=self.project_id,
341
- location=self.location,
342
- prompt=self.prompt,
343
- generation_config=self.generation_config,
344
- safety_settings=self.safety_settings,
345
- pretrained_model=self.pretrained_model,
346
- media_gcs_path=self.media_gcs_path,
347
- mime_type=self.mime_type,
348
- )
349
-
350
- self.log.info("Model response: %s", response)
351
- self.xcom_push(context, key="prompt_response", value=response)
352
-
353
- return response
354
-
355
-
356
34
  @deprecated(
357
35
  planned_removal_date="April 09, 2025",
358
36
  use_instead="GenerativeModelGenerateContentOperator",
@@ -23,10 +23,6 @@ from __future__ import annotations
23
23
  from collections.abc import Sequence
24
24
  from typing import TYPE_CHECKING, Any
25
25
 
26
- from google.api_core.exceptions import NotFound
27
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
28
- from google.cloud.aiplatform_v1 import types
29
-
30
26
  from airflow.configuration import conf
31
27
  from airflow.exceptions import AirflowException
32
28
  from airflow.providers.google.cloud.hooks.vertex_ai.hyperparameter_tuning_job import (
@@ -38,13 +34,15 @@ from airflow.providers.google.cloud.links.vertex_ai import (
38
34
  )
39
35
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
40
36
  from airflow.providers.google.cloud.triggers.vertex_ai import CreateHyperparameterTuningJobTrigger
37
+ from google.api_core.exceptions import NotFound
38
+ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
39
+ from google.cloud.aiplatform_v1 import types
41
40
 
42
41
  if TYPE_CHECKING:
42
+ from airflow.utils.context import Context
43
43
  from google.api_core.retry import Retry
44
44
  from google.cloud.aiplatform import HyperparameterTuningJob, gapic, hyperparameter_tuning
45
45
 
46
- from airflow.utils.context import Context
47
-
48
46
 
49
47
  class CreateHyperparameterTuningJobOperator(GoogleCloudBaseOperator):
50
48
  """
@@ -22,10 +22,6 @@ from __future__ import annotations
22
22
  from collections.abc import Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
- from google.api_core.exceptions import NotFound
26
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
27
- from google.cloud.aiplatform_v1.types import Model, model_service
28
-
29
25
  from airflow.providers.google.cloud.hooks.vertex_ai.model_service import ModelServiceHook
30
26
  from airflow.providers.google.cloud.links.vertex_ai import (
31
27
  VertexAIModelExportLink,
@@ -33,11 +29,13 @@ from airflow.providers.google.cloud.links.vertex_ai import (
33
29
  VertexAIModelListLink,
34
30
  )
35
31
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
32
+ from google.api_core.exceptions import NotFound
33
+ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
34
+ from google.cloud.aiplatform_v1.types import Model, model_service
36
35
 
37
36
  if TYPE_CHECKING:
38
- from google.api_core.retry import Retry
39
-
40
37
  from airflow.utils.context import Context
38
+ from google.api_core.retry import Retry
41
39
 
42
40
 
43
41
  class DeleteModelOperator(GoogleCloudBaseOperator):
@@ -23,10 +23,6 @@ from collections.abc import Sequence
23
23
  from functools import cached_property
24
24
  from typing import TYPE_CHECKING, Any
25
25
 
26
- from google.api_core.exceptions import NotFound
27
- from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
28
- from google.cloud.aiplatform_v1 import types
29
-
30
26
  from airflow.configuration import conf
31
27
  from airflow.exceptions import AirflowException
32
28
  from airflow.providers.google.cloud.hooks.vertex_ai.pipeline_job import PipelineJobHook
@@ -36,14 +32,16 @@ from airflow.providers.google.cloud.links.vertex_ai import (
36
32
  )
37
33
  from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
38
34
  from airflow.providers.google.cloud.triggers.vertex_ai import RunPipelineJobTrigger
35
+ from google.api_core.exceptions import NotFound
36
+ from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
37
+ from google.cloud.aiplatform_v1 import types
39
38
 
40
39
  if TYPE_CHECKING:
40
+ from airflow.utils.context import Context
41
41
  from google.api_core.retry import Retry
42
42
  from google.cloud.aiplatform import PipelineJob
43
43
  from google.cloud.aiplatform.metadata import experiment_resources
44
44
 
45
- from airflow.utils.context import Context
46
-
47
45
 
48
46
  class RunPipelineJobOperator(GoogleCloudBaseOperator):
49
47
  """
@@ -22,17 +22,15 @@ from __future__ import annotations
22
22
  from collections.abc import Sequence
23
23
  from typing import TYPE_CHECKING
24
24
 
25
+ from airflow.providers.google.cloud.hooks.video_intelligence import CloudVideoIntelligenceHook
26
+ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
25
27
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
26
28
  from google.cloud.videointelligence_v1 import Feature, VideoContext
27
29
  from google.protobuf.json_format import MessageToDict
28
30
 
29
- from airflow.providers.google.cloud.hooks.video_intelligence import CloudVideoIntelligenceHook
30
- from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
31
-
32
31
  if TYPE_CHECKING:
33
- from google.api_core.retry import Retry
34
-
35
32
  from airflow.utils.context import Context
33
+ from google.api_core.retry import Retry
36
34
 
37
35
 
38
36
  class CloudVideoIntelligenceDetectVideoLabelsOperator(GoogleCloudBaseOperator):
@@ -23,6 +23,9 @@ from collections.abc import Sequence
23
23
  from copy import deepcopy
24
24
  from typing import TYPE_CHECKING, Any
25
25
 
26
+ from airflow.providers.google.cloud.hooks.vision import CloudVisionHook
27
+ from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
28
+ from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
26
29
  from google.api_core.exceptions import AlreadyExists
27
30
  from google.api_core.gapic_v1.method import DEFAULT, _MethodDefault
28
31
  from google.cloud.vision_v1 import (
@@ -33,16 +36,11 @@ from google.cloud.vision_v1 import (
33
36
  ReferenceImage,
34
37
  )
35
38
 
36
- from airflow.providers.google.cloud.hooks.vision import CloudVisionHook
37
- from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator
38
- from airflow.providers.google.common.hooks.base_google import PROVIDE_PROJECT_ID
39
-
40
39
  if TYPE_CHECKING:
40
+ from airflow.utils.context import Context
41
41
  from google.api_core.retry import Retry
42
42
  from google.protobuf.field_mask_pb2 import FieldMask
43
43
 
44
- from airflow.utils.context import Context
45
-
46
44
 
47
45
  MetaData = Sequence[tuple[str, str]]
48
46