tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-07-16 14:17:10 +00:00
Code Issues Packages Projects Releases Wiki Activity

Merge pull request #9223 from BerriAI/litellm_dev_03_13_2025_p3

Browse Source 
Emit audit logs on All user + model Create/Update/Delete endpoints
This commit is contained in:
Krish Dholakia
2025-03-13 20:12:24 -07:00
committed by GitHub
parent e8c67f25e3 f89dbb8ab3
commit 7b229bd7ae
8 changed files with 585 additions and 330 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/_new_secret_config.yaml
+1
View File
@@ -4,3 +4,4 @@ model_list:
model: azure/chatgpt-v-2
api_key: os.environ/AZURE_API_KEY
api_base: os.environ/AZURE_API_BASE
litellm/proxy/_types.py
+9 -10
View File
@@ -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):
@@ -1584,11 +1584,11 @@ 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] = {}
user_email: Optional[str]
user_email: Optional[str] = None
models: list = []
tpm_limit: Optional[int] = None
rpm_limit: Optional[int] = None
@@ -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
@@ -1677,18 +1678,16 @@ 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"]
table_name: Literal[
LitellmTableNames.TEAM_TABLE_NAME,
LitellmTableNames.USER_TABLE_NAME,
LitellmTableNames.KEY_TABLE_NAME,
LitellmTableNames.PROXY_MODEL_TABLE_NAME,
]
action: AUDIT_ACTIONS
table_name: LitellmTableNames
object_id: str
before_value: Optional[Json] = None
updated_values: Optional[Json] = None
litellm/proxy/common_utils/encrypt_decrypt_utils.py
+4 -1
View File
@@ -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
litellm/proxy/management_endpoints/internal_user_endpoints.py
+119 -7
View File
@@ -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
@@ -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),
@@ -649,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()
@@ -662,11 +732,18 @@ 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 {}
existing_metadata = cast(Dict, getattr(existing_user_row, "metadata", {}) or {})
non_default_values = prepare_metadata_fields(
data=data,
@@ -711,10 +788,45 @@ 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:
verbose_proxy_logger.error(
verbose_proxy_logger.exception(
"litellm.proxy.proxy_server.user_update(): Exception occured - {}".format(
str(e)
)
litellm/proxy/management_endpoints/model_management_endpoints.py
+389 -1
View File
@@ -10,6 +10,7 @@ model/{model_id}/update - PATCH endpoint for model update.
#### MODEL MANAGEMENT ####
import asyncio
import json
import uuid
from typing import Optional, cast
@@ -22,7 +23,9 @@ from litellm.constants import LITELLM_PROXY_ADMIN_NAME
from litellm.proxy._types import (
CommonProxyErrors,
LiteLLM_ProxyModelTable,
LitellmTableNames,
LitellmUserRoles,
ModelInfoDelete,
PrismaCompatibleUpdateDBModel,
ProxyErrorTypes,
ProxyException,
@@ -36,6 +39,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,
@@ -264,7 +268,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,
@@ -329,3 +333,387 @@ 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),
):
from litellm.proxy.proxy_server import llm_router
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,
)
#### [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,
premium_user,
prisma_client,
proxy_config,
proxy_logging_obj,
store_model_in_db,
)
try:
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,
llm_router,
prisma_client,
store_model_in_db,
)
try:
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,
)
litellm/proxy/management_helpers/audit_logs.py
+56 -1
View File
@@ -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:
litellm/proxy/proxy_server.py
+1 -308
View File
@@ -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,
@@ -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)",
@@ -6592,84 +6363,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",
tests/local_testing/test_add_update_models.py
+6 -2
View File
@@ -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