mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-12 23:01:41 +00:00
feat(proxy): denormalize aliases into audit log rows
Adds object_alias, object_team_id, object_team_alias, changed_by_user_email and changed_by_key_alias columns to LiteLLM_AuditLog, stamped at write time from the audited object's blob and the caller's auth context with single lookups as fallback. The migration backfills existing rows idempotently and adds indexes on (table_name, object_id), object_team_id and updated_at. The read path returns the columns verbatim and object_team filters on them; the previous read-time enrichment is removed. Key rows resolve object_alias via the key table because blob key_alias is masked
This commit is contained in:
ryan-crabbe-berri
parent
c2a03a137e
commit
eff4e5f826
11 changed files with 583 additions and 442 deletions
|
|
@ -7,7 +7,7 @@ GET - /audit/{id} - Get audit log by id
|
|||
GET - /audit - Get all audit logs
|
||||
"""
|
||||
|
||||
from typing import TYPE_CHECKING, Any, Dict, Final, List, NamedTuple, Optional, Sequence, Tuple
|
||||
from typing import Any, Dict, Optional
|
||||
|
||||
#### AUDIT LOGGING ####
|
||||
from fastapi import APIRouter, Depends, HTTPException, Query
|
||||
|
|
@ -16,157 +16,11 @@ from litellm_enterprise.types.proxy.audit_logging_endpoints import (
|
|||
PaginatedAuditLogResponse,
|
||||
)
|
||||
|
||||
from litellm.proxy._types import CommonProxyErrors, LitellmTableNames, UserAPIKeyAuth
|
||||
from litellm.proxy._types import CommonProxyErrors, UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
router = APIRouter()
|
||||
|
||||
_KEY_TABLE: Final[str] = LitellmTableNames.KEY_TABLE_NAME.value
|
||||
_TEAM_TABLE: Final[str] = LitellmTableNames.TEAM_TABLE_NAME.value
|
||||
_USER_TABLE: Final[str] = LitellmTableNames.USER_TABLE_NAME.value
|
||||
_ORG_TABLE: Final[str] = "LiteLLM_OrganizationTable"
|
||||
_MODEL_TABLE: Final[str] = LitellmTableNames.PROXY_MODEL_TABLE_NAME.value
|
||||
|
||||
_BLOB_ALIAS_KEYS: Final[Dict[str, Tuple[str, ...]]] = {
|
||||
_KEY_TABLE: ("key_alias",),
|
||||
_TEAM_TABLE: ("team_alias",),
|
||||
_USER_TABLE: ("user_alias", "user_email"),
|
||||
_ORG_TABLE: ("organization_alias",),
|
||||
_MODEL_TABLE: ("model_name",),
|
||||
}
|
||||
|
||||
|
||||
class _AliasMaps(NamedTuple):
|
||||
key_alias_by_token: Dict[str, str]
|
||||
team_alias_by_id: Dict[str, str]
|
||||
user_alias_by_id: Dict[str, str]
|
||||
user_email_by_id: Dict[str, str]
|
||||
org_alias_by_id: Dict[str, str]
|
||||
model_name_by_id: Dict[str, str]
|
||||
|
||||
|
||||
def _object_ids_for_table(audit_logs: Sequence[AuditLogResponse], table_name: str) -> frozenset:
|
||||
return frozenset(log.object_id for log in audit_logs if log.table_name == table_name and log.object_id)
|
||||
|
||||
|
||||
async def _fetch_alias_maps(prisma_client: "PrismaClient", audit_logs: Sequence[AuditLogResponse]) -> _AliasMaps:
|
||||
tokens: Final = _object_ids_for_table(audit_logs, _KEY_TABLE) | frozenset(
|
||||
log.changed_by_api_key for log in audit_logs if log.changed_by_api_key
|
||||
)
|
||||
user_ids: Final = _object_ids_for_table(audit_logs, _USER_TABLE) | frozenset(
|
||||
log.changed_by for log in audit_logs if log.changed_by
|
||||
)
|
||||
team_ids: Final = _object_ids_for_table(audit_logs, _TEAM_TABLE)
|
||||
org_ids: Final = _object_ids_for_table(audit_logs, _ORG_TABLE)
|
||||
model_ids: Final = _object_ids_for_table(audit_logs, _MODEL_TABLE)
|
||||
|
||||
key_rows: Final = (
|
||||
await prisma_client.db.litellm_verificationtoken.find_many(where={"token": {"in": list(tokens)}})
|
||||
if tokens
|
||||
else []
|
||||
)
|
||||
user_rows: Final = (
|
||||
await prisma_client.db.litellm_usertable.find_many(where={"user_id": {"in": list(user_ids)}})
|
||||
if user_ids
|
||||
else []
|
||||
)
|
||||
team_rows: Final = (
|
||||
await prisma_client.db.litellm_teamtable.find_many(where={"team_id": {"in": list(team_ids)}})
|
||||
if team_ids
|
||||
else []
|
||||
)
|
||||
org_rows: Final = (
|
||||
await prisma_client.db.litellm_organizationtable.find_many(where={"organization_id": {"in": list(org_ids)}})
|
||||
if org_ids
|
||||
else []
|
||||
)
|
||||
model_rows: Final = (
|
||||
await prisma_client.db.litellm_proxymodeltable.find_many(where={"model_id": {"in": list(model_ids)}})
|
||||
if model_ids
|
||||
else []
|
||||
)
|
||||
|
||||
return _AliasMaps(
|
||||
key_alias_by_token={row.token: row.key_alias for row in key_rows if row.key_alias},
|
||||
team_alias_by_id={row.team_id: row.team_alias for row in team_rows if row.team_alias},
|
||||
user_alias_by_id={row.user_id: row.user_alias for row in user_rows if row.user_alias},
|
||||
user_email_by_id={row.user_id: row.user_email for row in user_rows if row.user_email},
|
||||
org_alias_by_id={row.organization_id: row.organization_alias for row in org_rows if row.organization_alias},
|
||||
model_name_by_id={row.model_id: row.model_name for row in model_rows if row.model_name},
|
||||
)
|
||||
|
||||
|
||||
def _db_object_alias(log: AuditLogResponse, aliases: _AliasMaps) -> str | None:
|
||||
if log.table_name == _KEY_TABLE:
|
||||
return aliases.key_alias_by_token.get(log.object_id)
|
||||
if log.table_name == _TEAM_TABLE:
|
||||
return aliases.team_alias_by_id.get(log.object_id)
|
||||
if log.table_name == _USER_TABLE:
|
||||
return aliases.user_alias_by_id.get(log.object_id) or aliases.user_email_by_id.get(log.object_id)
|
||||
if log.table_name == _ORG_TABLE:
|
||||
return aliases.org_alias_by_id.get(log.object_id)
|
||||
if log.table_name == _MODEL_TABLE:
|
||||
return aliases.model_name_by_id.get(log.object_id)
|
||||
return None
|
||||
|
||||
|
||||
def _alias_from_blobs(log: AuditLogResponse, blob_keys: Tuple[str, ...]) -> str | None:
|
||||
for blob in (log.updated_values, log.before_value):
|
||||
if not isinstance(blob, dict):
|
||||
continue
|
||||
for blob_key in blob_keys:
|
||||
value = blob.get(blob_key)
|
||||
if isinstance(value, str) and value:
|
||||
return value
|
||||
return None
|
||||
|
||||
|
||||
def _enrich_audit_log(log: AuditLogResponse, aliases: _AliasMaps) -> AuditLogResponse:
|
||||
object_alias: Final = _db_object_alias(log, aliases) or _alias_from_blobs(
|
||||
log, _BLOB_ALIAS_KEYS.get(log.table_name, ())
|
||||
)
|
||||
return log.model_copy(
|
||||
update={
|
||||
"object_alias": object_alias,
|
||||
"changed_by_user_email": aliases.user_email_by_id.get(log.changed_by),
|
||||
"changed_by_key_alias": aliases.key_alias_by_token.get(log.changed_by_api_key),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def _enrich_audit_logs(
|
||||
prisma_client: "PrismaClient", audit_logs: Sequence[AuditLogResponse]
|
||||
) -> List[AuditLogResponse]:
|
||||
if not audit_logs:
|
||||
return []
|
||||
aliases: Final = await _fetch_alias_maps(prisma_client, audit_logs)
|
||||
return [_enrich_audit_log(log, aliases) for log in audit_logs]
|
||||
|
||||
|
||||
_TEAM_ALIAS_CANDIDATE_LIMIT: Final[int] = 100
|
||||
_TEAM_ALIAS_CANDIDATE_SQL: Final[str] = (
|
||||
f'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_alias LIKE $1 LIMIT {_TEAM_ALIAS_CANDIDATE_LIMIT}'
|
||||
)
|
||||
|
||||
|
||||
def _contains_like_pattern(value: str) -> str:
|
||||
escaped: Final = value.replace("\\", "\\\\").replace("%", "\\%").replace("_", "\\_")
|
||||
return f"%{escaped}%"
|
||||
|
||||
|
||||
async def _build_object_team_condition(prisma_client: "PrismaClient", object_team: str) -> Dict[str, Any]:
|
||||
team_rows: Final = await prisma_client.db.query_raw(_TEAM_ALIAS_CANDIDATE_SQL, _contains_like_pattern(object_team))
|
||||
match_values: Final = dict.fromkeys([object_team, *(row["team_id"] for row in team_rows)])
|
||||
return {
|
||||
"OR": [
|
||||
_build_json_field_or_condition("team_alias", object_team),
|
||||
*(_build_json_field_or_condition("team_id", value) for value in match_values),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]:
|
||||
"""
|
||||
|
|
@ -189,6 +43,15 @@ def _build_json_field_or_condition(json_key: str, value: str) -> Dict[str, Any]:
|
|||
}
|
||||
|
||||
|
||||
def _build_object_team_condition(object_team: str) -> Dict[str, Any]:
|
||||
return {
|
||||
"OR": [
|
||||
{"object_team_id": object_team},
|
||||
{"object_team_alias": {"contains": object_team}},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
@router.get(
|
||||
"/audit",
|
||||
tags=["Audit Logging"],
|
||||
|
|
@ -213,8 +76,8 @@ async def get_audit_logs(
|
|||
object_team: str | None = Query(
|
||||
None,
|
||||
description=(
|
||||
"Filter by team id or alias: matches team_id or team_alias present in before_value or "
|
||||
"updated_values JSON, or teams whose team_alias contains this value (PostgreSQL only)"
|
||||
"Filter by team: matches the row's object_team_id exactly "
|
||||
"or rows whose object_team_alias contains this value"
|
||||
),
|
||||
),
|
||||
object_key_hash: Optional[str] = Query(
|
||||
|
|
@ -233,10 +96,9 @@ async def get_audit_logs(
|
|||
|
||||
Returns a paginated response of audit logs matching the specified filters.
|
||||
|
||||
Note: object_team_id, object_team and object_key_hash use Prisma JSON path
|
||||
filtering, which requires PostgreSQL. object_team matches a team_id or
|
||||
team_alias in the audit blobs, or any team whose team_alias contains the
|
||||
value.
|
||||
Note: object_team_id and object_key_hash use Prisma JSON path filtering,
|
||||
which requires PostgreSQL. object_team filters on the denormalized
|
||||
object_team_id and object_team_alias columns instead.
|
||||
"""
|
||||
from litellm.proxy.proxy_server import prisma_client
|
||||
|
||||
|
|
@ -277,9 +139,7 @@ async def get_audit_logs(
|
|||
_build_json_field_or_condition("token", object_key_hash)
|
||||
]
|
||||
if object_team:
|
||||
where_conditions["AND"] = where_conditions.get("AND", []) + [
|
||||
await _build_object_team_condition(prisma_client, object_team)
|
||||
]
|
||||
where_conditions["AND"] = where_conditions.get("AND", []) + [_build_object_team_condition(object_team)]
|
||||
|
||||
# Build sort conditions
|
||||
order_by: Dict[str, Any] = {}
|
||||
|
|
@ -300,14 +160,9 @@ async def get_audit_logs(
|
|||
total_count = await prisma_client.db.litellm_auditlog.count(where=where_conditions)
|
||||
total_pages = -(-total_count // page_size) # Ceiling division
|
||||
|
||||
enriched_logs: Final = await _enrich_audit_logs(
|
||||
prisma_client,
|
||||
[AuditLogResponse(**audit_log.model_dump()) for audit_log in audit_logs] if audit_logs else [],
|
||||
)
|
||||
|
||||
# Return paginated response
|
||||
return PaginatedAuditLogResponse(
|
||||
audit_logs=enriched_logs,
|
||||
audit_logs=[AuditLogResponse(**audit_log.model_dump()) for audit_log in audit_logs] if audit_logs else [],
|
||||
total=total_count,
|
||||
page=page,
|
||||
page_size=page_size,
|
||||
|
|
@ -352,5 +207,5 @@ async def get_audit_log_by_id(id: str, user_api_key_dict: UserAPIKeyAuth = Depen
|
|||
if audit_log is None:
|
||||
raise HTTPException(status_code=404, detail={"message": f"Audit log with ID {id} not found"})
|
||||
|
||||
enriched_logs: Final = await _enrich_audit_logs(prisma_client, [AuditLogResponse(**audit_log.model_dump())])
|
||||
return enriched_logs[0]
|
||||
# Convert to response model
|
||||
return AuditLogResponse(**audit_log.model_dump())
|
||||
|
|
|
|||
|
|
@ -17,6 +17,8 @@ class AuditLogResponse(BaseModel):
|
|||
before_value: Optional[Dict[str, Any]] = None
|
||||
updated_values: Optional[Dict[str, Any]] = None
|
||||
object_alias: str | None = None
|
||||
object_team_id: str | None = None
|
||||
object_team_alias: str | None = None
|
||||
changed_by_user_email: str | None = None
|
||||
changed_by_key_alias: str | None = None
|
||||
|
||||
|
|
|
|||
|
|
@ -0,0 +1,89 @@
|
|||
-- AlterTable
|
||||
ALTER TABLE "LiteLLM_AuditLog" ADD COLUMN IF NOT EXISTS "object_alias" TEXT,
|
||||
ADD COLUMN IF NOT EXISTS "object_team_id" TEXT,
|
||||
ADD COLUMN IF NOT EXISTS "object_team_alias" TEXT,
|
||||
ADD COLUMN IF NOT EXISTS "changed_by_user_email" TEXT,
|
||||
ADD COLUMN IF NOT EXISTS "changed_by_key_alias" TEXT;
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_AuditLog_table_name_object_id_idx" ON "LiteLLM_AuditLog"("table_name", "object_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_AuditLog_object_team_id_idx" ON "LiteLLM_AuditLog"("object_team_id");
|
||||
|
||||
-- CreateIndex
|
||||
CREATE INDEX IF NOT EXISTS "LiteLLM_AuditLog_updated_at_idx" ON "LiteLLM_AuditLog"("updated_at");
|
||||
|
||||
-- Backfill object_alias from the JSON blobs captured at change time (updated_values wins over before_value).
|
||||
-- Key rows are excluded here: the audit writer masks key_alias inside the blobs, so they resolve via join below
|
||||
UPDATE "LiteLLM_AuditLog"
|
||||
SET "object_alias" = CASE "table_name"
|
||||
WHEN 'LiteLLM_TeamTable' THEN COALESCE(
|
||||
NULLIF("updated_values"->>'team_alias', ''),
|
||||
NULLIF("before_value"->>'team_alias', '')
|
||||
)
|
||||
WHEN 'LiteLLM_UserTable' THEN COALESCE(
|
||||
NULLIF("updated_values"->>'user_alias', ''),
|
||||
NULLIF("updated_values"->>'user_email', ''),
|
||||
NULLIF("before_value"->>'user_alias', ''),
|
||||
NULLIF("before_value"->>'user_email', '')
|
||||
)
|
||||
WHEN 'LiteLLM_OrganizationTable' THEN COALESCE(
|
||||
NULLIF("updated_values"->>'organization_alias', ''),
|
||||
NULLIF("before_value"->>'organization_alias', '')
|
||||
)
|
||||
WHEN 'LiteLLM_ProxyModelTable' THEN COALESCE(
|
||||
NULLIF("updated_values"->>'model_name', ''),
|
||||
NULLIF("before_value"->>'model_name', '')
|
||||
)
|
||||
END
|
||||
WHERE "object_alias" IS NULL;
|
||||
|
||||
-- Backfill object_alias for key rows via the current key table (blob key_alias is masked; deleted keys stay NULL)
|
||||
UPDATE "LiteLLM_AuditLog" a
|
||||
SET "object_alias" = v."key_alias"
|
||||
FROM "LiteLLM_VerificationToken" v
|
||||
WHERE a."object_alias" IS NULL
|
||||
AND a."table_name" = 'LiteLLM_VerificationToken'
|
||||
AND a."object_id" = v."token"
|
||||
AND v."key_alias" IS NOT NULL;
|
||||
|
||||
-- Backfill object_team_id from the JSON blobs
|
||||
UPDATE "LiteLLM_AuditLog"
|
||||
SET "object_team_id" = COALESCE(
|
||||
NULLIF("updated_values"->>'team_id', ''),
|
||||
NULLIF("before_value"->>'team_id', '')
|
||||
)
|
||||
WHERE "object_team_id" IS NULL;
|
||||
|
||||
-- Backfill object_team_alias from the JSON blobs (team rows carry it even after the team is deleted)
|
||||
UPDATE "LiteLLM_AuditLog"
|
||||
SET "object_team_alias" = COALESCE(
|
||||
NULLIF("updated_values"->>'team_alias', ''),
|
||||
NULLIF("before_value"->>'team_alias', '')
|
||||
)
|
||||
WHERE "object_team_alias" IS NULL;
|
||||
|
||||
-- Backfill object_team_alias for remaining rows via the current team table
|
||||
UPDATE "LiteLLM_AuditLog" a
|
||||
SET "object_team_alias" = t."team_alias"
|
||||
FROM "LiteLLM_TeamTable" t
|
||||
WHERE a."object_team_alias" IS NULL
|
||||
AND a."object_team_id" = t."team_id"
|
||||
AND t."team_alias" IS NOT NULL;
|
||||
|
||||
-- Backfill changed_by_user_email via the current user table (deleted actors stay NULL)
|
||||
UPDATE "LiteLLM_AuditLog" a
|
||||
SET "changed_by_user_email" = u."user_email"
|
||||
FROM "LiteLLM_UserTable" u
|
||||
WHERE a."changed_by_user_email" IS NULL
|
||||
AND a."changed_by" = u."user_id"
|
||||
AND u."user_email" IS NOT NULL;
|
||||
|
||||
-- Backfill changed_by_key_alias via the current key table (deleted keys stay NULL)
|
||||
UPDATE "LiteLLM_AuditLog" a
|
||||
SET "changed_by_key_alias" = v."key_alias"
|
||||
FROM "LiteLLM_VerificationToken" v
|
||||
WHERE a."changed_by_key_alias" IS NULL
|
||||
AND a."changed_by_api_key" = v."token"
|
||||
AND v."key_alias" IS NOT NULL;
|
||||
|
|
@ -732,6 +732,15 @@ model LiteLLM_AuditLog {
|
|||
object_id String // id of the object being audited. This can be the key id, team id, user id, model id
|
||||
before_value Json? // value of the row
|
||||
updated_values Json? // value of the row after change
|
||||
object_alias String?
|
||||
object_team_id String?
|
||||
object_team_alias String?
|
||||
changed_by_user_email String?
|
||||
changed_by_key_alias String?
|
||||
|
||||
@@index([table_name, object_id])
|
||||
@@index([object_team_id])
|
||||
@@index([updated_at])
|
||||
}
|
||||
|
||||
// Track daily user spend metrics per model and key
|
||||
|
|
|
|||
|
|
@ -3148,6 +3148,11 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase):
|
|||
object_id: str
|
||||
before_value: Json | None = None
|
||||
updated_values: Json | None = None
|
||||
object_alias: str | None = None
|
||||
object_team_id: str | None = None
|
||||
object_team_alias: str | None = None
|
||||
changed_by_user_email: str | None = None
|
||||
changed_by_key_alias: str | None = None
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
|
|
|
|||
|
|
@ -5,7 +5,9 @@ Functions to create audit logs for LiteLLM Proxy
|
|||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from typing import Final
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
from pydantic import BaseModel, ValidationError
|
||||
|
||||
import litellm
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -20,6 +22,9 @@ from litellm.proxy._types import (
|
|||
from litellm.repositories.table_repositories import AuditLogRepository
|
||||
from litellm.types.utils import StandardAuditLogPayload
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
|
||||
_audit_log_callback_cache: Final[dict[str, CustomLogger]] = {}
|
||||
ALLOW_LITELLM_CHANGED_BY_HEADER_METADATA_KEY: Final = "allow_litellm_changed_by_header"
|
||||
|
||||
|
|
@ -143,6 +148,128 @@ async def _dispatch_audit_log_to_callbacks(
|
|||
verbose_proxy_logger.error("Failed dispatching audit log to callback: %s", e)
|
||||
|
||||
|
||||
class _AuditBlobAliasFields(BaseModel):
|
||||
team_alias: str | None = None
|
||||
user_alias: str | None = None
|
||||
user_email: str | None = None
|
||||
organization_alias: str | None = None
|
||||
model_name: str | None = None
|
||||
team_id: str | None = None
|
||||
|
||||
|
||||
def _blob_alias_fields(value: object) -> _AuditBlobAliasFields:
|
||||
try:
|
||||
if isinstance(value, dict):
|
||||
return _AuditBlobAliasFields.model_validate(value)
|
||||
if isinstance(value, str):
|
||||
return _AuditBlobAliasFields.model_validate_json(value)
|
||||
except ValidationError:
|
||||
return _AuditBlobAliasFields()
|
||||
return _AuditBlobAliasFields()
|
||||
|
||||
|
||||
def _derive_object_alias(table_name: str, updated: _AuditBlobAliasFields, before: _AuditBlobAliasFields) -> str | None:
|
||||
if table_name == LitellmTableNames.TEAM_TABLE_NAME.value:
|
||||
return updated.team_alias or before.team_alias or None
|
||||
if table_name == LitellmTableNames.USER_TABLE_NAME.value:
|
||||
return updated.user_alias or updated.user_email or before.user_alias or before.user_email or None
|
||||
if table_name == "LiteLLM_OrganizationTable":
|
||||
return updated.organization_alias or before.organization_alias or None
|
||||
if table_name == LitellmTableNames.PROXY_MODEL_TABLE_NAME.value:
|
||||
return updated.model_name or before.model_name or None
|
||||
return None
|
||||
|
||||
|
||||
async def _lookup_team_alias(prisma_client: "PrismaClient", team_id: str) -> str | None:
|
||||
try:
|
||||
row: Final = await prisma_client.db.litellm_teamtable.find_unique(where={"team_id": team_id})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("audit log team alias lookup failed: %s", e)
|
||||
return None
|
||||
alias: Final = None if row is None else row.team_alias
|
||||
return alias if isinstance(alias, str) and alias else None
|
||||
|
||||
|
||||
async def _lookup_user_email(prisma_client: "PrismaClient", user_id: str) -> str | None:
|
||||
try:
|
||||
row: Final = await prisma_client.db.litellm_usertable.find_unique(where={"user_id": user_id})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("audit log user email lookup failed: %s", e)
|
||||
return None
|
||||
email: Final = None if row is None else row.user_email
|
||||
return email if isinstance(email, str) and email else None
|
||||
|
||||
|
||||
async def _lookup_key_alias(prisma_client: "PrismaClient", token: str) -> str | None:
|
||||
try:
|
||||
row: Final = await prisma_client.db.litellm_verificationtoken.find_unique(where={"token": token})
|
||||
except Exception as e:
|
||||
verbose_proxy_logger.debug("audit log key alias lookup failed: %s", e)
|
||||
return None
|
||||
alias: Final = None if row is None else row.key_alias
|
||||
return alias if isinstance(alias, str) and alias else None
|
||||
|
||||
|
||||
def _serialized_blob(value: object) -> object:
|
||||
return json.dumps(value) if isinstance(value, dict) else value
|
||||
|
||||
|
||||
async def _with_denormalized_aliases(
|
||||
request_data: LiteLLM_AuditLogs, prisma_client: "PrismaClient | None"
|
||||
) -> LiteLLM_AuditLogs:
|
||||
updated: Final = _blob_alias_fields(request_data.updated_values)
|
||||
before: Final = _blob_alias_fields(request_data.before_value)
|
||||
table_name: Final = (
|
||||
request_data.table_name.value
|
||||
if isinstance(request_data.table_name, LitellmTableNames)
|
||||
else str(request_data.table_name)
|
||||
)
|
||||
is_team_row: Final = table_name == LitellmTableNames.TEAM_TABLE_NAME.value
|
||||
is_key_row: Final = table_name == LitellmTableNames.KEY_TABLE_NAME.value
|
||||
changed_by: Final = request_data.changed_by if isinstance(request_data.changed_by, str) else None
|
||||
object_alias: Final = request_data.object_alias or (
|
||||
await _lookup_key_alias(prisma_client, request_data.object_id)
|
||||
if is_key_row and prisma_client is not None
|
||||
else _derive_object_alias(table_name, updated, before)
|
||||
)
|
||||
object_team_id: Final = (
|
||||
request_data.object_team_id
|
||||
or updated.team_id
|
||||
or before.team_id
|
||||
or (request_data.object_id if is_team_row else None)
|
||||
or None
|
||||
)
|
||||
object_team_alias: Final = (
|
||||
request_data.object_team_alias
|
||||
or updated.team_alias
|
||||
or before.team_alias
|
||||
or (
|
||||
await _lookup_team_alias(prisma_client, object_team_id)
|
||||
if prisma_client is not None and object_team_id
|
||||
else None
|
||||
)
|
||||
)
|
||||
changed_by_user_email: Final = request_data.changed_by_user_email or (
|
||||
await _lookup_user_email(prisma_client, changed_by) if prisma_client is not None and changed_by else None
|
||||
)
|
||||
changed_by_key_alias: Final = request_data.changed_by_key_alias or (
|
||||
await _lookup_key_alias(prisma_client, request_data.changed_by_api_key)
|
||||
if prisma_client is not None and request_data.changed_by_api_key
|
||||
else None
|
||||
)
|
||||
return request_data.model_copy(
|
||||
update={
|
||||
"object_alias": object_alias,
|
||||
"object_team_id": object_team_id,
|
||||
"object_team_alias": object_team_alias,
|
||||
"changed_by_user_email": changed_by_user_email,
|
||||
"changed_by_key_alias": changed_by_key_alias,
|
||||
"updated_values": _serialized_blob(request_data.updated_values),
|
||||
"before_value": _serialized_blob(request_data.before_value),
|
||||
}
|
||||
)
|
||||
|
||||
|
||||
async def create_object_audit_log(
|
||||
object_id: str,
|
||||
action: AUDIT_ACTIONS,
|
||||
|
|
@ -188,6 +315,8 @@ async def create_object_audit_log(
|
|||
action=action,
|
||||
updated_values=after_value,
|
||||
before_value=before_value,
|
||||
changed_by_user_email=(user_api_key_dict.user_email if _changed_by == user_api_key_dict.user_id else None),
|
||||
changed_by_key_alias=user_api_key_dict.key_alias,
|
||||
)
|
||||
)
|
||||
|
||||
|
|
@ -209,20 +338,16 @@ async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
|
|||
|
||||
verbose_proxy_logger.debug("creating audit log for %s", request_data)
|
||||
|
||||
if isinstance(request_data.updated_values, dict):
|
||||
request_data.updated_values = json.dumps(request_data.updated_values)
|
||||
|
||||
if isinstance(request_data.before_value, dict):
|
||||
request_data.before_value = json.dumps(request_data.before_value)
|
||||
enriched: Final = await _with_denormalized_aliases(request_data, prisma_client)
|
||||
|
||||
# Dispatch to external audit log callbacks regardless of DB availability
|
||||
await _dispatch_audit_log_to_callbacks(request_data)
|
||||
await _dispatch_audit_log_to_callbacks(enriched)
|
||||
|
||||
if prisma_client is None:
|
||||
verbose_proxy_logger.error("prisma_client is None, cannot write audit log to DB")
|
||||
return
|
||||
|
||||
_request_data: Final = request_data.model_dump(exclude_none=True)
|
||||
_request_data: Final = enriched.model_dump(exclude_none=True)
|
||||
|
||||
try:
|
||||
await AuditLogRepository(prisma_client).table.create(
|
||||
|
|
|
|||
|
|
@ -732,6 +732,15 @@ model LiteLLM_AuditLog {
|
|||
object_id String // id of the object being audited. This can be the key id, team id, user id, model id
|
||||
before_value Json? // value of the row
|
||||
updated_values Json? // value of the row after change
|
||||
object_alias String?
|
||||
object_team_id String?
|
||||
object_team_alias String?
|
||||
changed_by_user_email String?
|
||||
changed_by_key_alias String?
|
||||
|
||||
@@index([table_name, object_id])
|
||||
@@index([object_team_id])
|
||||
@@index([updated_at])
|
||||
}
|
||||
|
||||
// Track daily user spend metrics per model and key
|
||||
|
|
|
|||
|
|
@ -732,6 +732,15 @@ model LiteLLM_AuditLog {
|
|||
object_id String // id of the object being audited. This can be the key id, team id, user id, model id
|
||||
before_value Json? // value of the row
|
||||
updated_values Json? // value of the row after change
|
||||
object_alias String?
|
||||
object_team_id String?
|
||||
object_team_alias String?
|
||||
changed_by_user_email String?
|
||||
changed_by_key_alias String?
|
||||
|
||||
@@index([table_name, object_id])
|
||||
@@index([object_team_id])
|
||||
@@index([updated_at])
|
||||
}
|
||||
|
||||
// Track daily user spend metrics per model and key
|
||||
|
|
|
|||
|
|
@ -34,11 +34,6 @@ def mock_prisma_client():
|
|||
mock.db.litellm_auditlog.find_many = AsyncMock()
|
||||
mock.db.litellm_auditlog.find_unique = AsyncMock()
|
||||
mock.db.litellm_auditlog.count = AsyncMock()
|
||||
mock.db.litellm_verificationtoken.find_many = AsyncMock(return_value=[])
|
||||
mock.db.litellm_usertable.find_many = AsyncMock(return_value=[])
|
||||
mock.db.litellm_teamtable.find_many = AsyncMock(return_value=[])
|
||||
mock.db.litellm_organizationtable.find_many = AsyncMock(return_value=[])
|
||||
mock.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
|
||||
yield mock
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -1,9 +1,8 @@
|
|||
"""
|
||||
Tests for audit log alias enrichment and the combined object_team filter (LIT-4997).
|
||||
Tests for the denormalized audit log alias columns and the object_team filter (LIT-4997).
|
||||
"""
|
||||
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import patch
|
||||
|
||||
from fastapi import FastAPI
|
||||
|
|
@ -11,7 +10,6 @@ from fastapi.testclient import TestClient
|
|||
from litellm_enterprise.proxy.audit_logging_endpoints import (
|
||||
_build_json_field_or_condition,
|
||||
_build_object_team_condition,
|
||||
_enrich_audit_logs,
|
||||
)
|
||||
from litellm_enterprise.proxy.audit_logging_endpoints import router as audit_router
|
||||
from litellm_enterprise.types.proxy.audit_logging_endpoints import AuditLogResponse
|
||||
|
|
@ -20,7 +18,7 @@ from litellm.proxy._types import UserAPIKeyAuth
|
|||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
|
||||
|
||||
class FakeTable:
|
||||
class FakeAuditLogTable:
|
||||
def __init__(self, rows=()):
|
||||
self.rows = list(rows)
|
||||
self.find_many_calls = []
|
||||
|
|
@ -29,14 +27,7 @@ class FakeTable:
|
|||
self.find_many_calls.append(where)
|
||||
return self.rows
|
||||
|
||||
|
||||
class FakeAuditLogTable(FakeTable):
|
||||
def __init__(self, rows=()):
|
||||
super().__init__(rows)
|
||||
self.count_calls = []
|
||||
|
||||
async def count(self, where=None):
|
||||
self.count_calls.append(where)
|
||||
return len(self.rows)
|
||||
|
||||
async def find_unique(self, where):
|
||||
|
|
@ -44,28 +35,8 @@ class FakeAuditLogTable(FakeTable):
|
|||
|
||||
|
||||
class FakeDb:
|
||||
def __init__(
|
||||
self,
|
||||
audit_logs=(),
|
||||
keys=(),
|
||||
users=(),
|
||||
teams=(),
|
||||
orgs=(),
|
||||
models=(),
|
||||
team_id_rows=(),
|
||||
):
|
||||
def __init__(self, audit_logs=()):
|
||||
self.litellm_auditlog = FakeAuditLogTable(audit_logs)
|
||||
self.litellm_verificationtoken = FakeTable(keys)
|
||||
self.litellm_usertable = FakeTable(users)
|
||||
self.litellm_teamtable = FakeTable(teams)
|
||||
self.litellm_organizationtable = FakeTable(orgs)
|
||||
self.litellm_proxymodeltable = FakeTable(models)
|
||||
self.team_id_rows = list(team_id_rows)
|
||||
self.query_raw_calls = []
|
||||
|
||||
async def query_raw(self, sql, *args):
|
||||
self.query_raw_calls.append((sql, *args))
|
||||
return self.team_id_rows
|
||||
|
||||
|
||||
class FakePrismaClient:
|
||||
|
|
@ -88,185 +59,6 @@ def make_log(**overrides) -> AuditLogResponse:
|
|||
return AuditLogResponse(**{**defaults, **overrides})
|
||||
|
||||
|
||||
async def test_enrichment_resolves_object_alias_per_table_name():
|
||||
"""object_alias comes from the right table per table_name; changed_by fields resolve too."""
|
||||
logs = [
|
||||
make_log(id="l1", table_name="LiteLLM_VerificationToken", object_id="hash-1"),
|
||||
make_log(id="l2", table_name="LiteLLM_TeamTable", object_id="team-1"),
|
||||
make_log(id="l3", table_name="LiteLLM_UserTable", object_id="user-1"),
|
||||
make_log(id="l4", table_name="LiteLLM_UserTable", object_id="user-2"),
|
||||
make_log(id="l5", table_name="LiteLLM_OrganizationTable", object_id="org-1"),
|
||||
make_log(id="l6", table_name="LiteLLM_ProxyModelTable", object_id="model-1"),
|
||||
make_log(
|
||||
id="l7",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="team-1",
|
||||
changed_by="admin-user",
|
||||
changed_by_api_key="hash-admin",
|
||||
),
|
||||
]
|
||||
db = FakeDb(
|
||||
keys=[
|
||||
SimpleNamespace(token="hash-1", key_alias="prod-key"),
|
||||
SimpleNamespace(token="hash-admin", key_alias="admin-key"),
|
||||
],
|
||||
users=[
|
||||
SimpleNamespace(user_id="user-1", user_alias="Alice", user_email="alice@example.com"),
|
||||
SimpleNamespace(user_id="user-2", user_alias=None, user_email="bob@example.com"),
|
||||
SimpleNamespace(user_id="admin-user", user_alias=None, user_email="admin@example.com"),
|
||||
],
|
||||
teams=[SimpleNamespace(team_id="team-1", team_alias="ml-team")],
|
||||
orgs=[SimpleNamespace(organization_id="org-1", organization_alias="acme-org")],
|
||||
models=[SimpleNamespace(model_id="model-1", model_name="gpt-5.2")],
|
||||
)
|
||||
|
||||
enriched = await _enrich_audit_logs(FakePrismaClient(db), logs)
|
||||
by_id = {log.id: log for log in enriched}
|
||||
|
||||
assert by_id["l1"].object_alias == "prod-key"
|
||||
assert by_id["l2"].object_alias == "ml-team"
|
||||
assert by_id["l3"].object_alias == "Alice"
|
||||
assert by_id["l4"].object_alias == "bob@example.com"
|
||||
assert by_id["l5"].object_alias == "acme-org"
|
||||
assert by_id["l6"].object_alias == "gpt-5.2"
|
||||
assert by_id["l7"].changed_by_user_email == "admin@example.com"
|
||||
assert by_id["l7"].changed_by_key_alias == "admin-key"
|
||||
assert by_id["l1"].changed_by_user_email is None
|
||||
assert by_id["l1"].changed_by_key_alias is None
|
||||
|
||||
|
||||
async def test_enrichment_runs_one_query_per_entity_type():
|
||||
"""A page with many rows triggers at most one find_many per entity table, ids batched via `in`."""
|
||||
logs = [
|
||||
make_log(id="l1", table_name="LiteLLM_VerificationToken", object_id="hash-1", changed_by="u1"),
|
||||
make_log(id="l2", table_name="LiteLLM_VerificationToken", object_id="hash-2", changed_by="u2"),
|
||||
make_log(
|
||||
id="l3",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="team-1",
|
||||
changed_by="u1",
|
||||
changed_by_api_key="hash-caller",
|
||||
),
|
||||
]
|
||||
db = FakeDb()
|
||||
|
||||
await _enrich_audit_logs(FakePrismaClient(db), logs)
|
||||
|
||||
assert len(db.litellm_verificationtoken.find_many_calls) == 1
|
||||
assert len(db.litellm_usertable.find_many_calls) == 1
|
||||
assert len(db.litellm_teamtable.find_many_calls) == 1
|
||||
assert len(db.litellm_organizationtable.find_many_calls) == 0
|
||||
assert len(db.litellm_proxymodeltable.find_many_calls) == 0
|
||||
assert set(db.litellm_verificationtoken.find_many_calls[0]["token"]["in"]) == {
|
||||
"hash-1",
|
||||
"hash-2",
|
||||
"hash-caller",
|
||||
}
|
||||
assert set(db.litellm_usertable.find_many_calls[0]["user_id"]["in"]) == {"u1", "u2"}
|
||||
assert db.litellm_teamtable.find_many_calls[0] == {"team_id": {"in": ["team-1"]}}
|
||||
|
||||
|
||||
async def test_enrichment_falls_back_to_blobs_for_deleted_objects():
|
||||
"""When DB lookups miss (deleted objects), aliases come from updated_values then before_value."""
|
||||
logs = [
|
||||
make_log(
|
||||
id="l1",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="gone-team",
|
||||
action="deleted",
|
||||
before_value={"team_id": "gone-team", "team_alias": "old-team"},
|
||||
),
|
||||
make_log(
|
||||
id="l2",
|
||||
table_name="LiteLLM_VerificationToken",
|
||||
object_id="gone-hash",
|
||||
before_value={"key_alias": "old-alias"},
|
||||
updated_values={"key_alias": "new-alias"},
|
||||
),
|
||||
make_log(
|
||||
id="l3",
|
||||
table_name="LiteLLM_UserTable",
|
||||
object_id="gone-user",
|
||||
updated_values={"user_email": "gone@example.com"},
|
||||
),
|
||||
make_log(
|
||||
id="l4",
|
||||
table_name="LiteLLM_OrganizationTable",
|
||||
object_id="gone-org",
|
||||
before_value={"organization_alias": "old-org"},
|
||||
),
|
||||
make_log(id="l5", table_name="SomeUnknownTable", object_id="x", updated_values={"team_alias": "nope"}),
|
||||
]
|
||||
|
||||
enriched = await _enrich_audit_logs(FakePrismaClient(FakeDb()), logs)
|
||||
by_id = {log.id: log for log in enriched}
|
||||
|
||||
assert by_id["l1"].object_alias == "old-team"
|
||||
assert by_id["l2"].object_alias == "new-alias"
|
||||
assert by_id["l3"].object_alias == "gone@example.com"
|
||||
assert by_id["l4"].object_alias == "old-org"
|
||||
assert by_id["l5"].object_alias is None
|
||||
|
||||
|
||||
async def test_enrichment_db_lookup_wins_over_blob():
|
||||
"""A live DB row beats a stale alias captured in the audit blobs."""
|
||||
logs = [
|
||||
make_log(
|
||||
id="l1",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
object_id="team-1",
|
||||
before_value={"team_alias": "stale-alias"},
|
||||
)
|
||||
]
|
||||
db = FakeDb(teams=[SimpleNamespace(team_id="team-1", team_alias="current-alias")])
|
||||
|
||||
enriched = await _enrich_audit_logs(FakePrismaClient(db), logs)
|
||||
|
||||
assert enriched[0].object_alias == "current-alias"
|
||||
|
||||
|
||||
async def test_build_object_team_condition_matches_id_and_alias():
|
||||
"""object_team ORs the raw value with every team_id whose team_alias contains it,
|
||||
via a projected and capped query so one request cannot load the whole team table."""
|
||||
db = FakeDb(team_id_rows=[{"team_id": "team-1"}, {"team_id": "team-2"}])
|
||||
|
||||
condition = await _build_object_team_condition(FakePrismaClient(db), "prod")
|
||||
|
||||
assert db.query_raw_calls == [
|
||||
('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_alias LIKE $1 LIMIT 100', "%prod%")
|
||||
]
|
||||
assert db.litellm_teamtable.find_many_calls == []
|
||||
assert condition == {
|
||||
"OR": [
|
||||
_build_json_field_or_condition("team_alias", "prod"),
|
||||
_build_json_field_or_condition("team_id", "prod"),
|
||||
_build_json_field_or_condition("team_id", "team-1"),
|
||||
_build_json_field_or_condition("team_id", "team-2"),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
async def test_build_object_team_condition_escapes_like_wildcards():
|
||||
"""LIKE wildcards in the user-supplied value are escaped, not treated as patterns."""
|
||||
db = FakeDb()
|
||||
|
||||
await _build_object_team_condition(FakePrismaClient(db), "pr_od%te\\am")
|
||||
|
||||
assert db.query_raw_calls[0][1] == "%pr\\_od\\%te\\\\am%"
|
||||
|
||||
|
||||
async def test_build_object_team_condition_deleted_team_matches_blob_alias():
|
||||
"""With no live team rows the condition still matches blob team_alias and the raw value as team_id."""
|
||||
condition = await _build_object_team_condition(FakePrismaClient(FakeDb()), "gone-team")
|
||||
|
||||
assert condition == {
|
||||
"OR": [
|
||||
_build_json_field_or_condition("team_alias", "gone-team"),
|
||||
_build_json_field_or_condition("team_id", "gone-team"),
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def _client_for(db: FakeDb) -> TestClient:
|
||||
app = FastAPI()
|
||||
app.include_router(audit_router)
|
||||
|
|
@ -274,81 +66,116 @@ def _client_for(db: FakeDb) -> TestClient:
|
|||
return TestClient(app)
|
||||
|
||||
|
||||
def test_get_audit_logs_object_team_filter_and_enrichment():
|
||||
"""GET /audit?object_team=... ANDs in the combined id/alias condition and returns enriched rows."""
|
||||
def test_build_object_team_condition_matches_id_and_alias_columns():
|
||||
"""object_team ORs an exact object_team_id match with an object_team_alias contains match."""
|
||||
assert _build_object_team_condition("prod") == {
|
||||
"OR": [
|
||||
{"object_team_id": "prod"},
|
||||
{"object_team_alias": {"contains": "prod"}},
|
||||
]
|
||||
}
|
||||
|
||||
|
||||
def test_get_audit_logs_returns_denormalized_columns_verbatim():
|
||||
"""GET /audit passes the alias columns straight through from the DB row."""
|
||||
audit_row = make_log(
|
||||
id="l1",
|
||||
table_name="LiteLLM_TeamTable",
|
||||
action="deleted",
|
||||
object_id="team-1",
|
||||
object_alias="ml-team",
|
||||
object_team_id="team-1",
|
||||
object_team_alias="ml-team",
|
||||
changed_by="admin-user",
|
||||
updated_values={"team_id": "team-1"},
|
||||
changed_by_api_key="hash-admin",
|
||||
changed_by_user_email="admin@example.com",
|
||||
changed_by_key_alias="admin-key",
|
||||
)
|
||||
db = FakeDb(
|
||||
audit_logs=[audit_row],
|
||||
users=[SimpleNamespace(user_id="admin-user", user_alias=None, user_email="admin@example.com")],
|
||||
teams=[SimpleNamespace(team_id="team-1", team_alias="prod-team")],
|
||||
team_id_rows=[{"team_id": "team-1"}],
|
||||
)
|
||||
client = _client_for(db)
|
||||
db = FakeDb(audit_logs=[audit_row])
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", FakePrismaClient(db)):
|
||||
response = client.get("/audit?object_team=prod")
|
||||
response = _client_for(db).get("/audit")
|
||||
|
||||
assert response.status_code == 200
|
||||
log = response.json()["audit_logs"][0]
|
||||
assert log["object_alias"] == "ml-team"
|
||||
assert log["object_team_id"] == "team-1"
|
||||
assert log["object_team_alias"] == "ml-team"
|
||||
assert log["changed_by_user_email"] == "admin@example.com"
|
||||
assert log["changed_by_key_alias"] == "admin-key"
|
||||
|
||||
|
||||
def test_get_audit_logs_object_team_filter_uses_columns():
|
||||
"""GET /audit?object_team=... ANDs in the column-based id-or-alias condition."""
|
||||
db = FakeDb()
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", FakePrismaClient(db)):
|
||||
response = _client_for(db).get("/audit?object_team=prod")
|
||||
|
||||
assert response.status_code == 200
|
||||
where = db.litellm_auditlog.find_many_calls[0]
|
||||
assert where["AND"] == [
|
||||
{
|
||||
"OR": [
|
||||
_build_json_field_or_condition("team_alias", "prod"),
|
||||
_build_json_field_or_condition("team_id", "prod"),
|
||||
_build_json_field_or_condition("team_id", "team-1"),
|
||||
{"object_team_id": "prod"},
|
||||
{"object_team_alias": {"contains": "prod"}},
|
||||
]
|
||||
}
|
||||
]
|
||||
log = response.json()["audit_logs"][0]
|
||||
assert log["object_alias"] == "prod-team"
|
||||
assert log["changed_by_user_email"] == "admin@example.com"
|
||||
assert log["changed_by_key_alias"] is None
|
||||
|
||||
|
||||
def test_get_audit_logs_object_team_id_filter_unchanged():
|
||||
"""The pre-existing object_team_id param still builds its exact condition, no alias lookup."""
|
||||
"""The pre-existing object_team_id param still builds its exact JSON-blob condition."""
|
||||
db = FakeDb()
|
||||
client = _client_for(db)
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", FakePrismaClient(db)):
|
||||
response = client.get("/audit?object_team_id=team-1")
|
||||
response = _client_for(db).get("/audit?object_team_id=team-1")
|
||||
|
||||
assert response.status_code == 200
|
||||
where = db.litellm_auditlog.find_many_calls[0]
|
||||
assert where["AND"] == [_build_json_field_or_condition("team_id", "team-1")]
|
||||
assert db.litellm_teamtable.find_many_calls == []
|
||||
assert db.query_raw_calls == []
|
||||
assert where["AND"] == [
|
||||
{
|
||||
"OR": [
|
||||
{"before_value": {"path": ["team_id"], "string_contains": "team-1"}},
|
||||
{"updated_values": {"path": ["team_id"], "string_contains": "team-1"}},
|
||||
]
|
||||
}
|
||||
]
|
||||
|
||||
|
||||
def test_get_audit_log_by_id_is_enriched():
|
||||
"""GET /audit/{id} carries the same alias enrichment as the list endpoint."""
|
||||
def test_get_audit_log_by_id_returns_denormalized_columns():
|
||||
"""GET /audit/{id} carries the same alias columns as the list endpoint."""
|
||||
audit_row = make_log(
|
||||
id="l1",
|
||||
table_name="LiteLLM_VerificationToken",
|
||||
object_id="gone-hash",
|
||||
action="deleted",
|
||||
before_value={"key_alias": "deleted-key"},
|
||||
changed_by="admin-user",
|
||||
changed_by_api_key="hash-admin",
|
||||
object_alias="deleted-key",
|
||||
changed_by_user_email="admin@example.com",
|
||||
changed_by_key_alias="admin-key",
|
||||
)
|
||||
db = FakeDb(
|
||||
audit_logs=[audit_row],
|
||||
keys=[SimpleNamespace(token="hash-admin", key_alias="admin-key")],
|
||||
users=[SimpleNamespace(user_id="admin-user", user_alias=None, user_email="admin@example.com")],
|
||||
)
|
||||
client = _client_for(db)
|
||||
db = FakeDb(audit_logs=[audit_row])
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", FakePrismaClient(db)):
|
||||
response = client.get("/audit/l1")
|
||||
response = _client_for(db).get("/audit/l1")
|
||||
|
||||
assert response.status_code == 200
|
||||
body = response.json()
|
||||
assert body["object_alias"] == "deleted-key"
|
||||
assert body["changed_by_user_email"] == "admin@example.com"
|
||||
assert body["changed_by_key_alias"] == "admin-key"
|
||||
|
||||
|
||||
def test_alias_columns_default_to_none_for_legacy_rows():
|
||||
"""Rows written before the migration serialize with null alias columns, not errors."""
|
||||
db = FakeDb(audit_logs=[make_log(id="l1")])
|
||||
|
||||
with patch("litellm.proxy.proxy_server.prisma_client", FakePrismaClient(db)):
|
||||
response = _client_for(db).get("/audit")
|
||||
|
||||
log = response.json()["audit_logs"][0]
|
||||
assert log["object_alias"] is None
|
||||
assert log["object_team_id"] is None
|
||||
assert log["object_team_alias"] is None
|
||||
assert log["changed_by_user_email"] is None
|
||||
assert log["changed_by_key_alias"] is None
|
||||
|
|
|
|||
|
|
@ -7,18 +7,20 @@ Tests the flow: create_audit_log_for_update -> _dispatch_audit_log_to_callbacks
|
|||
import asyncio
|
||||
import json
|
||||
from datetime import datetime, timezone
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames
|
||||
from litellm.proxy._types import LiteLLM_AuditLogs, LitellmTableNames, UserAPIKeyAuth
|
||||
from litellm.proxy.management_helpers.audit_logs import (
|
||||
_audit_log_task_done_callback,
|
||||
_build_audit_log_payload,
|
||||
_dispatch_audit_log_to_callbacks,
|
||||
create_audit_log_for_update,
|
||||
create_object_audit_log,
|
||||
)
|
||||
from litellm.types.utils import StandardAuditLogPayload
|
||||
|
||||
|
|
@ -130,9 +132,7 @@ class TestDispatchAuditLogToCallbacks:
|
|||
async def test_nonblocking_on_callback_failure(self):
|
||||
"""Callback errors should not propagate."""
|
||||
mock_logger = MagicMock(spec=CustomLogger)
|
||||
mock_logger.async_log_audit_log_event = AsyncMock(
|
||||
side_effect=RuntimeError("boom")
|
||||
)
|
||||
mock_logger.async_log_audit_log_event = AsyncMock(side_effect=RuntimeError("boom"))
|
||||
litellm.audit_log_callbacks = [mock_logger]
|
||||
|
||||
audit_log = _make_audit_log()
|
||||
|
|
@ -236,9 +236,7 @@ class TestCreateAuditLogForUpdateWithCallbacks:
|
|||
patch("litellm.store_audit_logs", True),
|
||||
patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma,
|
||||
):
|
||||
mock_prisma.db.litellm_auditlog.create = AsyncMock(
|
||||
side_effect=RuntimeError("DB connection lost")
|
||||
)
|
||||
mock_prisma.db.litellm_auditlog.create = AsyncMock(side_effect=RuntimeError("DB connection lost"))
|
||||
|
||||
audit_log = _make_audit_log()
|
||||
await create_audit_log_for_update(audit_log)
|
||||
|
|
@ -254,9 +252,7 @@ class TestAuditLogTaskDoneCallback:
|
|||
mock_task = MagicMock(spec=asyncio.Task)
|
||||
mock_task.exception.return_value = RuntimeError("callback failed")
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger"
|
||||
) as mock_logger:
|
||||
with patch("litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger") as mock_logger:
|
||||
_audit_log_task_done_callback(mock_task)
|
||||
mock_logger.error.assert_called_once()
|
||||
assert "callback failed" in str(mock_logger.error.call_args)
|
||||
|
|
@ -266,9 +262,7 @@ class TestAuditLogTaskDoneCallback:
|
|||
mock_task = MagicMock(spec=asyncio.Task)
|
||||
mock_task.exception.return_value = None
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger"
|
||||
) as mock_logger:
|
||||
with patch("litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger") as mock_logger:
|
||||
_audit_log_task_done_callback(mock_task)
|
||||
mock_logger.error.assert_not_called()
|
||||
|
||||
|
|
@ -277,9 +271,7 @@ class TestAuditLogTaskDoneCallback:
|
|||
mock_task = MagicMock(spec=asyncio.Task)
|
||||
mock_task.exception.side_effect = asyncio.CancelledError()
|
||||
|
||||
with patch(
|
||||
"litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger"
|
||||
) as mock_logger:
|
||||
with patch("litellm.proxy.management_helpers.audit_logs.verbose_proxy_logger") as mock_logger:
|
||||
_audit_log_task_done_callback(mock_task)
|
||||
mock_logger.error.assert_not_called()
|
||||
|
||||
|
|
@ -424,7 +416,6 @@ class TestS3AuditCallbackParamsDecoupling:
|
|||
def test_empty_dict_opts_in(self):
|
||||
"""`s3_audit_callback_params = {}` is opt-in (truthy-by-presence) and
|
||||
produces a separate instance with no bucket configured (env/IAM-only)."""
|
||||
from litellm.integrations.s3_v2 import S3Logger
|
||||
from litellm.litellm_core_utils.litellm_logging import (
|
||||
_init_custom_logger_compatible_class,
|
||||
)
|
||||
|
|
@ -469,3 +460,228 @@ class TestS3AuditCallbackParamsDecoupling:
|
|||
assert second is not None
|
||||
assert id(second) != id(first)
|
||||
assert second.s3_bucket_name == "second"
|
||||
|
||||
|
||||
class _FakeFindUniqueTable:
|
||||
def __init__(self, row=None):
|
||||
self.row = row
|
||||
self.where_calls = []
|
||||
|
||||
async def find_unique(self, where):
|
||||
self.where_calls.append(where)
|
||||
return self.row
|
||||
|
||||
|
||||
class _FakeAuditCreateTable:
|
||||
def __init__(self):
|
||||
self.created = []
|
||||
|
||||
async def create(self, data):
|
||||
self.created.append(data)
|
||||
|
||||
|
||||
class _FakeDb:
|
||||
def __init__(self, team_row=None, user_row=None, key_row=None):
|
||||
self.litellm_auditlog = _FakeAuditCreateTable()
|
||||
self.litellm_teamtable = _FakeFindUniqueTable(team_row)
|
||||
self.litellm_usertable = _FakeFindUniqueTable(user_row)
|
||||
self.litellm_verificationtoken = _FakeFindUniqueTable(key_row)
|
||||
|
||||
|
||||
class _FakePrismaClient:
|
||||
def __init__(self, db):
|
||||
self.db = db
|
||||
|
||||
|
||||
def _gates(fake_prisma):
|
||||
return (
|
||||
patch("litellm.proxy.proxy_server.premium_user", True),
|
||||
patch("litellm.store_audit_logs", True),
|
||||
patch("litellm.proxy.proxy_server.prisma_client", fake_prisma),
|
||||
)
|
||||
|
||||
|
||||
class TestAuditLogAliasDenormalization:
|
||||
@pytest.mark.asyncio
|
||||
async def test_team_delete_denormalizes_from_blob_without_lookups(self):
|
||||
"""A team delete row gets object_alias/object_team_id/object_team_alias straight from its blob."""
|
||||
db = _FakeDb(
|
||||
user_row=SimpleNamespace(user_email="admin@example.com"),
|
||||
key_row=SimpleNamespace(key_alias="admin-key"),
|
||||
)
|
||||
p1, p2, p3 = _gates(_FakePrismaClient(db))
|
||||
with p1, p2, p3:
|
||||
await create_audit_log_for_update(
|
||||
LiteLLM_AuditLogs(
|
||||
id="a1",
|
||||
updated_at=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
changed_by="admin-user",
|
||||
changed_by_api_key="hash-admin",
|
||||
action="deleted",
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
object_id="team-1",
|
||||
before_value=json.dumps({"team_id": "team-1", "team_alias": "ml-team"}),
|
||||
)
|
||||
)
|
||||
|
||||
data = db.litellm_auditlog.created[0]
|
||||
assert data["object_alias"] == "ml-team"
|
||||
assert data["object_team_id"] == "team-1"
|
||||
assert data["object_team_alias"] == "ml-team"
|
||||
assert data["changed_by_user_email"] == "admin@example.com"
|
||||
assert data["changed_by_key_alias"] == "admin-key"
|
||||
assert db.litellm_teamtable.where_calls == []
|
||||
assert db.litellm_usertable.where_calls == [{"user_id": "admin-user"}]
|
||||
assert db.litellm_verificationtoken.where_calls == [{"token": "hash-admin"}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_key_create_resolves_aliases_via_lookup_not_masked_blob(self):
|
||||
"""Key rows resolve object_alias through the key table because the blob's key_alias is masked,
|
||||
and resolve object_team_alias through the team table from the blob's team_id."""
|
||||
db = _FakeDb(
|
||||
team_row=SimpleNamespace(team_alias="team-nine"),
|
||||
key_row=SimpleNamespace(key_alias="prod-key"),
|
||||
)
|
||||
p1, p2, p3 = _gates(_FakePrismaClient(db))
|
||||
with p1, p2, p3:
|
||||
await create_audit_log_for_update(
|
||||
LiteLLM_AuditLogs(
|
||||
id="a2",
|
||||
updated_at=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
action="created",
|
||||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||||
object_id="hash-1",
|
||||
updated_values=json.dumps({"key_alias": "prod-key", "team_id": "team-9"}),
|
||||
)
|
||||
)
|
||||
|
||||
data = db.litellm_auditlog.created[0]
|
||||
assert json.loads(data["updated_values"])["key_alias"] == "********"
|
||||
assert data["object_alias"] == "prod-key"
|
||||
assert data["object_team_id"] == "team-9"
|
||||
assert data["object_team_alias"] == "team-nine"
|
||||
assert db.litellm_verificationtoken.where_calls == [{"token": "hash-1"}]
|
||||
assert db.litellm_teamtable.where_calls == [{"team_id": "team-9"}]
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleted_key_object_alias_absent_never_masked_junk(self):
|
||||
"""A key delete whose row is already gone leaves object_alias unset instead of writing the masked blob value."""
|
||||
db = _FakeDb()
|
||||
p1, p2, p3 = _gates(_FakePrismaClient(db))
|
||||
with p1, p2, p3:
|
||||
await create_audit_log_for_update(
|
||||
LiteLLM_AuditLogs(
|
||||
id="a2b",
|
||||
updated_at=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
action="deleted",
|
||||
table_name=LitellmTableNames.KEY_TABLE_NAME,
|
||||
object_id="gone-hash",
|
||||
before_value=json.dumps({"key_alias": "prod-key"}),
|
||||
)
|
||||
)
|
||||
|
||||
data = db.litellm_auditlog.created[0]
|
||||
assert "object_alias" not in data
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_user_row_prefers_alias_then_email_updated_blob_first(self):
|
||||
"""User rows derive object_alias as user_alias then user_email, updated_values before before_value."""
|
||||
db = _FakeDb()
|
||||
p1, p2, p3 = _gates(_FakePrismaClient(db))
|
||||
with p1, p2, p3:
|
||||
await create_audit_log_for_update(
|
||||
LiteLLM_AuditLogs(
|
||||
id="a3",
|
||||
updated_at=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
action="updated",
|
||||
table_name=LitellmTableNames.USER_TABLE_NAME,
|
||||
object_id="user-1",
|
||||
before_value=json.dumps({"user_alias": "Old Alias", "user_email": "old@example.com"}),
|
||||
updated_values=json.dumps({"user_email": "new@example.com"}),
|
||||
)
|
||||
)
|
||||
|
||||
assert db.litellm_auditlog.created[0]["object_alias"] == "new@example.com"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_caller_provided_actor_fields_skip_lookups(self):
|
||||
"""Actor email and key alias passed by the caller are kept verbatim with no DB lookups."""
|
||||
db = _FakeDb(
|
||||
user_row=SimpleNamespace(user_email="wrong@example.com"),
|
||||
key_row=SimpleNamespace(key_alias="wrong-key"),
|
||||
)
|
||||
p1, p2, p3 = _gates(_FakePrismaClient(db))
|
||||
with p1, p2, p3:
|
||||
await create_audit_log_for_update(
|
||||
LiteLLM_AuditLogs(
|
||||
id="a4",
|
||||
updated_at=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
changed_by="admin-user",
|
||||
changed_by_api_key="hash-admin",
|
||||
action="updated",
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
object_id="team-1",
|
||||
updated_values=json.dumps({"team_alias": "ml-team"}),
|
||||
changed_by_user_email="right@example.com",
|
||||
changed_by_key_alias="right-key",
|
||||
)
|
||||
)
|
||||
|
||||
data = db.litellm_auditlog.created[0]
|
||||
assert data["changed_by_user_email"] == "right@example.com"
|
||||
assert data["changed_by_key_alias"] == "right-key"
|
||||
assert db.litellm_usertable.where_calls == []
|
||||
assert db.litellm_verificationtoken.where_calls == []
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deleted_actor_stays_absent(self):
|
||||
"""When actor lookups miss, the actor columns are left unset rather than written as empty."""
|
||||
db = _FakeDb()
|
||||
p1, p2, p3 = _gates(_FakePrismaClient(db))
|
||||
with p1, p2, p3:
|
||||
await create_audit_log_for_update(
|
||||
LiteLLM_AuditLogs(
|
||||
id="a5",
|
||||
updated_at=datetime(2026, 8, 1, tzinfo=timezone.utc),
|
||||
changed_by="gone-user",
|
||||
changed_by_api_key="gone-hash",
|
||||
action="deleted",
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
object_id="team-1",
|
||||
before_value=json.dumps({"team_id": "team-1", "team_alias": "ml-team"}),
|
||||
)
|
||||
)
|
||||
|
||||
data = db.litellm_auditlog.created[0]
|
||||
assert "changed_by_user_email" not in data
|
||||
assert "changed_by_key_alias" not in data
|
||||
assert data["object_alias"] == "ml-team"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_create_object_audit_log_uses_auth_context(self):
|
||||
"""create_object_audit_log stamps actor email and key alias from UserAPIKeyAuth without lookups."""
|
||||
db = _FakeDb()
|
||||
p1, p2, p3 = _gates(_FakePrismaClient(db))
|
||||
with p1, p2, p3:
|
||||
await create_object_audit_log(
|
||||
object_id="team-1",
|
||||
action="updated",
|
||||
litellm_changed_by=None,
|
||||
user_api_key_dict=UserAPIKeyAuth(
|
||||
api_key="hash-admin",
|
||||
user_id="admin-user",
|
||||
user_email="admin@example.com",
|
||||
key_alias="admin-key",
|
||||
),
|
||||
litellm_proxy_admin_name="default_user_id",
|
||||
table_name=LitellmTableNames.TEAM_TABLE_NAME,
|
||||
after_value=json.dumps({"team_id": "team-1", "team_alias": "ml-team"}),
|
||||
)
|
||||
|
||||
data = db.litellm_auditlog.created[0]
|
||||
assert data["changed_by"] == "admin-user"
|
||||
assert data["changed_by_user_email"] == "admin@example.com"
|
||||
assert data["changed_by_key_alias"] == "admin-key"
|
||||
assert data["object_alias"] == "ml-team"
|
||||
assert db.litellm_usertable.where_calls == []
|
||||
assert db.litellm_verificationtoken.where_calls == []
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue