mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
35176fa64a
The daily guardrail metrics and usage-unit upserts are non-idempotent increments, but the retry loop re-sent every failed row on any exception. An ambiguous post-send failure such as a read timeout after the write had already committed therefore stacked a second increment and inflated the billable unit totals served by the guardrail usage endpoints. Retry only DB_RETRY_SAFE_ERROR_TYPES (httpx.ConnectError), the same rule the spend writer and autorouter rollup use for increment upserts, and log any other failure once as terminal for that row while the rest of the batch still lands. Follows up #37225
243 lines
8.9 KiB
Python
243 lines
8.9 KiB
Python
import json
|
|
from datetime import datetime, timezone
|
|
from typing import Any
|
|
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import httpx
|
|
import pytest
|
|
|
|
from litellm.proxy.guardrails.usage_tracking import process_spend_logs_guardrail_usage
|
|
|
|
|
|
def _prisma() -> MagicMock:
|
|
client = MagicMock()
|
|
db = client.db
|
|
db.litellm_dailyguardrailmetrics.upsert = AsyncMock()
|
|
db.litellm_dailyguardrailusageunits.upsert = AsyncMock()
|
|
db.litellm_spendlogguardrailindex.create_many = AsyncMock()
|
|
return client
|
|
|
|
|
|
def _payload(
|
|
request_id: str,
|
|
*,
|
|
team_id: str | None = "team-a",
|
|
api_key: str = "hashed-key-1",
|
|
usage: dict[str, Any] | None = None,
|
|
guardrail_status: str = "success",
|
|
) -> dict[str, Any]:
|
|
entry: dict[str, Any] = {
|
|
"guardrail_id": "bedrock-guard",
|
|
"guardrail_status": guardrail_status,
|
|
}
|
|
if usage is not None:
|
|
entry["guardrail_usage"] = usage
|
|
return {
|
|
"request_id": request_id,
|
|
"startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc),
|
|
"team_id": team_id,
|
|
"api_key": api_key,
|
|
"metadata": json.dumps({"guardrail_information": [entry]}),
|
|
}
|
|
|
|
|
|
def _units_upserts(prisma: MagicMock) -> dict[tuple, int]:
|
|
calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list
|
|
out: dict[tuple, int] = {}
|
|
for c in calls:
|
|
where = c.kwargs["where"]["guardrail_id_date_team_id_api_key_usage_unit"]
|
|
create = c.kwargs["data"]["create"]
|
|
assert create["units"] == c.kwargs["data"]["update"]["units"]["increment"]
|
|
assert {k: create[k] for k in where} == where
|
|
out[tuple(where[k] for k in ("guardrail_id", "date", "team_id", "api_key", "usage_unit"))] = create["units"]
|
|
return out
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_usage_units_rolled_up_by_guardrail_team_key_and_date():
|
|
"""
|
|
LIT-5650: billable units must aggregate per (guardrail, date, team, key,
|
|
counter): same-key payloads sum into one upsert, a team-less payload gets
|
|
its own empty-string-team row, and blocked invocations (which Bedrock
|
|
still bills for) count exactly like passed ones.
|
|
"""
|
|
prisma = _prisma()
|
|
logs = [
|
|
_payload("r1", usage={"topicPolicyUnits": 1, "contentPolicyUnits": 1}),
|
|
_payload(
|
|
"r2",
|
|
usage={"topicPolicyUnits": 1, "contentPolicyUnits": 2},
|
|
guardrail_status="guardrail_intervened",
|
|
),
|
|
_payload("r3", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 2,
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3,
|
|
("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits"): 1,
|
|
}
|
|
|
|
|
|
def _fake_sleep() -> tuple[AsyncMock, list[float]]:
|
|
delays: list[float] = []
|
|
sleep = AsyncMock(side_effect=lambda delay: delays.append(delay))
|
|
return sleep, delays
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_failing_upsert_does_not_drop_remaining_writes():
|
|
"""
|
|
A DB error on one daily-metrics or usage-unit upsert must not cancel the
|
|
remaining upserts in the flushed batch, or the usage endpoints would
|
|
permanently under-report billable counters.
|
|
"""
|
|
prisma = _prisma()
|
|
prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
|
|
prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [httpx.ConnectError("db down"), None, None]
|
|
sleep, _ = _fake_sleep()
|
|
logs = [
|
|
_payload("r1", usage={"topicPolicyUnits": 1}),
|
|
_payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep)
|
|
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1,
|
|
("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits"): 1,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_transient_upsert_failure_is_retried_with_backoff_for_failed_rows_only():
|
|
"""
|
|
A connection error (the write provably never reached the database) must
|
|
not permanently drop billed units from the aggregates: only the rows that
|
|
failed are re-sent, after exponential backoff, and the batch ends once
|
|
every row has landed.
|
|
"""
|
|
prisma = _prisma()
|
|
prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [httpx.ConnectError("blip"), None, None]
|
|
sleep, delays = _fake_sleep()
|
|
logs = [
|
|
_payload("r1", usage={"topicPolicyUnits": 1}),
|
|
_payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep)
|
|
|
|
calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list
|
|
assert len(calls) == 3
|
|
assert calls[2].kwargs["where"] == calls[0].kwargs["where"]
|
|
assert delays == [1]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_persistent_upsert_failure_stops_after_three_retries():
|
|
prisma = _prisma()
|
|
prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
|
|
sleep, delays = _fake_sleep()
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep)
|
|
|
|
assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 4
|
|
assert delays == [1, 2, 4]
|
|
assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1
|
|
|
|
|
|
def _units_upsert_wheres(prisma: MagicMock) -> list[tuple]:
|
|
return [
|
|
tuple(
|
|
c.kwargs["where"]["guardrail_id_date_team_id_api_key_usage_unit"][k]
|
|
for k in ("guardrail_id", "date", "team_id", "api_key", "usage_unit")
|
|
)
|
|
for c in prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_post_send_failure_is_never_retried_so_increments_cannot_double_count():
|
|
"""
|
|
Follow-up to #37225: the units upsert is a non-idempotent increment, so an
|
|
ambiguous post-send failure (read timeout after the statement may have
|
|
committed) must be attempted exactly once. Re-sending it stacks a second
|
|
increment and inflates billable unit totals. Only a connection error proves
|
|
the write never reached the database and may be retried; the other rows in
|
|
the batch still land either way.
|
|
"""
|
|
prisma = _prisma()
|
|
prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = [
|
|
httpx.ReadTimeout("read timed out"),
|
|
httpx.ConnectError("refused"),
|
|
None,
|
|
]
|
|
sleep, delays = _fake_sleep()
|
|
logs = [
|
|
_payload("r1", usage={"topicPolicyUnits": 1}),
|
|
_payload("r2", team_id=None, api_key="hashed-key-2", usage={"topicPolicyUnits": 1}),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs, sleep=sleep)
|
|
|
|
timed_out_row = ("bedrock-guard", "2026-08-17", "", "hashed-key-2", "topicPolicyUnits")
|
|
refused_row = ("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits")
|
|
assert _units_upsert_wheres(prisma) == [timed_out_row, refused_row, refused_row]
|
|
assert delays == [1]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_generic_upsert_exception_is_terminal_for_that_row_only():
|
|
prisma = _prisma()
|
|
prisma.db.litellm_dailyguardrailmetrics.upsert.side_effect = RuntimeError("constraint violation")
|
|
sleep, delays = _fake_sleep()
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep)
|
|
|
|
assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 1
|
|
assert delays == []
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_zero_and_non_int_usage_counters_are_skipped():
|
|
prisma = _prisma()
|
|
logs = [
|
|
_payload(
|
|
"r1",
|
|
usage={
|
|
"topicPolicyUnits": 1,
|
|
"wordPolicyUnits": 0,
|
|
"contentPolicyImageUnits": 0,
|
|
"oddball": "not-an-int",
|
|
"boolish": True,
|
|
},
|
|
),
|
|
_payload("r2", usage=None),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_payload_without_request_id_is_skipped_like_the_metrics_path():
|
|
prisma = _prisma()
|
|
logs = [
|
|
{**_payload("ignored", usage={"topicPolicyUnits": 5}), "request_id": None},
|
|
_payload("r2", usage={"topicPolicyUnits": 1}),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 1,
|
|
}
|
|
assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]["requests_evaluated"] == 1
|