From 8822e086ec5650e0ae485a60db272a35c35f4196 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Wed, 14 Oct 2026 06:58:22 +0000 Subject: [PATCH] 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 Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../management/test_budget_updates.py | 184 ++++++++ .../pricing/test_configured_prices.py | 99 ++++ .../pricing/test_cost_margin_discount.py | 320 +++++++++++++ .../spend/test_end_user_budget_enforcement.py | 438 ++++++++++++++++++ .../spend/test_global_spend_reset.py | 159 +++++++ .../spend/test_key_budget_lockout.py | 186 +++++++- .../spend/test_scoped_spend_reports.py | 228 +++++++++ .../spend/test_spend_log_read_scope.py | 92 +++- 8 files changed, 1703 insertions(+), 3 deletions(-) create mode 100644 tests/integration/pricing/test_cost_margin_discount.py create mode 100644 tests/integration/spend/test_end_user_budget_enforcement.py create mode 100644 tests/integration/spend/test_global_spend_reset.py create mode 100644 tests/integration/spend/test_scoped_spend_reports.py diff --git a/tests/integration/management/test_budget_updates.py b/tests/integration/management/test_budget_updates.py index 25d525bcaaa..7a79e35d2af 100644 --- a/tests/integration/management/test_budget_updates.py +++ b/tests/integration/management/test_budget_updates.py @@ -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} diff --git a/tests/integration/pricing/test_configured_prices.py b/tests/integration/pricing/test_configured_prices.py index 93290a404fc..b7eaadc05f9 100644 --- a/tests/integration/pricing/test_configured_prices.py +++ b/tests/integration/pricing/test_configured_prices.py @@ -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 diff --git a/tests/integration/pricing/test_cost_margin_discount.py b/tests/integration/pricing/test_cost_margin_discount.py new file mode 100644 index 00000000000..bfd2002f871 --- /dev/null +++ b/tests/integration/pricing/test_cost_margin_discount.py @@ -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) diff --git a/tests/integration/spend/test_end_user_budget_enforcement.py b/tests/integration/spend/test_end_user_budget_enforcement.py new file mode 100644 index 00000000000..de0bfcd5194 --- /dev/null +++ b/tests/integration/spend/test_end_user_budget_enforcement.py @@ -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}, + ) diff --git a/tests/integration/spend/test_global_spend_reset.py b/tests/integration/spend/test_global_spend_reset.py new file mode 100644 index 00000000000..fad5c14175a --- /dev/null +++ b/tests/integration/spend/test_global_spend_reset.py @@ -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 diff --git a/tests/integration/spend/test_key_budget_lockout.py b/tests/integration/spend/test_key_budget_lockout.py index 02831cd59cd..7108978f6e5 100644 --- a/tests/integration/spend/test_key_budget_lockout.py +++ b/tests/integration/spend/test_key_budget_lockout.py @@ -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) == [] diff --git a/tests/integration/spend/test_scoped_spend_reports.py b/tests/integration/spend/test_scoped_spend_reports.py new file mode 100644 index 00000000000..229e5ac8a8b --- /dev/null +++ b/tests/integration/spend/test_scoped_spend_reports.py @@ -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 diff --git a/tests/integration/spend/test_spend_log_read_scope.py b/tests/integration/spend/test_spend_log_read_scope.py index f9034371e35..8a208d61aed 100644 --- a/tests/integration/spend/test_spend_log_read_scope.py +++ b/tests/integration/spend/test_spend_log_read_scope.py @@ -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