diff --git a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py index 049336905ef..f507d2dfb97 100644 --- a/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py +++ b/litellm/llms/vertex_ai/context_caching/vertex_ai_context_caching.py @@ -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"], diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index 51de4fa3f69..ddf1c9af0d4 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -51,6 +51,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, @@ -142,15 +143,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 {}), ) @@ -2505,8 +2508,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)) @@ -3269,8 +3270,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)) diff --git a/litellm/types/llms/vertex_ai.py b/litellm/types/llms/vertex_ai.py index 5ff792d6310..ca74741af16 100644 --- a/litellm/types/llms/vertex_ai.py +++ b/litellm/types/llms/vertex_ai.py @@ -576,6 +576,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"