mirror of
https://github.com/BerriAI/litellm.git
synced 2025-04-24 18:24:20 +00:00
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:
parent
b5851769fc
commit
d66db2207b
3 changed files with 81 additions and 15 deletions
|
@ -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,
|
||||
)
|
||||
|
|
|
@ -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) =>
|
||||
|
|
|
@ -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>
|
||||
);
|
||||
|
|
Loading…
Add table
Add a link
Reference in a new issue