feat: Vertex AI Agent Engine GA SDK

PiperOrigin-RevId: 722713179
googleapis · Feb 4, 2025 · a45cb9a · a45cb9a
1 parent 58fbabd
commit a45cb9a
Show file tree

Hide file tree

Showing 6 changed files with 3,548 additions and 0 deletions.
diff --git a/google/cloud/aiplatform/compat/services/__init__.py b/google/cloud/aiplatform/compat/services/__init__.py
@@ -167,6 +167,12 @@
 from google.cloud.aiplatform_v1.services.prediction_service import (
     async_client as prediction_service_async_client_v1,
 )
+from google.cloud.aiplatform_v1.services.reasoning_engine_service import (
+    client as reasoning_engine_service_client_v1,
+)
+from google.cloud.aiplatform_v1.services.reasoning_engine_execution_service import (
+    client as reasoning_engine_execution_service_client_v1,
+)
 from google.cloud.aiplatform_v1.services.schedule_service import (
     client as schedule_service_client_v1,
 )
@@ -209,6 +215,8 @@
     pipeline_service_client_v1,
     prediction_service_client_v1,
     prediction_service_async_client_v1,
+    reasoning_engine_execution_service_client_v1,
+    reasoning_engine_service_client_v1,
     schedule_service_client_v1,
     specialist_pool_service_client_v1,
     tensorboard_service_client_v1,

diff --git a/google/cloud/aiplatform/utils/__init__.py b/google/cloud/aiplatform/utils/__init__.py
@@ -86,6 +86,8 @@
     pipeline_service_client_v1,
     prediction_service_client_v1,
     prediction_service_async_client_v1,
+    reasoning_engine_service_client_v1,
+    reasoning_engine_execution_service_client_v1,
     schedule_service_client_v1,
     tensorboard_service_client_v1,
     vizier_service_client_v1,
@@ -971,6 +973,28 @@ class ReasoningEngineExecutionClientWithOverride(ClientWithOverride):
     )
 
 
+class AgentEngineClientWithOverride(ClientWithOverride):
+    _is_temporary = True
+    _default_version = compat.V1
+    _version_map = (
+        (
+            compat.V1,
+            reasoning_engine_service_client_v1.ReasoningEngineServiceClient,
+        ),
+    )
+
+
+class AgentEngineExecutionClientWithOverride(ClientWithOverride):
+    _is_temporary = True
+    _default_version = compat.V1
+    _version_map = (
+        (
+            compat.V1,
+            reasoning_engine_execution_service_client_v1.ReasoningEngineExecutionServiceClient,
+        ),
+    )
+
+
 class VertexRagDataClientWithOverride(ClientWithOverride):
     _is_temporary = True
     _default_version = compat.DEFAULT_VERSION
@@ -1026,6 +1050,8 @@ class VertexRagClientWithOverride(ClientWithOverride):
     PersistentResourceClientWithOverride,
     ReasoningEngineClientWithOverride,
     ReasoningEngineExecutionClientWithOverride,
+    AgentEngineClientWithOverride,
+    AgentEngineExecutionClientWithOverride,
     ModelMonitoringClientWithOverride,
 )