From 0de298786c43427cb1a20b91cbabd1ce921c16da Mon Sep 17 00:00:00 2001 From: Alexey Volkov Date: Tue, 17 Sep 2024 14:31:02 -0700 Subject: [PATCH] fix: GenAI - Fixed `GenerativeModel.compute_tokens` for v1 API PiperOrigin-RevId: 675711678 --- vertexai/generative_models/_generative_models.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/vertexai/generative_models/_generative_models.py b/vertexai/generative_models/_generative_models.py index 30073e2ae4..74b03d51fc 100644 --- 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,