diff --git a/onetwo/backends/vertexai_api.py b/onetwo/backends/vertexai_api.py
index 8cdb831..1191975 100644
--- a/onetwo/backends/vertexai_api.py
+++ b/onetwo/backends/vertexai_api.py
@@ -24,9 +24,9 @@
 from typing import Any, Final, Iterable, cast
 
 from absl import logging
-from google.cloud.aiplatform import vertexai
-from google.cloud.aiplatform.vertexai import generative_models
-from google.cloud.aiplatform.vertexai import language_models
+import vertexai
+from vertexai import generative_models
+from vertexai import language_models
 import immutabledict
 from onetwo.backends import backends_base
 from onetwo.builtins import formatting
diff --git a/onetwo/backends/vertexai_api_test.py b/onetwo/backends/vertexai_api_test.py
index 5fbb9e4..27e987d 100644
--- a/onetwo/backends/vertexai_api_test.py
+++ b/onetwo/backends/vertexai_api_test.py
@@ -21,12 +21,12 @@
 
 from absl.testing import absltest
 from absl.testing import parameterized
-from google.cloud.aiplatform import vertexai
-from google.cloud.aiplatform.aiplatform_v1beta1.types import (
+import vertexai
+from google.cloud.aiplatform_v1beta1.types import (
     prediction_service as gapic_prediction_service_types,
 )
-from google.cloud.aiplatform.vertexai import generative_models
-from google.cloud.aiplatform.vertexai import language_models
+from vertexai import generative_models
+from vertexai import language_models
 from onetwo.backends import vertexai_api
 from onetwo.builtins import llm
 from onetwo.core import caching