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:
yucheng 2026-09-14 20:45:47 +00:00
parent c3f52fe0d5
commit 7ebb169a4d
2 changed files with 15 additions and 6 deletions

View file

@ -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

View file

@ -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)