mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-17 23:51:30 +00:00
6c81a5c423
A row that received both priced and unpriced increments used to collapse to cost NULL, throwing away the priced subtotal and making every unit on it read as untracked. The rollup now carries a second column, untracked_units, that the aggregator increments for units with no known price while cost keeps accruing for the rest, so cost covers exactly units - untracked_units. Rows written before the migration keep cost NULL and still read as untracked in full The endpoints read untracked units off the column (or the whole row for a legacy NULL) rather than from a NULL filter, and the policies overview now fills totalUntrackedUsageUnits, which the previous commit missed Claude-Session: https://claude.ai/code/session_01EX13mWex6RaBo9PYnkAtFW
579 lines
25 KiB
Python
579 lines
25 KiB
Python
"""
|
|
Tests for the /guardrails/usage/* endpoints backing the dashboard Guardrail Monitor.
|
|
|
|
Regression (LIT-2529): guardrails defined in config.yaml live only in
|
|
IN_MEMORY_GUARDRAIL_HANDLER, so the monitor's overview/detail/logs endpoints —
|
|
which read the litellm_guardrailstable Prisma table — could not see them:
|
|
detail 404'd, overview omitted them (or rendered them as Custom/Guardrail
|
|
orphans), and logs missed their logical-name alias.
|
|
"""
|
|
|
|
from datetime import datetime
|
|
from typing import Any, Optional
|
|
from unittest.mock import AsyncMock, MagicMock, patch
|
|
|
|
import pytest
|
|
|
|
|
|
from fastapi import HTTPException
|
|
from prisma.errors import TableNotFoundError
|
|
|
|
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
|
|
from litellm.proxy.guardrails.guardrail_registry import InMemoryGuardrailHandler
|
|
from litellm.proxy.guardrails.usage_endpoints import (
|
|
guardrails_usage_detail,
|
|
guardrails_usage_logs,
|
|
guardrails_usage_overview,
|
|
policies_usage_overview,
|
|
)
|
|
from litellm.types.guardrails import Guardrail, LitellmParams
|
|
|
|
ADMIN = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN)
|
|
# Query() defaults don't resolve to None when the handler is called directly.
|
|
START, END = "2026-04-20", "2026-04-27"
|
|
|
|
|
|
def _config_handler(*guardrails: Guardrail) -> InMemoryGuardrailHandler:
|
|
"""A real handler seeded with config-sourced YAML guardrails (no callbacks)."""
|
|
handler = InMemoryGuardrailHandler()
|
|
for g in guardrails:
|
|
gid = g["guardrail_id"]
|
|
handler.IN_MEMORY_GUARDRAILS[gid] = g
|
|
handler._sources[gid] = "config"
|
|
return handler
|
|
|
|
|
|
def _yaml_guardrail(
|
|
guardrail_id: str = "yaml-1",
|
|
name: str = "yaml-pii",
|
|
provider: str = "presidio",
|
|
info: Optional[dict] = None,
|
|
) -> Guardrail:
|
|
return Guardrail(
|
|
guardrail_id=guardrail_id,
|
|
guardrail_name=name,
|
|
litellm_params=LitellmParams(guardrail=provider, mode="pre_call"),
|
|
guardrail_info=info if info is not None else {"type": "PII", "description": "yaml-defined"},
|
|
)
|
|
|
|
|
|
def _db_row(guardrail_id: str = "db-1", name: str = "db-guard", provider: str = "aim") -> Any:
|
|
"""A Prisma-style row: attribute access, litellm_params/guardrail_info as plain dicts."""
|
|
row = MagicMock(spec=["guardrail_id", "guardrail_name", "litellm_params", "guardrail_info"])
|
|
row.guardrail_id = guardrail_id
|
|
row.guardrail_name = name
|
|
row.litellm_params = {"guardrail": provider, "mode": "pre_call"}
|
|
row.guardrail_info = {"type": "ContentSafety", "description": "db-defined"}
|
|
return row
|
|
|
|
|
|
def _metric(guardrail_id: str, date: str = "2026-04-25", requests: int = 10, passed: int = 8, blocked: int = 2) -> Any:
|
|
m = MagicMock()
|
|
m.guardrail_id = guardrail_id
|
|
m.date = date
|
|
m.requests_evaluated = requests
|
|
m.passed_count = passed
|
|
m.blocked_count = blocked
|
|
m.flagged_count = 0
|
|
return m
|
|
|
|
|
|
def _units_row(
|
|
guardrail_id: str,
|
|
date: str = "2026-04-25",
|
|
team_id: str = "",
|
|
api_key: str = "",
|
|
usage_unit: str = "contentPolicyUnits",
|
|
units: int = 1,
|
|
cost: float | None = None,
|
|
untracked_units: int = 0,
|
|
) -> Any:
|
|
"""cost=None is a row written before the cost column existed (untracked in full)."""
|
|
r = MagicMock()
|
|
r.guardrail_id = guardrail_id
|
|
r.date = date
|
|
r.team_id = team_id
|
|
r.api_key = api_key
|
|
r.usage_unit = usage_unit
|
|
r.units = units
|
|
r.cost = cost
|
|
r.untracked_units = untracked_units
|
|
return r
|
|
|
|
|
|
def _prisma(
|
|
*,
|
|
find_many=None,
|
|
find_unique=None,
|
|
metrics=None,
|
|
index_find_many=None,
|
|
units=None,
|
|
) -> MagicMock:
|
|
client = MagicMock()
|
|
db = client.db
|
|
db.litellm_guardrailstable.find_many = AsyncMock(return_value=find_many or [])
|
|
db.litellm_guardrailstable.find_unique = AsyncMock(return_value=find_unique)
|
|
db.litellm_dailyguardrailmetrics.find_many = AsyncMock(return_value=metrics or [])
|
|
db.litellm_dailyguardrailusageunits.find_many = AsyncMock(return_value=units or [])
|
|
db.litellm_spendlogguardrailindex.find_many = AsyncMock(return_value=index_find_many or [])
|
|
db.litellm_spendlogguardrailindex.count = AsyncMock(return_value=0)
|
|
db.litellm_spendlogs.find_many = AsyncMock(return_value=[])
|
|
return client
|
|
|
|
|
|
def _patches(prisma: MagicMock, handler: InMemoryGuardrailHandler):
|
|
return (
|
|
patch("litellm.proxy.proxy_server.prisma_client", prisma),
|
|
patch("litellm.proxy.guardrails.guardrail_registry.IN_MEMORY_GUARDRAIL_HANDLER", handler),
|
|
)
|
|
|
|
|
|
# ---- detail -----------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_returns_yaml_guardrail_when_db_misses():
|
|
prisma = _prisma(find_unique=None)
|
|
handler = _config_handler(_yaml_guardrail())
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_detail(
|
|
guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
|
|
)
|
|
assert resp.guardrail_id == "yaml-1"
|
|
assert resp.guardrail_name == "yaml-pii"
|
|
assert resp.provider == "presidio" # coerced from the LitellmParams pydantic model
|
|
assert resp.type == "PII" # from guardrail_info
|
|
assert resp.description == "yaml-defined"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_404_when_neither_db_nor_config():
|
|
prisma = _prisma(find_unique=None)
|
|
handler = _config_handler() # empty
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2, pytest.raises(HTTPException) as exc:
|
|
await guardrails_usage_detail(guardrail_id="ghost", start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
assert exc.value.status_code == 404
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_does_not_surface_db_sourced_in_memory_entry():
|
|
"""A stale in-memory entry (source=db, gone from DB) must 404, not resurface."""
|
|
prisma = _prisma(find_unique=None)
|
|
handler = InMemoryGuardrailHandler()
|
|
stale = _yaml_guardrail(guardrail_id="stale-1", name="stale")
|
|
handler.IN_MEMORY_GUARDRAILS["stale-1"] = stale
|
|
handler._sources["stale-1"] = "db"
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2, pytest.raises(HTTPException) as exc:
|
|
await guardrails_usage_detail(guardrail_id="stale-1", start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
assert exc.value.status_code == 404
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_db_row_still_resolves():
|
|
prisma = _prisma(find_unique=_db_row(guardrail_id="db-1", provider="aim"))
|
|
handler = _config_handler()
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_detail(
|
|
guardrail_id="db-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
|
|
)
|
|
assert resp.provider == "aim"
|
|
assert resp.type == "ContentSafety"
|
|
|
|
|
|
# ---- overview ---------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_includes_yaml_guardrail_with_no_metrics():
|
|
"""The core bug: a YAML guardrail with zero metrics must still appear as a row."""
|
|
prisma = _prisma(find_many=[]) # no DB guardrails
|
|
handler = _config_handler(_yaml_guardrail())
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
rows = [r for r in resp.rows if r.id == "yaml-1"]
|
|
assert len(rows) == 1
|
|
assert rows[0].name == "yaml-pii"
|
|
assert rows[0].provider == "presidio"
|
|
assert rows[0].type == "PII"
|
|
assert rows[0].requestsEvaluated == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_yaml_metrics_matched_by_logical_name():
|
|
"""Daily metrics are keyed by logical name; the YAML row must pick them up."""
|
|
prisma = _prisma(
|
|
find_many=[],
|
|
metrics=[_metric("yaml-pii", requests=10, blocked=2)], # keyed by name, not uuid
|
|
)
|
|
handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
rows = [r for r in resp.rows if r.id == "yaml-uuid"]
|
|
assert len(rows) == 1
|
|
assert rows[0].requestsEvaluated == 10
|
|
assert rows[0].failRate == 20.0
|
|
# must not also emit an orphan row keyed by the logical name
|
|
assert [r for r in resp.rows if r.id == "yaml-pii"] == []
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_excludes_db_sourced_in_memory_entry():
|
|
"""union must not resurrect a stale db-sourced in-memory guardrail."""
|
|
prisma = _prisma(find_many=[])
|
|
handler = InMemoryGuardrailHandler()
|
|
handler.IN_MEMORY_GUARDRAILS["cfg"] = _yaml_guardrail(guardrail_id="cfg", name="cfg-guard")
|
|
handler._sources["cfg"] = "config"
|
|
handler.IN_MEMORY_GUARDRAILS["stale"] = _yaml_guardrail(guardrail_id="stale", name="stale-guard")
|
|
handler._sources["stale"] = "db"
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
ids = {r.id for r in resp.rows}
|
|
assert "cfg" in ids
|
|
assert "stale" not in ids
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_reports_usage_units_per_row_and_total():
|
|
"""LIT-5650: billable units must surface per guardrail row (matched by
|
|
logical name like the daily metrics) and as a response-level total."""
|
|
prisma = _prisma(
|
|
find_many=[],
|
|
metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)],
|
|
units=[
|
|
_units_row("yaml-pii", usage_unit="topicPolicyUnits", units=4),
|
|
_units_row("yaml-pii", usage_unit="contentPolicyUnits", units=3),
|
|
_units_row("yaml-pii", team_id="team-a", usage_unit="contentPolicyUnits", units=2),
|
|
_units_row("other-guard", usage_unit="topicPolicyUnits", units=7),
|
|
],
|
|
)
|
|
handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
row = next(r for r in resp.rows if r.id == "yaml-uuid")
|
|
assert row.usageUnits == {"topicPolicyUnits": 4, "contentPolicyUnits": 5}
|
|
assert resp.totalUsageUnits == {"topicPolicyUnits": 11, "contentPolicyUnits": 5}
|
|
units_where = prisma.db.litellm_dailyguardrailusageunits.find_many.call_args.kwargs["where"]
|
|
assert units_where == {"date": {"gte": START, "lte": END}}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_breaks_units_down_by_day_team_and_key():
|
|
prisma = _prisma(
|
|
find_unique=None,
|
|
units=[
|
|
_units_row("yaml-pii", date="2026-04-25", team_id="team-a", api_key="hash-1", units=2),
|
|
_units_row("yaml-pii", date="2026-04-25", team_id="", api_key="hash-2", units=1),
|
|
_units_row(
|
|
"yaml-pii", date="2026-04-24", team_id="team-a", api_key="hash-1", usage_unit="topicPolicyUnits"
|
|
),
|
|
],
|
|
)
|
|
handler = _config_handler(_yaml_guardrail())
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_detail(
|
|
guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
|
|
)
|
|
assert resp.usage_units == {"contentPolicyUnits": 3, "topicPolicyUnits": 1}
|
|
assert [p.model_dump() for p in resp.usage_units_daily] == [
|
|
{"date": "2026-04-24", "units": {"topicPolicyUnits": 1}, "cost": None},
|
|
{"date": "2026-04-25", "units": {"contentPolicyUnits": 3}, "cost": None},
|
|
]
|
|
assert resp.usage_units_by_team == {
|
|
"team-a": {"contentPolicyUnits": 2, "topicPolicyUnits": 1},
|
|
"": {"contentPolicyUnits": 1},
|
|
}
|
|
assert resp.usage_units_by_key == {
|
|
"hash-1": {"contentPolicyUnits": 2, "topicPolicyUnits": 1},
|
|
"hash-2": {"contentPolicyUnits": 1},
|
|
}
|
|
units_where = prisma.db.litellm_dailyguardrailusageunits.find_many.call_args.kwargs["where"]
|
|
assert units_where == {"guardrail_id": {"in": ["yaml-pii", "yaml-1"]}, "date": {"gte": START, "lte": END}}
|
|
|
|
|
|
def _units_table_missing() -> TableNotFoundError:
|
|
return TableNotFoundError(
|
|
data={"user_facing_error": {"meta": {"table": "public.LiteLLM_DailyGuardrailUsageUnits"}}}
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_degrades_units_to_empty_when_units_table_is_missing():
|
|
prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)])
|
|
prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing())
|
|
handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
row = next(r for r in resp.rows if r.id == "yaml-uuid")
|
|
assert (row.requestsEvaluated, row.usageUnits) == (4, {})
|
|
assert (resp.totalRequests, resp.totalBlocked, resp.totalUsageUnits) == (4, 1, {})
|
|
assert (row.cost, resp.totalCost) == (None, None)
|
|
assert (row.untrackedUsageUnits, resp.totalUntrackedUsageUnits) == ({}, {})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_reports_cost_per_row_and_total_summing_only_tracked_days():
|
|
"""LIT-5652: cost rides the units rollup. Rows written before the cost column
|
|
carry NULL and rows whose every unit was unpriced carry 0.0 with
|
|
untracked_units == units; both must drop out of the sum rather than read as
|
|
$0, and a guardrail with only such rows reports None, not 0.0."""
|
|
prisma = _prisma(
|
|
find_many=[],
|
|
metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)],
|
|
units=[
|
|
_units_row("yaml-pii", usage_unit="contentPolicyUnits", units=1000, cost=0.15),
|
|
_units_row("yaml-pii", team_id="team-a", usage_unit="contentPolicyUnits", units=2000, cost=0.3),
|
|
_units_row("yaml-pii", date="2026-04-24", usage_unit="contentPolicyUnits", units=5000, cost=None),
|
|
_units_row(
|
|
"yaml-pii", date="2026-04-23", usage_unit="topicPolicyUnits", units=9, cost=0.0, untracked_units=9
|
|
),
|
|
_units_row("legacy-guard", usage_unit="topicPolicyUnits", units=7, cost=None),
|
|
],
|
|
)
|
|
handler = _config_handler(
|
|
_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"),
|
|
_yaml_guardrail(guardrail_id="legacy-uuid", name="legacy-guard"),
|
|
)
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
by_id = {r.id: r for r in resp.rows}
|
|
assert by_id["yaml-uuid"].cost == pytest.approx(0.45)
|
|
assert by_id["legacy-uuid"].cost is None
|
|
assert resp.totalCost == pytest.approx(0.45)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_reports_the_units_its_cost_leaves_out_per_row_and_total():
|
|
"""A row's cost covers only the units that had a price, so the response must
|
|
say exactly which units (per counter) that cost excludes: the row's own
|
|
untracked_units, or all of its units when it predates the cost column. A
|
|
guardrail whose rows are all priced reports none, one whose rows are all
|
|
unpriced reports all of its units, and a mixed row keeps its priced subtotal
|
|
while reporting just the unpriced share."""
|
|
prisma = _prisma(
|
|
find_many=[],
|
|
metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)],
|
|
units=[
|
|
_units_row("yaml-pii", usage_unit="contentPolicyUnits", units=1000, cost=0.15, untracked_units=200),
|
|
_units_row("yaml-pii", date="2026-04-24", usage_unit="contentPolicyUnits", units=5000, cost=None),
|
|
_units_row(
|
|
"yaml-pii", date="2026-04-24", usage_unit="topicPolicyUnits", units=40, cost=0.0, untracked_units=40
|
|
),
|
|
_units_row("yaml-pii", usage_unit="wordPolicyUnits", units=9, cost=0.0),
|
|
_units_row("legacy-guard", usage_unit="topicPolicyUnits", units=7, cost=None),
|
|
_units_row("priced-guard", usage_unit="contentPolicyUnits", units=3, cost=0.0003),
|
|
],
|
|
)
|
|
handler = _config_handler(
|
|
_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"),
|
|
_yaml_guardrail(guardrail_id="legacy-uuid", name="legacy-guard"),
|
|
_yaml_guardrail(guardrail_id="priced-uuid", name="priced-guard"),
|
|
)
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
by_id = {r.id: r for r in resp.rows}
|
|
assert by_id["yaml-uuid"].usageUnits == {"contentPolicyUnits": 6000, "topicPolicyUnits": 40, "wordPolicyUnits": 9}
|
|
assert by_id["yaml-uuid"].cost == pytest.approx(0.15)
|
|
assert by_id["yaml-uuid"].untrackedUsageUnits == {"contentPolicyUnits": 5200, "topicPolicyUnits": 40}
|
|
assert by_id["legacy-uuid"].untrackedUsageUnits == {"topicPolicyUnits": 7}
|
|
assert by_id["priced-uuid"].untrackedUsageUnits == {}
|
|
assert resp.totalUntrackedUsageUnits == {"contentPolicyUnits": 5200, "topicPolicyUnits": 47}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_breaks_cost_down_by_unit_day_team_and_key():
|
|
"""Every cost breakdown keeps the same keys as its units twin so the UI can
|
|
render them side by side, with None where that group has no tracked cost."""
|
|
prisma = _prisma(
|
|
find_unique=None,
|
|
units=[
|
|
_units_row("yaml-pii", date="2026-04-25", team_id="team-a", api_key="hash-1", units=1000, cost=0.15),
|
|
_units_row(
|
|
"yaml-pii", date="2026-04-25", team_id="", api_key="hash-2", units=200, cost=0.03, untracked_units=50
|
|
),
|
|
_units_row(
|
|
"yaml-pii",
|
|
date="2026-04-24",
|
|
team_id="team-a",
|
|
api_key="hash-1",
|
|
usage_unit="topicPolicyUnits",
|
|
units=10,
|
|
cost=None,
|
|
),
|
|
],
|
|
)
|
|
handler = _config_handler(_yaml_guardrail())
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_detail(
|
|
guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
|
|
)
|
|
assert resp.cost == pytest.approx(0.18)
|
|
assert resp.cost_by_unit == {"contentPolicyUnits": pytest.approx(0.18), "topicPolicyUnits": None}
|
|
assert [p.model_dump() for p in resp.usage_units_daily] == [
|
|
{"date": "2026-04-24", "units": {"topicPolicyUnits": 10}, "cost": None},
|
|
{"date": "2026-04-25", "units": {"contentPolicyUnits": 1200}, "cost": pytest.approx(0.18)},
|
|
]
|
|
assert resp.cost_by_team == {"team-a": pytest.approx(0.15), "": pytest.approx(0.03)}
|
|
assert resp.cost_by_key == {"hash-1": pytest.approx(0.15), "hash-2": pytest.approx(0.03)}
|
|
assert resp.cost_by_team.keys() == resp.usage_units_by_team.keys()
|
|
assert resp.cost_by_key.keys() == resp.usage_units_by_key.keys()
|
|
assert resp.untracked_usage_units == {"contentPolicyUnits": 50, "topicPolicyUnits": 10}
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_degrades_units_to_empty_when_units_table_is_missing():
|
|
prisma = _prisma(metrics=[_metric("yaml-pii", requests=4, passed=3, blocked=1)])
|
|
prisma.db.litellm_dailyguardrailusageunits.find_many = AsyncMock(side_effect=_units_table_missing())
|
|
handler = _config_handler(_yaml_guardrail())
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_detail(
|
|
guardrail_id="yaml-1", start_date=START, end_date=END, user_api_key_dict=ADMIN
|
|
)
|
|
assert (resp.requestsEvaluated, resp.failRate) == (4, 25.0)
|
|
assert (resp.usage_units, list(resp.usage_units_daily), resp.usage_units_by_team, resp.usage_units_by_key) == (
|
|
{},
|
|
[],
|
|
{},
|
|
{},
|
|
)
|
|
assert (resp.cost, resp.cost_by_unit, resp.cost_by_team, resp.cost_by_key) == (None, {}, {}, {})
|
|
assert resp.untracked_usage_units == {}
|
|
|
|
|
|
# ---- logs -------------------------------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_logs_resolves_config_guardrail_logical_name():
|
|
"""The index query must include the YAML guardrail's logical name alias."""
|
|
prisma = _prisma(find_unique=None)
|
|
handler = _config_handler(_yaml_guardrail(guardrail_id="yaml-uuid", name="yaml-pii"))
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
await guardrails_usage_logs(
|
|
guardrail_id="yaml-uuid",
|
|
policy_id=None,
|
|
page=1,
|
|
page_size=50,
|
|
action=None,
|
|
start_date=START,
|
|
end_date=END,
|
|
user_api_key_dict=ADMIN,
|
|
)
|
|
where = prisma.db.litellm_spendlogguardrailindex.find_many.call_args.kwargs["where"]
|
|
assert where["guardrail_id"] == {"in": ["yaml-uuid", "yaml-pii"]}
|
|
|
|
|
|
# ---- date window cap (LIT-5762) ---------------------------------------------
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_rejects_range_over_max_days():
|
|
prisma = _prisma()
|
|
handler = _config_handler()
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2, pytest.raises(HTTPException) as exc:
|
|
await guardrails_usage_overview(start_date="2020-01-01", end_date=END, user_api_key_dict=ADMIN)
|
|
assert exc.value.status_code == 400
|
|
assert "366" in str(exc.value.detail)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_accepts_range_at_exactly_max_days():
|
|
prisma = _prisma()
|
|
handler = _config_handler()
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
resp = await guardrails_usage_overview(start_date="2025-04-26", end_date="2026-04-27", user_api_key_dict=ADMIN)
|
|
assert resp.totalRequests == 0
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_rejects_malformed_dates():
|
|
prisma = _prisma()
|
|
handler = _config_handler()
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2, pytest.raises(HTTPException) as exc:
|
|
await guardrails_usage_overview(start_date="not-a-date", end_date=END, user_api_key_dict=ADMIN)
|
|
assert exc.value.status_code == 400
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_overview_rejects_non_canonical_date_format():
|
|
prisma = _prisma()
|
|
handler = _config_handler()
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2, pytest.raises(HTTPException) as exc:
|
|
await guardrails_usage_overview(start_date="20260420", end_date=END, user_api_key_dict=ADMIN)
|
|
assert exc.value.status_code == 400
|
|
assert "YYYY-MM-DD" in str(exc.value.detail)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_rejects_reversed_dates():
|
|
prisma = _prisma(find_unique=_db_row())
|
|
handler = _config_handler()
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2, pytest.raises(HTTPException) as exc:
|
|
await guardrails_usage_detail(guardrail_id="db-1", start_date=END, end_date=START, user_api_key_dict=ADMIN)
|
|
assert exc.value.status_code == 400
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_policies_overview_returns_a_full_row_and_totals():
|
|
"""Regression: the policies overview shares the guardrail response model, so
|
|
every field added there (usage units, cost, untracked units) must be filled
|
|
here too or the endpoint 500s on model validation."""
|
|
policy = MagicMock(spec=["policy_id", "policy_name"])
|
|
policy.policy_id = "pol-1"
|
|
policy.policy_name = "block-pii"
|
|
metric = _metric("pol-1", requests=10, passed=8, blocked=2)
|
|
metric.policy_id = "pol-1"
|
|
prisma = _prisma()
|
|
prisma.db.litellm_policytable.find_many = AsyncMock(return_value=[policy])
|
|
prisma.db.litellm_dailypolicymetrics.find_many = AsyncMock(return_value=[metric])
|
|
p1, p2 = _patches(prisma, _config_handler())
|
|
with p1, p2:
|
|
resp = await policies_usage_overview(start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
row = next(r for r in resp.rows if r.id == "pol-1")
|
|
assert (row.name, row.type, row.requestsEvaluated, row.failRate) == ("block-pii", "Policy", 10, 20.0)
|
|
assert (row.usageUnits, row.cost, row.untrackedUsageUnits) == ({}, None, {})
|
|
assert (resp.totalRequests, resp.totalBlocked, resp.passRate) == (10, 2, 80.0)
|
|
assert (resp.totalUsageUnits, resp.totalCost, resp.totalUntrackedUsageUnits) == ({}, None, {})
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_policies_overview_rejects_range_over_max_days():
|
|
prisma = _prisma()
|
|
handler = _config_handler()
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2, pytest.raises(HTTPException) as exc:
|
|
await policies_usage_overview(start_date="2020-01-01", end_date=END, user_api_key_dict=ADMIN)
|
|
assert exc.value.status_code == 400
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_detail_prev_trend_query_is_bounded():
|
|
"""Regression: the trend query scanned every metrics row before start_date."""
|
|
prisma = _prisma(find_unique=_db_row())
|
|
handler = _config_handler()
|
|
p1, p2 = _patches(prisma, handler)
|
|
with p1, p2:
|
|
await guardrails_usage_detail(guardrail_id="db-1", start_date=START, end_date=END, user_api_key_dict=ADMIN)
|
|
wheres = [c.kwargs["where"] for c in prisma.db.litellm_dailyguardrailmetrics.find_many.await_args_list]
|
|
prev_wheres = [w for w in wheres if "lt" in w.get("date", {})]
|
|
assert prev_wheres
|
|
assert all("gte" in w["date"] for w in prev_wheres)
|