scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

refactor(vertex_ai): tidy cache creation usage merge

Browse source 
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Devin AI 2026-09-18 20:14:06 +00:00 • committed by kerry
parent e04305025e
commit 0223492cc2
3 changed files with 17 additions and 10 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py
View file

@ -1,6 +1,7 @@
from typing import Final, Literal
import httpx
from pydantic import TypeAdapter
import litellm
from litellm._logging import verbose_logger
@ -14,6 +15,7 @@ from litellm.llms.custom_httpx.http_handler import (
)
from litellm.llms.openai.openai import AllMessageValues
from litellm.types.llms.vertex_ai import (
VERTEX_AI_CACHED_CONTENT_KEY,
CachedContentListAllResponseBody,
VertexAICachedContentCreation,
VertexAICachedContentResponseObject,
@ -31,7 +33,6 @@ from .transformation import (
local_cache_obj: Final = Cache(type=LiteLLMCacheType.LOCAL) # only used for calling 'get_cache_key' function
MAX_PAGINATION_PAGES: Final = 100 # Reasonable upper bound for pagination
VERTEX_AI_CACHED_CONTENT_KEY: Final = "vertex_ai_cached_content"
class ContextCachingEndpoints(VertexBase):
@ -426,7 +427,9 @@ class ContextCachingEndpoints(VertexBase):
raise VertexAIError(status_code=408, message="Timeout error occurred.")
raw_response_cached: Final = response.json()
cached_content_response_obj: Final = VertexAICachedContentResponseObject(**raw_response_cached)
cached_content_response_obj: Final = TypeAdapter(VertexAICachedContentResponseObject).validate_python(
raw_response_cached
)
usage_metadata: Final = cached_content_response_obj.get("usageMetadata", {})
cached_content_creation: Final = VertexAICachedContentCreation(
name=cached_content_response_obj["name"],
@ -588,7 +591,9 @@ class ContextCachingEndpoints(VertexBase):
raise VertexAIError(status_code=408, message="Timeout error occurred.")
raw_response_cached: Final = response.json()
cached_content_response_obj: Final = VertexAICachedContentResponseObject(**raw_response_cached)
cached_content_response_obj: Final = TypeAdapter(VertexAICachedContentResponseObject).validate_python(
raw_response_cached
)
usage_metadata: Final = cached_content_response_obj.get("usageMetadata", {})
cached_content_creation: Final = VertexAICachedContentCreation(
name=cached_content_response_obj["name"],

13
litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py
View file

@ -52,6 +52,7 @@ from litellm.types.llms.openai import (
OpenAIChatCompletionFinishReason,
)
from litellm.types.llms.vertex_ai import (
VERTEX_AI_CACHED_CONTENT_KEY,
VERTEX_AI_PROVIDER_METADATA_FIELDS,
VERTEX_CREDENTIALS_TYPES,
Candidates,
@ -144,15 +145,17 @@ def _add_cache_creation_usage(usage: Usage, creation: VertexAICachedContentCreat
if usage.prompt_tokens_details is not None
else PromptTokensDetailsWrapper()
)
cache_read_tokens: Final = getattr(usage, "_cache_read_input_tokens", 0) or 0
cache_read_tokens: Final = prompt_tokens_details.cached_tokens or 0
server_tool_use: Final = usage.server_tool_use if hasattr(usage, "server_tool_use") else None
cost: Final = usage.cost if hasattr(usage, "cost") else None
return Usage(
prompt_tokens=usage.prompt_tokens + creation_tokens,
completion_tokens=usage.completion_tokens,
total_tokens=usage.total_tokens + creation_tokens,
prompt_tokens_details=prompt_tokens_details,
completion_tokens_details=usage.completion_tokens_details,
server_tool_use=getattr(usage, "server_tool_use", None),
cost=getattr(usage, "cost", None),
server_tool_use=server_tool_use,
cost=cost,
cache_creation_input_tokens=creation_tokens,
**({"cache_read_input_tokens": cache_read_tokens} if cache_read_tokens > 0 else {}),
)
@ -2485,8 +2488,6 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
) = VertexGeminiConfig._process_candidates(_candidates, model_response, logging_obj.optional_params)
base_usage: Final = VertexGeminiConfig._calculate_usage(completion_response=completion_response)
from ..context_caching.vertex_ai_context_caching import VERTEX_AI_CACHED_CONTENT_KEY
cached_content_creation: Final = logging_obj.model_call_details.get(VERTEX_AI_CACHED_CONTENT_KEY)
usage: Final = (
_add_cache_creation_usage(base_usage, cast(VertexAICachedContentCreation, cached_content_creation))
@ -3247,8 +3248,6 @@ class ModelResponseIterator:
base_usage: Final = VertexGeminiConfig._calculate_usage(
completion_response=processed_chunk,
)
from ..context_caching.vertex_ai_context_caching import VERTEX_AI_CACHED_CONTENT_KEY
cached_content_creation: Final = self.logging_obj.model_call_details.get(VERTEX_AI_CACHED_CONTENT_KEY)
usage: Final = (
_add_cache_creation_usage(base_usage, cast(VertexAICachedContentCreation, cached_content_creation))

3
litellm/types/llms/vertex_ai.py
View file

@ -589,6 +589,9 @@ class VertexAICachedContentCreation(TypedDict):
expire_time: ReadOnly[str | None]
VERTEX_AI_CACHED_CONTENT_KEY: Final = "vertex_ai_cached_content"
class TaskTypeEnum(Enum):
TASK_TYPE_UNSPECIFIED = "TASK_TYPE_UNSPECIFIED"
RETRIEVAL_QUERY = "RETRIEVAL_QUERY"