diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 100b0bf6db..acc6b6175e 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -5333,6 +5333,67 @@ async def _check_if_model_is_user_added( return filtered_models +def _check_if_model_is_team_model( + models: List[DeploymentTypedDict], user_row: LiteLLM_UserTable +) -> List[Dict]: + """ + Check if model is a team model + + Check if user is a member of the team that the model belongs to + """ + + user_team_models: List[Dict] = [] + for model in models: + model_team_id = model.get("model_info", {}).get("team_id", None) + + if model_team_id is not None: + if model_team_id in user_row.teams: + user_team_models.append(cast(Dict, model)) + + return user_team_models + + +async def non_admin_all_models( + all_models: List[Dict], + llm_router: Router, + user_api_key_dict: UserAPIKeyAuth, + prisma_client: Optional[PrismaClient], +): + """ + 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}, + ) + + all_models = await _check_if_model_is_user_added( + models=all_models, + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + ) + + if user_api_key_dict.user_id: + try: + user_row = await prisma_client.db.litellm_usertable.find_unique( + where={"user_id": user_api_key_dict.user_id} + ) + except Exception: + raise HTTPException(status_code=400, detail={"error": "User not found"}) + + all_models += _check_if_model_is_team_model( + models=llm_router.get_model_list() or [], + user_row=user_row, + ) + + return all_models + + @router.get( "/v2/model/info", 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", @@ -5378,16 +5439,10 @@ 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, + if user_models_only: + all_models = await non_admin_all_models( + all_models=all_models, + llm_router=llm_router, user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) diff --git a/ui/litellm-dashboard/src/components/model_dashboard.tsx b/ui/litellm-dashboard/src/components/model_dashboard.tsx index e87da27c24..d3dda3a96d 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard.tsx @@ -1126,13 +1126,15 @@ const ModelDashboard: React.FC = ({ diff --git a/ui/litellm-dashboard/src/components/model_dashboard/columns.tsx b/ui/litellm-dashboard/src/components/model_dashboard/columns.tsx index b1acf9ea5d..2602761a5b 100644 --- a/ui/litellm-dashboard/src/components/model_dashboard/columns.tsx +++ b/ui/litellm-dashboard/src/components/model_dashboard/columns.tsx @@ -7,6 +7,8 @@ import { TrashIcon, PencilIcon, PencilAltIcon } from "@heroicons/react/outline"; import DeleteModelButton from "../delete_model_button"; export const columns = ( + userRole: string, + userID: string, premiumUser: boolean, setSelectedModelId: (id: string) => void, setSelectedTeamId: (id: string) => void, @@ -226,23 +228,30 @@ export const columns = ( header: "", cell: ({ row }) => { const model = row.original; + const canEditModel = userRole === "Admin" || model.model_info?.created_by === userID; return (
{ - setSelectedModelId(model.model_info.id); - setEditModel(true); + if (canEditModel) { + setSelectedModelId(model.model_info.id); + setEditModel(true); + } }} + className={!canEditModel ? "opacity-50 cursor-not-allowed" : "cursor-pointer"} /> { - setSelectedModelId(model.model_info.id); - setEditModel(false); + if (canEditModel) { + setSelectedModelId(model.model_info.id); + setEditModel(false); + } }} + className={!canEditModel ? "opacity-50 cursor-not-allowed" : "cursor-pointer"} />
);