mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-26 01:12:21 +00:00
107 lines
4.5 KiB
Python
107 lines
4.5 KiB
Python
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
|