mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-05 02:41:56 +00:00
fix(guardrails): coalesce usage index rows per request and guardrail, keeping policy linkage
Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
parent
c3f52fe0d5
commit
7ebb169a4d
2 changed files with 15 additions and 6 deletions
|
|
@ -356,7 +356,7 @@ async def process_spend_logs_guardrail_usage(
|
|||
"flagged_count": 0,
|
||||
}
|
||||
)
|
||||
index_rows: Final[list[dict[str, object]]] = []
|
||||
index_rows_by_key: Final[dict[tuple[object, object], dict[str, object]]] = {}
|
||||
|
||||
for payload in logs_to_process:
|
||||
request_id = payload.get("request_id")
|
||||
|
|
@ -389,14 +389,15 @@ async def process_spend_logs_guardrail_usage(
|
|||
else:
|
||||
daily_guardrail[key]["flagged_count"] += 1
|
||||
policy_id = entry.get("policy_id")
|
||||
index_rows.append(
|
||||
{
|
||||
prior = index_rows_by_key.get((request_id, guardrail_id))
|
||||
if prior is None or (prior["policy_id"] is None and policy_id is not None):
|
||||
index_rows_by_key[(request_id, guardrail_id)] = {
|
||||
"request_id": request_id,
|
||||
"guardrail_id": guardrail_id,
|
||||
"policy_id": policy_id,
|
||||
"start_time": start_time,
|
||||
}
|
||||
)
|
||||
index_rows: Final = tuple(index_rows_by_key.values())
|
||||
|
||||
async with pending.lock:
|
||||
pending_metrics: Final = pending.metrics
|
||||
|
|
|
|||
|
|
@ -383,7 +383,12 @@ async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_nam
|
|||
{
|
||||
"guardrail_information": [
|
||||
{"guardrail_name": "cf", "guardrail_status": "not_run"},
|
||||
{"guardrail_name": "cf", "guardrail_id": "cf-uuid", "guardrail_status": "success"},
|
||||
{
|
||||
"guardrail_name": "cf",
|
||||
"guardrail_id": "cf-uuid",
|
||||
"policy_id": "pol-1",
|
||||
"guardrail_status": "success",
|
||||
},
|
||||
{"guardrail_name": "other", "guardrail_status": "not_run"},
|
||||
]
|
||||
}
|
||||
|
|
@ -392,7 +397,10 @@ async def test_not_run_entry_shares_index_key_with_evaluated_sibling_of_same_nam
|
|||
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"] for row in index_rows) == ["cf-uuid", "cf-uuid", "other"]
|
||||
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)
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue