mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-19 00:01:29 +00:00
656 lines
26 KiB
Python
656 lines
26 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 import usage_tracking
|
|
from litellm.proxy.guardrails.usage_tracking import (
|
|
_MAX_PENDING_ROWS,
|
|
PendingRollups,
|
|
_capped,
|
|
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",
|
|
cost_by_unit: dict[str, Any] | None = None,
|
|
cost_in_spend: bool | None = None,
|
|
) -> dict[str, Any]:
|
|
entry: dict[str, Any] = {
|
|
"guardrail_id": "bedrock-guard",
|
|
"guardrail_status": guardrail_status,
|
|
}
|
|
if usage is not None:
|
|
entry["guardrail_usage"] = usage
|
|
if cost_by_unit is not None:
|
|
entry["guardrail_cost_by_unit"] = cost_by_unit
|
|
if cost_in_spend is not None:
|
|
entry["guardrail_cost_in_spend"] = cost_in_spend
|
|
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
|
|
|
|
|
|
def _cost_upserts(prisma: MagicMock) -> dict[str, tuple[float, int]]:
|
|
"""usage_unit -> (cost, untracked_units) written on create; the update path must increment by the same."""
|
|
calls = prisma.db.litellm_dailyguardrailusageunits.upsert.call_args_list
|
|
out: dict[str, tuple[float, int]] = {}
|
|
for c in calls:
|
|
create = c.kwargs["data"]["create"]
|
|
update = c.kwargs["data"]["update"]
|
|
assert update["cost"] == {"increment": create["cost"]}
|
|
assert update["untracked_units"] == {"increment": create["untracked_units"]}
|
|
out[create["usage_unit"]] = (create["cost"], create["untracked_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,
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_flagged_status_counts_as_flagged_not_passed_or_blocked():
|
|
"""LIT-6894: a custom code flag() verdict lands in flagged_count on the Monitor rollup."""
|
|
prisma = _prisma()
|
|
logs = [
|
|
_payload("r1", guardrail_status="success"),
|
|
_payload("r2", guardrail_status="guardrail_flagged"),
|
|
_payload("r3", guardrail_status="guardrail_intervened"),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]
|
|
assert (create["requests_evaluated"], create["passed_count"], create["flagged_count"], create["blocked_count"]) == (
|
|
3,
|
|
1,
|
|
1,
|
|
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, pending=PendingRollups())
|
|
|
|
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()
|
|
pending = PendingRollups()
|
|
|
|
await process_spend_logs_guardrail_usage(
|
|
prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep, pending=pending
|
|
)
|
|
|
|
assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 4
|
|
assert delays == [1, 2, 4]
|
|
assert prisma.db.litellm_dailyguardrailusageunits.upsert.call_count == 1
|
|
assert dict(pending.metrics) == {
|
|
("bedrock-guard", "2026-08-17"): {
|
|
"requests_evaluated": 1,
|
|
"passed_count": 1,
|
|
"blocked_count": 0,
|
|
"flagged_count": 0,
|
|
}
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_retry_exhausted_rows_are_requeued_and_land_on_the_next_flush():
|
|
"""
|
|
LIT-5761: rollup rows whose connection-error retries exhaust must not be
|
|
silently lost. They are requeued and merged into the next flushed batch,
|
|
so the aggregates catch up once the database is reachable again.
|
|
"""
|
|
pending = PendingRollups()
|
|
down = _prisma()
|
|
down.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
|
|
down.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ConnectError("db down")
|
|
sleep, _ = _fake_sleep()
|
|
|
|
await process_spend_logs_guardrail_usage(
|
|
down, [_payload("r1", usage={"topicPolicyUnits": 2})], sleep=sleep, pending=pending
|
|
)
|
|
|
|
assert dict(pending.units) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): (2, 0.0, 2)
|
|
}
|
|
|
|
recovered = _prisma()
|
|
await process_spend_logs_guardrail_usage(
|
|
recovered, [_payload("r2", usage={"topicPolicyUnits": 3})], sleep=sleep, pending=pending
|
|
)
|
|
|
|
assert _units_upserts(recovered) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 5,
|
|
}
|
|
metrics_create = recovered.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]
|
|
assert metrics_create["requests_evaluated"] == 2
|
|
assert not pending.units
|
|
assert not pending.metrics
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_ambiguous_failures_are_never_requeued():
|
|
"""
|
|
A post-send failure (the increment may have committed) must stay dropped:
|
|
requeueing it would re-send a possibly applied increment and double-count.
|
|
"""
|
|
pending = PendingRollups()
|
|
prisma = _prisma()
|
|
prisma.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ReadTimeout("maybe committed")
|
|
sleep, delays = _fake_sleep()
|
|
|
|
await process_spend_logs_guardrail_usage(
|
|
prisma, [_payload("r1", usage={"topicPolicyUnits": 1})], sleep=sleep, pending=pending
|
|
)
|
|
|
|
assert delays == []
|
|
assert not pending.units
|
|
assert not pending.metrics
|
|
|
|
|
|
def test_pending_requeue_is_capped_dropping_oldest_rows():
|
|
rows = {index: index for index in range(_MAX_PENDING_ROWS + 5)}
|
|
|
|
capped = _capped(rows, "usage unit")
|
|
|
|
assert len(capped) == _MAX_PENDING_ROWS
|
|
assert 4 not in capped
|
|
assert _MAX_PENDING_ROWS + 4 in capped
|
|
|
|
|
|
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_not_run_entries_are_indexed_but_not_counted_as_evaluations():
|
|
"""
|
|
LIT-6314 records a not_run entry when message scoping leaves a guardrail
|
|
nothing to scan. The guardrail never evaluated the request, so counting it
|
|
as a passed evaluation would inflate daily pass rates; it still gets an
|
|
index row so per-request drill-down finds the spend log.
|
|
"""
|
|
prisma = _prisma()
|
|
logs = [_payload("r1", guardrail_status="not_run"), _payload("r2")]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]
|
|
assert metrics_create["requests_evaluated"] == 1
|
|
assert metrics_create["passed_count"] == 1
|
|
index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"]
|
|
assert sorted(row["request_id"] for row in index_rows) == ["r1", "r2"]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_name():
|
|
"""
|
|
The not_run entry from the shared base guardrail carries only guardrail_name,
|
|
while the evaluated entry from the same guardrail (e.g. content filter on the
|
|
output of a logging_only run) carries its guardrail_id. Keying them differently
|
|
lists one request twice in the monitor, once as not_run and once as passed.
|
|
"""
|
|
prisma = _prisma()
|
|
payload = _payload("r1")
|
|
payload["metadata"] = json.dumps(
|
|
{
|
|
"guardrail_information": [
|
|
{"guardrail_name": "cf", "guardrail_status": "not_run"},
|
|
{
|
|
"guardrail_name": "cf",
|
|
"guardrail_id": "cf-uuid",
|
|
"policy_id": "pol-1",
|
|
"guardrail_status": "success",
|
|
},
|
|
{"guardrail_name": "other", "guardrail_status": "not_run"},
|
|
]
|
|
}
|
|
)
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, [payload])
|
|
|
|
index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"]
|
|
assert sorted((row["guardrail_id"], row["policy_id"]) for row in index_rows) == [
|
|
("cf-uuid", "pol-1"),
|
|
("other", None),
|
|
]
|
|
metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]
|
|
assert (metrics_create["guardrail_id"], metrics_create["requests_evaluated"]) == ("cf-uuid", 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_malformed_not_run_entry_does_not_drop_the_batch():
|
|
prisma = _prisma()
|
|
payload = _payload("r1")
|
|
payload["metadata"] = json.dumps(
|
|
{
|
|
"guardrail_information": [
|
|
{"guardrail_name": ["not", "a", "string"], "guardrail_status": "success"},
|
|
{"guardrail_name": "", "guardrail_id": "cf-uuid", "guardrail_status": "success"},
|
|
{"guardrail_status": "success"},
|
|
]
|
|
}
|
|
)
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, [payload])
|
|
|
|
index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"]
|
|
assert [row["guardrail_id"] for row in index_rows] == ["cf-uuid"]
|
|
metrics_create = prisma.db.litellm_dailyguardrailmetrics.upsert.call_args.kwargs["data"]["create"]
|
|
assert (metrics_create["guardrail_id"], metrics_create["requests_evaluated"]) == ("cf-uuid", 1)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_batch_of_only_not_run_entries_writes_no_metrics_row():
|
|
prisma = _prisma()
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, [_payload("r1", guardrail_status="not_run")])
|
|
|
|
assert prisma.db.litellm_dailyguardrailmetrics.upsert.call_count == 0
|
|
index_rows = prisma.db.litellm_spendlogguardrailindex.create_many.call_args.kwargs["data"]
|
|
assert [row["request_id"] for row in index_rows] == ["r1"]
|
|
|
|
|
|
@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
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_cost_rolled_up_per_counter_alongside_units():
|
|
"""LIT-5652: the hook's per-counter cost lands on the same daily row as the
|
|
units it priced, summed across payloads exactly like the units are, and the
|
|
update path increments it so a second flush on the same day keeps adding."""
|
|
prisma = _prisma()
|
|
logs = [
|
|
_payload(
|
|
"r1",
|
|
usage={"contentPolicyUnits": 1000, "wordPolicyUnits": 50},
|
|
cost_by_unit={"contentPolicyUnits": 0.15, "wordPolicyUnits": 0.0},
|
|
),
|
|
_payload(
|
|
"r2",
|
|
usage={"contentPolicyUnits": 2000, "wordPolicyUnits": 10},
|
|
cost_by_unit={"contentPolicyUnits": 0.3, "wordPolicyUnits": 0.0},
|
|
),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000,
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "wordPolicyUnits"): 60,
|
|
}
|
|
costs = _cost_upserts(prisma)
|
|
assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0)
|
|
assert costs["wordPolicyUnits"] == (0.0, 0)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_counter_the_hook_could_not_price_is_stored_as_untracked_units_not_free():
|
|
"""A counter the cost map does not list arrives stamped as None. Its units
|
|
must land in untracked_units with no cost, so the row never reads as free,
|
|
while the priced counter on the same request keeps its cost."""
|
|
prisma = _prisma()
|
|
logs = [
|
|
_payload(
|
|
"r1",
|
|
usage={"contentPolicyUnits": 1000, "someFutureCounter": 3},
|
|
cost_by_unit={"contentPolicyUnits": 0.15, "someFutureCounter": None},
|
|
)
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 1000,
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "someFutureCounter"): 3,
|
|
}
|
|
costs = _cost_upserts(prisma)
|
|
assert costs["contentPolicyUnits"] == (pytest.approx(0.15), 0)
|
|
assert costs["someFutureCounter"] == (0.0, 3)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_mixed_priced_and_unpriced_increments_keep_the_subtotal_and_count_the_rest_untracked():
|
|
"""Priced and unpriced increments on the same row (a hook without pricing,
|
|
a pre-upgrade proxy in a mixed fleet) must keep the priced subtotal and
|
|
count exactly the unpriced units as untracked. Nulling the cost would throw
|
|
away a known number; keeping it alone would look exact while understating."""
|
|
prisma = _prisma()
|
|
logs = [
|
|
_payload("r1", usage={"contentPolicyUnits": 1000}, cost_by_unit={"contentPolicyUnits": 0.15}),
|
|
_payload("r2", usage={"contentPolicyUnits": 700}),
|
|
_payload("r3", usage={"contentPolicyUnits": 300}, cost_by_unit={"contentPolicyUnits": None}),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 2000,
|
|
}
|
|
assert _cost_upserts(prisma) == {"contentPolicyUnits": (pytest.approx(0.15), 1000)}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_report_only_and_forged_costs_are_not_rolled_up_but_units_are():
|
|
"""guardrail_cost_in_spend=False (Azure Prompt Shield) keeps its cost out of
|
|
spend, so the rollup must not record it either or the dashboard would show
|
|
a number the budget never charged. A negative or non-finite per-counter cost
|
|
is treated the same way rather than subtracting from the day."""
|
|
prisma = _prisma()
|
|
logs = [
|
|
_payload("r1", usage={"text_records": 3}, cost_by_unit={"text_records": 0.5}, cost_in_spend=False),
|
|
_payload("r2", usage={"contentPolicyUnits": 10}, cost_by_unit={"contentPolicyUnits": -0.5}),
|
|
_payload("r3", usage={"topicPolicyUnits": 10}, cost_by_unit={"topicPolicyUnits": float("inf")}),
|
|
]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs)
|
|
|
|
assert _units_upserts(prisma) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "text_records"): 3,
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 10,
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "topicPolicyUnits"): 10,
|
|
}
|
|
assert _cost_upserts(prisma) == {
|
|
"text_records": (0.0, 3),
|
|
"contentPolicyUnits": (0.0, 10),
|
|
"topicPolicyUnits": (0.0, 10),
|
|
}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_requeued_cost_is_added_to_the_next_flush():
|
|
"""Cost and untracked units must survive the connection-error requeue the
|
|
same way units do, or a DB blip would silently drop dollars (or the record
|
|
that some units had no price) while keeping the units themselves."""
|
|
pending = PendingRollups()
|
|
down = _prisma()
|
|
down.db.litellm_dailyguardrailmetrics.upsert.side_effect = httpx.ConnectError("db down")
|
|
down.db.litellm_dailyguardrailusageunits.upsert.side_effect = httpx.ConnectError("db down")
|
|
sleep, _ = _fake_sleep()
|
|
|
|
await process_spend_logs_guardrail_usage(
|
|
down,
|
|
[
|
|
_payload(
|
|
"r1",
|
|
usage={"contentPolicyUnits": 1000, "someFutureCounter": 3},
|
|
cost_by_unit={"contentPolicyUnits": 0.15, "someFutureCounter": None},
|
|
)
|
|
],
|
|
sleep=sleep,
|
|
pending=pending,
|
|
)
|
|
recovered = _prisma()
|
|
await process_spend_logs_guardrail_usage(
|
|
recovered,
|
|
[
|
|
_payload(
|
|
"r2",
|
|
usage={"contentPolicyUnits": 2000, "someFutureCounter": 4},
|
|
cost_by_unit={"contentPolicyUnits": 0.3, "someFutureCounter": None},
|
|
)
|
|
],
|
|
sleep=sleep,
|
|
pending=pending,
|
|
)
|
|
|
|
assert _units_upserts(recovered) == {
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "contentPolicyUnits"): 3000,
|
|
("bedrock-guard", "2026-08-17", "team-a", "hashed-key-1", "someFutureCounter"): 7,
|
|
}
|
|
costs = _cost_upserts(recovered)
|
|
assert costs["contentPolicyUnits"] == (pytest.approx(0.45), 0)
|
|
assert costs["someFutureCounter"] == (0.0, 7)
|
|
|
|
|
|
def _fan_out_payload(request_id: str, guardrail_ids: tuple[str, ...]) -> dict[str, Any]:
|
|
return {
|
|
"request_id": request_id,
|
|
"startTime": datetime(2026, 8, 17, 12, 0, tzinfo=timezone.utc),
|
|
"team_id": "team-a",
|
|
"api_key": "hashed-key-1",
|
|
"metadata": json.dumps(
|
|
{"guardrail_information": [{"guardrail_id": gid, "guardrail_status": "success"} for gid in guardrail_ids]}
|
|
),
|
|
}
|
|
|
|
|
|
def _index_rows_written(prisma: MagicMock) -> list[tuple[str, str]]:
|
|
return [
|
|
(row["request_id"], row["guardrail_id"])
|
|
for call in prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list
|
|
for row in call.kwargs["data"]
|
|
]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_index_rows_are_written_in_row_bounded_statements(monkeypatch):
|
|
"""
|
|
LIT-5931: the drain caps logs, not logs x guardrails, so a fan-out must be
|
|
split into statements the query engine can afford instead of one create_many.
|
|
"""
|
|
monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100)
|
|
prisma = _prisma()
|
|
guardrail_ids = tuple(f"guard-{i}" for i in range(50))
|
|
logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs, pending=PendingRollups())
|
|
|
|
statements = prisma.db.litellm_spendlogguardrailindex.create_many.call_args_list
|
|
assert [len(call.kwargs["data"]) for call in statements] == [100, 100, 50]
|
|
assert all(call.kwargs["skip_duplicates"] is True for call in statements)
|
|
assert _index_rows_written(prisma) == [(f"r{i}", gid) for i in range(5) for gid in guardrail_ids]
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_one_failing_index_statement_does_not_drop_the_others_or_the_rollup(monkeypatch):
|
|
monkeypatch.setattr(usage_tracking, "SPEND_LOG_WRITE_BATCH_MAX_ROWS", 100)
|
|
prisma = _prisma()
|
|
prisma.db.litellm_spendlogguardrailindex.create_many.side_effect = [None, httpx.ReadTimeout("ambiguous"), None]
|
|
guardrail_ids = tuple(f"guard-{i}" for i in range(50))
|
|
logs = [_fan_out_payload(f"r{i}", guardrail_ids) for i in range(5)]
|
|
|
|
await process_spend_logs_guardrail_usage(prisma, logs, pending=PendingRollups())
|
|
|
|
assert prisma.db.litellm_spendlogguardrailindex.create_many.await_count == 3
|
|
assert prisma.db.litellm_dailyguardrailmetrics.upsert.await_count == len(guardrail_ids)
|