SameerMesiah97 commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3700285584
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
Review Comment:
I would rename this to `RunAgentQueryOperator` to make the pairing with
`RunQueryJobOperator` more explicit.
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :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 input_data: Optional. Input for the Agent Engine class method in
JSON object format.
+ :param class_method: Optional. The Agent Engine class method to invoke.
Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request as
metadata.
+ :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",
+ "input_data",
+ "class_method",
+ "gcp_conn_id",
+ "impersonation_chain",
+ )
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ location: str,
+ agent_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | _MethodDefault = DEFAULT,
+ timeout: float | None = None,
+ metadata: Sequence[tuple[str, 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.input_data = input_data
+ self.class_method = class_method
+ self.retry = retry
+ self.timeout = timeout
+ self.metadata = metadata
+ 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) -> Any:
Review Comment:
Could you be more specific with the return type? `Any` is not ideal.
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :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 input_data: Optional. Input for the Agent Engine class method in
JSON object format.
+ :param class_method: Optional. The Agent Engine class method to invoke.
Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request as
metadata.
+ :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.
Review Comment:
1) Add the default values for each of the parameters to the docstring.
2) Indicate whether each parameter is templated by adding `(templated)` if
it is.
##########
providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py:
##########
@@ -141,6 +164,43 @@ def get_agent_engine(
name = self.build_agent_engine_name(project_id, location,
agent_engine_id)
return client.get(name=name, config=config)
+ @GoogleBaseHook.fallback_to_default_project_id
+ def query_agent_engine(
+ self,
+ location: str,
+ agent_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | _MethodDefault = DEFAULT,
+ timeout: float | None = None,
+ metadata: Sequence[tuple[str, str]] = (),
+ project_id: str = PROVIDE_PROJECT_ID,
+ ) -> QueryReasoningEngineResponse:
+ """
+ Query an Agent Engine synchronously.
+
+ :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 input_data: Optional. Input for the Agent Engine class method
in JSON object format.
+ :param class_method: Optional. The Agent Engine class method to
invoke. Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request
as metadata.
+ :param project_id: Optional. The ID of the Google Cloud project.
Defaults to the project
+ configured in the connection.
Review Comment:
I think some of the docstring entries do not have the default values
mentioned. Please include them.
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +169,82 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class QueryAgentEngineOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :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 input_data: Optional. Input for the Agent Engine class method in
JSON object format.
+ :param class_method: Optional. The Agent Engine class method to invoke.
Defaults to ``query``.
+ :param retry: Designation of what errors, if any, should be retried.
+ :param timeout: The timeout for this request.
+ :param metadata: Strings which should be sent along with the request as
metadata.
+ :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",
+ "input_data",
+ "class_method",
+ "gcp_conn_id",
+ "impersonation_chain",
+ )
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ location: str,
+ agent_engine_id: str,
+ input_data: dict[str, Any] | None = None,
+ class_method: str = "query",
+ retry: Retry | _MethodDefault = DEFAULT,
+ timeout: float | None = None,
+ metadata: Sequence[tuple[str, 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.input_data = input_data
+ self.class_method = class_method
+ self.retry = retry
+ self.timeout = timeout
+ self.metadata = metadata
+ 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) -> Any:
+ self.log.info("Querying Agent Engine %s.", self.agent_engine_id)
+ response = self.hook.query_agent_engine(
+ project_id=self.project_id,
+ location=self.location,
+ agent_engine_id=self.agent_engine_id,
+ input_data=self.input_data,
+ class_method=self.class_method,
+ retry=self.retry,
+ timeout=self.timeout,
+ metadata=self.metadata,
+ )
+ self.log.info("Agent Engine %s returned a response.",
self.agent_engine_id)
+ return QueryReasoningEngineResponse.to_dict(response).get("output")
Review Comment:
Is there a reason we're returning only the output field rather than the full
serialized response?
--
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]
