fix: GenAI - Fixed GenerativeModel.compute_tokens for v1 API

PiperOrigin-RevId: 675711678
googleapis · Sep 17, 2024 · 0de2987 · 0de2987
1 parent f882657
commit 0de2987
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py
@@ -2960,7 +2960,7 @@ def _gapic_compute_tokens(
         prediction_resource_name: str,
         contents: List[gapic_content_types.Content],
     ) -> gapic_prediction_service_types.CountTokensResponse:
-        request = gapic_llm_utility_service_types.ComputeTokensRequest(
+        request = types_v1.ComputeTokensRequest(
             endpoint=prediction_resource_name,
             model=prediction_resource_name,
             contents=contents,
@@ -2972,7 +2972,7 @@ async def _gapic_compute_tokens_async(
         prediction_resource_name: str,
         contents: List[gapic_content_types.Content],
     ) -> gapic_prediction_service_types.CountTokensResponse:
-        request = gapic_llm_utility_service_types.ComputeTokensRequest(
+        request = types_v1.ComputeTokensRequest(
             endpoint=prediction_resource_name,
             model=prediction_resource_name,
             contents=contents,