diff --git a/litellm/proxy/_new_secret_config.yaml b/litellm/proxy/_new_secret_config.yaml index 169e8c355b..76317f050f 100644 --- a/litellm/proxy/_new_secret_config.yaml +++ b/litellm/proxy/_new_secret_config.yaml @@ -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 + \ No newline at end of file diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index a9fe6517ea..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): @@ -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 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..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 @@ -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) ) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 42ee2a5a40..2a2b7eae24 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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, + ) 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..250551b019 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, @@ -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", 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