mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-25 01:02:15 +00:00
245 lines
12 KiB
Python
245 lines
12 KiB
Python
import uuid
|
|
from contextlib import ExitStack
|
|
from hashlib import sha256
|
|
from typing import Final
|
|
|
|
import httpx
|
|
import pytest
|
|
from hypothesis import strategies as st
|
|
from hypothesis.stateful import RuleBasedStateMachine, rule, run_state_machine_as_test
|
|
|
|
from integration._support.client import Gateway, eventually
|
|
from integration._support.database import read_rows
|
|
from integration._support.generation import LIFECYCLE_SETTINGS, bounded_http_requests
|
|
|
|
|
|
@pytest.mark.covers("quota_management.response_cache.generated_sequences_preserve_content_and_accounting")
|
|
@pytest.mark.timeout(180)
|
|
def test_generated_cache_sequences_preserve_content_usage_and_zero_hit_cost(gateway: Gateway) -> None:
|
|
class CacheRequests(RuleBasedStateMachine):
|
|
def __init__(self) -> None:
|
|
super().__init__()
|
|
self.resources = ExitStack()
|
|
try:
|
|
self.scenario = self.resources.enter_context(gateway.scenario())
|
|
self.upstream = self.resources.enter_context(
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False)
|
|
)
|
|
self.model = self.scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
self.key = self.scenario.key(models=[self.model])
|
|
self.prefix = uuid.uuid4().hex
|
|
self.seen: frozenset[int] = frozenset()
|
|
self.requests = 0
|
|
self.paid = 0
|
|
self.failed = False
|
|
self.identities: dict[int, str] = {}
|
|
except BaseException:
|
|
with budget.cleanup():
|
|
self.resources.close()
|
|
raise
|
|
|
|
@rule(marker=st.integers(min_value=0, max_value=2))
|
|
def request(self, marker: int) -> None:
|
|
try:
|
|
self.perform_request(marker)
|
|
except BaseException:
|
|
self.failed = True
|
|
raise
|
|
|
|
def perform_request(self, marker: int) -> None:
|
|
self.upstream.get("/__observations").raise_for_status()
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": self.model,
|
|
"messages": [{"role": "user", "content": f"{self.prefix}-{marker}"}],
|
|
},
|
|
key=self.key,
|
|
)
|
|
assert response.status_code == 200, response.text
|
|
self.requests += 1
|
|
body: Final = response.json()
|
|
assert (
|
|
body["choices"][0]["message"]["content"]
|
|
== "Hello! This is a mock response from the fake OpenAI endpoint."
|
|
)
|
|
assert body["usage"]["total_tokens"] == 40
|
|
observed: Final = self.upstream.get("/__observations").json()["requests"]
|
|
expected_calls: Final = 0 if marker in self.seen else 1
|
|
assert len(observed) == expected_calls, observed
|
|
if marker not in self.seen:
|
|
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(0.06)
|
|
if marker in self.identities:
|
|
assert body["id"] == self.identities[marker]
|
|
else:
|
|
assert body["id"] not in self.identities.values()
|
|
self.identities = {**self.identities, marker: body["id"]}
|
|
self.paid += expected_calls
|
|
self.seen = self.seen.union((marker,))
|
|
|
|
def teardown(self) -> None:
|
|
try:
|
|
if self.requests and not self.failed:
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
"SELECT request_id, spend, cache_hit, prompt_tokens, "
|
|
'completion_tokens FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
|
|
(sha256(self.key.encode()).hexdigest(),),
|
|
),
|
|
lambda values: len(values) == self.requests,
|
|
seconds=70,
|
|
)
|
|
assert len({row["request_id"] for row in rows}) == self.requests
|
|
assert sum(float(row["spend"]) for row in rows) == pytest.approx(self.paid * 0.06)
|
|
assert sum(row["cache_hit"] == "True" for row in rows) == self.requests - self.paid
|
|
for row in rows:
|
|
assert row["prompt_tokens"] == 20 and row["completion_tokens"] == 20
|
|
if row["cache_hit"] == "True":
|
|
assert float(row["spend"]) == 0 and "_cache_hit" in row["request_id"]
|
|
assert any(
|
|
row["request_id"].startswith(identity + "_cache_hit")
|
|
for identity in self.identities.values()
|
|
)
|
|
else:
|
|
assert row["request_id"] in self.identities.values()
|
|
assert float(row["spend"]) == pytest.approx(0.06)
|
|
finally:
|
|
with budget.cleanup():
|
|
self.resources.close()
|
|
|
|
with bounded_http_requests((gateway,), limit=2000) as budget:
|
|
run_state_machine_as_test(CacheRequests, settings=LIFECYCLE_SETTINGS)
|
|
|
|
|
|
@pytest.mark.covers("quota_management.response_cache.repeated_hits_preserve_identity_and_single_charge")
|
|
def test_repeated_hits_keep_response_identity_and_create_distinct_zero_cost_rows(gateway: Gateway) -> None:
|
|
with (
|
|
gateway.scenario() as scenario,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model])
|
|
prompt: Final = f"repeated cache {uuid.uuid4().hex}"
|
|
upstream.get("/__observations").raise_for_status()
|
|
results: Final = tuple(
|
|
gateway.post(
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "user", "content": prompt}],
|
|
"metadata": {"integration_marker": f"{prompt}-{index}"},
|
|
},
|
|
key=key,
|
|
)
|
|
for index in range(3)
|
|
)
|
|
assert len(upstream.get("/__observations").json()["requests"]) == 1
|
|
assert len({result["id"] for result in results}) == 1
|
|
for result in results:
|
|
assert (
|
|
result["choices"][0]["message"]["content"]
|
|
== "Hello! This is a mock response from the fake OpenAI endpoint."
|
|
)
|
|
assert result["usage"]["total_tokens"] == 40
|
|
rows: Final = eventually(
|
|
lambda: read_rows(
|
|
'SELECT request_id, spend, cache_hit FROM "LiteLLM_SpendLogs" WHERE api_key=%s',
|
|
(sha256(key.encode()).hexdigest(),),
|
|
),
|
|
lambda values: len(values) == 3,
|
|
seconds=70,
|
|
)
|
|
assert len({row["request_id"] for row in rows}) == 3
|
|
assert sorted(float(row["spend"]) for row in rows) == [0, 0, 0.06]
|
|
for row in rows:
|
|
if row["cache_hit"] == "True":
|
|
assert float(row["spend"]) == 0
|
|
assert row["request_id"].startswith(results[0]["id"] + "_cache_hit")
|
|
else:
|
|
assert row["request_id"] == results[0]["id"] and float(row["spend"]) == pytest.approx(0.06)
|
|
|
|
|
|
@pytest.mark.covers("quota_management.budget.key.boundary_blocks_before_provider_and_reset_restores")
|
|
def test_key_budget_at_boundary_blocks_provider_then_explicit_reset_restores(gateway: Gateway) -> None:
|
|
with (
|
|
gateway.scenario() as scenario,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
|
key: Final = scenario.key(models=[model], max_budget=0.06)
|
|
control: Final = scenario.key(models=[model])
|
|
first: Final = gateway.chat(model, key=key, text=f"budget {uuid.uuid4().hex}")
|
|
assert first["usage"]["total_tokens"] == 40
|
|
digest: Final = sha256(key.encode()).hexdigest()
|
|
spent: Final = eventually(
|
|
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
|
|
lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06,
|
|
seconds=70,
|
|
)
|
|
assert float(spent[0]["spend"]) == pytest.approx(0.06)
|
|
upstream.get("/__observations").raise_for_status()
|
|
denied: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": f"over budget {uuid.uuid4().hex}"}]},
|
|
key=key,
|
|
)
|
|
assert denied.status_code == 422 and denied.json()["error"]["type"] == "budget_exceeded", denied.text
|
|
assert upstream.get("/__observations").json()["requests"] == []
|
|
assert gateway.chat(model, key=control, text=f"control {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
|
|
gateway.post("/key/update", {"key": key, "spend": 0})
|
|
assert read_rows('SELECT spend, max_budget FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)) == [
|
|
{"spend": 0.0, "max_budget": 0.06}
|
|
]
|
|
assert gateway.chat(model, key=key, text=f"reset {uuid.uuid4().hex}")["usage"]["total_tokens"] == 40
|
|
eventually(
|
|
lambda: read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (digest,)),
|
|
lambda values: len(values) == 1 and float(values[0]["spend"]) >= 0.06,
|
|
seconds=70,
|
|
)
|
|
upstream.get("/__observations").raise_for_status()
|
|
denied_again: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{"model": model, "messages": [{"role": "user", "content": f"boundary again {uuid.uuid4().hex}"}]},
|
|
key=key,
|
|
)
|
|
assert denied_again.status_code == 422 and denied_again.json()["error"]["type"] == "budget_exceeded", (
|
|
denied_again.text
|
|
)
|
|
assert upstream.get("/__observations").json()["requests"] == []
|
|
|
|
|
|
@pytest.mark.covers("quota_management.response_cache.system_messages_partition_cache_identity")
|
|
def test_different_system_messages_do_not_share_a_cached_response(gateway: Gateway) -> None:
|
|
with (
|
|
gateway.scenario() as scenario,
|
|
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
|
):
|
|
model: Final = scenario.model()
|
|
prompt: Final = uuid.uuid4().hex
|
|
identities: dict[str, str] = {}
|
|
for system, expected_calls in (("first policy", 1), ("second policy", 1), ("first policy", 0)):
|
|
upstream.get("/__observations").raise_for_status()
|
|
response: Final = gateway.request(
|
|
"POST",
|
|
"/v1/chat/completions",
|
|
{
|
|
"model": model,
|
|
"messages": [{"role": "system", "content": system}, {"role": "user", "content": prompt}],
|
|
},
|
|
)
|
|
assert response.status_code == 200 and response.json()["usage"]["total_tokens"] == 40, response.text
|
|
calls: Final = upstream.get("/__observations").json()["requests"]
|
|
assert len(calls) == expected_calls
|
|
if system in identities:
|
|
assert response.json()["id"] == identities[system]
|
|
else:
|
|
assert response.json()["id"] not in identities.values()
|
|
identities = {**identities, system: response.json()["id"]}
|
|
if calls:
|
|
assert calls[0]["body"]["messages"] == [
|
|
{"role": "system", "content": system},
|
|
{"role": "user", "content": prompt},
|
|
]
|