diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 994e46b50bd..fb294df9740 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -9,6 +9,7 @@ import os from typing import TYPE_CHECKING, Any, Dict, Literal, Optional, Tuple from litellm._logging import verbose_logger +from litellm.caching.in_memory_cache import InMemoryCache from litellm.litellm_core_utils.asyncify import asyncify from litellm.llms.base import BaseLLM from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler @@ -28,10 +29,7 @@ class VertexBase(BaseLLM): self.access_token: Optional[str] = None self.refresh_token: Optional[str] = None self._credentials: Optional[GoogleCredentialsObject] = None - self._credentials_project_mapping: Dict[ - Tuple[Optional[VERTEX_CREDENTIALS_TYPES], Optional[str]], - GoogleCredentialsObject, - ] = {} + self._credentials_project_mapping: InMemoryCache = InMemoryCache() self.project_id: Optional[str] = None self.async_handler: Optional[AsyncHTTPHandler] = None @@ -268,11 +266,13 @@ class VertexBase(BaseLLM): f"Checking cached credentials for project_id: {project_id}" ) - if credential_cache_key in self._credentials_project_mapping: + + credential_cached_object: Optional[GoogleCredentialsObject] = self._credentials_project_mapping.get_cache(credential_cache_key) + if credential_cached_object: verbose_logger.debug( f"Cached credentials found for project_id: {project_id}." ) - _credentials = self._credentials_project_mapping[credential_cache_key] + _credentials = credential_cached_object verbose_logger.debug("Using cached credentials") credential_project_id = _credentials.quota_project_id or getattr( _credentials, "project_id", None @@ -300,7 +300,10 @@ class VertexBase(BaseLLM): ) ) - self._credentials_project_mapping[credential_cache_key] = _credentials + self._credentials_project_mapping.set_cache( + key=credential_cache_key, + value=_credentials, + ) ## VALIDATE CREDENTIALS verbose_logger.debug(f"Validating credentials for project_id: {project_id}")