chore: initial public snapshot for github upload
This commit is contained in:
@@ -0,0 +1,176 @@
|
||||
"""
|
||||
Track guardrail and policy usage for the dashboard: upsert daily metrics and
|
||||
insert into SpendLogGuardrailIndex when spend logs are written.
|
||||
"""
|
||||
|
||||
import json
|
||||
from collections import defaultdict
|
||||
from datetime import datetime, timezone
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
|
||||
def _guardrail_status_to_action(status: Optional[str]) -> str:
|
||||
"""Map StandardLogging guardrail_status to blocked/passed/flagged."""
|
||||
if not status:
|
||||
return "passed"
|
||||
s = (status or "").lower()
|
||||
if "intervened" in s or "block" in s:
|
||||
return "blocked"
|
||||
if "fail" in s or "error" in s:
|
||||
return "flagged"
|
||||
return "passed"
|
||||
|
||||
|
||||
def _parse_guardrail_info_from_payload(payload: Dict[str, Any]) -> List[Dict[str, Any]]:
|
||||
"""Extract guardrail_information from spend log payload metadata."""
|
||||
meta = payload.get("metadata")
|
||||
if not meta:
|
||||
return []
|
||||
if isinstance(meta, str):
|
||||
try:
|
||||
meta = json.loads(meta)
|
||||
except (json.JSONDecodeError, TypeError):
|
||||
return []
|
||||
if not isinstance(meta, dict):
|
||||
return []
|
||||
info = meta.get("guardrail_information") or meta.get(
|
||||
"standard_logging_guardrail_information"
|
||||
)
|
||||
if not isinstance(info, list):
|
||||
return []
|
||||
return info
|
||||
|
||||
|
||||
def _date_str(dt: datetime) -> str:
|
||||
"""YYYY-MM-DD in UTC."""
|
||||
if dt.tzinfo is None:
|
||||
dt = dt.replace(tzinfo=timezone.utc)
|
||||
return dt.astimezone(timezone.utc).strftime("%Y-%m-%d")
|
||||
|
||||
|
||||
async def process_spend_logs_guardrail_usage(
|
||||
prisma_client: PrismaClient,
|
||||
logs_to_process: List[Dict[str, Any]],
|
||||
) -> None:
|
||||
"""
|
||||
After spend logs are written: update DailyGuardrailMetrics and insert
|
||||
SpendLogGuardrailIndex rows from guardrail_information in each payload.
|
||||
"""
|
||||
if not logs_to_process:
|
||||
return
|
||||
# Aggregate daily metrics by (guardrail_id, date). Latency/score metrics dropped.
|
||||
daily_guardrail: Dict[tuple, Dict[str, Any]] = defaultdict(
|
||||
lambda: {
|
||||
"requests_evaluated": 0,
|
||||
"passed_count": 0,
|
||||
"blocked_count": 0,
|
||||
"flagged_count": 0,
|
||||
}
|
||||
)
|
||||
index_rows: List[Dict[str, Any]] = []
|
||||
|
||||
for payload in logs_to_process:
|
||||
request_id = payload.get("request_id")
|
||||
start_time = payload.get("startTime")
|
||||
if not request_id or not start_time:
|
||||
continue
|
||||
if isinstance(start_time, str):
|
||||
try:
|
||||
start_time = datetime.fromisoformat(start_time.replace("Z", "+00:00"))
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
date_key = _date_str(start_time)
|
||||
|
||||
for entry in _parse_guardrail_info_from_payload(payload):
|
||||
guardrail_id = (
|
||||
entry.get("guardrail_id") or entry.get("guardrail_name") or ""
|
||||
)
|
||||
if not guardrail_id:
|
||||
continue
|
||||
key = (guardrail_id, date_key)
|
||||
daily_guardrail[key]["requests_evaluated"] += 1
|
||||
action = _guardrail_status_to_action(entry.get("guardrail_status"))
|
||||
if action == "passed":
|
||||
daily_guardrail[key]["passed_count"] += 1
|
||||
elif action == "blocked":
|
||||
daily_guardrail[key]["blocked_count"] += 1
|
||||
else:
|
||||
daily_guardrail[key]["flagged_count"] += 1
|
||||
policy_id = entry.get("policy_id")
|
||||
index_rows.append(
|
||||
{
|
||||
"request_id": request_id,
|
||||
"guardrail_id": guardrail_id,
|
||||
"policy_id": policy_id,
|
||||
"start_time": start_time,
|
||||
}
|
||||
)
|
||||
|
||||
if not daily_guardrail and not index_rows:
|
||||
return
|
||||
|
||||
try:
|
||||
# Insert index rows (skip duplicates by request_id + guardrail_id)
|
||||
if index_rows:
|
||||
index_data = []
|
||||
for r in index_rows:
|
||||
st = r["start_time"]
|
||||
if isinstance(st, str):
|
||||
try:
|
||||
st = datetime.fromisoformat(st.replace("Z", "+00:00"))
|
||||
except (ValueError, TypeError):
|
||||
continue
|
||||
index_data.append(
|
||||
{
|
||||
"request_id": r["request_id"],
|
||||
"guardrail_id": r["guardrail_id"],
|
||||
"policy_id": r.get("policy_id"),
|
||||
"start_time": st,
|
||||
}
|
||||
)
|
||||
try:
|
||||
await prisma_client.db.litellm_spendlogguardrailindex.create_many(
|
||||
data=index_data,
|
||||
skip_duplicates=True,
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug(
|
||||
"Guardrail usage tracking: index create_many skipped: %s", e
|
||||
)
|
||||
|
||||
# Upsert daily guardrail metrics (counts only; latency/score dropped)
|
||||
for (guardrail_id, date_key), agg in daily_guardrail.items():
|
||||
n = int(agg["requests_evaluated"])
|
||||
if n == 0:
|
||||
continue
|
||||
await prisma_client.db.litellm_dailyguardrailmetrics.upsert(
|
||||
where={
|
||||
"guardrail_id_date": {
|
||||
"guardrail_id": guardrail_id,
|
||||
"date": date_key,
|
||||
}
|
||||
},
|
||||
data={
|
||||
"create": {
|
||||
"guardrail_id": guardrail_id,
|
||||
"date": date_key,
|
||||
"requests_evaluated": n,
|
||||
"passed_count": int(agg["passed_count"]),
|
||||
"blocked_count": int(agg["blocked_count"]),
|
||||
"flagged_count": int(agg["flagged_count"]),
|
||||
},
|
||||
"update": {
|
||||
"requests_evaluated": {"increment": n},
|
||||
"passed_count": {"increment": int(agg["passed_count"])},
|
||||
"blocked_count": {"increment": int(agg["blocked_count"])},
|
||||
"flagged_count": {"increment": int(agg["flagged_count"])},
|
||||
},
|
||||
},
|
||||
)
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.warning(
|
||||
"Guardrail usage tracking failed (non-fatal): %s", e
|
||||
)
|
||||
Reference in New Issue
Block a user