feat: Add v1beta1 service clients for reasoning engine

yeesian · copybara-github · commit 13ec7e0c6364 · 2024-03-21T07:19:31.000-07:00
PiperOrigin-RevId: 617833044
diff --git a/google/cloud/aiplatform/compat/services/__init__.py b/google/cloud/aiplatform/compat/services/__init__.py
@@ -69,6 +69,12 @@
 from google.cloud.aiplatform_v1beta1.services.prediction_service import (
     async_client as prediction_service_async_client_v1beta1,
 )
+from google.cloud.aiplatform_v1beta1.services.reasoning_engine_service import (
+    client as reasoning_engine_service_client_v1beta1,
+)
+from google.cloud.aiplatform_v1beta1.services.reasoning_engine_execution_service import (
+    client as reasoning_engine_execution_service_client_v1beta1,
+)
 from google.cloud.aiplatform_v1beta1.services.schedule_service import (
     client as schedule_service_client_v1beta1,
 )
diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py
@@ -57,6 +57,8 @@
     vizier_service_client_v1beta1,
     model_garden_service_client_v1beta1,
     persistent_resource_service_client_v1beta1,
+    reasoning_engine_service_client_v1beta1,
+    reasoning_engine_execution_service_client_v1beta1,
 )
 from google.cloud.aiplatform.compat.services import (
     dataset_service_client_v1,
@@ -746,6 +748,28 @@ class PersistentResourceClientWithOverride(ClientWithOverride):
     )
 
 
+class ReasoningEngineClientWithOverride(ClientWithOverride):
+    _is_temporary = True
+    _default_version = compat.V1BETA1
+    _version_map = (
+        (
+            compat.V1BETA1,
+            reasoning_engine_service_client_v1beta1.ReasoningEngineServiceClient,
+        ),
+    )
+
+
+class ReasoningEngineExecutionClientWithOverride(ClientWithOverride):
+    _is_temporary = True
+    _default_version = compat.V1BETA1
+    _version_map = (
+        (
+            compat.V1BETA1,
+            reasoning_engine_execution_service_client_v1beta1.ReasoningEngineExecutionServiceClient,
+        ),
+    )
+
+
 VertexAiServiceClientWithOverride = TypeVar(
     "VertexAiServiceClientWithOverride",
     DatasetClientWithOverride,
@@ -763,6 +787,8 @@ class PersistentResourceClientWithOverride(ClientWithOverride):
     VizierClientWithOverride,
     ModelGardenClientWithOverride,
     PersistentResourceClientWithOverride,
+    ReasoningEngineClientWithOverride,
+    ReasoningEngineExecutionClientWithOverride,
 )
 
 

Original file line number	Diff line number	Diff line change
`@@ -69,6 +69,12 @@`
`69`	`69`	`from google.cloud.aiplatform_v1beta1.services.prediction_service import (`
`70`	`70`	`async_client as prediction_service_async_client_v1beta1,`
`71`	`71`	`)`
	`72`	`+from google.cloud.aiplatform_v1beta1.services.reasoning_engine_service import (`
	`73`	`+ client as reasoning_engine_service_client_v1beta1,`
	`74`	`+)`
	`75`	`+from google.cloud.aiplatform_v1beta1.services.reasoning_engine_execution_service import (`
	`76`	`+ client as reasoning_engine_execution_service_client_v1beta1,`
	`77`	`+)`
`72`	`78`	`from google.cloud.aiplatform_v1beta1.services.schedule_service import (`
`73`	`79`	`client as schedule_service_client_v1beta1,`
`74`	`80`	`)`