MaksYermak commented on code in PR #68479: URL: https://github.com/apache/airflow/pull/68479#discussion_r3481809794
########## providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py: ########## @@ -0,0 +1,332 @@ +# +# 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. +"""This module contains a Google Cloud Vertex AI Agent Engine hook.""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import google.auth.transport.requests +from asgiref.sync import sync_to_async +from vertexai import Client + +from airflow.providers.google.common.hooks.base_google import ( + PROVIDE_PROJECT_ID, + GoogleBaseAsyncHook, + GoogleBaseHook, +) + +if TYPE_CHECKING: + from vertexai._genai import types + + +def _serialize_value(value: Any) -> Any: + """Recursively convert SDK model objects to JSON-serializable types.""" + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if isinstance(value, dict): + return {key: _serialize_value(item) for key, item in value.items()} + if isinstance(value, list): + return [_serialize_value(item) for item in value] + if isinstance(value, tuple): + return tuple(_serialize_value(item) for item in value) + return value + + +class AgentEngineHook(GoogleBaseHook): + """Hook for Google Cloud Vertex AI Agent Engine APIs.""" + + def __init__( + self, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__( + gcp_conn_id=gcp_conn_id, + impersonation_chain=impersonation_chain, + **kwargs, + ) + + def get_agent_engine_client(self, project_id: str, location: str): + """Return the Vertex AI Agent Engine client.""" + return Client( + project=project_id, + location=location, + credentials=self.get_credentials(), + ).agent_engines + + @staticmethod + def build_agent_engine_name(project_id: str, location: str, agent_engine_id: str) -> str: + """Build a fully qualified Agent Engine resource name.""" + return f"projects/{project_id}/locations/{location}/reasoningEngines/{agent_engine_id}" + + @GoogleBaseHook.fallback_to_default_project_id + def create_agent_engine( + self, + location: str, + agent: Any | None = None, + config: types.AgentEngineConfigOrDict | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Create an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent: Optional. The agent object to deploy. + :param config: Optional. Configuration for the Agent Engine. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + return client.create(agent=agent, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def get_agent_engine( + self, + location: str, + agent_engine_id: str, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Get an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + name = self.build_agent_engine_name(project_id, location, agent_engine_id) + return client.get(name=name) + + @GoogleBaseHook.fallback_to_default_project_id + def query_agent_engine( Review Comment: @AlejandroMorgante could we rename this function from `query_agent_engine` to `run_query_job` for being consistent with what this function does under the hood? ########## providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py: ########## @@ -0,0 +1,332 @@ +# +# 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. +"""This module contains a Google Cloud Vertex AI Agent Engine hook.""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import google.auth.transport.requests +from asgiref.sync import sync_to_async +from vertexai import Client + +from airflow.providers.google.common.hooks.base_google import ( + PROVIDE_PROJECT_ID, + GoogleBaseAsyncHook, + GoogleBaseHook, +) + +if TYPE_CHECKING: + from vertexai._genai import types + + +def _serialize_value(value: Any) -> Any: + """Recursively convert SDK model objects to JSON-serializable types.""" + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if isinstance(value, dict): + return {key: _serialize_value(item) for key, item in value.items()} + if isinstance(value, list): + return [_serialize_value(item) for item in value] + if isinstance(value, tuple): + return tuple(_serialize_value(item) for item in value) + return value + + +class AgentEngineHook(GoogleBaseHook): + """Hook for Google Cloud Vertex AI Agent Engine APIs.""" + + def __init__( + self, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__( + gcp_conn_id=gcp_conn_id, + impersonation_chain=impersonation_chain, + **kwargs, + ) + + def get_agent_engine_client(self, project_id: str, location: str): + """Return the Vertex AI Agent Engine client.""" + return Client( + project=project_id, + location=location, + credentials=self.get_credentials(), + ).agent_engines + + @staticmethod + def build_agent_engine_name(project_id: str, location: str, agent_engine_id: str) -> str: + """Build a fully qualified Agent Engine resource name.""" + return f"projects/{project_id}/locations/{location}/reasoningEngines/{agent_engine_id}" + + @GoogleBaseHook.fallback_to_default_project_id + def create_agent_engine( + self, + location: str, + agent: Any | None = None, + config: types.AgentEngineConfigOrDict | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Create an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent: Optional. The agent object to deploy. + :param config: Optional. Configuration for the Agent Engine. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + return client.create(agent=agent, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def get_agent_engine( + self, + location: str, + agent_engine_id: str, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Get an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + name = self.build_agent_engine_name(project_id, location, agent_engine_id) + return client.get(name=name) + + @GoogleBaseHook.fallback_to_default_project_id + def query_agent_engine( + self, + location: str, + agent_engine_id: str, + config: types.RunQueryJobAgentEngineConfigOrDict, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.RunQueryJobResult: + """ + Run a query job on an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param config: Required. Configuration for the query job (``query``, ``output_gcs_uri``). + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + name = self.build_agent_engine_name(project_id, location, agent_engine_id) + return client.run_query_job(name=name, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def check_query_agent_engine_job( + self, + location: str, + operation_name: str, Review Comment: @AlejandroMorgante could you please divide the `operation_name` by parts as you did for other methods? ########## providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py: ########## @@ -0,0 +1,332 @@ +# +# 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. +"""This module contains a Google Cloud Vertex AI Agent Engine hook.""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import google.auth.transport.requests +from asgiref.sync import sync_to_async +from vertexai import Client + +from airflow.providers.google.common.hooks.base_google import ( + PROVIDE_PROJECT_ID, + GoogleBaseAsyncHook, + GoogleBaseHook, +) + +if TYPE_CHECKING: + from vertexai._genai import types + + +def _serialize_value(value: Any) -> Any: + """Recursively convert SDK model objects to JSON-serializable types.""" + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if isinstance(value, dict): + return {key: _serialize_value(item) for key, item in value.items()} + if isinstance(value, list): + return [_serialize_value(item) for item in value] + if isinstance(value, tuple): + return tuple(_serialize_value(item) for item in value) + return value + + +class AgentEngineHook(GoogleBaseHook): + """Hook for Google Cloud Vertex AI Agent Engine APIs.""" + + def __init__( + self, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__( + gcp_conn_id=gcp_conn_id, + impersonation_chain=impersonation_chain, + **kwargs, + ) + + def get_agent_engine_client(self, project_id: str, location: str): + """Return the Vertex AI Agent Engine client.""" + return Client( + project=project_id, + location=location, + credentials=self.get_credentials(), + ).agent_engines + + @staticmethod + def build_agent_engine_name(project_id: str, location: str, agent_engine_id: str) -> str: + """Build a fully qualified Agent Engine resource name.""" + return f"projects/{project_id}/locations/{location}/reasoningEngines/{agent_engine_id}" + + @GoogleBaseHook.fallback_to_default_project_id + def create_agent_engine( + self, + location: str, + agent: Any | None = None, + config: types.AgentEngineConfigOrDict | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Create an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent: Optional. The agent object to deploy. + :param config: Optional. Configuration for the Agent Engine. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + return client.create(agent=agent, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def get_agent_engine( + self, + location: str, + agent_engine_id: str, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Get an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + name = self.build_agent_engine_name(project_id, location, agent_engine_id) + return client.get(name=name) + + @GoogleBaseHook.fallback_to_default_project_id + def query_agent_engine( + self, + location: str, + agent_engine_id: str, + config: types.RunQueryJobAgentEngineConfigOrDict, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.RunQueryJobResult: + """ + Run a query job on an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param config: Required. Configuration for the query job (``query``, ``output_gcs_uri``). + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + name = self.build_agent_engine_name(project_id, location, agent_engine_id) + return client.run_query_job(name=name, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def check_query_agent_engine_job( + self, + location: str, + operation_name: str, + config: types.CheckQueryJobAgentEngineConfigOrDict | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.CheckQueryJobResult: + """ + Check a query job on an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param operation_name: Required. The query job operation name. + :param config: Optional. Configuration for checking the query job. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + return client.check_query_job(name=operation_name, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def wait_for_query_agent_engine_job( + self, + location: str, + operation_name: str, + config: types.CheckQueryJobAgentEngineConfigOrDict | None = None, + poll_interval: float = 30, + timeout: float | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.CheckQueryJobResult: + """ + Wait until an Agent Engine query job completes. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param operation_name: Required. The query job operation name. + :param config: Optional. Configuration for checking the query job. + :param poll_interval: Time, in seconds, to wait between checks. + :param timeout: Optional timeout, in seconds. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + start_time = time.monotonic() + while True: + query_job = self.check_query_agent_engine_job( + project_id=project_id, + location=location, + operation_name=operation_name, + config=config, + ) + status = getattr(query_job, "status", None) + if status == "SUCCESS": + return query_job + if status == "FAILED": + raise RuntimeError(f"Agent Engine query job {operation_name} failed.") + if status not in (None, "RUNNING"): + raise RuntimeError( + f"Agent Engine query job {operation_name} completed with unexpected status {status}." + ) + if timeout is not None and time.monotonic() - start_time >= timeout: + raise TimeoutError(f"Timed out waiting for Agent Engine query job {operation_name}") + self.log.info("Waiting for Agent Engine query job %s to complete.", operation_name) + time.sleep(poll_interval) + + @GoogleBaseHook.fallback_to_default_project_id + def update_agent_engine( + self, + location: str, + agent_engine_id: str, + config: types.AgentEngineConfigOrDict, + agent: Any | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Update an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param config: Required. Configuration for the Agent Engine update. + :param agent: Optional. The updated agent object to deploy. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + name = self.build_agent_engine_name(project_id, location, agent_engine_id) + return client.update(name=name, agent=agent, config=config) Review Comment: I see that Client has `agent_engine` [parameter](https://github.com/googleapis/python-aiplatform/blob/main/vertexai/_genai/agent_engines.py#L2682) could you please add this parameter to the Hook method? Because hook's method and client's method should be consistent in parameters ########## providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py: ########## @@ -0,0 +1,494 @@ +# +# 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. +"""This module contains Google Vertex AI Agent Engine operators.""" + +from __future__ import annotations + +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import conf +from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import AgentEngineHook, _serialize_value +from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator +from airflow.providers.google.cloud.triggers.vertex_ai import ( + AgentEngineDeleteTrigger, + AgentEngineQueryJobTrigger, +) + +if TYPE_CHECKING: + from vertexai._genai import types + + from airflow.providers.common.compat.sdk import Context + + +def _serialize_agent_engine(agent_engine: types.AgentEngine) -> dict[str, Any]: + api_resource = getattr(agent_engine, "api_resource", None) + if api_resource is not None: + return _serialize_value(api_resource) + return _serialize_value(agent_engine) + + +class CreateAgentEngineOperator(GoogleCloudBaseOperator): + """ + Create a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent: Optional. The agent object to deploy. + :param config: Optional. Configuration for the Agent Engine. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ( + "project_id", + "location", + "agent", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent: Any | None = None, + config: types.AgentEngineConfigOrDict | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent = agent + self.config = config + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Creating Agent Engine.") + agent_engine = self.hook.create_agent_engine( + project_id=self.project_id, + location=self.location, + agent=self.agent, + config=self.config, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine was created.") + return result + + +class GetAgentEngineOperator(GoogleCloudBaseOperator): + """ + Get a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ("project_id", "location", "agent_engine_id", "gcp_conn_id", "impersonation_chain") + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Getting Agent Engine %s.", self.agent_engine_id) + agent_engine = self.hook.get_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine %s was retrieved.", self.agent_engine_id) + return result + + +class QueryAgentEngineOperator(GoogleCloudBaseOperator): + """ + Run a query job on a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param config: Required. Configuration for the query job (``query``, ``output_gcs_uri``). + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ( + "project_id", + "location", + "agent_engine_id", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + config: types.RunQueryJobAgentEngineConfigOrDict, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.config = config + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Running query job on Agent Engine %s.", self.agent_engine_id) + query_job = self.hook.query_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + config=self.config, + ) + self.log.info("Query job was started on Agent Engine %s.", self.agent_engine_id) + return _serialize_value(query_job) Review Comment: @AlejandroMorgante could we discuss it from airflow users perspective. Because in my opinion right now this Operator looks useless, because user can use Hook for run query job. And this operator will not work well via XCom with `CheckQueryAgentEngineOperator` because you serialize `RunQueryJobResult` object. And also `CheckQueryAgentEngineOperator` doesn't have option to take object. In my opinion we should have Operator which runs `QueryJob` on the agent then waiting for `QueryJob` results and then return the result to XCom. For doing it you have `query_agent_engine` and `wait_for_query_agent_engine_job` ########## providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py: ########## @@ -0,0 +1,494 @@ +# +# 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. +"""This module contains Google Vertex AI Agent Engine operators.""" + +from __future__ import annotations + +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import conf +from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import AgentEngineHook, _serialize_value +from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator +from airflow.providers.google.cloud.triggers.vertex_ai import ( + AgentEngineDeleteTrigger, + AgentEngineQueryJobTrigger, +) + +if TYPE_CHECKING: + from vertexai._genai import types + + from airflow.providers.common.compat.sdk import Context + + +def _serialize_agent_engine(agent_engine: types.AgentEngine) -> dict[str, Any]: + api_resource = getattr(agent_engine, "api_resource", None) + if api_resource is not None: + return _serialize_value(api_resource) + return _serialize_value(agent_engine) + + +class CreateAgentEngineOperator(GoogleCloudBaseOperator): + """ + Create a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent: Optional. The agent object to deploy. + :param config: Optional. Configuration for the Agent Engine. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ( + "project_id", + "location", + "agent", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent: Any | None = None, + config: types.AgentEngineConfigOrDict | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent = agent + self.config = config + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Creating Agent Engine.") + agent_engine = self.hook.create_agent_engine( + project_id=self.project_id, + location=self.location, + agent=self.agent, + config=self.config, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine was created.") + return result + + +class GetAgentEngineOperator(GoogleCloudBaseOperator): + """ + Get a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ("project_id", "location", "agent_engine_id", "gcp_conn_id", "impersonation_chain") + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Getting Agent Engine %s.", self.agent_engine_id) + agent_engine = self.hook.get_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine %s was retrieved.", self.agent_engine_id) + return result + + +class QueryAgentEngineOperator(GoogleCloudBaseOperator): + """ + Run a query job on a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param config: Required. Configuration for the query job (``query``, ``output_gcs_uri``). + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ( + "project_id", + "location", + "agent_engine_id", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + config: types.RunQueryJobAgentEngineConfigOrDict, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.config = config + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Running query job on Agent Engine %s.", self.agent_engine_id) + query_job = self.hook.query_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + config=self.config, + ) + self.log.info("Query job was started on Agent Engine %s.", self.agent_engine_id) + return _serialize_value(query_job) + + +class CheckQueryAgentEngineOperator(GoogleCloudBaseOperator): + """ + Check a query job on a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param operation_name: Required. The query job operation name (e.g. from the ``job_name`` field of the result of ``QueryAgentEngineOperator``). + :param config: Optional. Configuration for checking the query job. + :param poll_interval: Time, in seconds, to wait between checks. + :param timeout: Optional timeout, in seconds. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + :param deferrable: Run operator in the deferrable mode. + """ + + template_fields = ( + "project_id", + "location", + "operation_name", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + operation_name: str, + config: types.CheckQueryJobAgentEngineConfigOrDict | None = None, + poll_interval: float = 30, + timeout: float | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.operation_name = operation_name + self.config = config + self.poll_interval = poll_interval + self.timeout = timeout + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + self.deferrable = deferrable + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Checking Agent Engine query job %s.", self.operation_name) + if self.deferrable: + self.defer( + trigger=AgentEngineQueryJobTrigger( + project_id=self.project_id, + location=self.location, + operation_name=self.operation_name, + config=self.config, + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + poll_interval=self.poll_interval, + timeout=self.timeout, + ), + method_name="execute_complete", + ) + + query_job = self.hook.wait_for_query_agent_engine_job( + project_id=self.project_id, + location=self.location, + operation_name=self.operation_name, + config=self.config, + poll_interval=self.poll_interval, + timeout=self.timeout, + ) + result = _serialize_value(query_job) + self.log.info("Agent Engine query job %s completed.", self.operation_name) + return result + + def execute_complete(self, context: Context, event: dict[str, Any] | None = None) -> dict[str, Any]: + if event is None: + raise RuntimeError("No event received in trigger callback") + if event["status"] == "success": + self.log.info("Agent Engine query job completed.") + return event["query_job"] + if event["status"] == "timeout": + raise TimeoutError(event["message"]) + raise RuntimeError(event["message"]) + + +class UpdateAgentEngineOperator(GoogleCloudBaseOperator): + """ + Update a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param agent: Optional. The updated agent object to deploy. + :param config: Required. Configuration for the Agent Engine update. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ( + "project_id", + "location", + "agent_engine_id", + "agent", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + config: types.AgentEngineConfigOrDict, + agent: Any | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.agent = agent + self.config = config + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Updating Agent Engine %s.", self.agent_engine_id) + agent_engine = self.hook.update_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + agent=self.agent, + config=self.config, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine %s was updated.", self.agent_engine_id) + return result + + +class DeleteAgentEngineOperator(GoogleCloudBaseOperator): + """ + Delete a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param force: Optional. Whether to delete child resources. + :param config: Optional. Additional deletion configuration. + :param wait_for_completion: Whether to wait until the delete operation completes. + :param poll_interval: Time, in seconds, to wait between checks. + :param timeout: Optional timeout, in seconds. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + :param deferrable: Run operator in the deferrable mode. + """ + + template_fields = ( + "project_id", + "location", + "agent_engine_id", + "force", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + force: bool | None = None, + config: types.DeleteAgentEngineConfigOrDict | None = None, + wait_for_completion: bool = True, + poll_interval: float = 30, + timeout: float | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.force = force + self.config = config + self.wait_for_completion = wait_for_completion + self.poll_interval = poll_interval + self.timeout = timeout + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + self.deferrable = deferrable + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Deleting Agent Engine %s.", self.agent_engine_id) + operation = self.hook.delete_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + force=self.force, + config=self.config, + ) + result = _serialize_value(operation) + if not self.wait_for_completion: + return result + + operation_name = getattr(operation, "name", None) + if not operation_name: + raise RuntimeError("Delete Agent Engine operation did not include an operation name.") + + if getattr(operation, "done", False): + self.log.info("Agent Engine %s was deleted.", self.agent_engine_id) + return result + + if self.deferrable: + self.defer( + trigger=AgentEngineDeleteTrigger( + location=self.location, + agent_engine_id=self.agent_engine_id, + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + poll_interval=self.poll_interval, + timeout=self.timeout, + operation_name=operation_name, + ), + method_name="execute_complete", + kwargs={"operation": result}, + ) Review Comment: @AlejandroMorgante how long the delete operation run? We do not have any `deferable` mode for delete operation in google provider, because waiting takes second. Mostly we add `deferrable` mode only for long-running operation for free `worker` resources and start waiting on `trigerrer`. @VladaZakharova what do you think about this? ########## providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py: ########## @@ -0,0 +1,494 @@ +# +# 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. +"""This module contains Google Vertex AI Agent Engine operators.""" + +from __future__ import annotations + +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import conf +from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import AgentEngineHook, _serialize_value +from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator +from airflow.providers.google.cloud.triggers.vertex_ai import ( + AgentEngineDeleteTrigger, + AgentEngineQueryJobTrigger, +) + +if TYPE_CHECKING: + from vertexai._genai import types + + from airflow.providers.common.compat.sdk import Context + + +def _serialize_agent_engine(agent_engine: types.AgentEngine) -> dict[str, Any]: + api_resource = getattr(agent_engine, "api_resource", None) + if api_resource is not None: + return _serialize_value(api_resource) + return _serialize_value(agent_engine) + + +class CreateAgentEngineOperator(GoogleCloudBaseOperator): + """ + Create a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent: Optional. The agent object to deploy. + :param config: Optional. Configuration for the Agent Engine. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ( + "project_id", + "location", + "agent", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent: Any | None = None, + config: types.AgentEngineConfigOrDict | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent = agent + self.config = config + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Creating Agent Engine.") + agent_engine = self.hook.create_agent_engine( + project_id=self.project_id, + location=self.location, + agent=self.agent, + config=self.config, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine was created.") + return result + + +class GetAgentEngineOperator(GoogleCloudBaseOperator): + """ + Get a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ("project_id", "location", "agent_engine_id", "gcp_conn_id", "impersonation_chain") + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Getting Agent Engine %s.", self.agent_engine_id) + agent_engine = self.hook.get_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine %s was retrieved.", self.agent_engine_id) + return result + + +class QueryAgentEngineOperator(GoogleCloudBaseOperator): Review Comment: @AlejandroMorgante could we rename it from `QueryAgentEngineOperator` to `RunQueryJobOperator`? ########## providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py: ########## @@ -0,0 +1,332 @@ +# +# 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. +"""This module contains a Google Cloud Vertex AI Agent Engine hook.""" + +from __future__ import annotations + +import time +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import google.auth.transport.requests +from asgiref.sync import sync_to_async +from vertexai import Client + +from airflow.providers.google.common.hooks.base_google import ( + PROVIDE_PROJECT_ID, + GoogleBaseAsyncHook, + GoogleBaseHook, +) + +if TYPE_CHECKING: + from vertexai._genai import types + + +def _serialize_value(value: Any) -> Any: + """Recursively convert SDK model objects to JSON-serializable types.""" + if hasattr(value, "model_dump"): + return value.model_dump(mode="json") + if isinstance(value, dict): + return {key: _serialize_value(item) for key, item in value.items()} + if isinstance(value, list): + return [_serialize_value(item) for item in value] + if isinstance(value, tuple): + return tuple(_serialize_value(item) for item in value) + return value + + +class AgentEngineHook(GoogleBaseHook): + """Hook for Google Cloud Vertex AI Agent Engine APIs.""" + + def __init__( + self, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__( + gcp_conn_id=gcp_conn_id, + impersonation_chain=impersonation_chain, + **kwargs, + ) + + def get_agent_engine_client(self, project_id: str, location: str): + """Return the Vertex AI Agent Engine client.""" + return Client( + project=project_id, + location=location, + credentials=self.get_credentials(), + ).agent_engines + + @staticmethod + def build_agent_engine_name(project_id: str, location: str, agent_engine_id: str) -> str: + """Build a fully qualified Agent Engine resource name.""" + return f"projects/{project_id}/locations/{location}/reasoningEngines/{agent_engine_id}" + + @GoogleBaseHook.fallback_to_default_project_id + def create_agent_engine( + self, + location: str, + agent: Any | None = None, + config: types.AgentEngineConfigOrDict | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Create an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent: Optional. The agent object to deploy. + :param config: Optional. Configuration for the Agent Engine. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + return client.create(agent=agent, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def get_agent_engine( + self, + location: str, + agent_engine_id: str, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.AgentEngine: + """ + Get an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + name = self.build_agent_engine_name(project_id, location, agent_engine_id) + return client.get(name=name) + + @GoogleBaseHook.fallback_to_default_project_id + def query_agent_engine( + self, + location: str, + agent_engine_id: str, + config: types.RunQueryJobAgentEngineConfigOrDict, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.RunQueryJobResult: + """ + Run a query job on an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param config: Required. Configuration for the query job (``query``, ``output_gcs_uri``). + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + name = self.build_agent_engine_name(project_id, location, agent_engine_id) + return client.run_query_job(name=name, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def check_query_agent_engine_job( + self, + location: str, + operation_name: str, + config: types.CheckQueryJobAgentEngineConfigOrDict | None = None, + project_id: str = PROVIDE_PROJECT_ID, + ) -> types.CheckQueryJobResult: + """ + Check a query job on an Agent Engine. + + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param operation_name: Required. The query job operation name. + :param config: Optional. Configuration for checking the query job. + :param project_id: Optional. The ID of the Google Cloud project. Defaults to the project + configured in the connection. + """ + client = self.get_agent_engine_client(project_id=project_id, location=location) + return client.check_query_job(name=operation_name, config=config) + + @GoogleBaseHook.fallback_to_default_project_id + def wait_for_query_agent_engine_job( + self, + location: str, + operation_name: str, Review Comment: @AlejandroMorgante could you please divide the `operation_name` by parts as you did for other methods? ########## providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py: ########## @@ -0,0 +1,494 @@ +# +# 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. +"""This module contains Google Vertex AI Agent Engine operators.""" + +from __future__ import annotations + +from collections.abc import Sequence +from functools import cached_property +from typing import TYPE_CHECKING, Any + +from airflow.providers.common.compat.sdk import conf +from airflow.providers.google.cloud.hooks.vertex_ai.agent_engine import AgentEngineHook, _serialize_value +from airflow.providers.google.cloud.operators.cloud_base import GoogleCloudBaseOperator +from airflow.providers.google.cloud.triggers.vertex_ai import ( + AgentEngineDeleteTrigger, + AgentEngineQueryJobTrigger, +) + +if TYPE_CHECKING: + from vertexai._genai import types + + from airflow.providers.common.compat.sdk import Context + + +def _serialize_agent_engine(agent_engine: types.AgentEngine) -> dict[str, Any]: + api_resource = getattr(agent_engine, "api_resource", None) + if api_resource is not None: + return _serialize_value(api_resource) + return _serialize_value(agent_engine) + + +class CreateAgentEngineOperator(GoogleCloudBaseOperator): + """ + Create a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent: Optional. The agent object to deploy. + :param config: Optional. Configuration for the Agent Engine. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ( + "project_id", + "location", + "agent", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent: Any | None = None, + config: types.AgentEngineConfigOrDict | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent = agent + self.config = config + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Creating Agent Engine.") + agent_engine = self.hook.create_agent_engine( + project_id=self.project_id, + location=self.location, + agent=self.agent, + config=self.config, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine was created.") + return result + + +class GetAgentEngineOperator(GoogleCloudBaseOperator): + """ + Get a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ("project_id", "location", "agent_engine_id", "gcp_conn_id", "impersonation_chain") + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Getting Agent Engine %s.", self.agent_engine_id) + agent_engine = self.hook.get_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + ) + result = _serialize_agent_engine(agent_engine) + self.log.info("Agent Engine %s was retrieved.", self.agent_engine_id) + return result + + +class QueryAgentEngineOperator(GoogleCloudBaseOperator): + """ + Run a query job on a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param agent_engine_id: Required. The Agent Engine ID. + :param config: Required. Configuration for the query job (``query``, ``output_gcs_uri``). + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + """ + + template_fields = ( + "project_id", + "location", + "agent_engine_id", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + agent_engine_id: str, + config: types.RunQueryJobAgentEngineConfigOrDict, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.agent_engine_id = agent_engine_id + self.config = config + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Running query job on Agent Engine %s.", self.agent_engine_id) + query_job = self.hook.query_agent_engine( + project_id=self.project_id, + location=self.location, + agent_engine_id=self.agent_engine_id, + config=self.config, + ) + self.log.info("Query job was started on Agent Engine %s.", self.agent_engine_id) + return _serialize_value(query_job) + + +class CheckQueryAgentEngineOperator(GoogleCloudBaseOperator): + """ + Check a query job on a Vertex AI Agent Engine. + + :param project_id: Required. The ID of the Google Cloud project that the service belongs to. + :param location: Required. The ID of the Google Cloud location that the service belongs to. + :param operation_name: Required. The query job operation name (e.g. from the ``job_name`` field of the result of ``QueryAgentEngineOperator``). + :param config: Optional. Configuration for checking the query job. + :param poll_interval: Time, in seconds, to wait between checks. + :param timeout: Optional timeout, in seconds. + :param gcp_conn_id: The connection ID to use connecting to Google Cloud. + :param impersonation_chain: Optional service account to impersonate using short-term credentials. + :param deferrable: Run operator in the deferrable mode. + """ + + template_fields = ( + "project_id", + "location", + "operation_name", + "config", + "gcp_conn_id", + "impersonation_chain", + ) + + def __init__( + self, + *, + project_id: str, + location: str, + operation_name: str, + config: types.CheckQueryJobAgentEngineConfigOrDict | None = None, + poll_interval: float = 30, + timeout: float | None = None, + gcp_conn_id: str = "google_cloud_default", + impersonation_chain: str | Sequence[str] | None = None, + deferrable: bool = conf.getboolean("operators", "default_deferrable", fallback=False), + **kwargs, + ) -> None: + super().__init__(**kwargs) + self.project_id = project_id + self.location = location + self.operation_name = operation_name + self.config = config + self.poll_interval = poll_interval + self.timeout = timeout + self.gcp_conn_id = gcp_conn_id + self.impersonation_chain = impersonation_chain + self.deferrable = deferrable + + @cached_property + def hook(self) -> AgentEngineHook: + return AgentEngineHook( + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + ) + + def execute(self, context: Context) -> dict[str, Any]: + self.log.info("Checking Agent Engine query job %s.", self.operation_name) + if self.deferrable: + self.defer( + trigger=AgentEngineQueryJobTrigger( + project_id=self.project_id, + location=self.location, + operation_name=self.operation_name, + config=self.config, + gcp_conn_id=self.gcp_conn_id, + impersonation_chain=self.impersonation_chain, + poll_interval=self.poll_interval, + timeout=self.timeout, + ), + method_name="execute_complete", + ) + + query_job = self.hook.wait_for_query_agent_engine_job( + project_id=self.project_id, + location=self.location, + operation_name=self.operation_name, + config=self.config, + poll_interval=self.poll_interval, + timeout=self.timeout, + ) + result = _serialize_value(query_job) + self.log.info("Agent Engine query job %s completed.", self.operation_name) + return result Review Comment: @AlejandroMorgante I do not see any reason for having this Operator. We can wait for `QueryJob` results inside `RunQueryJobOperator` in `deferrable` on `non-deferable` mode. It is how we do for the most of google operators which works with Jobs. -- 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]
