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

fix(vertex_ai): bill context cache creation tokens and surface cache metadata

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:07:34 +00:00 • committed by kerry
parent 327515a3ba
commit e04305025e
5 changed files with 348 additions and 11 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

@ -15,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 (
CachedContentListAllResponseBody,
VertexAICachedContentCreation,
VertexAICachedContentResponseObject,
)
from litellm.utils import is_prompt_caching_valid_prompt
@ -30,6 +31,7 @@ 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):
@ -424,9 +426,16 @@ class ContextCachingEndpoints(VertexBase):
raise VertexAIError(status_code=408, message="Timeout error occurred.")
raw_response_cached: Final = response.json()
cached_content_response_obj: Final = VertexAICachedContentResponseObject(
name=raw_response_cached.get("name"), model=raw_response_cached.get("model")
cached_content_response_obj: Final = VertexAICachedContentResponseObject(**raw_response_cached)
usage_metadata: Final = cached_content_response_obj.get("usageMetadata", {})
cached_content_creation: Final = VertexAICachedContentCreation(
name=cached_content_response_obj["name"],
model=cached_content_response_obj["model"],
total_token_count=usage_metadata.get("totalTokenCount", 0),
create_time=cached_content_response_obj.get("createTime"),
expire_time=cached_content_response_obj.get("expireTime"),
)
logging_obj.model_call_details[VERTEX_AI_CACHED_CONTENT_KEY] = cached_content_creation
return (
non_cached_messages,
optional_params,
@ -579,9 +588,16 @@ class ContextCachingEndpoints(VertexBase):
raise VertexAIError(status_code=408, message="Timeout error occurred.")
raw_response_cached: Final = response.json()
cached_content_response_obj: Final = VertexAICachedContentResponseObject(
name=raw_response_cached.get("name"), model=raw_response_cached.get("model")
cached_content_response_obj: Final = VertexAICachedContentResponseObject(**raw_response_cached)
usage_metadata: Final = cached_content_response_obj.get("usageMetadata", {})
cached_content_creation: Final = VertexAICachedContentCreation(
name=cached_content_response_obj["name"],
model=cached_content_response_obj["model"],
total_token_count=usage_metadata.get("totalTokenCount", 0),
create_time=cached_content_response_obj.get("createTime"),
expire_time=cached_content_response_obj.get("expireTime"),
)
logging_obj.model_call_details[VERTEX_AI_CACHED_CONTENT_KEY] = cached_content_creation
return (
non_cached_messages,
optional_params,

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

@ -66,6 +66,7 @@ from litellm.types.llms.vertex_ai import (
ToolConfig,
Tools,
UsageMetadata,
VertexAICachedContentCreation,
VertexToolName,
)
from litellm.types.utils import (
@ -133,6 +134,30 @@ def _served_model_name(model_version: object) -> str | None:
return model_version.split("@", 1)[0]
def _add_cache_creation_usage(usage: Usage, creation: VertexAICachedContentCreation) -> Usage:
creation_tokens: Final = creation["total_token_count"]
if creation_tokens <= 0:
return usage
prompt_tokens_details: Final = (
PromptTokensDetailsWrapper(**usage.prompt_tokens_details.model_dump())
if usage.prompt_tokens_details is not None
else PromptTokensDetailsWrapper()
)
cache_read_tokens: Final = getattr(usage, "_cache_read_input_tokens", 0) or 0
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),
cache_creation_input_tokens=creation_tokens,
**({"cache_read_input_tokens": cache_read_tokens} if cache_read_tokens > 0 else {}),
)
class VertexAIBaseConfig:
def get_mapped_special_auth_params(self) -> dict:
"""
@ -2459,7 +2484,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig):
_, # cumulative_tool_call_index not needed in non-streaming
) = VertexGeminiConfig._process_candidates(_candidates, model_response, logging_obj.optional_params)
usage: Final = VertexGeminiConfig._calculate_usage(completion_response=completion_response)
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))
if isinstance(cached_content_creation, dict) and "total_token_count" in cached_content_creation
else base_usage
)
VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata)
@ -3211,9 +3244,17 @@ class ModelResponseIterator:
if "usageMetadata" not in processed_chunk:
return None
usage: Final = VertexGeminiConfig._calculate_usage(
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))
if isinstance(cached_content_creation, dict) and "total_token_count" in cached_content_creation
else base_usage
)
VertexGeminiConfig._set_grounding_usage_counters(usage, grounding_metadata)

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

@ -2,6 +2,8 @@ from enum import Enum
from typing import Any, Final, Literal, Protocol
from typing_extensions import (
NotRequired,
ReadOnly,
Required,
TypedDict,
)
@ -567,9 +569,24 @@ class MultimodalPredictions(TypedDict):
predictions: list[MultimodalPrediction]
class VertexAICachedContentUsageMetadata(TypedDict):
totalTokenCount: ReadOnly[int]
class VertexAICachedContentResponseObject(TypedDict):
name: str
model: str
name: ReadOnly[str]
model: ReadOnly[str]
usageMetadata: NotRequired[ReadOnly[VertexAICachedContentUsageMetadata]]
createTime: NotRequired[ReadOnly[str]]
expireTime: NotRequired[ReadOnly[str]]
class VertexAICachedContentCreation(TypedDict):
name: ReadOnly[str]
model: ReadOnly[str]
total_token_count: ReadOnly[int]
create_time: ReadOnly[str | None]
expire_time: ReadOnly[str | None]
class TaskTypeEnum(Enum):

96
tests/test_litellm/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py
View file

@ -15,11 +15,12 @@ from litellm.llms.anthropic.experimental_pass_through.messages import handler as
from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler
from litellm.llms.gemini.chat.transformation import GoogleAIStudioGeminiConfig
from litellm.llms.vertex_ai.common_utils import VertexAIError
from litellm.llms.vertex_ai.gemini import vertex_and_google_ai_studio_gemini as vertex_gemini_module
from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import (
VertexGeminiConfig,
)
from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata
from litellm.types.utils import ChoiceLogprobs, Usage
from litellm.types.llms.vertex_ai import GeminiFinishReason, UsageMetadata, VertexAICachedContentCreation
from litellm.types.utils import ChoiceLogprobs, PromptTokensDetailsWrapper, Usage
from litellm.utils import CustomStreamWrapper
@ -6340,3 +6341,94 @@ def test_gemini_multi_candidate_messages_do_not_share_state():
assert resp.choices[1].message.tool_calls is None
assert getattr(resp.choices[1].message, "reasoning_content", None) is None
assert resp.choices[1].provider_specific_fields["native_finish_reason"] == "STOP"
def test_add_cache_creation_usage_preserves_cache_read_and_bills_creation_tokens():
base_usage = Usage(
prompt_tokens=10010,
completion_tokens=1,
total_tokens=10011,
prompt_tokens_details=PromptTokensDetailsWrapper(cached_tokens=10000),
)
creation: VertexAICachedContentCreation = {
"name": "cached-content",
"model": "gemini-2.0-flash",
"total_token_count": 10000,
"create_time": None,
"expire_time": None,
}
merged_usage = vertex_gemini_module._add_cache_creation_usage(base_usage, creation)
assert merged_usage.prompt_tokens == 20010
assert merged_usage.total_tokens == 20011
assert merged_usage.cache_creation_input_tokens == 10000
assert merged_usage.prompt_tokens_details.cached_tokens == 10000
assert merged_usage.prompt_tokens_details.cache_creation_tokens == 10000
model = "gemini-2.0-flash"
base_response = ModelResponse(model=model, usage=base_usage)
merged_response = ModelResponse(model=model, usage=merged_usage)
base_cost = litellm.completion_cost(
completion_response=base_response,
model=model,
custom_llm_provider="vertex_ai",
)
merged_cost = litellm.completion_cost(
completion_response=merged_response,
model=model,
custom_llm_provider="vertex_ai",
)
model_info = litellm.get_model_info(model=model, custom_llm_provider="vertex_ai")
creation_rate = model_info.get("cache_creation_input_token_cost") or model_info["input_cost_per_token"]
assert merged_cost > base_cost
assert merged_cost == pytest.approx(base_cost + 10000 * creation_rate)
@pytest.mark.parametrize("include_creation", [True, False])
def test_transform_response_applies_cache_creation_usage(include_creation):
model = "gemini-2.0-flash"
logging_obj = MagicMock()
logging_obj.model_call_details = (
{
"vertex_ai_cached_content": {
"name": "cached-content",
"model": model,
"total_token_count": 10000,
"create_time": None,
"expire_time": None,
}
}
if include_creation
else {}
)
raw_response = MagicMock()
raw_response.json.return_value = {
"candidates": [{"content": {"parts": [{"text": "Hello"}]}, "finishReason": "STOP"}],
"usageMetadata": {
"promptTokenCount": 10010,
"cachedContentTokenCount": 10000,
"candidatesTokenCount": 1,
"totalTokenCount": 10011,
},
}
result = VertexGeminiConfig().transform_response(
model=model,
raw_response=raw_response,
model_response=ModelResponse(),
logging_obj=logging_obj,
request_data={},
messages=[],
optional_params={},
litellm_params={},
encoding=None,
)
if include_creation:
assert result.usage.prompt_tokens == 20010
assert result.usage.cache_creation_input_tokens == 10000
else:
assert result.usage.prompt_tokens == 10010
assert not hasattr(result.usage, "cache_creation_input_tokens")

173
tests/unit/llms/vertex_ai/context_caching/test_vertex_ai_context_caching.py
View file

@ -1,4 +1,3 @@
from typing import List
from unittest.mock import AsyncMock, MagicMock, patch
import httpx
@ -40,6 +39,7 @@ class TestContextCachingEndpoints:
"""Setup for each test method"""
self.context_caching = ContextCachingEndpoints()
self.mock_logging = MagicMock(spec=Logging)
self.mock_logging.model_call_details = {}
self.mock_client = MagicMock(spec=HTTPHandler)
self.mock_async_client = MagicMock(spec=AsyncHTTPHandler)
@ -2119,6 +2119,28 @@ class TestContextCachingMultiRegionUrls:
def setup_method(self):
self.caching = ContextCachingEndpoints()
self.context_caching = self.caching
self.mock_logging = MagicMock(spec=Logging)
self.mock_logging.model_call_details = {}
self.mock_client = MagicMock(spec=HTTPHandler)
self.mock_async_client = MagicMock(spec=AsyncHTTPHandler)
self.sample_messages = [
{
"role": "system",
"content": "You are a helpful assistant",
"cache_control": {"type": "ephemeral"},
},
{"role": "user", "content": "Hello, how are you?"},
]
self.sample_optional_params = {}
self._token_check_patcher = patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.is_prompt_caching_valid_prompt",
return_value=True,
)
self._token_check_patcher.start()
def teardown_method(self):
self._token_check_patcher.stop()
@pytest.mark.parametrize("location", ["eu", "us"])
def test_vertex_ai_multi_region_uses_rep_host(self, location):
@ -2161,3 +2183,152 @@ class TestContextCachingMultiRegionUrls:
assert url.startswith("https://aiplatform.googleapis.com/")
assert "/locations/global/cachedContents" in url
def test_check_and_create_cache_stashes_creation_metadata(self):
self.mock_logging.model_call_details = {}
cached_messages = [self.sample_messages[0]]
non_cached_messages = [self.sample_messages[1]]
response = MagicMock()
response.json.return_value = {
"name": "new_cache_name",
"model": "gemini-1.5-pro",
"usageMetadata": {"totalTokenCount": 10000},
"createTime": "2025-01-01T00:00:00Z",
"expireTime": "2025-01-02T00:00:00Z",
}
with (
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages",
return_value=(cached_messages, non_cached_messages),
),
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj.get_cache_key",
return_value="test_cache_key",
),
patch.object(self.context_caching, "check_cache", return_value=None),
patch.object(
self.context_caching,
"_get_token_and_url_context_caching",
return_value=("token", "https://test-url.com"),
),
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching",
return_value={"model": "gemini-1.5-pro", "contents": []},
),
):
self.mock_client.post.return_value = response
self.context_caching.check_and_create_cache(
messages=self.sample_messages,
optional_params=self.sample_optional_params.copy(),
api_key="test_key",
api_base=None,
model="gemini-1.5-pro",
client=self.mock_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider="vertex_ai",
vertex_project="test_project",
vertex_location="us-central1",
vertex_auth_header="token",
)
assert self.mock_logging.model_call_details["vertex_ai_cached_content"] == {
"name": "new_cache_name",
"model": "gemini-1.5-pro",
"total_token_count": 10000,
"create_time": "2025-01-01T00:00:00Z",
"expire_time": "2025-01-02T00:00:00Z",
}
@pytest.mark.asyncio
async def test_async_check_and_create_cache_stashes_creation_metadata(self):
self.mock_logging.model_call_details = {}
cached_messages = [self.sample_messages[0]]
non_cached_messages = [self.sample_messages[1]]
response = MagicMock()
response.json.return_value = {
"name": "new_cache_name",
"model": "gemini-1.5-pro",
"usageMetadata": {"totalTokenCount": 10000},
"createTime": "2025-01-01T00:00:00Z",
"expireTime": "2025-01-02T00:00:00Z",
}
with (
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages",
return_value=(cached_messages, non_cached_messages),
),
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj.get_cache_key",
return_value="test_cache_key",
),
patch.object(self.context_caching, "async_check_cache", return_value=None),
patch.object(
self.context_caching,
"_get_token_and_url_context_caching",
return_value=("token", "https://test-url.com"),
),
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.transform_openai_messages_to_gemini_context_caching",
return_value={"model": "gemini-1.5-pro", "contents": []},
),
):
self.mock_async_client.post = AsyncMock(return_value=response)
await self.context_caching.async_check_and_create_cache(
messages=self.sample_messages,
optional_params=self.sample_optional_params.copy(),
api_key="test_key",
api_base=None,
model="gemini-1.5-pro",
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider="vertex_ai",
vertex_project="test_project",
vertex_location="us-central1",
vertex_auth_header="token",
)
assert self.mock_logging.model_call_details["vertex_ai_cached_content"] == {
"name": "new_cache_name",
"model": "gemini-1.5-pro",
"total_token_count": 10000,
"create_time": "2025-01-01T00:00:00Z",
"expire_time": "2025-01-02T00:00:00Z",
}
@pytest.mark.asyncio
async def test_async_check_and_create_cache_reuse_does_not_stash_creation_metadata(self):
self.mock_logging.model_call_details = {}
cached_messages = [self.sample_messages[0]]
non_cached_messages = [self.sample_messages[1]]
with (
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.separate_cached_messages",
return_value=(cached_messages, non_cached_messages),
),
patch(
"litellm.llms.vertex_ai.context_caching.vertex_ai_context_caching.local_cache_obj.get_cache_key",
return_value="test_cache_key",
),
patch.object(self.context_caching, "async_check_cache", return_value="existing_cache"),
):
await self.context_caching.async_check_and_create_cache(
messages=self.sample_messages,
optional_params=self.sample_optional_params.copy(),
api_key="test_key",
api_base=None,
model="gemini-1.5-pro",
client=self.mock_async_client,
timeout=30.0,
logging_obj=self.mock_logging,
custom_llm_provider="vertex_ai",
vertex_project="test_project",
vertex_location="us-central1",
vertex_auth_header="token",
)
assert "vertex_ai_cached_content" not in self.mock_logging.model_call_details