From 5cfae0e98a50ffe3929400f9c636fa5c01cad78d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 13 Mar 2025 16:47:58 -0700 Subject: [PATCH 1/7] feat(internal_user_endpoints.py): emit audit log on `/user/new` event --- litellm/proxy/_new_secret_config.yaml | 12 +--- litellm/proxy/_types.py | 7 +- .../common_utils/encrypt_decrypt_utils.py | 5 +- .../internal_user_endpoints.py | 70 +++++++++++++++++++ 4 files changed, 81 insertions(+), 13 deletions(-) diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 9436e8af0f..718e22806d 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -5,13 +5,5 @@ model_list: aws_region_name: "us-east-1" litellm_credential_name: "azure" -credential_list: - - credential_name: azure - credential_values: - api_key: os.environ/AZURE_API_KEY - api_base: os.environ/AZURE_API_BASE - credential_info: - description: "Azure API Key and Base URL" - type: "azure" - required: true - default: "azure" \ No newline at end of file +litellm_settings: + store_audit_logs: true diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a9fe6517ea..8611fbd716 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1584,7 +1584,7 @@ class NewOrganizationResponse(LiteLLM_OrganizationTable): class LiteLLM_UserTable(LiteLLMPydanticObjectBase): user_id: str - max_budget: Optional[float] + max_budget: Optional[float] = None spend: float = 0.0 model_max_budget: Optional[Dict] = {} model_spend: Optional[Dict] = {} @@ -1677,12 +1677,15 @@ class LiteLLM_ErrorLogs(LiteLLMPydanticObjectBase): endTime: Union[str, datetime, None] +AUDIT_ACTIONS = Literal["created", "updated", "deleted", "blocked"] + + class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): id: str updated_at: datetime changed_by: Optional[Any] = None changed_by_api_key: Optional[str] = None - action: Literal["created", "updated", "deleted", "blocked"] + action: AUDIT_ACTIONS table_name: Literal[ LitellmTableNames.TEAM_TABLE_NAME, LitellmTableNames.USER_TABLE_NAME, diff --git a/litellm/proxy/common_utils/encrypt_decrypt_utils.py b/litellm/proxy/common_utils/encrypt_decrypt_utils.py index 7be60d9dc6..ec9279a089 100644 --- a/litellm/proxy/common_utils/encrypt_decrypt_utils.py +++ b/litellm/proxy/common_utils/encrypt_decrypt_utils.py @@ -53,8 +53,11 @@ def decrypt_value_helper(value: str): # if it's not str - do not decrypt it, return the value return value except Exception as e: + import traceback + + traceback.print_stack() verbose_proxy_logger.error( - f"Error decrypting value, Did your master_key/salt key change recently? : {value}\nError: {str(e)}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" + f"Error decrypting value, Did your master_key/salt key change recently? \nError: {str(e)}\nSet permanent salt key - https://docs.litellm.ai/docs/proxy/prod#5-set-litellm-salt-key" ) # [Non-Blocking Exception. - this should not block decrypting other values] pass diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 85b58ba596..aa4c3c343b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -29,12 +29,53 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( generate_key_helper_fn, prepare_metadata_fields, ) +from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update from litellm.proxy.management_helpers.utils import management_endpoint_wrapper from litellm.proxy.utils import handle_exception_on_proxy router = APIRouter() +async def create_internal_user_audit_log( + user_id: str, + action: AUDIT_ACTIONS, + litellm_changed_by: Optional[str], + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: Optional[str], + before_value: Optional[str] = None, + after_value: Optional[str] = None, +): + """ + Create an audit log for an internal user. + + Parameters: + - user_id: str - The id of the user to create the audit log for. + - action: AUDIT_ACTIONS - The action to create the audit log for. + - user_row: LiteLLM_UserTable - The user row to create the audit log for. + - litellm_changed_by: Optional[str] - The user id of the user who is changing the user. + - user_api_key_dict: UserAPIKeyAuth - The user api key dictionary. + - litellm_proxy_admin_name: Optional[str] - The name of the proxy admin. + """ + if not litellm.store_audit_logs: + return + + await create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by + or user_api_key_dict.user_id + or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=LitellmTableNames.USER_TABLE_NAME, + object_id=user_id, + action=action, + updated_values=after_value, + before_value=before_value, + ) + ) + + def _update_internal_new_user_params(data_json: dict, data: NewUserRequest) -> dict: if "user_id" in data_json and data_json["user_id"] is None: data_json["user_id"] = str(uuid.uuid4()) @@ -169,6 +210,7 @@ async def new_user( try: from litellm.proxy.proxy_server import ( general_settings, + litellm_proxy_admin_name, prisma_client, proxy_logging_obj, ) @@ -254,6 +296,34 @@ async def new_user( ) ) + try: + if prisma_client is None: + raise Exception(CommonProxyErrors.db_not_connected_error.value) + user_row: BaseModel = await prisma_client.db.litellm_usertable.find_first( + where={"user_id": response["user_id"]} + ) + + user_row_litellm_typed = LiteLLM_UserTable( + **user_row.model_dump(exclude_none=True) + ) + asyncio.create_task( + create_internal_user_audit_log( + user_id=user_row_litellm_typed.user_id, + action="created", + litellm_changed_by=user_api_key_dict.user_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + before_value=None, + after_value=user_row_litellm_typed.model_dump_json( + exclude_none=True + ), + ) + ) + except Exception as e: + verbose_proxy_logger.warning( + "Unable to create audit log for user on `/user/new` - {}".format(str(e)) + ) + return NewUserResponse( key=response.get("token", ""), expires=response.get("expires", None), From e90b3d9c4cab34c563b46fe4a8008237a1d2e30b Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 13 Mar 2025 18:17:05 -0700 Subject: [PATCH 2/7] feat(internal_user_endpoints.py): add audit logs on /user/update --- litellm/proxy/_types.py | 1 + .../internal_user_endpoints.py | 50 +++++++++++++++++-- 2 files changed, 47 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 8611fbd716..c071195aa4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1598,6 +1598,7 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase): sso_user_id: Optional[str] = None budget_duration: Optional[str] = None budget_reset_at: Optional[datetime] = None + metadata: Optional[dict] = None @model_validator(mode="before") @classmethod diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index aa4c3c343b..309f451dd9 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -719,7 +719,7 @@ async def user_update( """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import litellm_proxy_admin_name, prisma_client try: data_json: dict = data.json() @@ -732,9 +732,16 @@ async def user_update( data_json=data_json, data=data ) - existing_user_row = await prisma_client.get_data( - user_id=data.user_id, table_name="user", query_type="find_unique" - ) + existing_user_row: Optional[BaseModel] = None + if data.user_id is not None: + existing_user_row = await prisma_client.db.litellm_usertable.find_first( + where={"user_id": data.user_id} + ) + if existing_user_row is None: + raise Exception(f"User not found, passed user_id={data.user_id}") + existing_user_row = LiteLLM_UserTable( + **existing_user_row.model_dump(exclude_none=True) + ) existing_metadata = existing_user_row.metadata if existing_user_row else {} @@ -781,6 +788,41 @@ async def user_update( data=non_default_values, table_name="user", ) + + if response is not None: # emit audit log + try: + user_row: BaseModel = ( + await prisma_client.db.litellm_usertable.find_first( + where={"user_id": response["user_id"]} + ) + ) + + user_row_litellm_typed = LiteLLM_UserTable( + **user_row.model_dump(exclude_none=True) + ) + asyncio.create_task( + create_internal_user_audit_log( + user_id=user_row_litellm_typed.user_id, + action="updated", + litellm_changed_by=user_api_key_dict.user_id, + user_api_key_dict=user_api_key_dict, + litellm_proxy_admin_name=litellm_proxy_admin_name, + before_value=( + existing_user_row.model_dump_json(exclude_none=True) + if existing_user_row + else None + ), + after_value=user_row_litellm_typed.model_dump_json( + exclude_none=True + ), + ) + ) + except Exception as e: + verbose_proxy_logger.warning( + "Unable to create audit log for user on `/user/update` - {}".format( + str(e) + ) + ) return response # type: ignore # update based on remaining passed in values except Exception as e: From 37b30395c9facc3778e40d111e904d20a4fd199d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 13 Mar 2025 18:48:38 -0700 Subject: [PATCH 3/7] feat(model_management_endpoints.py): emit audit logs on model delete --- litellm/proxy/_types.py | 11 +- .../management_endpoints/common_utils.py | 9 +- .../internal_user_endpoints.py | 2 +- .../model_management_endpoints.py | 106 ++++++++++++++++++ .../proxy/management_helpers/audit_logs.py | 57 +++++++++- litellm/proxy/proxy_server.py | 80 +------------ 6 files changed, 175 insertions(+), 90 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index c071195aa4..7d8c394cf0 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -133,7 +133,7 @@ class LitellmTableNames(str, enum.Enum): TEAM_TABLE_NAME = "LiteLLM_TeamTable" USER_TABLE_NAME = "LiteLLM_UserTable" KEY_TABLE_NAME = "LiteLLM_VerificationToken" - PROXY_MODEL_TABLE_NAME = "LiteLLM_ModelTable" + PROXY_MODEL_TABLE_NAME = "LiteLLM_ProxyModelTable" def hash_token(token: str): @@ -1588,7 +1588,7 @@ class LiteLLM_UserTable(LiteLLMPydanticObjectBase): spend: float = 0.0 model_max_budget: Optional[Dict] = {} model_spend: Optional[Dict] = {} - user_email: Optional[str] + user_email: Optional[str] = None models: list = [] tpm_limit: Optional[int] = None rpm_limit: Optional[int] = None @@ -1687,12 +1687,7 @@ class LiteLLM_AuditLogs(LiteLLMPydanticObjectBase): changed_by: Optional[Any] = None changed_by_api_key: Optional[str] = None action: AUDIT_ACTIONS - table_name: Literal[ - LitellmTableNames.TEAM_TABLE_NAME, - LitellmTableNames.USER_TABLE_NAME, - LitellmTableNames.KEY_TABLE_NAME, - LitellmTableNames.PROXY_MODEL_TABLE_NAME, - ] + table_name: LitellmTableNames object_id: str before_value: Optional[Json] = None updated_values: Optional[Json] = None diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index d80a06c597..a09db75301 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,11 +1,18 @@ -from typing import Any, Union +import uuid +from datetime import datetime, timezone +from typing import Any, Optional, Union +import litellm from litellm.proxy._types import ( + AUDIT_ACTIONS, GenerateKeyRequest, + LiteLLM_AuditLogs, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_TeamTable, + LitellmTableNames, UserAPIKeyAuth, ) +from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update from litellm.proxy.utils import _premium_user_check diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 309f451dd9..fd627a4f5b 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -826,7 +826,7 @@ async def user_update( return response # type: ignore # update based on remaining passed in values except Exception as e: - verbose_proxy_logger.error( + verbose_proxy_logger.exception( "litellm.proxy.proxy_server.user_update(): Exception occured - {}".format( str(e) ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 42ee2a5a40..861869c145 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -10,19 +10,26 @@ model/{model_id}/update - PATCH endpoint for model update. #### MODEL MANAGEMENT #### +import asyncio import json import uuid +from datetime import datetime, timezone from typing import Optional, cast from fastapi import APIRouter, Depends, HTTPException, Request, status from pydantic import BaseModel +import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._types import ( + AUDIT_ACTIONS, CommonProxyErrors, + LiteLLM_AuditLogs, LiteLLM_ProxyModelTable, + LitellmTableNames, LitellmUserRoles, + ModelInfoDelete, PrismaCompatibleUpdateDBModel, ProxyErrorTypes, ProxyException, @@ -36,6 +43,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( team_model_add, update_team, ) +from litellm.proxy.management_helpers.audit_logs import create_object_audit_log from litellm.proxy.utils import PrismaClient from litellm.types.router import ( Deployment, @@ -329,3 +337,101 @@ def check_if_team_id_matches_key( if user_api_key_dict.team_id != team_id: can_make_call = False return can_make_call + + +#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 +@router.post( + "/model/delete", + description="Allows deleting models in the model list in the config.yaml", + tags=["model management"], + dependencies=[Depends(user_api_key_auth)], +) +async def delete_model( + model_info: ModelInfoDelete, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config + try: + """ + [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 + + - Check if id in db + - Delete + """ + + from litellm.proxy.proxy_server import ( + llm_router, + prisma_client, + store_model_in_db, + ) + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={ + "error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys" + }, + ) + + # update DB + if store_model_in_db is True: + """ + - store model_list in db + - store keys separately + """ + # encrypt litellm params # + result = await prisma_client.db.litellm_proxymodeltable.delete( + where={"model_id": model_info.id} + ) + + if result is None: + raise HTTPException( + status_code=400, + detail={"error": f"Model with id={model_info.id} not found in db"}, + ) + + ## DELETE FROM ROUTER ## + if llm_router is not None: + llm_router.delete_deployment(id=model_info.id) + + ## CREATE AUDIT LOG ## + asyncio.create_task( + create_object_audit_log( + object_id=model_info.id, + action="deleted", + user_api_key_dict=user_api_key_dict, + table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME, + before_value=result.model_dump_json(exclude_none=True), + after_value=None, + litellm_changed_by=user_api_key_dict.user_id, + litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME, + ) + ) + return {"message": f"Model: {result.model_id} deleted successfully"} + else: + raise HTTPException( + status_code=500, + detail={ + "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." + }, + ) + + except Exception as e: + verbose_proxy_logger.exception( + f"Failed to delete model. Due to error - {str(e)}" + ) + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Authentication Error({str(e)})"), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=status.HTTP_400_BAD_REQUEST, + ) diff --git a/litellm/proxy/management_helpers/audit_logs.py b/litellm/proxy/management_helpers/audit_logs.py index b023e90965..d6c83c3856 100644 --- a/litellm/proxy/management_helpers/audit_logs.py +++ b/litellm/proxy/management_helpers/audit_logs.py @@ -3,13 +3,68 @@ Functions to create audit logs for LiteLLM Proxy """ import json +import uuid +from datetime import datetime, timezone import litellm from litellm._logging import verbose_proxy_logger -from litellm.proxy._types import LiteLLM_AuditLogs +from litellm.proxy._types import ( + AUDIT_ACTIONS, + LiteLLM_AuditLogs, + LitellmTableNames, + Optional, + UserAPIKeyAuth, +) + + +async def create_object_audit_log( + object_id: str, + action: AUDIT_ACTIONS, + litellm_changed_by: Optional[str], + user_api_key_dict: UserAPIKeyAuth, + litellm_proxy_admin_name: Optional[str], + table_name: LitellmTableNames, + before_value: Optional[str] = None, + after_value: Optional[str] = None, +): + """ + Create an audit log for an internal user. + + Parameters: + - user_id: str - The id of the user to create the audit log for. + - action: AUDIT_ACTIONS - The action to create the audit log for. + - user_row: LiteLLM_UserTable - The user row to create the audit log for. + - litellm_changed_by: Optional[str] - The user id of the user who is changing the user. + - user_api_key_dict: UserAPIKeyAuth - The user api key dictionary. + - litellm_proxy_admin_name: Optional[str] - The name of the proxy admin. + """ + if not litellm.store_audit_logs: + return + + await create_audit_log_for_update( + request_data=LiteLLM_AuditLogs( + id=str(uuid.uuid4()), + updated_at=datetime.now(timezone.utc), + changed_by=litellm_changed_by + or user_api_key_dict.user_id + or litellm_proxy_admin_name, + changed_by_api_key=user_api_key_dict.api_key, + table_name=table_name, + object_id=object_id, + action=action, + updated_values=after_value, + before_value=before_value, + ) + ) async def create_audit_log_for_update(request_data: LiteLLM_AuditLogs): + """ + Create an audit log for an object. + """ + if not litellm.store_audit_logs: + return + from litellm.proxy.proxy_server import premium_user, prisma_client if premium_user is not True: diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cfd722def6..cc4df723b9 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3182,7 +3182,7 @@ class ProxyStartupEvent: # add proxy budget to db in the user table asyncio.create_task( - generate_key_helper_fn( + generate_key_helper_fn( # type: ignore request_type="user", user_id=litellm_proxy_budget_name, duration=None, @@ -6592,84 +6592,6 @@ async def model_group_info( return {"data": model_groups} -#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 -@router.post( - "/model/delete", - description="Allows deleting models in the model list in the config.yaml", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], -) -async def delete_model(model_info: ModelInfoDelete): - global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config - try: - """ - [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 - - - Check if id in db - - Delete - """ - - global prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={ - "error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys" - }, - ) - - # update DB - if store_model_in_db is True: - """ - - store model_list in db - - store keys separately - """ - # encrypt litellm params # - result = await prisma_client.db.litellm_proxymodeltable.delete( - where={"model_id": model_info.id} - ) - - if result is None: - raise HTTPException( - status_code=400, - detail={"error": f"Model with id={model_info.id} not found in db"}, - ) - - ## DELETE FROM ROUTER ## - if llm_router is not None: - llm_router.delete_deployment(id=model_info.id) - - return {"message": f"Model: {result.model_id} deleted successfully"} - else: - raise HTTPException( - status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, - ) - - except Exception as e: - verbose_proxy_logger.exception( - f"Failed to delete model. Due to error - {str(e)}" - ) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) - - @router.get( "/model/settings", description="Returns provider name, description, and required parameters for each provider", From 9145e8db77f54b19e2fc89269217329d81f650dd Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 13 Mar 2025 19:00:27 -0700 Subject: [PATCH 4/7] feat: fix linting errors --- .../management_endpoints/internal_user_endpoints.py | 4 ++-- .../management_endpoints/model_management_endpoints.py | 9 ++++++++- 2 files changed, 10 insertions(+), 3 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index fd627a4f5b..a638f4c743 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -15,7 +15,7 @@ import asyncio import traceback import uuid from datetime import datetime, timedelta, timezone -from typing import Any, List, Optional, Union +from typing import Any, List, Optional, Union, cast import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -743,7 +743,7 @@ async def user_update( **existing_user_row.model_dump(exclude_none=True) ) - existing_metadata = existing_user_row.metadata if existing_user_row else {} + existing_metadata = cast(Dict, getattr(existing_user_row, "metadata", {}) or {}) non_default_values = prepare_metadata_fields( data=data, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 861869c145..cf4696fff1 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -350,7 +350,14 @@ async def delete_model( model_info: ModelInfoDelete, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config + from litellm.proxy.proxy_server import ( + general_settings, + llm_model_list, + llm_router, + proxy_config, + user_config_file_path, + ) + try: """ [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 From dc3b02920f53e6dd9f1add6a60eef9ca9f9beaa7 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 13 Mar 2025 19:17:40 -0700 Subject: [PATCH 5/7] feat(model_management_endpoints.py): support audit logs on `/model/add` and `/model/update` endpoints complete CUD endpoint audit logging on models + users --- .../model_management_endpoints.py | 299 +++++++++++++++++- litellm/proxy/proxy_server.py | 229 -------------- 2 files changed, 298 insertions(+), 230 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index cf4696fff1..f6e3cad1d6 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -272,7 +272,7 @@ async def _add_team_model_to_db( model_params: Deployment, user_api_key_dict: UserAPIKeyAuth, prisma_client: PrismaClient, -): +) -> Optional[LiteLLM_ProxyModelTable]: """ If 'team_id' is provided, @@ -442,3 +442,300 @@ async def delete_model( param=getattr(e, "param", "None"), code=status.HTTP_400_BAD_REQUEST, ) + + +#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 +@router.post( + "/model/new", + description="Allows adding new models to the model list in the config.yaml", + tags=["model management"], + dependencies=[Depends(user_api_key_auth)], +) +async def add_new_model( + model_params: Deployment, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + from litellm.proxy.proxy_server import ( + general_settings, + llm_model_list, + llm_router, + master_key, + premium_user, + prisma_client, + proxy_config, + proxy_logging_obj, + store_model_in_db, + user_config_file_path, + ) + + try: + import base64 + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={ + "error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys" + }, + ) + + if model_params.model_info.team_id is not None and premium_user is not True: + raise HTTPException( + status_code=403, + detail={"error": CommonProxyErrors.not_premium_user.value}, + ) + + if not check_if_team_id_matches_key( + team_id=model_params.model_info.team_id, user_api_key_dict=user_api_key_dict + ): + raise HTTPException( + status_code=403, + detail={"error": "Team ID does not match the API key's team ID"}, + ) + + model_response: Optional[LiteLLM_ProxyModelTable] = None + # update DB + if store_model_in_db is True: + """ + - store model_list in db + - store keys separately + """ + + try: + _original_litellm_model_name = model_params.model_name + if model_params.model_info.team_id is None: + model_response = await _add_model_to_db( + model_params=model_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + else: + model_response = await _add_team_model_to_db( + model_params=model_params, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + await proxy_config.add_deployment( + prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj + ) + # don't let failed slack alert block the /model/new response + _alerting = general_settings.get("alerting", []) or [] + if "slack" in _alerting: + # send notification - new model added + await proxy_logging_obj.slack_alerting_instance.model_added_alert( + model_name=model_params.model_name, + litellm_model_name=_original_litellm_model_name, + passed_model_info=model_params.model_info, + ) + except Exception as e: + verbose_proxy_logger.exception(f"Exception in add_new_model: {e}") + + else: + raise HTTPException( + status_code=500, + detail={ + "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." + }, + ) + + if model_response is None: + raise HTTPException( + status_code=500, + detail={ + "error": "Failed to add model to db. Check your server logs for more details." + }, + ) + + ## CREATE AUDIT LOG ## + asyncio.create_task( + create_object_audit_log( + object_id=model_response.model_id, + action="created", + user_api_key_dict=user_api_key_dict, + table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME, + before_value=None, + after_value=( + model_response.model_dump_json(exclude_none=True) + if isinstance(model_response, BaseModel) + else None + ), + litellm_changed_by=user_api_key_dict.user_id, + litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME, + ) + ) + + return model_response + + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.add_new_model(): Exception occured - {}".format( + str(e) + ) + ) + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Authentication Error({str(e)})"), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=status.HTTP_400_BAD_REQUEST, + ) + + +#### MODEL MANAGEMENT #### +@router.post( + "/model/update", + description="Edit existing model params", + tags=["model management"], + dependencies=[Depends(user_api_key_auth)], +) +async def update_model( + model_params: updateDeployment, + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), +): + """ + Old endpoint for model update. Makes a PUT request. + + Use `/model/{model_id}/update` to PATCH the stored model in db. + """ + from litellm.proxy.proxy_server import ( + LITELLM_PROXY_ADMIN_NAME, + general_settings, + llm_model_list, + llm_router, + master_key, + prisma_client, + proxy_config, + proxy_logging_obj, + store_model_in_db, + user_config_file_path, + ) + + try: + import base64 + + if prisma_client is None: + raise HTTPException( + status_code=500, + detail={ + "error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys" + }, + ) + # update DB + if store_model_in_db is True: + _model_id = None + _model_info = getattr(model_params, "model_info", None) + if _model_info is None: + raise Exception("model_info not provided") + + _model_id = _model_info.id + if _model_id is None: + raise Exception("model_info.id not provided") + _existing_litellm_params = ( + await prisma_client.db.litellm_proxymodeltable.find_unique( + where={"model_id": _model_id} + ) + ) + if _existing_litellm_params is None: + if ( + llm_router is not None + and llm_router.get_deployment(model_id=_model_id) is not None + ): + raise HTTPException( + status_code=400, + detail={ + "error": "Can't edit model. Model in config. Store model in db via `/model/new`. to edit." + }, + ) + raise Exception("model not found") + _existing_litellm_params_dict = dict( + _existing_litellm_params.litellm_params + ) + + if model_params.litellm_params is None: + raise Exception("litellm_params not provided") + + _new_litellm_params_dict = model_params.litellm_params.dict( + exclude_none=True + ) + + ### ENCRYPT PARAMS ### + for k, v in _new_litellm_params_dict.items(): + encrypted_value = encrypt_value_helper(value=v) + model_params.litellm_params[k] = encrypted_value + + ### MERGE WITH EXISTING DATA ### + merged_dictionary = {} + _mp = model_params.litellm_params.dict() + + for key, value in _mp.items(): + if value is not None: + merged_dictionary[key] = value + elif ( + key in _existing_litellm_params_dict + and _existing_litellm_params_dict[key] is not None + ): + merged_dictionary[key] = _existing_litellm_params_dict[key] + else: + pass + + _data: dict = { + "litellm_params": json.dumps(merged_dictionary), # type: ignore + "updated_by": user_api_key_dict.user_id or LITELLM_PROXY_ADMIN_NAME, + } + model_response = await prisma_client.db.litellm_proxymodeltable.update( + where={"model_id": _model_id}, + data=_data, # type: ignore + ) + + ## CREATE AUDIT LOG ## + asyncio.create_task( + create_object_audit_log( + object_id=_model_id, + action="updated", + user_api_key_dict=user_api_key_dict, + table_name=LitellmTableNames.PROXY_MODEL_TABLE_NAME, + before_value=( + _existing_litellm_params.model_dump_json(exclude_none=True) + if isinstance(_existing_litellm_params, BaseModel) + else None + ), + after_value=( + model_response.model_dump_json(exclude_none=True) + if isinstance(model_response, BaseModel) + else None + ), + litellm_changed_by=user_api_key_dict.user_id, + litellm_proxy_admin_name=LITELLM_PROXY_ADMIN_NAME, + ) + ) + + return model_response + except Exception as e: + verbose_proxy_logger.exception( + "litellm.proxy.proxy_server.update_model(): Exception occured - {}".format( + str(e) + ) + ) + if isinstance(e, HTTPException): + raise ProxyException( + message=getattr(e, "detail", f"Authentication Error({str(e)})"), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), + ) + elif isinstance(e, ProxyException): + raise e + raise ProxyException( + message="Authentication Error, " + str(e), + type=ProxyErrorTypes.auth_error, + param=getattr(e, "param", "None"), + code=status.HTTP_400_BAD_REQUEST, + ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index cc4df723b9..250551b019 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5438,235 +5438,6 @@ async def transform_request(request: TransformRequestBody): return return_raw_request(endpoint=request.call_type, kwargs=request.request_body) -#### [BETA] - This is a beta endpoint, format might change based on user feedback. - https://github.com/BerriAI/litellm/issues/964 -@router.post( - "/model/new", - description="Allows adding new models to the model list in the config.yaml", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], -) -async def add_new_model( - model_params: Deployment, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config, prisma_client, master_key, store_model_in_db, proxy_logging_obj, premium_user - try: - import base64 - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={ - "error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys" - }, - ) - - if model_params.model_info.team_id is not None and premium_user is not True: - raise HTTPException( - status_code=403, - detail={"error": CommonProxyErrors.not_premium_user.value}, - ) - - if not check_if_team_id_matches_key( - team_id=model_params.model_info.team_id, user_api_key_dict=user_api_key_dict - ): - raise HTTPException( - status_code=403, - detail={"error": "Team ID does not match the API key's team ID"}, - ) - - model_response = None - # update DB - if store_model_in_db is True: - """ - - store model_list in db - - store keys separately - """ - - try: - _original_litellm_model_name = model_params.model_name - if model_params.model_info.team_id is None: - model_response = await _add_model_to_db( - model_params=model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) - else: - model_response = await _add_team_model_to_db( - model_params=model_params, - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - ) - await proxy_config.add_deployment( - prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj - ) - # don't let failed slack alert block the /model/new response - _alerting = general_settings.get("alerting", []) or [] - if "slack" in _alerting: - # send notification - new model added - await proxy_logging_obj.slack_alerting_instance.model_added_alert( - model_name=model_params.model_name, - litellm_model_name=_original_litellm_model_name, - passed_model_info=model_params.model_info, - ) - except Exception as e: - verbose_proxy_logger.exception(f"Exception in add_new_model: {e}") - - else: - raise HTTPException( - status_code=500, - detail={ - "error": "Set `'STORE_MODEL_IN_DB='True'` in your env to enable this feature." - }, - ) - - return model_response - - except Exception as e: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.add_new_model(): Exception occured - {}".format( - str(e) - ) - ) - verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) - - -#### MODEL MANAGEMENT #### -@router.post( - "/model/update", - description="Edit existing model params", - tags=["model management"], - dependencies=[Depends(user_api_key_auth)], -) -async def update_model( - model_params: updateDeployment, - user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -): - """ - Old endpoint for model update. Makes a PUT request. - - Use `/model/{model_id}/update` to PATCH the stored model in db. - """ - global llm_router, llm_model_list, general_settings, user_config_file_path, proxy_config, prisma_client, master_key, store_model_in_db, proxy_logging_obj - try: - import base64 - - global prisma_client - - if prisma_client is None: - raise HTTPException( - status_code=500, - detail={ - "error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys" - }, - ) - # update DB - if store_model_in_db is True: - _model_id = None - _model_info = getattr(model_params, "model_info", None) - if _model_info is None: - raise Exception("model_info not provided") - - _model_id = _model_info.id - if _model_id is None: - raise Exception("model_info.id not provided") - _existing_litellm_params = ( - await prisma_client.db.litellm_proxymodeltable.find_unique( - where={"model_id": _model_id} - ) - ) - if _existing_litellm_params is None: - if ( - llm_router is not None - and llm_router.get_deployment(model_id=_model_id) is not None - ): - raise HTTPException( - status_code=400, - detail={ - "error": "Can't edit model. Model in config. Store model in db via `/model/new`. to edit." - }, - ) - raise Exception("model not found") - _existing_litellm_params_dict = dict( - _existing_litellm_params.litellm_params - ) - - if model_params.litellm_params is None: - raise Exception("litellm_params not provided") - - _new_litellm_params_dict = model_params.litellm_params.dict( - exclude_none=True - ) - - ### ENCRYPT PARAMS ### - for k, v in _new_litellm_params_dict.items(): - encrypted_value = encrypt_value_helper(value=v) - model_params.litellm_params[k] = encrypted_value - - ### MERGE WITH EXISTING DATA ### - merged_dictionary = {} - _mp = model_params.litellm_params.dict() - - for key, value in _mp.items(): - if value is not None: - merged_dictionary[key] = value - elif ( - key in _existing_litellm_params_dict - and _existing_litellm_params_dict[key] is not None - ): - merged_dictionary[key] = _existing_litellm_params_dict[key] - else: - pass - - _data: dict = { - "litellm_params": json.dumps(merged_dictionary), # type: ignore - "updated_by": user_api_key_dict.user_id or litellm_proxy_admin_name, - } - model_response = await prisma_client.db.litellm_proxymodeltable.update( - where={"model_id": _model_id}, - data=_data, # type: ignore - ) - - return model_response - except Exception as e: - verbose_proxy_logger.error( - "litellm.proxy.proxy_server.update_model(): Exception occured - {}".format( - str(e) - ) - ) - verbose_proxy_logger.debug(traceback.format_exc()) - if isinstance(e, HTTPException): - raise ProxyException( - message=getattr(e, "detail", f"Authentication Error({str(e)})"), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=getattr(e, "status_code", status.HTTP_400_BAD_REQUEST), - ) - elif isinstance(e, ProxyException): - raise e - raise ProxyException( - message="Authentication Error, " + str(e), - type=ProxyErrorTypes.auth_error, - param=getattr(e, "param", "None"), - code=status.HTTP_400_BAD_REQUEST, - ) - - @router.get( "/v2/model/info", description="v2 - returns all the models set on the config.yaml, shows 'user_access' = True if the user has access to the model. Provides more info about each model in /models, including config.yaml descriptions (except api key and api base)", From 81abb9c6a4ae0105a99050f534f43884203e1d1d Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 13 Mar 2025 19:26:10 -0700 Subject: [PATCH 6/7] test: fix tests --- tests/local_testing/test_add_update_models.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/tests/local_testing/test_add_update_models.py b/tests/local_testing/test_add_update_models.py index 4a04fcdf4a..1ea714d9b7 100644 --- a/tests/local_testing/test_add_update_models.py +++ b/tests/local_testing/test_add_update_models.py @@ -15,8 +15,12 @@ sys.path.insert( 0, os.path.abspath("../..") ) # Adds the parent directory to the system path import pytest, logging, asyncio -import litellm, asyncio -from litellm.proxy.proxy_server import add_new_model, update_model, LitellmUserRoles +import litellm +from litellm.proxy.management_endpoints.model_management_endpoints import ( + add_new_model, + update_model, +) +from litellm.proxy._types import LitellmUserRoles from litellm._logging import verbose_proxy_logger from litellm.proxy.utils import PrismaClient, ProxyLogging from litellm.proxy.management_endpoints.team_endpoints import new_team From f2d0aaacbcc3c0a8d6080f4f6e60bc256d87a126 Mon Sep 17 00:00:00 2001 From: Krrish Dholakia Date: Thu, 13 Mar 2025 19:26:46 -0700 Subject: [PATCH 7/7] fix: fix linting errors --- .../management_endpoints/common_utils.py | 9 +------ .../model_management_endpoints.py | 24 +------------------ 2 files changed, 2 insertions(+), 31 deletions(-) diff --git a/litellm/proxy/management_endpoints/common_utils.py b/litellm/proxy/management_endpoints/common_utils.py index a09db75301..d80a06c597 100644 --- a/litellm/proxy/management_endpoints/common_utils.py +++ b/litellm/proxy/management_endpoints/common_utils.py @@ -1,18 +1,11 @@ -import uuid -from datetime import datetime, timezone -from typing import Any, Optional, Union +from typing import Any, Union -import litellm from litellm.proxy._types import ( - AUDIT_ACTIONS, GenerateKeyRequest, - LiteLLM_AuditLogs, LiteLLM_ManagementEndpoint_MetadataFields_Premium, LiteLLM_TeamTable, - LitellmTableNames, UserAPIKeyAuth, ) -from litellm.proxy.management_helpers.audit_logs import create_audit_log_for_update from litellm.proxy.utils import _premium_user_check diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index f6e3cad1d6..2a2b7eae24 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -13,19 +13,15 @@ model/{model_id}/update - PATCH endpoint for model update. import asyncio import json import uuid -from datetime import datetime, timezone from typing import Optional, cast from fastapi import APIRouter, Depends, HTTPException, Request, status from pydantic import BaseModel -import litellm from litellm._logging import verbose_proxy_logger from litellm.constants import LITELLM_PROXY_ADMIN_NAME from litellm.proxy._types import ( - AUDIT_ACTIONS, CommonProxyErrors, - LiteLLM_AuditLogs, LiteLLM_ProxyModelTable, LitellmTableNames, LitellmUserRoles, @@ -350,13 +346,7 @@ async def delete_model( model_info: ModelInfoDelete, user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), ): - from litellm.proxy.proxy_server import ( - general_settings, - llm_model_list, - llm_router, - proxy_config, - user_config_file_path, - ) + from litellm.proxy.proxy_server import llm_router try: """ @@ -457,19 +447,14 @@ async def add_new_model( ): from litellm.proxy.proxy_server import ( general_settings, - llm_model_list, - llm_router, - master_key, premium_user, prisma_client, proxy_config, proxy_logging_obj, store_model_in_db, - user_config_file_path, ) try: - import base64 if prisma_client is None: raise HTTPException( @@ -607,19 +592,12 @@ async def update_model( """ from litellm.proxy.proxy_server import ( LITELLM_PROXY_ADMIN_NAME, - general_settings, - llm_model_list, llm_router, - master_key, prisma_client, - proxy_config, - proxy_logging_obj, store_model_in_db, - user_config_file_path, ) try: - import base64 if prisma_client is None: raise HTTPException(