MaksYermak commented on code in PR #70933:
URL: https://github.com/apache/airflow/pull/70933#discussion_r3734465685
##########
providers/google/docs/operators/cloud/vertex_ai.rst:
##########
@@ -50,6 +50,17 @@ To get an Agent Engine you can use
:start-after: [START how_to_cloud_vertex_ai_get_agent_engine_operator]
:end-before: [END how_to_cloud_vertex_ai_get_agent_engine_operator]
+To query an Agent Engine synchronously you can use
Review Comment:
@AlejandroMorgante could you please update documentation and mentioned that
is ReasoningEngine operator?
##########
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_reasoning_engine(
+ self,
+ location: str,
+ reasoning_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.
Review Comment:
@AlejandroMorgante could you also change the docstring here to mention that
it is Reasoning Engine?
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +170,90 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class RunReasoningEngineQueryOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :param project_id: Required (templated). The ID of the Google Cloud
project that the service
+ belongs to.
+ :param location: Required (templated). The ID of the Google Cloud location
that the service
+ belongs to.
+ :param reasoning_engine_id: Required (templated). The Reasoning Engine
resource ID for the
+ Agent Engine.
+ :param input_data: Optional (templated). Input for the Agent Engine class
method in JSON object
+ format. Defaults to ``None``.
+ :param class_method: Optional (templated). 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 gcp_conn_id: The connection ID to use connecting to Google Cloud
(templated). Defaults
+ to ``google_cloud_default``.
+ :param impersonation_chain: Optional service account to impersonate using
short-term credentials
+ (templated). Defaults to ``None``.
+ """
+
+ template_fields = (
+ "project_id",
+ "location",
+ "reasoning_engine_id",
+ "input_data",
+ "class_method",
+ "gcp_conn_id",
+ "impersonation_chain",
+ )
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ location: str,
+ reasoning_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.reasoning_engine_id = reasoning_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) -> dict[str, JsonValue]:
+ self.log.info("Querying Agent Engine %s.", self.reasoning_engine_id)
Review Comment:
@AlejandroMorgante could you update the log message to "Querying Reasoning
Engine"?
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +170,90 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class RunReasoningEngineQueryOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
Review Comment:
@AlejandroMorgante could you change that it is query to reasoning engine?
##########
providers/google/src/airflow/providers/google/cloud/operators/vertex_ai/agent_engine.py:
##########
@@ -165,6 +170,90 @@ def execute(self, context: Context) -> dict[str, Any]:
return result
+class RunReasoningEngineQueryOperator(GoogleCloudBaseOperator):
+ """
+ Query a Vertex AI Agent Engine synchronously.
+
+ :param project_id: Required (templated). The ID of the Google Cloud
project that the service
+ belongs to.
+ :param location: Required (templated). The ID of the Google Cloud location
that the service
+ belongs to.
+ :param reasoning_engine_id: Required (templated). The Reasoning Engine
resource ID for the
+ Agent Engine.
+ :param input_data: Optional (templated). Input for the Agent Engine class
method in JSON object
+ format. Defaults to ``None``.
+ :param class_method: Optional (templated). 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 gcp_conn_id: The connection ID to use connecting to Google Cloud
(templated). Defaults
+ to ``google_cloud_default``.
+ :param impersonation_chain: Optional service account to impersonate using
short-term credentials
+ (templated). Defaults to ``None``.
+ """
+
+ template_fields = (
+ "project_id",
+ "location",
+ "reasoning_engine_id",
+ "input_data",
+ "class_method",
+ "gcp_conn_id",
+ "impersonation_chain",
+ )
+
+ def __init__(
+ self,
+ *,
+ project_id: str,
+ location: str,
+ reasoning_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.reasoning_engine_id = reasoning_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) -> dict[str, JsonValue]:
+ self.log.info("Querying Agent Engine %s.", self.reasoning_engine_id)
+ response = self.hook.query_reasoning_engine(
+ project_id=self.project_id,
+ location=self.location,
+ reasoning_engine_id=self.reasoning_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.reasoning_engine_id)
Review Comment:
@AlejandroMorgante the same is here could you change "Agent" to "Reasoning"?
--
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]
