phoenix/litellm-mirror
0
0
Fork 1
mirror of https://github.com/BerriAI/litellm.git synced 2025-04-27 03:34:10 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Allow team admins to add/update/delete models on UI + show api base and model id on request logs (#9572)

Browse source 
* 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 2025-03-27 12:06:31 -07:00 • committed by GitHub
parent a5fbe50f04
commit ed8c63b51e
11 changed files with 483 additions and 144 deletions
Download patch file Download diff file Expand all files Collapse all files

4
litellm/proxy/_types.py
View file

@ -292,6 +292,7 @@ class LiteLLMRoutes(enum.Enum):
"/team/available",
"/user/info",
"/model/info",
"/v1/model/info",
"/v2/model/info",
"/v2/key/info",
"/model_group/info",
@ -386,6 +387,7 @@ class LiteLLMRoutes(enum.Enum):
"/global/predict/spend/logs",
"/global/activity",
"/health/services",
"/get/litellm_model_cost_map",
] + info_routes
internal_user_routes = [
@ -412,6 +414,8 @@ class LiteLLMRoutes(enum.Enum):
"/team/member_add",
"/team/member_delete",
"/model/new",
"/model/update",
"/model/delete",
] # routes that manage their own allowed/disallowed logic
## Org Admin Routes ##

170
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -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

51
litellm/proxy/proxy_server.py
View file

@ -215,7 +215,6 @@ from litellm.proxy.management_endpoints.key_management_endpoints import (
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
_add_team_model_to_db,
check_if_team_id_matches_key,
)
from litellm.proxy.management_endpoints.model_management_endpoints import (
router as model_management_router,
@ -5494,9 +5493,40 @@ async def transform_request(request: TransformRequestBody):
return return_raw_request(endpoint=request.call_type, kwargs=request.request_body)
async def _check_if_model_is_user_added(
models: List[Dict],
user_api_key_dict: UserAPIKeyAuth,
prisma_client: Optional[PrismaClient],
) -> List[Dict]:
"""
Check if model is in db
Check if db model is 'created_by' == user_api_key_dict.user_id
Only return models that match
"""
if prisma_client is None:
raise HTTPException(
status_code=500,
detail={"error": CommonProxyErrors.db_not_connected_error.value},
)
filtered_models = []
for model in models:
id = model.get("model_info", {}).get("id", None)
if id is None:
continue
db_model = await prisma_client.db.litellm_proxymodeltable.find_unique(
where={"model_id": id}
)
if db_model is not None:
if db_model.created_by == user_api_key_dict.user_id:
filtered_models.append(model)
return filtered_models
@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)",
description="v2 - returns models available to the user based on their API key permissions. Shows model info from config.yaml (except api key and api base). Filter to just user-added models with ?user_models_only=true",
tags=["model management"],
dependencies=[Depends(user_api_key_auth)],
include_in_schema=False,
@ -5506,6 +5536,9 @@ async def model_info_v2(
model: Optional[str] = fastapi.Query(
None, description="Specify the model name (optional)"
),
user_models_only: Optional[bool] = fastapi.Query(
False, description="Only return models added by this user"
),
debug: Optional[bool] = False,
):
"""
@ -5536,6 +5569,20 @@ async def model_info_v2(
if model is not None:
all_models = [m for m in all_models if m["model_name"] == model]
if user_models_only is True:
"""
Check if model is in db
Check if db model is 'created_by' == user_api_key_dict.user_id
Only return models that match
"""
all_models = await _check_if_model_is_user_added(
models=all_models,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
# fill in model info based on config.yaml and litellm model_prices_and_context_window.json
for _model in all_models:
# provided model_info in config.yaml