phoenix/litellm
1
0
Fork 0
forked from phoenix/litellm-mirror
Code Pull requests Activity

feat - working audit logs for create, update delete team

Browse source
This commit is contained in:
Ishaan Jaff 2024-06-05 17:50:27 -07:00
parent 970c3dfdc7
commit 5bd658493f
6 changed files with 128 additions and 32 deletions
Download patch file Download diff file Expand all files Collapse all files

1
litellm/__init__.py
View file

@ -60,6 +60,7 @@ _async_failure_callback: List[Callable] = (
pre_call_rules: List[Callable] = []
post_call_rules: List[Callable] = []
turn_off_message_logging: Optional[bool] = False
store_audit_logs = False # Enterprise feature, allow users to see audit logs
## end of callbacks #############
email: Optional[str] = (

15
litellm/proxy/_types.py
View file

@ -1278,6 +1278,21 @@ class LiteLLM_ErrorLogs(LiteLLMBase):
endTime: Union[str, datetime, None]
class LiteLLM_AuditLogs(LiteLLMBase):
id: str
updated_at: datetime
changed_by: str
action: Literal["created", "updated", "deleted"]
table_name: Literal[
LitellmTableNames.TEAM_TABLE_NAME,
LitellmTableNames.USER_TABLE_NAME,
LitellmTableNames.PROXY_MODEL_TABLE_NAME,
]
object_id: str
before_value: Optional[Json] = None
updated_values: Optional[Json] = None
class LiteLLM_SpendLogs_ResponseObject(LiteLLMBase):
response: Optional[List[Union[LiteLLM_SpendLogs, Any]]] = None

3
litellm/proxy/proxy_config.yaml
View file

@ -23,4 +23,5 @@ general_settings:
master_key: sk-1234
litellm_settings:
callbacks: ["otel"]
callbacks: ["otel"]
store_audit_logs: true

121
litellm/proxy/proxy_server.py
View file

@ -10357,45 +10357,62 @@ async def new_team(
}
},
)
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
if litellm.store_audit_logs is True:
_updated_values = complete_team_data.json(exclude_none=True)
_updated_values = json.dumps(_updated_values)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
table_name=LitellmTableNames.TEAM_TABLE_NAME,
object_id=data.team_id,
action="created",
updated_values=_updated_values,
before_value=None,
)
)
)
try:
return team_row.model_dump()
except Exception as e:
return team_row.dict()
async def create_audit_log_for_update(
action: Literal["create", "update", "delete"],
# fyi: pylint does not directly allow you to pass Literal["LiteLLM_TeamTable"]
# because LiteLLM_TeamTable is also defined in _types.py
table_name: Literal[
LitellmTableNames.TEAM_TABLE_NAME,
LitellmTableNames.USER_TABLE_NAME,
LitellmTableNames.PROXY_MODEL_TABLE_NAME,
],
object_id: str,
changed_by: str,
before_value: dict,
after_value: dict,
):
if not premium_user:
async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs):
if premium_user is not True:
return
try:
pass
if litellm.store_audit_logs is not True:
return
if prisma_client is None:
raise Exception("prisma_client is None, no DB connected")
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)
_request_data = request_data.dict(exclude_none=True)
try:
await prisma_client.db.litellm_auditlog.create(
data={
**_request_data, # type: ignore
}
)
except Exception as e:
# [Non-Blocking Exception. Do not allow blocking LLM API call]
verbose_proxy_logger.error(f"Failed Creating audit log {e}")
# await prisma_client.create_audit_log(
# data={
# "action": action,
# "object_id": object_id,
# "changed_by": changed_by,
# "before_value": json.dumps(before_value),
# "after_value": json.dumps(after_value),
# }
# )
return
@ -10471,7 +10488,28 @@ async def update_team(
team_id=data.team_id,
)
return team_row
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
if litellm.store_audit_logs is True:
_before_value = existing_team_row.json(exclude_none=True)
_before_value = json.dumps(_before_value)
_after_value: str = json.dumps(updated_kv)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=user_api_key_dict.user_id or litellm_proxy_admin_name,
table_name=LitellmTableNames.TEAM_TABLE_NAME,
object_id=data.team_id,
action="updated",
updated_values=_after_value,
before_value=_before_value,
)
)
)
return team_row
@router.post(
@ -10742,6 +10780,35 @@ async def delete_team(
detail={"error": f"Team not found, passed team_id={team_id}"},
)
# Enterprise Feature - Audit Logging. Enable with litellm.store_audit_logs = True
# we do this after the first for loop, since first for loop is for validation. we only want this inserted after validation passes
if litellm.store_audit_logs is True:
# make an audit log for each team deleted
for team_id in data.team_ids:
team_row = await prisma_client.get_data( # type: ignore
team_id=team_id, table_name="team", query_type="find_unique"
)
_team_row = team_row.json(exclude_none=True)
asyncio.create_task(
create_audit_log_for_update(
request_data=LiteLLM_AuditLogs(
id=str(uuid.uuid4()),
updated_at=datetime.now(timezone.utc),
changed_by=user_api_key_dict.user_id
or litellm_proxy_admin_name,
table_name=LitellmTableNames.TEAM_TABLE_NAME,
object_id=team_id,
action="deleted",
updated_values="{}",
before_value=_team_row,
)
)
)
# End of Audit logging
## DELETE ASSOCIATED KEYS
await prisma_client.delete_data(team_id_list=data.team_ids, table_name="key")
## DELETE TEAMS

12
litellm/proxy/schema.prisma
View file

@ -243,4 +243,16 @@ model LiteLLM_InvitationLink {
liteLLM_user_table_user LiteLLM_UserTable @relation("UserId", fields: [user_id], references: [user_id])
liteLLM_user_table_created LiteLLM_UserTable @relation("CreatedBy", fields: [created_by], references: [user_id])
liteLLM_user_table_updated LiteLLM_UserTable @relation("UpdatedBy", fields: [updated_by], references: [user_id])
}
model LiteLLM_AuditLog {
id String @id @default(uuid())
updated_at DateTime @default(now())
changed_by String // user or system that performed the action
action String // create, update, delete
table_name String // on of LitellmTableNames.TEAM_TABLE_NAME, LitellmTableNames.USER_TABLE_NAME, LitellmTableNames.PROXY_MODEL_TABLE_NAME,
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
}

8
schema.prisma
View file

@ -246,13 +246,13 @@ model LiteLLM_InvitationLink {
}
model AuditLog {
id Int @id @default(autoincrement())
model LiteLLM_AuditLog {
id String @id @default(uuid())
updated_at DateTime @default(now())
changed_by String // user or system that performed the action
action String // create, update, delete
object_type String // team, key, user, model
table_name String // on of LitellmTableNames.TEAM_TABLE_NAME, LitellmTableNames.USER_TABLE_NAME, LitellmTableNames.PROXY_MODEL_TABLE_NAME,
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
after_value Json? // value of the row after change
updated_values Json? // value of the row after change
}