diff --git a/tests/integration/contracts.json b/tests/integration/contracts.json index ee9b81f63dc..84f1ac03387 100644 --- a/tests/integration/contracts.json +++ b/tests/integration/contracts.json @@ -290,6 +290,9 @@ "tests/integration/spend/test_spend_calculate.py::test_spend_calculate_rejects_unpriced_model_with_400": [ "quota_management.spend_tracking.spend_calculate.rejects_unpriced_model" ], + "tests/integration/spend/test_vertex_context_cache.py::test_gemini_context_cache_creation_is_billed_once_at_its_own_rate": [ + "quota_management.spend_tracking.vertex_context_cache.creation_and_read_tokens_charged" + ], "tests/integration/management/test_partial_update_sequences.py::test_restricted_actor_cannot_detach_key_from_project": [ "mgmt.key.update.project_detach_denied_to_restricted_actor" ], diff --git a/tests/integration/spend/test_vertex_context_cache.py b/tests/integration/spend/test_vertex_context_cache.py new file mode 100644 index 00000000000..3f9a6e28e11 --- /dev/null +++ b/tests/integration/spend/test_vertex_context_cache.py @@ -0,0 +1,107 @@ +import uuid +from typing import Final + +import pytest +from integration._support.client import Gateway, eventually, object_value +from integration._support.database import read_rows +from integration._support.upstream import delete_scenario, register_scenario +from integration.cost_calculation.cost_tracking_case import JsonResponse, RoutedResponse + +_CACHE_TOKENS: Final = 2500 +_GEMINI_MODEL: Final = "gemini-2.5-pro" + + +@pytest.mark.covers("quota_management.spend_tracking.vertex_context_cache.creation_and_read_tokens_charged") +@pytest.mark.timeout(180) +def test_gemini_context_cache_creation_is_billed_once_at_its_own_rate(gateway: Gateway) -> None: + scenario_id: Final = f"sc-{uuid.uuid4().hex}" + handle: Final = register_scenario( + scenario_id, + RoutedResponse( + content_type="application/x-routed", + routes={ + f"GET /models/{_GEMINI_MODEL}:cachedContents": JsonResponse( + content_type="application/json", body={"cachedContents": []} + ), + f"POST /models/{_GEMINI_MODEL}:cachedContents": JsonResponse( + content_type="application/json", + body={ + "name": "cachedContents/created-cache", + "model": f"models/{_GEMINI_MODEL}", + "createTime": "2026-01-01T00:00:00Z", + "expireTime": "2026-01-01T01:00:00Z", + "usageMetadata": {"totalTokenCount": _CACHE_TOKENS}, + }, + ), + f"POST /models/{_GEMINI_MODEL}:generateContent": JsonResponse( + content_type="application/json", + body={ + "candidates": [ + { + "content": {"role": "model", "parts": [{"text": "cached answer"}]}, + "finishReason": "STOP", + "index": 0, + } + ], + "modelVersion": _GEMINI_MODEL, + "usageMetadata": { + "promptTokenCount": 2520, + "cachedContentTokenCount": _CACHE_TOKENS, + "candidatesTokenCount": 10, + "totalTokenCount": 2530, + }, + }, + ), + }, + ), + ) + with gateway.scenario() as scenario: + scenario.cleanups.callback(delete_scenario, handle) + model: Final = scenario.model( + model=f"gemini/{_GEMINI_MODEL}", + api_key="sk-scripted-provider", + api_base=handle.api_base(), + input_cost_per_token=0.001, + output_cost_per_token=0.002, + cache_read_input_token_cost=0.0004, + cache_creation_input_token_cost=0.0005, + ) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [ + { + "role": "system", + "content": [ + { + "type": "text", + "text": "cacheable prefix " * 4000, + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": "answer from the cache"}, + ], + }, + ) + assert response.status_code == 200, response.text + body: Final = object_value(response.json()) + usage: Final = object_value(body["usage"]) + assert usage["prompt_tokens"] == 2520 + _CACHE_TOKENS, body + assert usage["cache_creation_input_tokens"] == _CACHE_TOKENS, body + expected_spend: Final = ( + 20 * 0.001 + _CACHE_TOKENS * 0.0004 + _CACHE_TOKENS * 0.0005 + 10 * 0.002 + ) + assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected_spend) + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (body["id"],), + ), + lambda values: len(values) == 1, + seconds=70, + ) + assert rows[0]["prompt_tokens"] == 2520 + _CACHE_TOKENS, rows + assert float(rows[0]["spend"]) == pytest.approx(expected_spend), rows