Allow team members to see team models (#9742)

* fix(proxy_server.py): allow team member to see team models

* fix(model_dashboard.tsx): show edit + delete icons to be disabled if user is not admin and did not create models

* fix(proxy_server.py): fix ruff function size error

* fix(proxy_server.py): fix user model filter check
This commit is contained in:
Krish Dholakia 2025-04-04 20:36:48 -07:00 committed by GitHub
parent b5851769fc
commit d66db2207b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 81 additions and 15 deletions

View file

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

View file

@ -1126,13 +1126,15 @@ const ModelDashboard: React.FC<ModelDashboardProps> = ({
</div>
<ModelDataTable
columns={columns(
userRole,
userID,
premiumUser,
setSelectedModelId,
setSelectedTeamId,
getDisplayModelName,
handleEditClick,
handleRefreshClick,
setEditModel
setEditModel,
)}
data={modelData.data.filter(
(model: any) =>

View file

@ -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 (
<div className="flex items-center justify-end gap-2 pr-4">
<Icon
icon={PencilAltIcon}
size="sm"
onClick={() => {
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"}
/>
<Icon
icon={TrashIcon}
size="sm"
onClick={() => {
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"}
/>
</div>
);