MaksYermak commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3702978401
##########
providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py:
##########
@@ -141,6 +164,45 @@ 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.
+ Defaults to ``None``.
+ :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.
Defaults to ``DEFAULT``.
+ :param timeout: The timeout for this request. Defaults to ``None``.
+ :param metadata: Strings which should be sent along with the request
as metadata. Defaults
+ to an empty tuple.
+ :param project_id: Optional. The ID of the Google Cloud project.
Defaults to the project
+ configured in the connection.
+ """
+ client =
self.get_reasoning_engine_execution_service_client(location=location)
+ name = client.reasoning_engine_path(project_id, location,
agent_engine_id)
Review Comment:
@AlejandroMorgante as I see in the
[code](https://github.com/googleapis/python-aiplatform/blob/main/google/cloud/aiplatform_v1/services/reasoning_engine_execution_service/client.py#L269)
it is `reasoning_engine` parameter not the `agent_engine`. Could you please
change `agent_engine_id` to `reasoning_engine_id` to be consist with the Client
code?
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +170,89 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class RunAgentQueryOperator(GoogleCloudBaseOperator):
Review Comment:
@AlejandroMorgante the same question as for Hook method, why do you use the
`RunAgentQueryOperator` name when under the hood the Client will call
`query_reasoning_engine` method? I think it is better to use
`RunReasoningEngineQueryOperator` name
##########
providers/google/src/airflow/providers/google/cloud/hooks/vertex_ai/agent_engine.py:
##########
@@ -141,6 +164,45 @@ 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(
Review Comment:
@AlejandroMorgante could you please clarify for me why do you use
`query_agent_engine` name here but later in the code the `Client` call
`query_reasoning_engine` method? In my opinion it is make sense to use the
`query_reasoning_engine` name for the Hook method
--
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]
