fix(model_management_endpoints.py): fix allowing team admins to update team models (#9697)
All checks were successful
Read Version from pyproject.toml / read-version (push) Successful in 17s
Helm unit test / unit-test (push) Successful in 22s

* fix(model_management_endpoints.py): fix allowing team admins to update their models

* test(test_models.py): add e2e test to for team model flow

ensure team admin can always add / edit / delete team models
This commit is contained in:
Krish Dholakia 2025-04-01 22:28:15 -07:00 committed by GitHub
parent 3d0313b15b
commit 6c69ad4c89
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 115 additions and 33 deletions

View file

@ -394,7 +394,7 @@ class ModelManagementAuthChecks:
@staticmethod
async def can_user_make_model_call(
model_params: Union[Deployment, updateDeployment],
model_params: Deployment,
user_api_key_dict: UserAPIKeyAuth,
prisma_client: PrismaClient,
premium_user: bool,
@ -723,8 +723,38 @@ async def update_model(
},
)
_model_id = None
_model_info = getattr(model_params, "model_info", None)
if _model_info is None:
raise Exception("model_info not provided")
_model_id = _model_info.id
if _model_id is None:
raise Exception("model_info.id not provided")
_existing_litellm_params = (
await prisma_client.db.litellm_proxymodeltable.find_unique(
where={"model_id": _model_id}
)
)
if _existing_litellm_params is None:
if (
llm_router is not None
and llm_router.get_deployment(model_id=_model_id) is not None
):
raise HTTPException(
status_code=400,
detail={
"error": "Can't edit model. Model in config. Store model in db via `/model/new`. to edit."
},
)
else:
raise Exception("model not found")
deployment = Deployment(**_existing_litellm_params.model_dump())
await ModelManagementAuthChecks.can_user_make_model_call(
model_params=model_params,
model_params=deployment,
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
premium_user=premium_user,
@ -732,31 +762,6 @@ async def update_model(
# update DB
if store_model_in_db is True:
_model_id = None
_model_info = getattr(model_params, "model_info", None)
if _model_info is None:
raise Exception("model_info not provided")
_model_id = _model_info.id
if _model_id is None:
raise Exception("model_info.id not provided")
_existing_litellm_params = (
await prisma_client.db.litellm_proxymodeltable.find_unique(
where={"model_id": _model_id}
)
)
if _existing_litellm_params is None:
if (
llm_router is not None
and llm_router.get_deployment(model_id=_model_id) is not None
):
raise HTTPException(
status_code=400,
detail={
"error": "Can't edit model. Model in config. Store model in db via `/model/new`. to edit."
},
)
raise Exception("model not found")
_existing_litellm_params_dict = dict(
_existing_litellm_params.litellm_params
)