mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-25 18:54:30 +00:00
Allow team admins to add/update/delete models on UI + show api base and model id on request logs (#9572)
* feat(view_logs.tsx): show model id + api base in request logs easier debugging * fix(index.tsx): fix length of api base easier viewing * refactor(leftnav.tsx): show models tab to team admin * feat(model_dashboard.tsx): add explainer for what the 'models' page is for team admin helps them understand how they can use it * feat(model_management_endpoints.py): restrict model add by team to just team admin allow team admin to add models via non-team keys (e.g. ui token) * test(test_add_update_models.py): update unit testing for new behaviour * fix(model_dashboard.tsx): show user the models * feat(proxy_server.py): add new query param 'user_models_only' to `/v2/model/info` Allows user to retrieve just the models they've added Used in UI to show internal users just the models they've added * feat(model_dashboard.tsx): allow team admins to view their own models * fix: allow ui user to fetch model cost map * feat(add_model_tab.tsx): require team admins to specify team when onboarding models * fix(_types.py): add `/v1/model/info` to info route `/model/info` was already there * fix(model_info_view.tsx): allow user to edit a model they created * fix(model_management_endpoints.py): allow team admin to update team model * feat(model_managament_endpoints.py): allow team admin to delete team models * fix(model_management_endpoints.py): don't require team id to be set when adding a model * fix(proxy_server.py): fix linting error * fix: fix ui linting error * fix(model_management_endpoints.py): ensure consistent auth checks on all model calls * test: remove old test - function no longer exists in same form * test: add updated mock testing
This commit is contained in:
Krish Dholakia
GitHub
parent
3ca34a181c
commit
63c9f59373
11 changed files with 483 additions and 144 deletions
|
|
@ -13,7 +13,7 @@ model/{model_id}/update - PATCH endpoint for model update.
|
|||
import asyncio
|
||||
import json
|
||||
import uuid
|
||||
from typing import Optional, cast
|
||||
from typing import Literal, Optional, Union, cast
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException, Request, status
|
||||
from pydantic import BaseModel
|
||||
|
|
@ -23,6 +23,7 @@ from litellm.constants import LITELLM_PROXY_ADMIN_NAME
|
|||
from litellm.proxy._types import (
|
||||
CommonProxyErrors,
|
||||
LiteLLM_ProxyModelTable,
|
||||
LiteLLM_TeamTable,
|
||||
LitellmTableNames,
|
||||
LitellmUserRoles,
|
||||
ModelInfoDelete,
|
||||
|
|
@ -35,6 +36,7 @@ from litellm.proxy._types import (
|
|||
)
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.management_endpoints.common_utils import _is_user_team_admin
|
||||
from litellm.proxy.management_endpoints.team_endpoints import (
|
||||
team_model_add,
|
||||
update_team,
|
||||
|
|
@ -317,22 +319,111 @@ async def _add_team_model_to_db(
|
|||
return model_response
|
||||
|
||||
|
||||
def check_if_team_id_matches_key(
|
||||
team_id: Optional[str], user_api_key_dict: UserAPIKeyAuth
|
||||
) -> bool:
|
||||
can_make_call = True
|
||||
if (
|
||||
user_api_key_dict.user_role
|
||||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
):
|
||||
class ModelManagementAuthChecks:
|
||||
"""
|
||||
Common auth checks for model management endpoints
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def can_user_make_team_model_call(
|
||||
team_id: str,
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
team_obj: Optional[LiteLLM_TeamTable] = None,
|
||||
premium_user: bool = False,
|
||||
) -> Literal[True]:
|
||||
if premium_user is False:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={"error": CommonProxyErrors.not_premium_user.value},
|
||||
)
|
||||
if (
|
||||
user_api_key_dict.user_role
|
||||
and user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN
|
||||
):
|
||||
return True
|
||||
elif team_obj is None or not _is_user_team_admin(
|
||||
user_api_key_dict=user_api_key_dict, team_obj=team_obj
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "Team ID={} does not match the API key's team ID={}, OR you are not the admin for this team. Check `/user/info` to verify your team admin status.".format(
|
||||
team_id, user_api_key_dict.team_id
|
||||
)
|
||||
},
|
||||
)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
async def allow_team_model_action(
|
||||
model_params: Union[Deployment, updateDeployment],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
premium_user: bool,
|
||||
) -> Literal[True]:
|
||||
if model_params.model_info is None or model_params.model_info.team_id is None:
|
||||
return True
|
||||
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},
|
||||
)
|
||||
|
||||
_existing_team_row = await prisma_client.db.litellm_teamtable.find_unique(
|
||||
where={"team_id": model_params.model_info.team_id}
|
||||
)
|
||||
|
||||
if _existing_team_row is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": "Team id={} does not exist in db".format(
|
||||
model_params.model_info.team_id
|
||||
)
|
||||
},
|
||||
)
|
||||
existing_team_row = LiteLLM_TeamTable(**_existing_team_row.model_dump())
|
||||
|
||||
ModelManagementAuthChecks.can_user_make_team_model_call(
|
||||
team_id=model_params.model_info.team_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
team_obj=existing_team_row,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
return True
|
||||
|
||||
@staticmethod
|
||||
async def can_user_make_model_call(
|
||||
model_params: Union[Deployment, updateDeployment],
|
||||
user_api_key_dict: UserAPIKeyAuth,
|
||||
prisma_client: PrismaClient,
|
||||
premium_user: bool,
|
||||
) -> Literal[True]:
|
||||
|
||||
## Check team model auth
|
||||
if (
|
||||
model_params.model_info is not None
|
||||
and model_params.model_info.team_id is not None
|
||||
):
|
||||
return ModelManagementAuthChecks.can_user_make_team_model_call(
|
||||
team_id=model_params.model_info.team_id,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
## Check non-team model auth
|
||||
elif user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
raise HTTPException(
|
||||
status_code=403,
|
||||
detail={
|
||||
"error": "User does not have permission to make this model call. Your role={}. You can only make model calls if you are a PROXY_ADMIN or if you are a team admin, by specifying a team_id in the model_info.".format(
|
||||
user_api_key_dict.user_role
|
||||
)
|
||||
},
|
||||
)
|
||||
else:
|
||||
return True
|
||||
|
||||
return True
|
||||
if team_id is None:
|
||||
if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN:
|
||||
can_make_call = False
|
||||
else:
|
||||
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
|
||||
|
|
@ -358,6 +449,7 @@ async def delete_model(
|
|||
|
||||
from litellm.proxy.proxy_server import (
|
||||
llm_router,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
store_model_in_db,
|
||||
)
|
||||
|
|
@ -370,6 +462,23 @@ async def delete_model(
|
|||
},
|
||||
)
|
||||
|
||||
model_in_db = await prisma_client.db.litellm_proxymodeltable.find_unique(
|
||||
where={"model_id": model_info.id}
|
||||
)
|
||||
if model_in_db is None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"Model with id={model_info.id} not found in db"},
|
||||
)
|
||||
|
||||
model_params = Deployment(**model_in_db.model_dump())
|
||||
await ModelManagementAuthChecks.can_user_make_model_call(
|
||||
model_params=model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
# update DB
|
||||
if store_model_in_db is True:
|
||||
"""
|
||||
|
|
@ -464,19 +573,13 @@ async def add_new_model(
|
|||
},
|
||||
)
|
||||
|
||||
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"},
|
||||
)
|
||||
## Auth check
|
||||
await ModelManagementAuthChecks.can_user_make_model_call(
|
||||
model_params=model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
model_response: Optional[LiteLLM_ProxyModelTable] = None
|
||||
# update DB
|
||||
|
|
@ -593,6 +696,7 @@ async def update_model(
|
|||
from litellm.proxy.proxy_server import (
|
||||
LITELLM_PROXY_ADMIN_NAME,
|
||||
llm_router,
|
||||
premium_user,
|
||||
prisma_client,
|
||||
store_model_in_db,
|
||||
)
|
||||
|
|
@ -606,6 +710,14 @@ async def update_model(
|
|||
"error": "No DB Connected. Here's how to do it - https://docs.litellm.ai/docs/proxy/virtual_keys"
|
||||
},
|
||||
)
|
||||
|
||||
await ModelManagementAuthChecks.can_user_make_model_call(
|
||||
model_params=model_params,
|
||||
user_api_key_dict=user_api_key_dict,
|
||||
prisma_client=prisma_client,
|
||||
premium_user=premium_user,
|
||||
)
|
||||
|
||||
# update DB
|
||||
if store_model_in_db is True:
|
||||
_model_id = None
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue