scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-17 23:51:30 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 66
litellm/tests/test_litellm/proxy/guardrails/test_usage_endpoints.py

748 lines
31 KiB
Python
Raw Blame History

"""
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}
assert resp.untracked_usage_units_by_team == {"team-a": {"topicPolicyUnits": 10}, "": {"contentPolicyUnits": 50}}
assert resp.untracked_usage_units_by_key == {
"hash-1": {"topicPolicyUnits": 10},
"hash-2": {"contentPolicyUnits": 50},
}
assert resp.untracked_usage_units_by_team.keys() == resp.usage_units_by_team.keys()
assert resp.untracked_usage_units_by_key.keys() == resp.usage_units_by_key.keys()
@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 == {}
assert (resp.untracked_usage_units_by_team, resp.untracked_usage_units_by_key) == ({}, {})
# ---- 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"]}
def _index_row(request_id: str, guardrail_id: str = "cc-flag") -> Any:
r = MagicMock(spec=["request_id", "guardrail_id", "policy_id", "start_time"])
r.request_id = request_id
r.guardrail_id = guardrail_id
return r
def _spend_log(request_id: str, *guardrail_statuses: str, guardrail_id: str = "cc-flag") -> Any:
sl = MagicMock(spec=["request_id", "metadata", "startTime", "model", "messages", "response"])
sl.request_id = request_id
sl.startTime = datetime(2026, 4, 25, 12, 0)
sl.model = "gpt-4o-mini"
sl.messages = [{"role": "user", "content": "hi"}]
sl.response = "ok"
sl.metadata = {
"guardrail_information": [
{
"guardrail_name": guardrail_id,
"guardrail_status": status,
"guardrail_response": (
{"action": "flag", "reason": "audit hit"} if status == "guardrail_flagged" else "allow"
),
"duration": 0.002,
}
for status in guardrail_statuses
]
}
return sl
@pytest.mark.asyncio
async def test_logs_reports_flagged_action_for_guardrail_flagged_status():
"""LIT-6894: Request Logs surface a custom code flag() verdict as flagged with its reason."""
prisma = _prisma(index_find_many=[_index_row("r-flag"), _index_row("r-pass"), _index_row("r-block")])
prisma.db.litellm_spendlogs.find_many = AsyncMock(
return_value=[
_spend_log("r-flag", "guardrail_flagged"),
_spend_log("r-pass", "success"),
_spend_log("r-block", "guardrail_intervened"),
]
)
p1, p2 = _patches(prisma, _config_handler())
with p1, p2:
resp = await guardrails_usage_logs(
guardrail_id="cc-flag",
policy_id=None,
page=1,
page_size=50,
action=None,
start_date=START,
end_date=END,
user_api_key_dict=ADMIN,
)
flagged_only = await guardrails_usage_logs(
guardrail_id="cc-flag",
policy_id=None,
page=1,
page_size=50,
action="flagged",
start_date=START,
end_date=END,
user_api_key_dict=ADMIN,
)
assert [(log.id, log.action) for log in resp.logs] == [
("r-flag", "flagged"),
("r-pass", "passed"),
("r-block", "blocked"),
]
assert resp.logs[0].reason == "{'action': 'flag', 'reason': 'audit hit'}"
assert [log.id for log in flagged_only.logs] == ["r-flag"]
@pytest.mark.asyncio
async def test_logs_reports_post_call_flag_when_pre_call_allowed():
"""LIT-6894: a guardrail on mode [pre_call, post_call] that allows the request but flags the response
shows as flagged, not hidden behind the pre_call allow entry."""
prisma = _prisma(index_find_many=[_index_row("r-post-flag")])
prisma.db.litellm_spendlogs.find_many = AsyncMock(
return_value=[_spend_log("r-post-flag", "success", "guardrail_flagged")]
)
p1, p2 = _patches(prisma, _config_handler())
with p1, p2:
resp = await guardrails_usage_logs(
guardrail_id="cc-flag",
policy_id=None,
page=1,
page_size=50,
action=None,
start_date=START,
end_date=END,
user_api_key_dict=ADMIN,
)
assert [(log.id, log.action, log.reason) for log in resp.logs] == [
("r-post-flag", "flagged", "{'action': 'flag', 'reason': 'audit hit'}")
]
# ---- 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)
@pytest.mark.asyncio
async def test_logs_report_not_run_entries_as_not_run_not_passed():
"""LIT-6314: a guardrail that never scanned must not be reported as a pass in the drill-down."""
index_row = MagicMock()
index_row.request_id = "req-nr"
index_row.guardrail_id = "db-1"
index_row.start_time = datetime(2026, 4, 22)
spend_log = MagicMock()
spend_log.request_id = "req-nr"
spend_log.model = "gpt-4o-mini"
spend_log.startTime = datetime(2026, 4, 22)
spend_log.metadata = {
"guardrail_information": [
{"guardrail_name": "db-1", "guardrail_status": "not_run", "duration": 0.0},
]
}
prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row])
prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[spend_log])
handler = _config_handler()
p1, p2 = _patches(prisma, handler)
with p1, p2:
resp = await guardrails_usage_logs(
guardrail_id="db-1",
policy_id=None,
page=1,
page_size=50,
action=None,
start_date=START,
end_date=END,
user_api_key_dict=ADMIN,
)
assert [log.action for log in resp.logs] == ["not_run"]
@pytest.mark.asyncio
async def test_logs_action_passed_filter_excludes_not_run_entries():
"""LIT-6314: filtering the drill-down for passes must not return unscanned requests."""
index_row = MagicMock()
index_row.request_id = "req-nr"
index_row.guardrail_id = "db-1"
index_row.start_time = datetime(2026, 4, 22)
spend_log = MagicMock()
spend_log.request_id = "req-nr"
spend_log.model = "gpt-4o-mini"
spend_log.startTime = datetime(2026, 4, 22)
spend_log.metadata = {"guardrail_information": [{"guardrail_name": "db-1", "guardrail_status": "not_run"}]}
prisma = _prisma(find_unique=_db_row(), index_find_many=[index_row])
prisma.db.litellm_spendlogs.find_many = AsyncMock(return_value=[spend_log])
handler = _config_handler()
p1, p2 = _patches(prisma, handler)
with p1, p2:
resp = await guardrails_usage_logs(
guardrail_id="db-1",
policy_id=None,
page=1,
page_size=50,
action="passed",
start_date=START,
end_date=END,
user_api_key_dict=ADMIN,
)
assert resp.logs == []
Reference in a new issue View git blame Copy permalink