mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-16 04:29:31 +00:00
test(integration): cover budgets_spend endpoint contracts (#44736)
* test(integration): cover budgets_spend endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover budgets_spend endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover budgets_spend endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): cover budgets_spend endpoint contracts Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): harden budgets_spend contracts per review Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): make spend report window midnight safe and drop redundant comments Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): pin organization spend report scope Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): close budgets and spend review findings Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(integration): drop streamed cost header absence assertion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(spend): cover streamed spend on linked budget tiers Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(spend): wait for async spend logs Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: mateo <mateo@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
60ee966885
commit
8822e086ec
8 changed files with 1703 additions and 3 deletions
|
|
@ -1,10 +1,13 @@
|
|||
from datetime import datetime, timedelta, timezone
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import pytest
|
||||
from pydantic import JsonValue, TypeAdapter
|
||||
|
||||
from tests.integration._support.client import Gateway, string_value
|
||||
from tests.integration._support.database import read_rows
|
||||
from tests.integration._support.process import owned_proxy
|
||||
|
||||
|
||||
def _persisted_reset_at(budget_id: str) -> datetime:
|
||||
|
|
@ -27,3 +30,184 @@ def test_shortening_budget_duration_moves_reset_at_onto_the_new_schedule(gateway
|
|||
updated: Final = _persisted_reset_at(budget_id)
|
||||
assert updated < ten_day_reset_at, f"{updated} not before {ten_day_reset_at}"
|
||||
assert before < updated <= before + timedelta(days=1, minutes=5), f"{updated} not within 1d of {before}"
|
||||
|
||||
|
||||
BUDGET_COLUMNS: Final = (
|
||||
"budget_id, max_budget, soft_budget, max_parallel_requests, tpm_limit, rpm_limit, tpd_limit, model_max_budget, "
|
||||
"budget_duration, allowed_models, temp_budget_increase, temp_budget_expiry, created_by, updated_by"
|
||||
)
|
||||
VOLATILE_INFO_FIELDS: Final = frozenset({"created_at", "updated_at", "budget_reset_at"})
|
||||
UNLOADED_RELATIONS: Final[dict[str, JsonValue]] = dict.fromkeys(
|
||||
(
|
||||
"organization",
|
||||
"projects",
|
||||
"keys",
|
||||
"end_users",
|
||||
"tags",
|
||||
"model_access_groups",
|
||||
"team_membership",
|
||||
"organization_membership",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
def _budget_row(budget_id: str) -> dict[str, JsonValue]:
|
||||
rows: Final = read_rows(
|
||||
f'SELECT {BUDGET_COLUMNS}, budget_reset_at::text AS reset_at FROM "LiteLLM_BudgetTable" WHERE budget_id = %s',
|
||||
(budget_id,),
|
||||
)
|
||||
assert len(rows) == 1, rows
|
||||
return rows[0]
|
||||
|
||||
|
||||
def _budget_info(gateway: Gateway, budget_id: str) -> dict[str, JsonValue]:
|
||||
response: Final = gateway.request("POST", "/budget/info", {"budgets": [budget_id]})
|
||||
assert response.status_code == 200, response.text
|
||||
entries: Final = TypeAdapter(list[dict[str, JsonValue]]).validate_json(response.content)
|
||||
assert len(entries) == 1, entries
|
||||
return entries[0]
|
||||
|
||||
|
||||
def _reset_at(value: JsonValue) -> datetime | None:
|
||||
if value is None:
|
||||
return None
|
||||
parsed: Final = datetime.fromisoformat(string_value(value).replace("Z", "+00:00"))
|
||||
return parsed if parsed.tzinfo is not None else parsed.replace(tzinfo=timezone.utc)
|
||||
|
||||
|
||||
def _assert_budget(
|
||||
gateway: Gateway,
|
||||
budget_id: str,
|
||||
expected: dict[str, JsonValue],
|
||||
reset_after: datetime | None,
|
||||
reset_within: timedelta,
|
||||
) -> None:
|
||||
row: Final = _budget_row(budget_id)
|
||||
reset_at: Final = _reset_at(row.pop("reset_at"))
|
||||
assert row == expected, row
|
||||
info: Final = _budget_info(gateway, budget_id)
|
||||
assert {name: value for name, value in info.items() if name not in VOLATILE_INFO_FIELDS} == {
|
||||
**expected,
|
||||
**UNLOADED_RELATIONS,
|
||||
}, info
|
||||
assert VOLATILE_INFO_FIELDS <= info.keys(), info
|
||||
assert _reset_at(info["budget_reset_at"]) == reset_at, info
|
||||
if reset_after is None:
|
||||
assert reset_at is None, row
|
||||
else:
|
||||
assert reset_at is not None and reset_after < reset_at <= reset_after + reset_within, (reset_at, reset_after)
|
||||
|
||||
|
||||
def _update(gateway: Gateway, body: dict[str, JsonValue]) -> None:
|
||||
response: Final = gateway.request("POST", "/budget/update", body)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
|
||||
def test_partial_update_keeps_sibling_budget_fields(gateway: Gateway) -> None:
|
||||
model_budget: Final[dict[str, JsonValue]] = {"openai/gpt-4o-mini": {"max_budget": 1, "budget_duration": "1d"}}
|
||||
with gateway.scenario() as scenario:
|
||||
created_at: Final = datetime.now(timezone.utc)
|
||||
budget_id: Final = scenario.budget(
|
||||
max_budget=10.0,
|
||||
soft_budget=5.0,
|
||||
tpm_limit=1000,
|
||||
rpm_limit=60,
|
||||
tpd_limit=100000,
|
||||
model_max_budget=model_budget,
|
||||
budget_duration="30d",
|
||||
)
|
||||
creator: Final = string_value(_budget_row(budget_id)["created_by"])
|
||||
expected: Final[dict[str, JsonValue]] = {
|
||||
"budget_id": budget_id,
|
||||
"max_budget": 10.0,
|
||||
"soft_budget": 5.0,
|
||||
"max_parallel_requests": None,
|
||||
"tpm_limit": 1000,
|
||||
"rpm_limit": 60,
|
||||
"tpd_limit": 100000,
|
||||
"model_max_budget": {"openai/gpt-4o-mini": {"max_budget": 1.0, "budget_duration": "1d"}},
|
||||
"budget_duration": "30d",
|
||||
"allowed_models": [],
|
||||
"temp_budget_increase": None,
|
||||
"temp_budget_expiry": None,
|
||||
"created_by": creator,
|
||||
"updated_by": creator,
|
||||
}
|
||||
_assert_budget(gateway, budget_id, expected, created_at, timedelta(days=30, minutes=5))
|
||||
|
||||
before_partial: Final = datetime.now(timezone.utc)
|
||||
_update(gateway, {"budget_id": budget_id, "max_budget": 20.0})
|
||||
expected["max_budget"] = 20.0
|
||||
_assert_budget(gateway, budget_id, expected, created_at, timedelta(days=30, minutes=5))
|
||||
|
||||
_update(
|
||||
gateway,
|
||||
{
|
||||
"budget_id": budget_id,
|
||||
"model_max_budget": {"openai/gpt-4o-mini": {"max_budget": 2, "budget_duration": "1d"}},
|
||||
},
|
||||
)
|
||||
expected["model_max_budget"] = {"openai/gpt-4o-mini": {"max_budget": 2.0, "budget_duration": "1d"}}
|
||||
_assert_budget(gateway, budget_id, expected, created_at, timedelta(days=30, minutes=5))
|
||||
|
||||
_update(gateway, {"budget_id": budget_id, "soft_budget": None})
|
||||
expected["soft_budget"] = None
|
||||
_assert_budget(gateway, budget_id, expected, created_at, timedelta(days=30, minutes=5))
|
||||
|
||||
before_hourly: Final = datetime.now(timezone.utc)
|
||||
_update(gateway, {"budget_id": budget_id, "budget_duration": "1h"})
|
||||
expected["budget_duration"] = "1h"
|
||||
_assert_budget(gateway, budget_id, expected, before_hourly, timedelta(hours=1, minutes=5))
|
||||
assert before_partial <= before_hourly
|
||||
|
||||
_update(gateway, {"budget_id": budget_id, "budget_duration": None})
|
||||
expected["budget_duration"] = None
|
||||
_assert_budget(gateway, budget_id, expected, None, timedelta())
|
||||
|
||||
pinned: Final = datetime(2030, 1, 1, tzinfo=timezone.utc)
|
||||
_update(gateway, {"budget_id": budget_id, "budget_duration": None, "budget_reset_at": pinned.isoformat()})
|
||||
_assert_budget(gateway, budget_id, expected, pinned - timedelta(microseconds=1), timedelta(microseconds=1))
|
||||
|
||||
rejected: Final = gateway.request(
|
||||
"POST",
|
||||
"/budget/update",
|
||||
{"budget_id": budget_id, "model_max_budget": {"openai/gpt-4o-mini": {"max_budget": "lots"}}},
|
||||
)
|
||||
assert rejected.status_code == 422, rejected.text
|
||||
assert rejected.json() == {
|
||||
"detail": [
|
||||
{
|
||||
"type": "float_parsing",
|
||||
"loc": ["body", "model_max_budget", "openai/gpt-4o-mini", "max_budget"],
|
||||
"msg": "Input should be a valid number, unable to parse string as a number",
|
||||
}
|
||||
]
|
||||
}, rejected.text
|
||||
_assert_budget(gateway, budget_id, expected, pinned - timedelta(microseconds=1), timedelta(microseconds=1))
|
||||
|
||||
|
||||
def test_model_max_budget_update_is_refused_by_the_handler_without_a_license(gateway: Gateway, tmp_path: Path) -> None:
|
||||
model_budget: Final[dict[str, JsonValue]] = {"openai/gpt-4o-mini": {"max_budget": 1, "budget_duration": "1d"}}
|
||||
with gateway.scenario() as scenario:
|
||||
budget_id: Final = scenario.budget(max_budget=10.0, model_max_budget=model_budget)
|
||||
before: Final = _budget_row(budget_id)
|
||||
with owned_proxy(gateway, tmp_path, {}, remove_environment=("LITELLM_LICENSE",)) as unlicensed:
|
||||
rejected: Final = unlicensed.request(
|
||||
"POST",
|
||||
"/budget/update",
|
||||
{"budget_id": budget_id, "model_max_budget": {"openai/gpt-4o-mini": {"max_budget": 2}}},
|
||||
)
|
||||
assert rejected.status_code == 400, rejected.text
|
||||
assert rejected.json() == {
|
||||
"detail": {
|
||||
"error": "Invalid model_max_budget: You must have an enterprise license to set model_max_budget. "
|
||||
"You must be a LiteLLM Enterprise user to use this feature. If you have a license please set "
|
||||
"`LITELLM_LICENSE` in your env. Get a 7 day trial key here: https://www.litellm.ai/enterprise#trial. "
|
||||
"\nPricing: https://www.litellm.ai/#pricing. Example of valid model_max_budget: "
|
||||
"https://docs.litellm.ai/docs/proxy/users"
|
||||
}
|
||||
}, rejected.text
|
||||
assert _budget_row(budget_id) == before
|
||||
served: Final = unlicensed.request("POST", "/budget/update", {"budget_id": budget_id, "max_budget": 20.0})
|
||||
assert served.status_code == 200, served.text
|
||||
assert _budget_row(budget_id) == {**before, "max_budget": 20.0}
|
||||
|
|
|
|||
|
|
@ -7,6 +7,7 @@ from hashlib import sha256
|
|||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
import yaml
|
||||
from pydantic import JsonValue
|
||||
|
|
@ -358,3 +359,101 @@ def test_loaded_router_preserves_cached_defaults_during_real_requests(gateway: G
|
|||
assert info["output_cost_per_token"] == (0.0 if alias == aliases[2] else 0.0000006)
|
||||
finally:
|
||||
router.reset()
|
||||
|
||||
|
||||
def _bounded_chat(gateway: Gateway, model: str, key: str, content: str) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _observed(upstream: httpx.Client) -> list[JsonValue]:
|
||||
response: Final = upstream.get("/__observations")
|
||||
response.raise_for_status()
|
||||
requests: Final = object_value(response.json())["requests"]
|
||||
assert isinstance(requests, list), response.text
|
||||
return requests
|
||||
|
||||
|
||||
def _upstream_chat(upstream_model: str, content: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"path": "/v1/chat/completions",
|
||||
"authorization": "Bearer integration-provider-key",
|
||||
"body": {"model": upstream_model, "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
|
||||
"method": "POST",
|
||||
"api_key": "",
|
||||
}
|
||||
|
||||
|
||||
def _block_unpriced(gateway: Gateway, enabled: bool) -> httpx.Response:
|
||||
return gateway.request("PATCH", "/config/block_requests_for_models_without_pricing", {"enabled": enabled})
|
||||
|
||||
|
||||
def test_block_requests_for_models_without_pricing_refuses_unpriced_deployments(gateway: Gateway) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
):
|
||||
original: Final = gateway.get("/config/block_requests_for_models_without_pricing")["enabled"]
|
||||
scenario.cleanups.callback(_block_unpriced, gateway, bool(original))
|
||||
enabled: Final = _block_unpriced(gateway, True)
|
||||
assert enabled.status_code == 200, enabled.text
|
||||
assert enabled.json() == {"enabled": True}, enabled.text
|
||||
assert gateway.get("/config/block_requests_for_models_without_pricing") == {"enabled": True}
|
||||
settings_rows: Final = read_rows(
|
||||
'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("litellm_settings",)
|
||||
)
|
||||
settings: Final = object_value(settings_rows[0]["param_value"])
|
||||
assert settings["block_requests_for_models_without_pricing"] is True, settings
|
||||
|
||||
unpriced_upstream_model: Final = f"integration-unpriced-{uuid.uuid4().hex}"
|
||||
unpriced: Final = scenario.model(model=f"openai/{unpriced_upstream_model}")
|
||||
priced: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
key: Final = scenario.key(models=[unpriced, priced])
|
||||
entries: Final = gateway.get("/model/info")["data"]
|
||||
assert isinstance(entries, list)
|
||||
unpriced_info: Final = next(
|
||||
object_value(entry) for entry in entries if object_value(entry)["model_name"] == unpriced
|
||||
)
|
||||
unpriced_model_info: Final = object_value(unpriced_info["model_info"])
|
||||
assert unpriced_model_info["input_cost_per_token"] is None, unpriced_model_info
|
||||
assert unpriced_model_info["output_cost_per_token"] is None, unpriced_model_info
|
||||
|
||||
_observed(upstream)
|
||||
blocked: Final = _bounded_chat(gateway, unpriced, key, f"unpriced {uuid.uuid4().hex}")
|
||||
assert blocked.status_code == 403, blocked.text
|
||||
assert blocked.json() == {
|
||||
"error": {
|
||||
"message": f"Model '{unpriced}' has no pricing in the cost map, so litellm cannot price the request. "
|
||||
"Requests for unpriced models are blocked because 'block_requests_for_models_without_pricing' is "
|
||||
"enabled. Add pricing (input_cost_per_token/output_cost_per_token) to allow the request.",
|
||||
"type": "model_cost_map_missing",
|
||||
"param": "model",
|
||||
"code": "403",
|
||||
}
|
||||
}, blocked.text
|
||||
assert _observed(upstream) == [], blocked.text
|
||||
|
||||
priced_content: Final = f"priced {uuid.uuid4().hex}"
|
||||
served: Final = _bounded_chat(gateway, priced, key, priced_content)
|
||||
assert served.status_code == 200, served.text
|
||||
assert float(served.headers["x-litellm-response-cost"]) == pytest.approx(0.06), served.text
|
||||
assert _observed(upstream) == [_upstream_chat("gpt-4o-mini", priced_content)], served.text
|
||||
|
||||
disabled: Final = _block_unpriced(gateway, False)
|
||||
assert disabled.status_code == 200, disabled.text
|
||||
assert disabled.json() == {"enabled": False}, disabled.text
|
||||
assert gateway.get("/config/block_requests_for_models_without_pricing") == {"enabled": False}
|
||||
settings_rows: Final = read_rows(
|
||||
'SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("litellm_settings",)
|
||||
)
|
||||
disabled_settings: Final = object_value(settings_rows[0]["param_value"])
|
||||
assert disabled_settings["block_requests_for_models_without_pricing"] is False, disabled_settings
|
||||
_observed(upstream)
|
||||
unpriced_content: Final = f"unpriced {uuid.uuid4().hex}"
|
||||
now_served: Final = _bounded_chat(gateway, unpriced, key, unpriced_content)
|
||||
assert now_served.status_code == 200, now_served.text
|
||||
assert _observed(upstream) == [_upstream_chat(unpriced_upstream_model, unpriced_content)], now_served.text
|
||||
|
|
|
|||
320
tests/integration/pricing/test_cost_margin_discount.py
Normal file
320
tests/integration/pricing/test_cost_margin_discount.py
Normal file
|
|
@ -0,0 +1,320 @@
|
|||
import json
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
|
||||
class CostEstimate(BaseModel):
|
||||
model: str
|
||||
cost_per_request: float
|
||||
input_cost_per_request: float
|
||||
output_cost_per_request: float
|
||||
margin_cost_per_request: float
|
||||
|
||||
|
||||
def _settings_row() -> dict[str, JsonValue]:
|
||||
rows: Final = read_rows('SELECT param_value FROM "LiteLLM_Config" WHERE param_name = %s', ("litellm_settings",))
|
||||
assert len(rows) == 1, rows
|
||||
value: Final = rows[0]["param_value"]
|
||||
return object_value(json.loads(value) if isinstance(value, str) else value)
|
||||
|
||||
|
||||
def _patch(gateway: Gateway, path: str, body: Mapping[str, JsonValue], message: str) -> dict[str, JsonValue]:
|
||||
response: Final = gateway.request("PATCH", path, body)
|
||||
assert response.status_code == 200, response.text
|
||||
parsed: Final = object_value(response.json())
|
||||
assert parsed == {"message": message, "status": "success", "values": dict(body)}, response.text
|
||||
return parsed
|
||||
|
||||
|
||||
class ChatChunk(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
|
||||
|
||||
def _observed(upstream: httpx.Client) -> list[JsonValue]:
|
||||
response: Final = upstream.get("/__observations")
|
||||
response.raise_for_status()
|
||||
requests: Final = object_value(response.json())["requests"]
|
||||
assert isinstance(requests, list), response.text
|
||||
return requests
|
||||
|
||||
|
||||
def _upstream_chat(content: str, streamed: bool) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"path": "/v1/chat/completions",
|
||||
"authorization": "Bearer integration-provider-key",
|
||||
"body": {
|
||||
"model": "gpt-4o-mini",
|
||||
"max_tokens": 20,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
**({"stream": True, "stream_options": {"include_usage": True}} if streamed else {}),
|
||||
},
|
||||
"method": "POST",
|
||||
"api_key": "",
|
||||
}
|
||||
|
||||
|
||||
def _chat_cost(gateway: Gateway, upstream: httpx.Client, model: str, key: str, expected: float) -> str:
|
||||
content: Final = f"margin {uuid.uuid4().hex}"
|
||||
_observed(upstream)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
request_id: Final = string_value(object_value(response.json())["id"])
|
||||
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(expected), response.text
|
||||
assert _observed(upstream) == [_upstream_chat(content, streamed=False)], response.text
|
||||
return request_id
|
||||
|
||||
|
||||
def _streamed_chat(gateway: Gateway, upstream: httpx.Client, model: str, key: str) -> str:
|
||||
content: Final = f"margin stream {uuid.uuid4().hex}"
|
||||
_observed(upstream)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 20, "stream": True, "messages": [{"role": "user", "content": content}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
events: Final = [line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: ")]
|
||||
assert events[-1] == "[DONE]", response.text
|
||||
chunks: Final = [ChatChunk.model_validate_json(event) for event in events[:-1]]
|
||||
assert {chunk.object for chunk in chunks} == {"chat.completion.chunk"}, response.text
|
||||
assert len({chunk.id for chunk in chunks}) == 1, response.text
|
||||
assert _observed(upstream) == [_upstream_chat(content, streamed=True)], response.text
|
||||
return chunks[0].id
|
||||
|
||||
|
||||
def _assert_logged_spends(expectations: tuple[tuple[str, float], ...]) -> None:
|
||||
placeholders: Final = ", ".join("%s" for _ in expectations)
|
||||
request_ids: Final = tuple(request_id for request_id, _ in expectations)
|
||||
rows: Final = eventually(
|
||||
lambda: read_rows(
|
||||
f'SELECT request_id, spend FROM "LiteLLM_SpendLogs" WHERE request_id IN ({placeholders})', request_ids
|
||||
),
|
||||
lambda values: len(values) == len(expectations),
|
||||
seconds=40,
|
||||
)
|
||||
actual: Final = {string_value(row["request_id"]): float(str(row["spend"])) for row in rows}
|
||||
assert actual == pytest.approx(dict(expectations)), rows
|
||||
|
||||
|
||||
def _estimated(gateway: Gateway, model: str) -> CostEstimate:
|
||||
response: Final = gateway.request(
|
||||
"POST", "/cost/estimate", {"model": model, "input_tokens": 20, "output_tokens": 20}
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
return CostEstimate.model_validate(response.json())
|
||||
|
||||
|
||||
def test_discount_and_margin_reprice_a_request_and_survive_rejection(gateway: Gateway) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
):
|
||||
original_discount: Final = gateway.get("/config/cost_discount_config")["values"]
|
||||
original_margin: Final = gateway.get("/config/cost_margin_config")["values"]
|
||||
|
||||
def restore() -> None:
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_discount_config",
|
||||
original_discount,
|
||||
"Cost discount configuration updated successfully",
|
||||
)
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_margin_config",
|
||||
original_margin,
|
||||
"Cost margin configuration updated successfully",
|
||||
)
|
||||
|
||||
scenario.cleanups.callback(restore)
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
key: Final = scenario.key(models=[model])
|
||||
|
||||
_patch(gateway, "/config/cost_discount_config", {}, "Cost discount configuration updated successfully")
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_margin_config",
|
||||
{},
|
||||
"Cost margin configuration updated successfully",
|
||||
)
|
||||
base_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.06)
|
||||
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_discount_config",
|
||||
{"openai": 0.05},
|
||||
"Cost discount configuration updated successfully",
|
||||
)
|
||||
discounted_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.057)
|
||||
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_margin_config",
|
||||
{"openai": 0.05},
|
||||
"Cost margin configuration updated successfully",
|
||||
)
|
||||
bare_margin_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.06 * 0.95 * 1.05)
|
||||
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_margin_config",
|
||||
{"openai": {"fixed_amount": 0.001}},
|
||||
"Cost margin configuration updated successfully",
|
||||
)
|
||||
fixed_margin_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.057 + 0.001)
|
||||
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_margin_config",
|
||||
{"openai": {"percentage": 0.1, "fixed_amount": 0.001}, "global": 0.05},
|
||||
"Cost margin configuration updated successfully",
|
||||
)
|
||||
provider_margin_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.057 * 1.1 + 0.001)
|
||||
streamed_margin_request_id: Final = _streamed_chat(gateway, upstream, model, key)
|
||||
|
||||
assert gateway.get("/config/cost_discount_config")["values"] == {"openai": 0.05}
|
||||
assert gateway.get("/config/cost_margin_config")["values"] == {
|
||||
"openai": {"percentage": 0.1, "fixed_amount": 0.001},
|
||||
"global": 0.05,
|
||||
}
|
||||
persisted: Final = _settings_row()
|
||||
assert persisted["cost_discount_config"] == {"openai": 0.05}, persisted
|
||||
assert persisted["cost_margin_config"] == {
|
||||
"openai": {"percentage": 0.1, "fixed_amount": 0.001},
|
||||
"global": 0.05,
|
||||
}, persisted
|
||||
|
||||
expected: Final = 0.057 * 1.1 + 0.001
|
||||
estimate: Final = _estimated(gateway, model)
|
||||
assert estimate.model == model
|
||||
assert estimate.cost_per_request == pytest.approx(expected), estimate
|
||||
assert estimate.input_cost_per_request == pytest.approx(0.02), estimate
|
||||
assert estimate.output_cost_per_request == pytest.approx(0.04), estimate
|
||||
discounted_subtotal: Final = estimate.cost_per_request - estimate.margin_cost_per_request
|
||||
assert (
|
||||
estimate.input_cost_per_request + estimate.output_cost_per_request - discounted_subtotal
|
||||
== pytest.approx(0.06 * 0.05)
|
||||
), estimate
|
||||
assert estimate.margin_cost_per_request == pytest.approx(0.06 * 0.95 * 0.1 + 0.001), estimate
|
||||
|
||||
_patch(gateway, "/config/cost_discount_config", {}, "Cost discount configuration updated successfully")
|
||||
no_discount_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.06 * 1.1 + 0.001)
|
||||
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_margin_config",
|
||||
{"global": 1},
|
||||
"Cost margin configuration updated successfully",
|
||||
)
|
||||
global_margin_request_id: Final = _chat_cost(gateway, upstream, model, key, 0.12)
|
||||
streamed_global_request_id: Final = _streamed_chat(gateway, upstream, model, key)
|
||||
|
||||
_assert_logged_spends(
|
||||
(
|
||||
(base_request_id, 0.06),
|
||||
(discounted_request_id, 0.057),
|
||||
(bare_margin_request_id, 0.06 * 0.95 * 1.05),
|
||||
(fixed_margin_request_id, 0.057 + 0.001),
|
||||
(provider_margin_request_id, 0.057 * 1.1 + 0.001),
|
||||
(streamed_margin_request_id, 0.057 * 1.1 + 0.001),
|
||||
(no_discount_request_id, 0.06 * 1.1 + 0.001),
|
||||
(global_margin_request_id, 0.12),
|
||||
(streamed_global_request_id, 0.12),
|
||||
)
|
||||
)
|
||||
|
||||
rejections: Final = (
|
||||
(
|
||||
"/config/cost_margin_config",
|
||||
{"not_a_provider": 0.1},
|
||||
{
|
||||
"detail": {
|
||||
"error": "Invalid provider(s): not_a_provider. Must be valid LiteLLM providers or 'global'. "
|
||||
"See https://docs.litellm.ai/docs/providers for the full list."
|
||||
}
|
||||
},
|
||||
),
|
||||
(
|
||||
"/config/cost_discount_config",
|
||||
{"openai": 1.5},
|
||||
{"detail": "Discount for openai must be between 0 and 1 (0% to 100%)"},
|
||||
),
|
||||
(
|
||||
"/config/cost_discount_config",
|
||||
{"global": 0.1},
|
||||
{
|
||||
"detail": {
|
||||
"error": "Invalid provider(s): global. Must be valid LiteLLM providers. "
|
||||
"See https://docs.litellm.ai/docs/providers for the full list."
|
||||
}
|
||||
},
|
||||
),
|
||||
)
|
||||
for path, body, error in rejections:
|
||||
rejected = gateway.request("PATCH", path, body)
|
||||
assert rejected.status_code == 400, rejected.text
|
||||
assert rejected.json() == error, rejected.text
|
||||
assert gateway.get("/config/cost_discount_config")["values"] == {}
|
||||
assert gateway.get("/config/cost_margin_config")["values"] == {"global": 1}
|
||||
final_settings: Final = _settings_row()
|
||||
assert final_settings["cost_discount_config"] == {}, final_settings
|
||||
assert final_settings["cost_margin_config"] == {"global": 1}, final_settings
|
||||
|
||||
|
||||
def test_streamed_response_cost_header_matches_billed_spend(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: a streamed chat with a fixed cost margin returns x-litellm-response-cost equal to the fixed amount "
|
||||
"(computed before usage is known) while LiteLLM_SpendLogs bills the full margined cost"
|
||||
)
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
):
|
||||
original_margin: Final = gateway.get("/config/cost_margin_config")["values"]
|
||||
scenario.cleanups.callback(
|
||||
_patch,
|
||||
gateway,
|
||||
"/config/cost_margin_config",
|
||||
original_margin,
|
||||
"Cost margin configuration updated successfully",
|
||||
)
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
key: Final = scenario.key(models=[model])
|
||||
_patch(
|
||||
gateway,
|
||||
"/config/cost_margin_config",
|
||||
{"openai": {"percentage": 0.1, "fixed_amount": 0.001}},
|
||||
"Cost margin configuration updated successfully",
|
||||
)
|
||||
content: Final = f"margin stream header {uuid.uuid4().hex}"
|
||||
_observed(upstream)
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 20, "stream": True, "messages": [{"role": "user", "content": content}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
events: Final = [
|
||||
line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: ")
|
||||
]
|
||||
assert events[-1] == "[DONE]", response.text
|
||||
request_id: Final = ChatChunk.model_validate_json(events[0]).id
|
||||
assert _observed(upstream) == [_upstream_chat(content, streamed=True)], response.text
|
||||
_assert_logged_spends(((request_id, 0.06 * 1.1 + 0.001),))
|
||||
advertised: Final = response.headers.get("x-litellm-response-cost")
|
||||
assert advertised is None or float(advertised) == pytest.approx(0.06 * 1.1 + 0.001), dict(response.headers)
|
||||
438
tests/integration/spend/test_end_user_budget_enforcement.py
Normal file
438
tests/integration/spend/test_end_user_budget_enforcement.py
Normal file
|
|
@ -0,0 +1,438 @@
|
|||
import json
|
||||
import re
|
||||
import uuid
|
||||
from collections.abc import Callable, Mapping
|
||||
from dataclasses import dataclass
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, Scenario, eventually, object_value, string_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
|
||||
from pydantic import BaseModel, JsonValue
|
||||
|
||||
PROVIDER_AUTHORIZATION: Final = "Bearer integration-provider-key"
|
||||
CANNED_CONTENT: Final = "Hello! This is a mock response from the fake OpenAI endpoint."
|
||||
RESPONSES_REPLY: Final = JsonResponse(
|
||||
content_type="application/json",
|
||||
body={
|
||||
"id": "resp_$UNIQUE_ID",
|
||||
"object": "response",
|
||||
"created_at": 1,
|
||||
"status": "completed",
|
||||
"model": "gpt-4o-mini",
|
||||
"output": [
|
||||
{
|
||||
"type": "message",
|
||||
"id": "msg_$UNIQUE_ID",
|
||||
"status": "completed",
|
||||
"role": "assistant",
|
||||
"content": [{"type": "output_text", "text": "ok", "annotations": []}],
|
||||
}
|
||||
],
|
||||
"usage": {"input_tokens": 20, "output_tokens": 20, "total_tokens": 40},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
class ChatCompletion(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
|
||||
|
||||
class ChatChunkDelta(BaseModel):
|
||||
content: str | None = None
|
||||
|
||||
|
||||
class ChatChunkChoice(BaseModel):
|
||||
delta: ChatChunkDelta
|
||||
|
||||
|
||||
class ChatChunk(BaseModel):
|
||||
object: str
|
||||
choices: list[ChatChunkChoice]
|
||||
|
||||
|
||||
class ResponsesResult(BaseModel):
|
||||
id: str
|
||||
object: str
|
||||
status: str
|
||||
|
||||
|
||||
class AnthropicMessage(BaseModel):
|
||||
type: str
|
||||
role: str
|
||||
content: list[dict[str, JsonValue]]
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Spelling:
|
||||
route: str
|
||||
fields: Callable[[str], dict[str, JsonValue]]
|
||||
headers: Callable[[str], dict[str, str]]
|
||||
upstream_fields: Callable[[str], dict[str, JsonValue]]
|
||||
stream: bool = False
|
||||
|
||||
|
||||
def _no_fields(_: str) -> dict[str, JsonValue]:
|
||||
return {}
|
||||
|
||||
|
||||
def _no_headers(_: str) -> dict[str, str]:
|
||||
return {}
|
||||
|
||||
|
||||
SPELLINGS: Final[Mapping[str, Spelling]] = {
|
||||
"user": Spelling("/v1/chat/completions", lambda user: {"user": user}, _no_headers, lambda user: {"user": user}),
|
||||
"x-litellm-customer-id": Spelling(
|
||||
"/v1/chat/completions", _no_fields, lambda user: {"x-litellm-customer-id": user}, _no_fields
|
||||
),
|
||||
"x-litellm-end-user-id": Spelling(
|
||||
"/v1/chat/completions", _no_fields, lambda user: {"x-litellm-end-user-id": user}, _no_fields
|
||||
),
|
||||
"litellm_metadata": Spelling(
|
||||
"/v1/chat/completions", lambda user: {"litellm_metadata": {"user": user}}, _no_headers, _no_fields
|
||||
),
|
||||
"litellm_metadata_json": Spelling(
|
||||
"/v1/chat/completions",
|
||||
lambda user: {"litellm_metadata": json.dumps({"user": user})},
|
||||
_no_headers,
|
||||
_no_fields,
|
||||
),
|
||||
"metadata": Spelling("/v1/chat/completions", lambda user: {"metadata": {"user_id": user}}, _no_headers, _no_fields),
|
||||
"metadata_json": Spelling(
|
||||
"/v1/chat/completions", lambda user: {"metadata": json.dumps({"user_id": user})}, _no_headers, _no_fields
|
||||
),
|
||||
"safety_identifier": Spelling(
|
||||
"/v1/chat/completions",
|
||||
lambda user: {"safety_identifier": user},
|
||||
_no_headers,
|
||||
lambda user: {"safety_identifier": user},
|
||||
),
|
||||
"user_stream": Spelling(
|
||||
"/v1/chat/completions",
|
||||
lambda user: {"user": user, "stream": True},
|
||||
_no_headers,
|
||||
lambda user: {"user": user, "stream": True, "stream_options": {"include_usage": True}},
|
||||
stream=True,
|
||||
),
|
||||
"litellm_metadata_stream": Spelling(
|
||||
"/v1/chat/completions",
|
||||
lambda user: {"litellm_metadata": {"user": user}, "stream": True},
|
||||
_no_headers,
|
||||
lambda _: {"stream": True, "stream_options": {"include_usage": True}},
|
||||
stream=True,
|
||||
),
|
||||
"responses_user": Spelling("/v1/responses", lambda user: {"user": user}, _no_headers, lambda user: {"user": user}),
|
||||
"responses_safety_identifier": Spelling(
|
||||
"/v1/responses",
|
||||
lambda user: {"safety_identifier": user},
|
||||
_no_headers,
|
||||
lambda user: {"safety_identifier": user},
|
||||
),
|
||||
"messages_metadata_user_id": Spelling(
|
||||
"/v1/messages", lambda user: {"metadata": {"user_id": user}}, _no_headers, lambda user: {"user": user}
|
||||
),
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class Deployment:
|
||||
model: str
|
||||
upstream_prefix: str
|
||||
|
||||
|
||||
def _deployment(scenario: Scenario, spelling: Spelling) -> Deployment:
|
||||
if spelling.route == "/v1/chat/completions":
|
||||
return Deployment(scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002), "")
|
||||
scenario_id: Final = f"enduser{uuid.uuid4().hex[:12]}"
|
||||
handle: Final = register_scenario(scenario_id, RESPONSES_REPLY)
|
||||
scenario.cleanups.callback(delete_scenario, handle)
|
||||
model: Final = scenario.model(
|
||||
api_base=f"{handle.api_base()}/v1", input_cost_per_token=0.001, output_cost_per_token=0.002
|
||||
)
|
||||
return Deployment(model, f"/{scenario_id}")
|
||||
|
||||
|
||||
def _client_body(spelling: Spelling, model: str, user_id: str, content: str) -> dict[str, JsonValue]:
|
||||
match spelling.route:
|
||||
case "/v1/chat/completions":
|
||||
base: dict[str, JsonValue] = {
|
||||
"model": model,
|
||||
"max_tokens": 20,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
}
|
||||
case "/v1/responses":
|
||||
base = {"model": model, "max_output_tokens": 20, "input": content}
|
||||
case _:
|
||||
base = {"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": content}]}
|
||||
return {**base, **spelling.fields(user_id)}
|
||||
|
||||
|
||||
def _upstream_request(spelling: Spelling, prefix: str, user_id: str, content: str) -> dict[str, JsonValue]:
|
||||
match spelling.route:
|
||||
case "/v1/chat/completions":
|
||||
path = "/v1/chat/completions"
|
||||
body: dict[str, JsonValue] = {
|
||||
"model": "gpt-4o-mini",
|
||||
"max_tokens": 20,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
}
|
||||
case "/v1/responses":
|
||||
path = f"{prefix}/v1/responses"
|
||||
body = {"model": "gpt-4o-mini", "input": content, "max_output_tokens": 20}
|
||||
case _:
|
||||
path = f"{prefix}/v1/responses"
|
||||
body = {
|
||||
"model": "gpt-4o-mini",
|
||||
"input": [{"type": "message", "role": "user", "content": [{"type": "input_text", "text": content}]}],
|
||||
"include": ["reasoning.encrypted_content"],
|
||||
"max_output_tokens": 20,
|
||||
}
|
||||
return {
|
||||
"path": path,
|
||||
"authorization": PROVIDER_AUTHORIZATION,
|
||||
"body": {**body, **spelling.upstream_fields(user_id)},
|
||||
"method": "POST",
|
||||
"api_key": "",
|
||||
}
|
||||
|
||||
|
||||
def _send(
|
||||
gateway: Gateway, spelling: Spelling, deployment: Deployment, key: str, user_id: str, content: str
|
||||
) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
spelling.route,
|
||||
_client_body(spelling, deployment.model, user_id, content),
|
||||
key=key,
|
||||
headers=spelling.headers(user_id),
|
||||
)
|
||||
|
||||
|
||||
def _assert_served(spelling: Spelling, response: httpx.Response) -> None:
|
||||
assert response.status_code == 200, response.text
|
||||
if spelling.stream:
|
||||
events: Final = [
|
||||
line.removeprefix("data: ") for line in response.text.splitlines() if line.startswith("data: ")
|
||||
]
|
||||
assert events[-1] == "[DONE]", response.text
|
||||
chunks: Final = [ChatChunk.model_validate_json(event) for event in events[:-1]]
|
||||
assert {chunk.object for chunk in chunks} == {"chat.completion.chunk"}, response.text
|
||||
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices) == CANNED_CONTENT
|
||||
return
|
||||
assert float(response.headers["x-litellm-response-cost"]) == pytest.approx(0.06), dict(response.headers)
|
||||
match spelling.route:
|
||||
case "/v1/chat/completions":
|
||||
assert ChatCompletion.model_validate_json(response.content).object == "chat.completion", response.text
|
||||
case "/v1/responses":
|
||||
parsed: Final = ResponsesResult.model_validate_json(response.content)
|
||||
assert (parsed.object, parsed.status) == ("response", "completed"), response.text
|
||||
case _:
|
||||
message: Final = AnthropicMessage.model_validate_json(response.content)
|
||||
assert (message.type, message.role, message.content) == (
|
||||
"message",
|
||||
"assistant",
|
||||
[{"type": "text", "text": "ok"}],
|
||||
), response.text
|
||||
|
||||
|
||||
def _delete_customer(gateway: Gateway, user_id: str) -> None:
|
||||
if not _end_user_spend(user_id):
|
||||
return
|
||||
response: Final = gateway.request("POST", "/customer/delete", {"user_ids": [user_id]})
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
|
||||
def _end_user_spend(user_id: str) -> list[dict[str, JsonValue]]:
|
||||
return read_rows('SELECT spend FROM "LiteLLM_EndUserTable" WHERE user_id = %s', (user_id,))
|
||||
|
||||
|
||||
def _observed(upstream: httpx.Client) -> list[JsonValue]:
|
||||
response: Final = upstream.get("/__observations")
|
||||
response.raise_for_status()
|
||||
requests: Final = object_value(response.json())["requests"]
|
||||
assert isinstance(requests, list), response.text
|
||||
return requests
|
||||
|
||||
|
||||
def _denial_error(denied: httpx.Response) -> dict[str, JsonValue]:
|
||||
assert denied.status_code == 422, denied.text
|
||||
return object_value(object_value(denied.json())["error"])
|
||||
|
||||
|
||||
def _prove_end_user_budget_lockout(
|
||||
gateway: Gateway,
|
||||
upstream: httpx.Client,
|
||||
spelling: Spelling,
|
||||
deployment: Deployment,
|
||||
key: str,
|
||||
end_user_id: str,
|
||||
other_end_user_id: str,
|
||||
expected_budget: Mapping[str, JsonValue],
|
||||
info_path: str = "/customer/info",
|
||||
) -> None:
|
||||
served_content: Final = f"served {end_user_id}"
|
||||
_observed(upstream)
|
||||
served: Final = _send(gateway, spelling, deployment, key, end_user_id, served_content)
|
||||
_assert_served(spelling, served)
|
||||
assert _observed(upstream) == [
|
||||
_upstream_request(spelling, deployment.upstream_prefix, end_user_id, served_content)
|
||||
], served.text
|
||||
eventually(
|
||||
lambda: _end_user_spend(end_user_id),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0.05,
|
||||
seconds=30,
|
||||
)
|
||||
persisted_spend: Final = float(str(_end_user_spend(end_user_id)[0]["spend"]))
|
||||
assert persisted_spend == pytest.approx(0.06), persisted_spend
|
||||
customer_info: Final = gateway.request("GET", info_path, params={"end_user_id": end_user_id})
|
||||
assert customer_info.status_code == 200, customer_info.text
|
||||
customer: Final = object_value(customer_info.json())
|
||||
assert customer["user_id"] == end_user_id, customer_info.text
|
||||
assert float(str(customer["spend"])) == pytest.approx(persisted_spend), customer_info.text
|
||||
assert customer["budget_id"] == expected_budget["budget_id"], customer_info.text
|
||||
budget: Final = customer["litellm_budget_table"]
|
||||
expected_max_budget: Final = expected_budget["max_budget"]
|
||||
if expected_max_budget is None:
|
||||
assert budget is None, customer_info.text
|
||||
else:
|
||||
assert float(str(object_value(budget)["max_budget"])) == float(str(expected_max_budget)), customer_info.text
|
||||
|
||||
_observed(upstream)
|
||||
denied: Final = _send(gateway, spelling, deployment, key, end_user_id, f"denied {end_user_id}")
|
||||
assert _observed(upstream) == [], denied.text
|
||||
error: Final = _denial_error(denied)
|
||||
assert error["type"] == "budget_exceeded", denied.text
|
||||
matched: Final = re.fullmatch(
|
||||
rf"ExceededBudget: End User={re.escape(end_user_id)} over budget\. "
|
||||
rf"Spend=([-+]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][-+]?\d+)?), Budget=0\.05",
|
||||
string_value(error["message"]),
|
||||
)
|
||||
assert matched is not None, denied.text
|
||||
assert float(matched.group(1)) == pytest.approx(persisted_spend), denied.text
|
||||
|
||||
other_content: Final = f"served {other_end_user_id}"
|
||||
other: Final = _send(gateway, spelling, deployment, key, other_end_user_id, other_content)
|
||||
_assert_served(spelling, other)
|
||||
assert _observed(upstream) == [
|
||||
_upstream_request(spelling, deployment.upstream_prefix, other_end_user_id, other_content)
|
||||
], other.text
|
||||
other_spend: Final = eventually(
|
||||
lambda: _end_user_spend(other_end_user_id),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0.05,
|
||||
seconds=30,
|
||||
)
|
||||
assert float(str(other_spend[0]["spend"])) == pytest.approx(0.06), other_spend
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spelling", tuple(SPELLINGS))
|
||||
def test_customer_max_budget_refuses_second_request(gateway: Gateway, spelling: str) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
):
|
||||
variant: Final = SPELLINGS[spelling]
|
||||
deployment: Final = _deployment(scenario, variant)
|
||||
key: Final = scenario.key(models=[deployment.model])
|
||||
end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
|
||||
created: Final = gateway.post("/customer/new", {"user_id": end_user_id, "max_budget": 0.05})
|
||||
assert created["user_id"] == end_user_id
|
||||
scenario.cleanups.callback(_delete_customer, gateway, end_user_id)
|
||||
other_end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
|
||||
scenario.cleanups.callback(_delete_customer, gateway, other_end_user_id)
|
||||
_prove_end_user_budget_lockout(
|
||||
gateway,
|
||||
upstream,
|
||||
variant,
|
||||
deployment,
|
||||
key,
|
||||
end_user_id,
|
||||
other_end_user_id,
|
||||
{"budget_id": string_value(created["budget_id"]), "max_budget": 0.05},
|
||||
)
|
||||
|
||||
|
||||
def test_end_user_alias_routes_create_and_read_a_budgeted_customer(gateway: Gateway) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
):
|
||||
variant: Final = SPELLINGS["user"]
|
||||
deployment: Final = _deployment(scenario, variant)
|
||||
key: Final = scenario.key(models=[deployment.model])
|
||||
end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
|
||||
created: Final = gateway.post("/end_user/new", {"user_id": end_user_id, "max_budget": 0.05})
|
||||
assert created["user_id"] == end_user_id, created
|
||||
assert float(str(object_value(created["litellm_budget_table"])["max_budget"])) == 0.05, created
|
||||
scenario.cleanups.callback(_delete_customer, gateway, end_user_id)
|
||||
other_end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
|
||||
scenario.cleanups.callback(_delete_customer, gateway, other_end_user_id)
|
||||
_prove_end_user_budget_lockout(
|
||||
gateway,
|
||||
upstream,
|
||||
variant,
|
||||
deployment,
|
||||
key,
|
||||
end_user_id,
|
||||
other_end_user_id,
|
||||
{"budget_id": string_value(created["budget_id"]), "max_budget": 0.05},
|
||||
info_path="/end_user/info",
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("spelling", ("litellm_metadata", "litellm_metadata_stream"))
|
||||
def test_customer_tier_budget_refuses_second_request(gateway: Gateway, spelling: str) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
):
|
||||
variant: Final = SPELLINGS[spelling]
|
||||
deployment: Final = _deployment(scenario, variant)
|
||||
tier: Final = scenario.budget(max_budget=0.05)
|
||||
key: Final = scenario.key(models=[deployment.model])
|
||||
end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
|
||||
created: Final = gateway.post("/customer/new", {"user_id": end_user_id, "budget_id": tier})
|
||||
assert created["user_id"] == end_user_id
|
||||
scenario.cleanups.callback(_delete_customer, gateway, end_user_id)
|
||||
other_end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
|
||||
scenario.cleanups.callback(_delete_customer, gateway, other_end_user_id)
|
||||
_prove_end_user_budget_lockout(
|
||||
gateway,
|
||||
upstream,
|
||||
variant,
|
||||
deployment,
|
||||
key,
|
||||
end_user_id,
|
||||
other_end_user_id,
|
||||
{"budget_id": tier, "max_budget": 0.05},
|
||||
)
|
||||
|
||||
|
||||
def test_key_end_user_budget_id_refuses_second_request(gateway: Gateway) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
httpx.Client(base_url=gateway.upstream_url, timeout=5, trust_env=False) as upstream,
|
||||
):
|
||||
variant: Final = SPELLINGS["metadata"]
|
||||
deployment: Final = _deployment(scenario, variant)
|
||||
tier: Final = scenario.budget(max_budget=0.05)
|
||||
key: Final = scenario.key(models=[deployment.model], metadata={"end_user_budget_id": tier})
|
||||
key_info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
|
||||
assert object_value(key_info["metadata"]) == {"end_user_budget_id": tier}, key_info
|
||||
end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
|
||||
other_end_user_id: Final = f"end-user-{uuid.uuid4().hex}"
|
||||
scenario.cleanups.callback(_delete_customer, gateway, end_user_id)
|
||||
scenario.cleanups.callback(_delete_customer, gateway, other_end_user_id)
|
||||
_prove_end_user_budget_lockout(
|
||||
gateway,
|
||||
upstream,
|
||||
variant,
|
||||
deployment,
|
||||
key,
|
||||
end_user_id,
|
||||
other_end_user_id,
|
||||
{"budget_id": None, "max_budget": None},
|
||||
)
|
||||
159
tests/integration/spend/test_global_spend_reset.py
Normal file
159
tests/integration/spend/test_global_spend_reset.py
Normal file
|
|
@ -0,0 +1,159 @@
|
|||
import uuid
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, model: str, key: str, content: str | None = None) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 20,
|
||||
"messages": [{"role": "user", "content": content or f"reset {uuid.uuid4().hex}"}],
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _observed(upstream: httpx.Client) -> list[JsonValue]:
|
||||
response: Final = upstream.get("/__observations")
|
||||
response.raise_for_status()
|
||||
requests: Final = object_value(response.json())["requests"]
|
||||
assert isinstance(requests, list), response.text
|
||||
return requests
|
||||
|
||||
|
||||
def _upstream_chat(content: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"path": "/v1/chat/completions",
|
||||
"authorization": "Bearer integration-provider-key",
|
||||
"body": {"model": "gpt-4o-mini", "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
|
||||
"method": "POST",
|
||||
"api_key": "",
|
||||
}
|
||||
|
||||
|
||||
def _token_spend(key: str) -> list[dict[str, object]]:
|
||||
return read_rows(
|
||||
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token = %s', (sha256(key.encode()).hexdigest(),)
|
||||
) # fmt: skip
|
||||
|
||||
|
||||
def _team_spend(team_id: str) -> list[dict[str, object]]:
|
||||
return read_rows('SELECT spend FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,))
|
||||
|
||||
|
||||
def _spend_log_totals(keys: tuple[str, ...]) -> list[dict[str, object]]:
|
||||
digests: Final = tuple(sha256(key.encode()).hexdigest() for key in keys)
|
||||
return read_rows(
|
||||
'SELECT COUNT(*) AS rows, COALESCE(SUM(spend), 0) AS total_spend FROM "LiteLLM_SpendLogs" '
|
||||
"WHERE api_key = ANY(%s::text[])",
|
||||
("{" + ",".join(digests) + "}",),
|
||||
)
|
||||
|
||||
|
||||
def _seed_spend(gateway: Gateway, model: str, team_id: str, key_a: str, key_b: str) -> None:
|
||||
assert _chat(gateway, model, key_a).status_code == 200
|
||||
assert _chat(gateway, model, key_b).status_code == 200
|
||||
for key in (key_a, key_b):
|
||||
eventually(
|
||||
lambda key=key: _token_spend(key),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0,
|
||||
seconds=70,
|
||||
)
|
||||
eventually(
|
||||
lambda: _team_spend(team_id),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) > 0,
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
|
||||
def test_global_spend_reset_zeroes_token_and_team_spend_but_keeps_logs(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)
|
||||
team_id: Final = scenario.team()
|
||||
key_a: Final = scenario.key(team_id=team_id, models=[model], max_budget=0.06)
|
||||
key_b: Final = scenario.key(team_id=team_id, models=[model])
|
||||
internal_user: Final = scenario.user(user_role="internal_user")
|
||||
internal_key: Final = scenario.key(user_id=internal_user, models=[model])
|
||||
_seed_spend(gateway, model, team_id, key_a, key_b)
|
||||
before: Final = eventually(
|
||||
lambda: _spend_log_totals((key_a, key_b)),
|
||||
lambda rows: (
|
||||
len(rows) == 1
|
||||
and int(str(rows[0]["rows"])) == 2
|
||||
and float(str(rows[0]["total_spend"])) == pytest.approx(0.12)
|
||||
),
|
||||
seconds=70,
|
||||
)
|
||||
assert float(str(before[0]["total_spend"])) == pytest.approx(0.12), before
|
||||
|
||||
refused: Final = gateway.request("POST", "/global/spend/reset", key=internal_key)
|
||||
assert refused.status_code == 401, refused.text
|
||||
assert refused.json() == {
|
||||
"error": {
|
||||
"message": "Authentication Error, Tried to access route=/global/spend/reset, which is only for MASTER KEY",
|
||||
"type": "auth_error",
|
||||
"param": "None",
|
||||
"code": "401",
|
||||
}
|
||||
}, refused.text
|
||||
assert float(str(_token_spend(key_a)[0]["spend"])) == pytest.approx(0.06), _token_spend(key_a)
|
||||
assert float(str(_token_spend(key_b)[0]["spend"])) == pytest.approx(0.06), _token_spend(key_b)
|
||||
assert float(str(_team_spend(team_id)[0]["spend"])) == pytest.approx(0.12), _team_spend(team_id)
|
||||
|
||||
_observed(upstream)
|
||||
denied: Final = _chat(gateway, model, key_a)
|
||||
assert denied.status_code == 422, denied.text
|
||||
assert object_value(denied.json()["error"])["type"] == "budget_exceeded", denied.text
|
||||
assert _observed(upstream) == [], denied.text
|
||||
|
||||
reset: Final = gateway.request("POST", "/global/spend/reset")
|
||||
assert reset.status_code == 200, reset.text
|
||||
assert reset.json() == {
|
||||
"message": "Spend for all API Keys and Teams reset successfully",
|
||||
"status": "success",
|
||||
}, reset.text
|
||||
for key in (key_a, key_b):
|
||||
rows: Final = _token_spend(key)
|
||||
assert float(str(rows[0]["spend"])) == 0, rows
|
||||
assert float(str(_team_spend(team_id)[0]["spend"])) == 0
|
||||
assert _spend_log_totals((key_a, key_b)) == before
|
||||
|
||||
|
||||
def test_reset_exhausted_key_is_served_again(gateway: Gateway) -> None:
|
||||
pytest.skip(
|
||||
"BUG: /global/spend/reset zeroes DB spend but leaves Redis spend:key counters, "
|
||||
"so an exhausted key stays refused with budget_exceeded"
|
||||
)
|
||||
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)
|
||||
assert _chat(gateway, model, key).status_code == 200
|
||||
eventually(
|
||||
lambda: _token_spend(key),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.06,
|
||||
seconds=70,
|
||||
)
|
||||
denied: Final = _chat(gateway, model, key)
|
||||
assert denied.status_code == 422, denied.text
|
||||
reset: Final = gateway.request("POST", "/global/spend/reset")
|
||||
assert reset.status_code == 200, reset.text
|
||||
content: Final = f"restored {uuid.uuid4().hex}"
|
||||
_observed(upstream)
|
||||
restored: Final = _chat(gateway, model, key, content)
|
||||
assert restored.status_code == 200, restored.text
|
||||
assert _observed(upstream) == [_upstream_chat(content)], restored.text
|
||||
|
|
@ -6,6 +6,8 @@ import httpx
|
|||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from openai.types.chat import ChatCompletionChunk
|
||||
from pydantic import JsonValue
|
||||
|
||||
|
||||
def test_an_exhausted_key_is_refused_inference_but_can_still_read_its_own_info(gateway: Gateway) -> None:
|
||||
|
|
@ -43,19 +45,46 @@ def test_an_exhausted_key_is_refused_inference_but_can_still_read_its_own_info(g
|
|||
assert own["max_budget"] == 0.06
|
||||
|
||||
|
||||
def _bounded_chat(gateway: Gateway, model: str, key: str) -> httpx.Response:
|
||||
def _bounded_chat(gateway: Gateway, model: str, key: str, content: str | None = None) -> httpx.Response:
|
||||
return gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": model,
|
||||
"max_tokens": 20,
|
||||
"messages": [{"role": "user", "content": f"key recovery {uuid.uuid4().hex}"}],
|
||||
"messages": [{"role": "user", "content": content or f"key recovery {uuid.uuid4().hex}"}],
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
|
||||
|
||||
def _observed(upstream: httpx.Client) -> list[JsonValue]:
|
||||
response: Final = upstream.get("/__observations")
|
||||
response.raise_for_status()
|
||||
requests: Final = object_value(response.json())["requests"]
|
||||
assert isinstance(requests, list), response.text
|
||||
return requests
|
||||
|
||||
|
||||
def _upstream_chat(content: str) -> dict[str, JsonValue]:
|
||||
return {
|
||||
"path": "/v1/chat/completions",
|
||||
"authorization": "Bearer integration-provider-key",
|
||||
"body": {"model": "gpt-4o-mini", "max_tokens": 20, "messages": [{"role": "user", "content": content}]},
|
||||
"method": "POST",
|
||||
"api_key": "",
|
||||
}
|
||||
|
||||
|
||||
def _served_chat(gateway: Gateway, upstream: httpx.Client, model: str, key: str) -> None:
|
||||
content: Final = f"tier served {uuid.uuid4().hex}"
|
||||
_observed(upstream)
|
||||
served: Final = _bounded_chat(gateway, model, key, content)
|
||||
assert served.status_code == 200, served.text
|
||||
assert float(served.headers["x-litellm-response-cost"]) == 0.06, dict(served.headers)
|
||||
assert _observed(upstream) == [_upstream_chat(content)], served.text
|
||||
|
||||
|
||||
def test_raising_a_spent_keys_budget_restores_serving(gateway: Gateway) -> None:
|
||||
with (
|
||||
gateway.scenario() as scenario,
|
||||
|
|
@ -82,3 +111,156 @@ def test_raising_a_spent_keys_budget_restores_serving(gateway: Gateway) -> None:
|
|||
served: Final = tuple(_bounded_chat(gateway, model, key) for _ in range(3))
|
||||
assert [response.status_code for response in served] == [200, 200, 200], [response.text for response in served]
|
||||
assert len(upstream.get("/__observations").json()["requests"]) == 3
|
||||
|
||||
|
||||
def _denied_with_budget_exceeded(gateway: Gateway, model: str, key: str) -> None:
|
||||
denied: Final = _bounded_chat(gateway, model, key)
|
||||
assert denied.status_code == 422, denied.text
|
||||
assert object_value(denied.json()["error"])["type"] == "budget_exceeded"
|
||||
|
||||
|
||||
@pytest.mark.timeout(149)
|
||||
def test_raising_a_spent_tier_budget_restores_serving(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)
|
||||
budget_id: Final = scenario.budget(budget_id=f"tier-{uuid.uuid4().hex}", max_budget=0.06, budget_duration="30d")
|
||||
key: Final = scenario.key(models=[model], budget_id=budget_id)
|
||||
info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
|
||||
assert info["budget_id"] == budget_id, info
|
||||
tier_rows: Final = read_rows(
|
||||
'SELECT max_budget, budget_duration FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
|
||||
)
|
||||
assert len(tier_rows) == 1
|
||||
assert float(str(tier_rows[0]["max_budget"])) == 0.06
|
||||
assert tier_rows[0]["budget_duration"] == "30d"
|
||||
_served_chat(gateway, upstream, model, key)
|
||||
eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),)
|
||||
),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.06,
|
||||
seconds=70,
|
||||
)
|
||||
upstream.get("/__observations").raise_for_status()
|
||||
_denied_with_budget_exceeded(gateway, model, key)
|
||||
assert upstream.get("/__observations").json()["requests"] == []
|
||||
gateway.post("/budget/update", {"budget_id": budget_id, "max_budget": 1.0})
|
||||
updated_rows: Final = read_rows(
|
||||
'SELECT max_budget FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
|
||||
)
|
||||
assert float(str(updated_rows[0]["max_budget"])) == 1.0, updated_rows
|
||||
restored_content: Final = f"tier restored {uuid.uuid4().hex}"
|
||||
served: Final = eventually(
|
||||
lambda: _bounded_chat(gateway, model, key, restored_content),
|
||||
lambda response: response.status_code == 200,
|
||||
seconds=125,
|
||||
)
|
||||
assert served.status_code == 200, served.text
|
||||
assert _observed(upstream) == [_upstream_chat(restored_content)], served.text
|
||||
|
||||
|
||||
@pytest.mark.timeout(149)
|
||||
def test_lowering_a_tier_budget_below_spend_blocks_the_key(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)
|
||||
budget_id: Final = scenario.budget(
|
||||
budget_id=f"tier-{uuid.uuid4().hex}", max_budget=1000.0, budget_duration="30d"
|
||||
)
|
||||
key: Final = scenario.key(models=[model], budget_id=budget_id)
|
||||
_served_chat(gateway, upstream, model, key)
|
||||
eventually(
|
||||
lambda: read_rows(
|
||||
'SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (sha256(key.encode()).hexdigest(),)
|
||||
),
|
||||
lambda rows: len(rows) == 1 and float(str(rows[0]["spend"])) >= 0.06,
|
||||
seconds=30,
|
||||
)
|
||||
upstream.get("/__observations").raise_for_status()
|
||||
probe: Final = gateway.request("POST", "/v1/chat/completions", {"model": model}, key=key)
|
||||
assert probe.status_code == 400, probe.text
|
||||
assert object_value(probe.json()["error"]) == {
|
||||
"message": "/chat/completions: Missing required parameter: 'messages'.",
|
||||
"type": "invalid_request_error",
|
||||
"param": "messages",
|
||||
"code": "400",
|
||||
}, probe.text
|
||||
assert upstream.get("/__observations").json()["requests"] == []
|
||||
gateway.post("/budget/update", {"budget_id": budget_id, "max_budget": 0.01})
|
||||
updated_rows: Final = read_rows(
|
||||
'SELECT max_budget FROM "LiteLLM_BudgetTable" WHERE budget_id = %s', (budget_id,)
|
||||
)
|
||||
assert len(updated_rows) == 1
|
||||
assert float(str(updated_rows[0]["max_budget"])) == 0.01, updated_rows
|
||||
upstream.get("/__observations").raise_for_status()
|
||||
lowered_probe: Final = eventually(
|
||||
lambda: gateway.request("POST", "/v1/chat/completions", {"model": model}, key=key),
|
||||
lambda response: (
|
||||
response.status_code == 422 and object_value(response.json()["error"])["type"] == "budget_exceeded"
|
||||
),
|
||||
seconds=125,
|
||||
)
|
||||
assert "Max budget: 0.01" in lowered_probe.text, lowered_probe.text
|
||||
denied: Final = _bounded_chat(gateway, model, key)
|
||||
assert denied.status_code == 422, denied.text
|
||||
assert object_value(denied.json()["error"])["type"] == "budget_exceeded", denied.text
|
||||
assert "Max budget: 0.01" in denied.text, denied.text
|
||||
assert upstream.get("/__observations").json()["requests"] == []
|
||||
|
||||
|
||||
@pytest.mark.timeout(149)
|
||||
def test_key_linked_to_a_tier_by_update_is_refused_once_streamed_spend_exhausts_it(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)
|
||||
budget_id: Final = scenario.budget(budget_id=f"tier-{uuid.uuid4().hex}", max_budget=0.06, budget_duration="30d")
|
||||
key: Final = scenario.key(models=[model])
|
||||
gateway.post("/key/update", {"key": key, "budget_id": budget_id})
|
||||
info: Final = object_value(gateway.get("/key/info", {"key": key})["info"])
|
||||
assert info["budget_id"] == budget_id, info
|
||||
|
||||
content: Final = f"tier streamed {uuid.uuid4().hex}"
|
||||
body: Final = {
|
||||
"model": model,
|
||||
"max_tokens": 20,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"stream": True,
|
||||
}
|
||||
_observed(upstream)
|
||||
served: Final = gateway.request("POST", "/v1/chat/completions", body, key=key)
|
||||
assert served.status_code == 200, served.text
|
||||
events: Final = tuple(
|
||||
line.removeprefix("data: ") for line in served.text.splitlines() if line.startswith("data: ")
|
||||
)
|
||||
assert events[-1] == "[DONE]", served.text
|
||||
chunks: Final = tuple(ChatCompletionChunk.model_validate_json(event) for event in events[:-1])
|
||||
assert {chunk.object for chunk in chunks} == {"chat.completion.chunk"}, served.text
|
||||
assert "".join(choice.delta.content or "" for chunk in chunks for choice in chunk.choices), served.text
|
||||
observations: Final = _observed(upstream)
|
||||
assert len(observations) == 1, observations
|
||||
upstream_body: Final = object_value(object_value(observations[0])["body"])
|
||||
assert upstream_body["stream"] is True, observations
|
||||
|
||||
denied: Final = eventually(
|
||||
lambda: gateway.request("POST", "/v1/chat/completions", body, key=key),
|
||||
lambda response: (
|
||||
response.status_code == 422 and object_value(response.json()["error"])["type"] == "budget_exceeded"
|
||||
),
|
||||
seconds=125,
|
||||
)
|
||||
assert denied.status_code == 422, denied.text
|
||||
assert object_value(denied.json()["error"])["type"] == "budget_exceeded", denied.text
|
||||
assert "Max budget: 0.06" in denied.text, denied.text
|
||||
_observed(upstream)
|
||||
fresh_denied: Final = gateway.request("POST", "/v1/chat/completions", body, key=key)
|
||||
assert fresh_denied.status_code == 422, fresh_denied.text
|
||||
assert object_value(fresh_denied.json()["error"])["type"] == "budget_exceeded", fresh_denied.text
|
||||
assert "Max budget: 0.06" in fresh_denied.text, fresh_denied.text
|
||||
assert _observed(upstream) == []
|
||||
|
|
|
|||
228
tests/integration/spend/test_scoped_spend_reports.py
Normal file
228
tests/integration/spend/test_scoped_spend_reports.py
Normal file
|
|
@ -0,0 +1,228 @@
|
|||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from datetime import date, datetime, time, timedelta
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from integration._support.client import Gateway, eventually, object_value
|
||||
from integration._support.database import read_rows
|
||||
from pydantic import BaseModel
|
||||
|
||||
|
||||
class ModelDetail(BaseModel):
|
||||
model: str
|
||||
total_cost: float
|
||||
total_input_tokens: float
|
||||
total_output_tokens: float
|
||||
team_id: str | None = None
|
||||
|
||||
|
||||
class ReportRow(BaseModel):
|
||||
api_key: str
|
||||
total_cost: float
|
||||
total_input_tokens: float
|
||||
total_output_tokens: float
|
||||
model_details: list[ModelDetail]
|
||||
|
||||
|
||||
def _chat(gateway: Gateway, model: str, key: str) -> None:
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": f"report {uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
|
||||
def _report(
|
||||
gateway: Gateway, path: str, *, key: str | None = None, params: Mapping[str, str] | None = None
|
||||
) -> httpx.Response:
|
||||
return gateway.request("GET", path, key=key, params=params)
|
||||
|
||||
|
||||
def _rows(
|
||||
gateway: Gateway, path: str, *, key: str | None = None, params: Mapping[str, str] | None = None
|
||||
) -> list[ReportRow]:
|
||||
response: Final = _report(gateway, path, key=key, params=params)
|
||||
assert response.status_code == 200, response.text
|
||||
return [ReportRow.model_validate(row) for row in response.json()]
|
||||
|
||||
|
||||
def _logged(scope_column: str, scope_value: str, start: datetime, end: datetime) -> list[dict[str, object]]:
|
||||
return read_rows(
|
||||
f'SELECT api_key, model, spend, prompt_tokens, completion_tokens, team_id FROM "LiteLLM_SpendLogs" '
|
||||
f'WHERE "startTime" >= %s AND "startTime" < %s AND {scope_column} = %s',
|
||||
(start.isoformat(), end.isoformat(), scope_value),
|
||||
)
|
||||
|
||||
|
||||
def _logged_keys(keys: tuple[str, ...], start: datetime, end: datetime) -> list[dict[str, object]]:
|
||||
placeholders: Final = ", ".join("%s" for _ in keys)
|
||||
return read_rows(
|
||||
'SELECT api_key, model, spend, prompt_tokens, completion_tokens, team_id FROM "LiteLLM_SpendLogs" '
|
||||
f'WHERE "startTime" >= %s AND "startTime" < %s AND api_key IN ({placeholders})',
|
||||
(start.isoformat(), end.isoformat(), *keys),
|
||||
)
|
||||
|
||||
|
||||
def _assert_report_matches_db(report: list[ReportRow], logged: list[dict[str, object]]) -> None:
|
||||
assert len(report) == len({row["api_key"] for row in logged}), (report, logged)
|
||||
by_key: Final[dict[str, dict[str, list[float]]]] = {}
|
||||
for row in logged:
|
||||
bucket: Final = by_key.setdefault(str(row["api_key"]), {})
|
||||
model_bucket: Final = bucket.setdefault(str(row["model"]), [0.0, 0.0, 0.0])
|
||||
model_bucket[0] += float(str(row["spend"]))
|
||||
model_bucket[1] += int(str(row["prompt_tokens"]))
|
||||
model_bucket[2] += int(str(row["completion_tokens"]))
|
||||
for entry in report:
|
||||
assert entry.api_key in by_key, (entry, logged)
|
||||
models: Final = by_key[entry.api_key]
|
||||
assert entry.total_cost == pytest.approx(sum(v[0] for v in models.values())), entry
|
||||
assert entry.total_input_tokens == pytest.approx(sum(v[1] for v in models.values())), entry
|
||||
assert entry.total_output_tokens == pytest.approx(sum(v[2] for v in models.values())), entry
|
||||
details: Final = {detail.model: detail for detail in entry.model_details}
|
||||
assert set(details) == set(models), (entry, logged)
|
||||
for model, detail in details.items():
|
||||
cost, prompt, completion = models[model]
|
||||
assert detail.total_cost == pytest.approx(cost), detail
|
||||
assert detail.total_input_tokens == pytest.approx(prompt), detail
|
||||
assert detail.total_output_tokens == pytest.approx(completion), detail
|
||||
|
||||
|
||||
def test_scoped_spend_reports_match_sql_and_enforce_caller_scope(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
org1: Final = scenario.organization()
|
||||
org2: Final = scenario.organization()
|
||||
team1: Final = scenario.team(organization_id=org1)
|
||||
team2: Final = scenario.team(organization_id=org2)
|
||||
user1: Final = scenario.member(team1, role="user")
|
||||
user2: Final = scenario.member(team2, role="user")
|
||||
org_admin: Final = scenario.org_member(org1, "org_admin")
|
||||
k1: Final = scenario.key(user_id=user1, team_id=team1, models=[model])
|
||||
k2: Final = scenario.key(user_id=user2, team_id=team2, models=[model])
|
||||
direct_org_key: Final = scenario.key(organization_id=org2, models=[model])
|
||||
org_admin_key: Final = scenario.key(user_id=org_admin, organization_id=org1, models=[model])
|
||||
org_admin_info: Final = object_value(gateway.get("/key/info", {"key": org_admin_key})["info"])
|
||||
assert org_admin_info["organization_id"] == org1, org_admin_info
|
||||
_chat(gateway, model, k1)
|
||||
_chat(gateway, model, k2)
|
||||
_chat(gateway, model, direct_org_key)
|
||||
k1_hash: Final = sha256(k1.encode()).hexdigest()
|
||||
k2_hash: Final = sha256(k2.encode()).hexdigest()
|
||||
direct_org_key_hash: Final = sha256(direct_org_key.encode()).hexdigest()
|
||||
logged: Final = eventually(
|
||||
lambda: read_rows(
|
||||
"SELECT api_key, to_char(\"startTime\", 'YYYY-MM-DD') AS day "
|
||||
'FROM "LiteLLM_SpendLogs" WHERE api_key IN (%s, %s, %s)',
|
||||
(k1_hash, k2_hash, direct_org_key_hash),
|
||||
),
|
||||
lambda rows: len(rows) == 3,
|
||||
seconds=30,
|
||||
)
|
||||
dates: Final = {date.fromisoformat(str(row["day"])) for row in logged}
|
||||
start_date: Final[date] = min(dates)
|
||||
end_date: Final[date] = max(dates)
|
||||
start: Final = datetime.combine(start_date, time.min)
|
||||
end: Final = datetime.combine(end_date, time.min) + timedelta(days=1)
|
||||
window: Final = {"start_date": start_date.isoformat(), "end_date": end_date.isoformat()}
|
||||
|
||||
_assert_report_matches_db(
|
||||
_rows(gateway, "/key/spend/report", key=k1, params=window), _logged("api_key", k1_hash, start, end)
|
||||
)
|
||||
_assert_report_matches_db(
|
||||
_rows(gateway, "/user/spend/report", key=k1, params=window), _logged('"user"', user1, start, end)
|
||||
)
|
||||
_assert_report_matches_db(
|
||||
_rows(gateway, "/team/spend/report", key=k1, params=window), _logged("team_id", team1, start, end)
|
||||
)
|
||||
org1_report: Final = _rows(gateway, "/organization/spend/report", key=org_admin_key, params=window)
|
||||
_assert_report_matches_db(org1_report, _logged_keys((k1_hash,), start, end))
|
||||
refused_org2: Final = _report(
|
||||
gateway,
|
||||
"/organization/spend/report",
|
||||
key=org_admin_key,
|
||||
params={**window, "organization_id": org2},
|
||||
)
|
||||
assert refused_org2.status_code == 403, refused_org2.text
|
||||
assert refused_org2.json() == {"detail": "You do not have access to this organization"}, refused_org2.text
|
||||
|
||||
own_org: Final = _report(gateway, "/organization/spend/report", key=k1, params=window)
|
||||
assert own_org.status_code == 403, own_org.text
|
||||
assert own_org.json() == {"detail": "You do not have access to this organization"}, own_org.text
|
||||
|
||||
assert _rows(gateway, "/key/spend/report", key=k1, params={**window, "api_key": k1}) == _rows(
|
||||
gateway, "/key/spend/report", key=k1, params={**window, "api_key": k1_hash}
|
||||
)
|
||||
|
||||
for path, param, value, detail in (
|
||||
(
|
||||
"/key/spend/report",
|
||||
"api_key",
|
||||
k2_hash,
|
||||
"Not authorized to view spend for a api_key other than your own",
|
||||
),
|
||||
(
|
||||
"/user/spend/report",
|
||||
"internal_user_id",
|
||||
user2,
|
||||
"Not authorized to view spend for a internal_user_id other than your own",
|
||||
),
|
||||
(
|
||||
"/team/spend/report",
|
||||
"team_id",
|
||||
team2,
|
||||
"Not authorized to view spend for a team_id other than your own",
|
||||
),
|
||||
("/organization/spend/report", "organization_id", org2, "You do not have access to this organization"),
|
||||
):
|
||||
refused: Final = _report(gateway, path, key=k1, params={**window, param: value})
|
||||
assert refused.status_code == 403, refused.text
|
||||
assert refused.json() == {"detail": detail}, refused.text
|
||||
|
||||
_assert_report_matches_db(
|
||||
_rows(gateway, "/key/spend/report", params={**window, "api_key": k2_hash}),
|
||||
_logged("api_key", k2_hash, start, end),
|
||||
)
|
||||
_assert_report_matches_db(
|
||||
_rows(gateway, "/user/spend/report", params={**window, "internal_user_id": user2}),
|
||||
_logged('"user"', user2, start, end),
|
||||
)
|
||||
_assert_report_matches_db(
|
||||
_rows(gateway, "/team/spend/report", params={**window, "team_id": team2}),
|
||||
_logged("team_id", team2, start, end),
|
||||
)
|
||||
org_report: Final = _rows(gateway, "/organization/spend/report", params={**window, "organization_id": org2})
|
||||
_assert_report_matches_db(org_report, _logged_keys((k2_hash, direct_org_key_hash), start, end))
|
||||
assert {detail.team_id for entry in org_report for detail in entry.model_details} == {team2, ""}, org_report
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"path",
|
||||
(
|
||||
"/key/spend/report",
|
||||
"/user/spend/report",
|
||||
"/team/spend/report",
|
||||
"/organization/spend/report",
|
||||
),
|
||||
)
|
||||
@pytest.mark.parametrize(
|
||||
("date_kind", "params", "detail"),
|
||||
(
|
||||
("missing", None, "Please provide start_date and end_date"),
|
||||
(
|
||||
"malformed",
|
||||
{"start_date": "2026-13-45", "end_date": "2026-01-01"},
|
||||
"start_date and end_date must be in YYYY-MM-DD format",
|
||||
),
|
||||
),
|
||||
)
|
||||
def test_spend_report_rejects_missing_and_malformed_dates(
|
||||
gateway: Gateway, path: str, date_kind: str, params: Mapping[str, str] | None, detail: str
|
||||
) -> None:
|
||||
response: Final = _report(gateway, path, params=params)
|
||||
assert response.status_code == 400, response.text
|
||||
assert response.json() == {"detail": detail}, response.text
|
||||
|
|
@ -3,13 +3,15 @@ import uuid
|
|||
from collections.abc import AsyncIterator
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timedelta, timezone
|
||||
from hashlib import sha256
|
||||
from typing import Final
|
||||
from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
|
||||
|
||||
import psycopg
|
||||
import pytest
|
||||
import pytest_asyncio
|
||||
from integration._support.client import Gateway
|
||||
from integration._support.client import Gateway, eventually, object_value, string_value
|
||||
from integration._support.database import read_rows
|
||||
from prisma import Prisma
|
||||
from psycopg import sql
|
||||
from psycopg.types.json import Jsonb
|
||||
|
|
@ -224,3 +226,91 @@ def test_spend_log_routes_preserve_user_and_permitted_team_access(
|
|||
admin: Final = gateway.request("GET", f"/spend/logs/ui/{rows[2].request_id}")
|
||||
assert admin.status_code == 200, admin.text
|
||||
assert admin.json()["messages"] == [{"role": "user", "content": rows[2].request_id + " payload"}], admin.text
|
||||
|
||||
|
||||
def _chat_request_id(gateway: Gateway, model: str, key: str) -> str:
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": model, "max_tokens": 20, "messages": [{"role": "user", "content": f"scope {uuid.uuid4().hex}"}]},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
body: Final = TypeAdapter(dict[str, object]).validate_python(response.json())
|
||||
return str(body["id"])
|
||||
|
||||
|
||||
def _legacy_spend_rows(gateway: Gateway, key: str, params: dict[str, str]) -> list[dict[str, object]]:
|
||||
response: Final = gateway.request("GET", "/spend/logs", key=key, params=params)
|
||||
assert response.status_code == 200, response.text
|
||||
return TypeAdapter(list[dict[str, object]]).validate_python(response.json())
|
||||
|
||||
|
||||
def test_legacy_spend_logs_clamp_internal_user_filters_to_their_own_user(gateway: Gateway) -> None:
|
||||
with gateway.scenario() as scenario:
|
||||
model: Final = scenario.model(input_cost_per_token=0.001, output_cost_per_token=0.002)
|
||||
user_a: Final = scenario.user(user_role="internal_user")
|
||||
user_b: Final = scenario.user(user_role="internal_user")
|
||||
key_a: Final = scenario.key(user_id=user_a, models=[model])
|
||||
key_b: Final = scenario.key(user_id=user_b, models=[model])
|
||||
request_a: Final = _chat_request_id(gateway, model, key_a)
|
||||
request_b: Final = _chat_request_id(gateway, model, key_b)
|
||||
for request_id in (request_a, request_b):
|
||||
eventually(
|
||||
lambda request_id=request_id: read_rows(
|
||||
'SELECT request_id, "user" FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (request_id,)
|
||||
),
|
||||
lambda rows: len(rows) == 1,
|
||||
seconds=70,
|
||||
)
|
||||
|
||||
filtered: Final = _legacy_spend_rows(gateway, key_a, {"user_id": user_b})
|
||||
assert filtered != [], filtered
|
||||
assert {string_value(object_value(row)["user"]) for row in filtered} == {user_a}, filtered
|
||||
assert request_b not in {string_value(object_value(row)["request_id"]) for row in filtered}
|
||||
assert request_a in {string_value(object_value(row)["request_id"]) for row in filtered}
|
||||
|
||||
by_key: Final = _legacy_spend_rows(gateway, key_a, {"api_key": key_b})
|
||||
assert by_key == [], by_key
|
||||
hashed_key_b: Final = sha256(key_b.encode()).hexdigest()
|
||||
by_hashed_key: Final = _legacy_spend_rows(gateway, key_a, {"api_key": hashed_key_b})
|
||||
assert by_hashed_key == [], by_hashed_key
|
||||
|
||||
by_request: Final = _legacy_spend_rows(gateway, key_a, {"request_id": request_b})
|
||||
assert by_request == [], by_request
|
||||
|
||||
now: Final = datetime.now(timezone.utc)
|
||||
window: Final = {
|
||||
"start_date": (now - timedelta(days=1)).strftime("%Y-%m-%d"),
|
||||
"end_date": (now + timedelta(days=1)).strftime("%Y-%m-%d"),
|
||||
}
|
||||
summarized: Final = _legacy_spend_rows(gateway, key_a, {**window, "summarize": "true"})
|
||||
assert user_b not in {name for row in summarized for name in object_value(object_value(row)["users"])}, (
|
||||
summarized
|
||||
)
|
||||
assert sum(float(str(object_value(row)["spend"])) for row in summarized) == pytest.approx(0.06), summarized
|
||||
unsummarized: Final = _legacy_spend_rows(gateway, key_a, {**window, "summarize": "false"})
|
||||
assert unsummarized != [], unsummarized
|
||||
assert {string_value(object_value(row)["user"]) for row in unsummarized} == {user_a}, unsummarized
|
||||
|
||||
master: Final = _legacy_spend_rows(gateway, gateway.key, {"user_id": user_b, "request_id": request_b})
|
||||
assert [string_value(object_value(row)["request_id"]) for row in master] == [request_b], master
|
||||
for api_key in (key_b, hashed_key_b):
|
||||
admin_by_key = _legacy_spend_rows(gateway, gateway.key, {"api_key": api_key})
|
||||
assert [
|
||||
(string_value(object_value(row)["request_id"]), object_value(row)["user"]) for row in admin_by_key
|
||||
] == [(request_b, user_b)], admin_by_key
|
||||
summarized_by_key: Final = _legacy_spend_rows(gateway, key_a, {**window, "summarize": "true", "api_key": key_b})
|
||||
assert sum(float(str(object_value(row)["spend"])) for row in summarized_by_key) == 0, summarized_by_key
|
||||
admin_summarized: Final = _legacy_spend_rows(
|
||||
gateway, gateway.key, {**window, "summarize": "true", "api_key": key_b}
|
||||
)
|
||||
assert sum(float(str(object_value(row)["spend"])) for row in admin_summarized) == pytest.approx(0.06), (
|
||||
admin_summarized
|
||||
)
|
||||
admin_users: Final = {
|
||||
name: float(str(spend))
|
||||
for row in admin_summarized
|
||||
for name, spend in object_value(object_value(row)["users"]).items()
|
||||
}
|
||||
assert admin_users == pytest.approx({user_b: 0.06}), admin_summarized
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue