scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-10-08 03:08:45 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 115

fix vtx caching, use bounded in mem cache

Browse source
This commit is contained in:
Ishaan Jaff 2025-07-09 08:54:07 -07:00
parent af9db827fc
commit 7dc405248c
1 changed files with 10 additions and 7 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

17
litellm/llms/vertex_ai/vertex_llm_base.py
View file

@ -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}")