Skip to content

Commit

Permalink
feat: Vertex AI Agent Engine GA SDK
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 722713179
  • Loading branch information
yeesian authored and copybara-github committed Feb 5, 2025
1 parent 58fbabd commit cc4d88c
Show file tree
Hide file tree
Showing 7 changed files with 3,746 additions and 0 deletions.
8 changes: 8 additions & 0 deletions google/cloud/aiplatform/compat/services/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
)
Expand Down Expand Up @@ -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,
Expand Down
26 changes: 26 additions & 0 deletions google/cloud/aiplatform/utils/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -1026,6 +1050,8 @@ class VertexRagClientWithOverride(ClientWithOverride):
PersistentResourceClientWithOverride,
ReasoningEngineClientWithOverride,
ReasoningEngineExecutionClientWithOverride,
AgentEngineClientWithOverride,
AgentEngineExecutionClientWithOverride,
ModelMonitoringClientWithOverride,
)

Expand Down
9 changes: 9 additions & 0 deletions setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -148,6 +148,14 @@
"typing_extensions",
]

agent_engines_extra_require = [
"cloudpickle >= 3.0, < 4.0",
"google-cloud-trace < 2",
"google-cloud-logging < 4",
"pydantic >= 2.10, < 3",
"typing_extensions",
]

evaluation_extra_require = [
"pandas >= 1.0.0",
"tqdm>=4.23.0",
Expand Down Expand Up @@ -279,6 +287,7 @@
"ray": ray_extra_require,
"ray_testing": ray_testing_extra_require,
"reasoningengine": reasoning_engine_extra_require,
"agent_engines": agent_engines_extra_require,
"evaluation": evaluation_extra_require,
"langchain": langchain_extra_require,
"langchain_testing": langchain_testing_extra_require,
Expand Down
Loading

0 comments on commit cc4d88c

Please sign in to comment.