scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207
litellm/tests/unit/accounting/test_projection.py
joshua-berri 894552205b
feat(spend): report guardrail spend separately in daily activity (#51010) 
* feat(spend): report guardrail spend separately in daily activity

Model and key spend in /user/daily/activity/aggregated include the cost of
guardrails like Bedrock guardrails, and nothing in the response says how much.
A customer that bills tenants per key and per model could not get model cost
alone, so tenants that pay AWS for their own guardrails were charged for them
twice

Daily spend rows now carry a guardrail_spend column next to spend, and every
breakdown in the daily activity responses exposes it, so model cost is
spend - guardrail_spend. spend keeps its meaning, so budgets still count
guardrail cost as before

The value comes from the request's cost_breakdown.guardrail_cost, capped at
the row's spend because an allowance discount or a repeat job poll can charge
less than the guardrail cost the breakdown records. The delivery ledger
freezes it into the accounting payload next to accounting_savings, since the
stored payload does not keep cost_breakdown, and the legacy daily queue writes
the same value

The migration only adds the column with a 0 default, so rows written before it
read 0

* fix(spend): count guardrail spend on rows billed without a cost breakdown

A request blocked before the LLM call is billed for its guardrails straight from
guardrail_information, so it carries no cost_breakdown and its guardrail_spend
read 0. Both writers now price guardrail_spend from guardrail_information with
guardrail_information_cost, the same function success and failure rows use to add
guardrail cost to spend

Also fills guardrail_spend in the client usage view built from request logs,
qualifies the field description for rows recorded before it existed, and adds
the lock timeout the migration checks require

* fix(ui): avoid a nested ternary when reading guardrail entries from request logs

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
2026-11-10 13:02:26 -08:00

579 lines
24 KiB
Python
Raw Permalink Blame History

import json
from collections.abc import Mapping
from datetime import datetime, timezone
from typing import Final
import pytest
from pydantic import TypeAdapter
from litellm.accounting.projection import project
from litellm.accounting.snapshot import build_accounting_payload
from litellm.constants import BILLED_RETRY_ATTEMPT_METADATA_KEY, INTERNAL_CALL_ORIGIN_METADATA_KEY
from litellm.proxy._types import SpendLogsPayload
from litellm.proxy.db.db_spend_update_writer import DBSpendUpdateWriter
from litellm.proxy.proxy_server import proxy_logging_obj
from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload
from litellm.proxy.utils import PrismaClient
from litellm.types.utils import ModelResponse, Usage
_NOW: Final = datetime(2026, 11, 2, tzinfo=timezone.utc)
_FIELDS: Final = TypeAdapter(dict[str, object])
@pytest.mark.parametrize("status", ["success", "failure"])
@pytest.mark.parametrize("origin", [None, "auto_router", "billed_retry"])
@pytest.mark.parametrize("spend", [0.0, 0.25])
@pytest.mark.parametrize(
("guardrail_information", "billed_guardrail_cost"),
[
(None, 0.0),
([{"guardrail_name": "guard", "guardrail_cost": 0.1}], 0.1),
([{"guardrail_name": "guard", "guardrail_cost": 0.5}], 0.5),
([{"guardrail_name": "guard", "guardrail_cost": 0.1, "guardrail_cost_in_spend": False}], 0.0),
],
)
@pytest.mark.asyncio
async def test_daily_projection_matches_legacy_writer(
status: str,
origin: str | None,
spend: float,
guardrail_information: list[dict[str, object]] | None,
billed_guardrail_cost: float,
) -> None:
payload: Final = get_logging_payload(
kwargs={
"model": "projection-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "projection-request",
"litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user", "status": status}},
},
response_obj=ModelResponse(
id="projection-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=_NOW,
end_time=_NOW,
)
metadata: Final = {
**_FIELDS.validate_json(payload["metadata"]),
**({BILLED_RETRY_ATTEMPT_METADATA_KEY: True} if origin == "billed_retry" else {}),
**({INTERNAL_CALL_ORIGIN_METADATA_KEY: origin} if origin not in (None, "billed_retry") else {}),
**({"guardrail_information": guardrail_information} if guardrail_information is not None else {}),
}
row: Final = {**payload, "spend": spend, "metadata": json.dumps(metadata), "request_duration_ms": 25}
legacy: Final = await DBSpendUpdateWriter()._common_add_spend_log_transaction_to_daily_transaction(
row, PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj), "user"
)
assert legacy is not None
accounting_input: Final[SpendLogsPayload] = {**row}
accounting: Final = build_accounting_payload(
accounting_input,
entities={"api_key": payload["api_key"], "user": payload["user"]},
model_access_groups=(),
windows=(),
router=lambda: None,
)
restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str))
stored_daily: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend")
assert dict(stored_daily.increments) == {name: legacy[name] for name, _ in stored_daily.increments}
assert dict(stored_daily.increments)["guardrail_spend"] == min(billed_guardrail_cost, spend)
assert dict(stored_daily.key)["date"] == legacy["date"]
assert dict(stored_daily.key)["api_key"] == legacy["api_key"]
assert not {"messages", "response", "proxy_server_request", "requester_ip_address"}.intersection(accounting)
def test_projection_uses_frozen_savings_and_keeps_duplicate_entity_tags() -> None:
row: Final[Mapping[str, object]] = {
"request_id": "saved",
"startTime": _NOW,
"api_key": "key",
"user": "user",
"model": "unpriced-projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"request_tags": '["tag", "tag"]',
"spend": 0.5,
"metadata": json.dumps(
{
"accounting_savings": {
"compression": 1.0,
"prompt_caching": -2.0,
"gateway_injected_caching": 3.0,
"autorouter": 4.0,
},
"autorouter_savings": 9000,
}
),
}
contributions: Final = project(row)
tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_TagTable")
daily_tags: Final = tuple(item for item in contributions if item.table == "LiteLLM_DailyTagSpend")
assert sum(dict(item.increments)["spend"] for item in tags) == 1.0
assert len(daily_tags) == 1
assert {key: value for key, value in daily_tags[0].increments if key.endswith("savings_spend")} == {
"compression_savings_spend": 1.0,
"prompt_caching_savings_spend": -2.0,
"gateway_injected_caching_savings_spend": 3.0,
"autorouter_savings_spend": 4.0,
}
def test_projection_carries_windows_and_does_not_reprice_guardrails() -> None:
row: Final = {
"request_id": "guarded",
"startTime": _NOW,
"api_key": "key",
"spend": 0.75,
"metadata": json.dumps(
{
"accounting_windows": [
{
"entity_type": "key",
"entity_id": "key",
"window_duration": "1d",
"window_start": _NOW.isoformat(),
}
],
"guardrail_information": [
{
"guardrail_id": "guard",
"guardrail_usage": {"text": 3, "images": 2},
"guardrail_cost_by_unit": {"text": 0.3, "images": None},
}
],
}
),
}
contributions: Final = project(row)
window: Final = next(item for item in contributions if item.table == "LiteLLM_BudgetWindowSpend")
units: Final = {
dict(item.key)["usage_unit"]: dict(item.increments)
for item in contributions
if item.table == "LiteLLM_DailyGuardrailUsageUnits"
}
assert dict(window.increments) == {"spend": 0.75}
assert dict(window.values)["window_start"] == _NOW.isoformat()
assert units == {
"text": {"units": 3, "cost": 0.3, "untracked_units": 0},
"images": {"units": 2, "cost": 0.0, "untracked_units": 2},
}
@pytest.mark.asyncio
async def test_entity_projection_matches_real_request_queues(monkeypatch: pytest.MonkeyPatch) -> None:
from litellm.proxy import proxy_server
from litellm.responses.streaming_iterator import drain_logging_tasks
client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj)
monkeypatch.setattr(proxy_server, "prisma_client", client)
monkeypatch.setattr(proxy_server, "disable_spend_logs", False)
monkeypatch.setitem(proxy_server.general_settings, "use_spend_logs_ledger", False)
monkeypatch.setattr(client, "spend_log_transactions", [])
monkeypatch.setattr(client, "model_usage_transactions", [])
writer: Final = DBSpendUpdateWriter()
charged: Final = await writer.update_database(
token="sk-projection",
user_id="user",
end_user_id="end-user",
team_id="team",
org_id="org",
project_id="project",
kwargs={
"model": "projection-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "entity-request",
"agent_id": "agent",
"user": "end-user",
"litellm_params": {
"metadata": {
"user_api_key": "sk-projection",
"user_api_key_user_id": "user",
"user_api_key_team_id": "team",
"user_api_key_org_id": "org",
"user_api_key_project_id": "project",
"user_api_key_end_user_id": "end-user",
"tags": ["tag", "tag"],
},
"user": "end-user",
},
},
completion_response=ModelResponse(
id="entity-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=_NOW,
end_time=_NOW,
response_cost=0.25,
)
await drain_logging_tasks()
assert charged
assert len(client.spend_log_transactions) == 1
queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue()
logged: Final[SpendLogsPayload] = client.spend_log_transactions[0]
projected: Final = project(
_FIELDS.validate_json(
json.dumps(
build_accounting_payload(
logged,
entities={
"api_key": logged["api_key"],
"user": "user",
"end_user": "end-user",
"team_id": "team",
"organization_id": "org",
"project_id": "project",
"agent_id": "agent",
},
model_access_groups=(),
windows=(),
router=lambda: None,
),
default=str,
)
)
)
legacy: Final = {(item["entity_type"].value, item["entity_id"]): item["response_cost"] for item in queued}
assert legacy[("key", proxy_server.hash_token("sk-projection"))] == next(
dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_VerificationToken"
)
for table, entity, identity in (
("LiteLLM_UserTable", "user", "user"),
("LiteLLM_TeamTable", "team", "team"),
("LiteLLM_OrganizationTable", "organization", "org"),
("LiteLLM_EndUserTable", "end_user", "end-user"),
("LiteLLM_AgentsTable", "agent", "agent"),
("LiteLLM_ProjectTable", "project", "project"),
):
assert legacy[(entity, identity)] == next(
dict(item.increments)["spend"] for item in projected if item.table == table
)
assert sum(dict(item.increments)["spend"] for item in projected if item.table == "LiteLLM_TagTable") == sum(
item["response_cost"] or 0 for item in queued if item["entity_type"].value == "tag"
)
daily_queues: Final = (
("LiteLLM_DailyUserSpend", writer.daily_spend_update_queue),
("LiteLLM_DailyTeamSpend", writer.daily_team_spend_update_queue),
("LiteLLM_DailyOrganizationSpend", writer.daily_org_spend_update_queue),
("LiteLLM_DailyEndUserSpend", writer.daily_end_user_spend_update_queue),
("LiteLLM_DailyAgentSpend", writer.daily_agent_spend_update_queue),
("LiteLLM_DailyTagSpend", writer.daily_tag_spend_update_queue),
)
for table, queue in daily_queues:
transactions = await queue.flush_and_get_aggregated_daily_spend_update_transactions()
matching = tuple(item for item in projected if item.table == table)
assert matching, table
contribution = matching[0]
assert len(transactions) == 1
transaction = next(iter(transactions.values()))
assert dict(contribution.increments) == {name: transaction[name] for name, _ in contribution.increments}
assert len(client.model_usage_transactions) == 1
model: Final = client.model_usage_transactions[0]
model_projection: Final = next(item for item in projected if item.table == "LiteLLM_DailyModelUsage")
assert dict(model_projection.increments) == {
"spend": model.spend,
"prompt_tokens": model.prompt_tokens,
"completion_tokens": model.completion_tokens,
"request_count": int(model.counts_as_request),
"successful_requests": int(model.successful),
"failed_requests": int(not model.successful),
}
@pytest.mark.parametrize("served_id", ["requested", "fallback", "missing"])
@pytest.mark.asyncio
async def test_projection_preserves_fallback_narrowed_model_access_groups(served_id: str) -> None:
from litellm import Router
from litellm.proxy.db.db_spend_update_writer import debitable_model_access_groups
router: Final = Router(
model_list=[
{
"model_name": "first",
"litellm_params": {"model": "openai/projection-test", "api_key": "unused"},
"model_info": {"id": "requested", "access_groups": ["first-group"]},
},
{
"model_name": "second",
"litellm_params": {"model": "openai/projection-test", "api_key": "unused"},
"model_info": {"id": "fallback", "access_groups": ["second-group", "unattributed"]},
},
]
)
attributed: Final = ("first-group", "second-group", "first-group")
writer: Final = DBSpendUpdateWriter()
client: Final = PrismaClient(database_url="postgresql://unused", proxy_logging_obj=proxy_logging_obj)
await writer._update_model_access_group_db(0.5, attributed, served_id, client, router)
queued: Final = await writer.spend_update_queue.flush_all_updates_from_in_memory_queue()
frozen_groups: Final = debitable_model_access_groups(attributed, served_id, router)
projected: Final = project(
{"spend": 0.5, "metadata": json.dumps({"accounting_model_access_groups": frozen_groups})}
)
assert {dict(item.key)["access_group_name"]: dict(item.increments)["spend"] for item in projected} == {
item["entity_id"]: item["response_cost"] for item in queued
}
assert "unattributed" not in frozen_groups
def test_guardrail_projection_skips_unmeasured_counters() -> None:
row: Final = {
"request_id": "guarded",
"startTime": _NOW,
"metadata": json.dumps(
{
"guardrail_information": [
{"guardrail_usage": {"text": 3}},
{
"guardrail_id": "guard",
"guardrail_usage": {"zero": 0, "negative": -1, "bool": True, "unknown": None, "text": 2},
},
]
}
),
}
units: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyGuardrailUsageUnits")
assert len(units) == 1
assert dict(units[0].key)["usage_unit"] == "text"
assert dict(units[0].increments) == {"units": 2, "cost": 0.0, "untracked_units": 2}
def test_lean_row_without_model_does_not_fabricate_daily_entity_usage() -> None:
projected: Final = project(
{
"request_id": "lean",
"api_key": "key",
"startTime": _NOW,
"spend": 0.5,
"metadata": {"accounting_entities": {"api_key": "key"}},
}
)
assert tuple(item.table for item in projected) == ("LiteLLM_VerificationToken", "LiteLLM_DailyModelUsage")
@pytest.mark.parametrize(
"cost_fields",
[
{"guardrail_cost_in_spend": False, "guardrail_cost_by_unit": {"text": 1.0}},
{"guardrail_cost_by_unit": {"text": -1.0}},
],
)
def test_report_only_or_invalid_guardrail_prices_remain_untracked(cost_fields: Mapping[str, object]) -> None:
row: Final = {
"request_id": "guarded",
"startTime": _NOW,
"request_tags": '["task:debugging"]',
"metadata": json.dumps(
{
"guardrail_information": [
{"guardrail_id": "guard", "guardrail_usage": {"text": 2}, **cost_fields},
]
}
),
}
contributions: Final = project(row)
guardrail: Final = next(item for item in contributions if item.table == "LiteLLM_DailyGuardrailUsageUnits")
model: Final = next(item for item in contributions if item.table == "LiteLLM_DailyModelUsage")
assert dict(guardrail.increments) == {"units": 2, "cost": 0.0, "untracked_units": 2}
assert dict(model.key)["task_type"] == "debugging"
@pytest.mark.parametrize("organization_id", [None, "", "charged-org"])
def test_daily_organization_uses_frozen_charged_identity(organization_id: str | None) -> None:
row: Final[Mapping[str, object]] = {
"startTime": _NOW,
"api_key": "key",
"model": "projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"organization_id": "stored-org",
"spend": 0.25,
"metadata": {"accounting_entities": {"organization_id": organization_id}},
}
daily: Final = tuple(item for item in project(row) if item.table == "LiteLLM_DailyOrganizationSpend")
if organization_id is None:
assert daily == ()
else:
assert len(daily) == 1
assert dict(daily[0].key)["organization_id"] == organization_id
assert dict(daily[0].increments)["spend"] == 0.25
@pytest.mark.parametrize("start", [_NOW, "2026-11-02T12:00:00+00:00", "2026-11-02 12:00:00+00:00"])
def test_every_daily_row_uses_the_calendar_date_of_the_request(start: object) -> None:
row: Final[Mapping[str, object]] = {
"request_id": "dated",
"startTime": start,
"api_key": "key",
"user": "user",
"team_id": "team",
"model": "unpriced-projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"spend": 0.5,
"metadata": "{}",
}
dates: Final = {
item.table: dict(item.key)["date"] for item in project(row) if item.table.startswith("LiteLLM_Daily")
}
assert {"LiteLLM_DailyUserSpend", "LiteLLM_DailyModelUsage"} <= dates.keys()
assert set(dates.values()) == {"2026-11-02"}
def test_unparseable_start_time_skips_only_guardrail_usage() -> None:
projected: Final = project(
{
"request_id": "bad-time",
"startTime": "yesterday at noon",
"spend": 0.5,
"metadata": {
"accounting_entities": {"api_key": "key"},
"guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 2}}],
},
}
)
tables: Final = {item.table for item in projected}
assert "LiteLLM_VerificationToken" in tables
assert "LiteLLM_DailyGuardrailUsageUnits" not in tables
def test_guardrail_usage_is_dated_in_utc() -> None:
projected: Final = project(
{
"request_id": "offset",
"startTime": "2026-11-02T22:30:00-05:00",
"metadata": {"guardrail_information": [{"guardrail_id": "guard", "guardrail_usage": {"text": 1}}]},
}
)
assert [dict(item.key)["date"] for item in projected if item.table == "LiteLLM_DailyGuardrailUsageUnits"] == [
"2026-11-03"
]
@pytest.mark.parametrize(
("untrusted", "metadata"),
[
({"request_tags": json.dumps(["x" * 16384, "normal"])}, {}),
({"end_user": "x" * 3072}, {}),
({}, {"compression_savings": {"tokens_saved": 2**63}}),
],
)
def test_untrusted_analytics_values_never_block_the_charge(
untrusted: Mapping[str, str], metadata: Mapping[str, object]
) -> None:
row: Final[Mapping[str, object]] = {
"request_id": "untrusted",
"startTime": _NOW,
"api_key": "key",
"user": "user",
"team_id": "team",
"model": "unpriced-projection-test",
"model_group": "group",
"custom_llm_provider": "openai",
"spend": 0.5,
**untrusted,
"metadata": {
**metadata,
"accounting_entities": {
"api_key": "key",
"user": "user",
"team_id": "team",
"end_user": untrusted.get("end_user"),
},
},
}
contributions: Final = project(row)
texts: Final = tuple(value for item in contributions for _, value in (*item.key, *item.values) if value)
numbers: Final = tuple(value for item in contributions for _, value in item.increments)
assert {"LiteLLM_VerificationToken", "LiteLLM_TeamTable", "LiteLLM_UserTable"} <= {
item.table for item in contributions
}
assert "LiteLLM_DailyUserSpend" in {item.table for item in contributions}
assert all(len(text.encode()) <= 2000 for text in texts)
assert all(abs(number) <= 2**40 for number in numbers)
def test_stored_payload_keeps_guardrail_billing_but_not_guardrail_content() -> None:
payload: Final = get_logging_payload(
kwargs={
"model": "projection-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "guardrail-request",
"litellm_params": {"metadata": {"user_api_key": "key", "user_api_key_user_id": "user"}},
},
response_obj=ModelResponse(
id="guardrail-request", model="projection-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=_NOW,
end_time=_NOW,
)
guardrail: Final = {
"guardrail_name": "pii",
"guardrail_usage": {"requests": 1},
"guardrail_cost_by_unit": {"requests": 0.25},
"guardrail_response": "card 4111 1111 1111 1111",
"match_details": [{"matched": "4111 1111 1111 1111"}],
}
headroom: Final = {
"guardrail_name": "compress",
"guardrail_provider": "headroom",
"guardrail_response": {"tokens_saved": 7, "compressed_prompt": "card 4111 1111 1111 1111"},
}
metadata: Final = {**_FIELDS.validate_json(payload["metadata"]), "guardrail_information": [guardrail, headroom]}
accounting_input: Final[SpendLogsPayload] = {
**payload,
"spend": 0.5,
"metadata": json.dumps(metadata),
"request_tags": '["nul\x00tag"]',
}
accounting: Final = build_accounting_payload(
accounting_input, entities={"api_key": "key"}, model_access_groups=(), windows=(), router=lambda: None
)
stored: Final = json.dumps(accounting, default=str)
restored: Final = _FIELDS.validate_json(stored)
usage_rows: Final = tuple(item for item in project(restored) if item.table == "LiteLLM_DailyGuardrailUsageUnits")
assert "4111" not in stored
assert "\\u0000" not in stored
assert [dict(item.increments) for item in usage_rows] == [{"units": 1, "cost": 0.25, "untracked_units": 0}]
daily_user: Final = next(item for item in project(restored) if item.table == "LiteLLM_DailyUserSpend")
assert dict(daily_user.increments)["compression_saved_tokens"] == 7
def test_caller_supplied_savings_beyond_json_range_are_clamped_and_never_drop_the_charge(
monkeypatch: pytest.MonkeyPatch,
) -> None:
import litellm
monkeypatch.setitem(
litellm.model_cost,
"savings-overflow-test",
{"input_cost_per_token": 0.002, "output_cost_per_token": 0.004, "litellm_provider": "openai", "mode": "chat"},
)
payload: Final[SpendLogsPayload] = {
**get_logging_payload(
kwargs={
"model": "savings-overflow-test",
"custom_llm_provider": "openai",
"call_type": "acompletion",
"litellm_call_id": "overflow-request",
"litellm_params": {"metadata": {"user_api_key": "key"}},
},
response_obj=ModelResponse(
id="overflow-request", model="savings-overflow-test", usage=Usage(prompt_tokens=10, completion_tokens=5)
),
start_time=_NOW,
end_time=_NOW,
),
"spend": 0.5,
"metadata": json.dumps({"compression_savings": {"tokens_saved": 10**400}}),
}
accounting: Final = build_accounting_payload(
payload, entities={"api_key": "key"}, model_access_groups=(), windows=(), router=lambda: None
)
restored: Final = _FIELDS.validate_json(json.dumps(accounting, default=str))
assert dict(restored["metadata"])["accounting_savings"]["compression"] == 0.0
assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1}
assert dict(restored["metadata"])["compression_savings"] == {"tokens_saved": 2**53 - 1}
assert [
dict(item.increments)["spend"] for item in project(restored) if item.table == "LiteLLM_VerificationToken"
] == [0.5]
Reference in a new issue View git blame Copy permalink