refac
This commit is contained in:
@@ -55,9 +55,7 @@ def is_valid_model_id(model_id: str) -> bool:
|
||||
PAGE_ITEM_COUNT = 30
|
||||
|
||||
|
||||
@router.get(
|
||||
"/list", response_model=ModelAccessListResponse
|
||||
) # do NOT use "/" as path, conflicts with main.py
|
||||
@router.get('/list', response_model=ModelAccessListResponse) # do NOT use "/" as path, conflicts with main.py
|
||||
async def get_models(
|
||||
query: Optional[str] = None,
|
||||
view_option: Optional[str] = None,
|
||||
@@ -68,7 +66,6 @@ async def get_models(
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
|
||||
limit = PAGE_ITEM_COUNT
|
||||
|
||||
page = max(1, page)
|
||||
@@ -76,25 +73,25 @@ async def get_models(
|
||||
|
||||
filter = {}
|
||||
if query:
|
||||
filter["query"] = query
|
||||
filter['query'] = query
|
||||
if view_option:
|
||||
filter["view_option"] = view_option
|
||||
filter['view_option'] = view_option
|
||||
if tag:
|
||||
filter["tag"] = tag
|
||||
filter['tag'] = tag
|
||||
if order_by:
|
||||
filter["order_by"] = order_by
|
||||
filter['order_by'] = order_by
|
||||
if direction:
|
||||
filter["direction"] = direction
|
||||
filter['direction'] = direction
|
||||
|
||||
# Pre-fetch user group IDs once - used for both filter and write_access check
|
||||
groups = Groups.get_groups_by_member_id(user.id, db=db)
|
||||
user_group_ids = {group.id for group in groups}
|
||||
|
||||
if not user.role == "admin" or not BYPASS_ADMIN_ACCESS_CONTROL:
|
||||
if not user.role == 'admin' or not BYPASS_ADMIN_ACCESS_CONTROL:
|
||||
if groups:
|
||||
filter["group_ids"] = [group.id for group in groups]
|
||||
filter['group_ids'] = [group.id for group in groups]
|
||||
|
||||
filter["user_id"] = user.id
|
||||
filter['user_id'] = user.id
|
||||
|
||||
result = Models.search_models(user.id, filter=filter, skip=skip, limit=limit, db=db)
|
||||
|
||||
@@ -102,9 +99,9 @@ async def get_models(
|
||||
model_ids = [model.id for model in result.items]
|
||||
writable_model_ids = AccessGrants.get_accessible_resource_ids(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_ids=model_ids,
|
||||
permission="write",
|
||||
permission='write',
|
||||
user_group_ids=user_group_ids,
|
||||
db=db,
|
||||
)
|
||||
@@ -114,7 +111,7 @@ async def get_models(
|
||||
ModelAccessResponse(
|
||||
**model.model_dump(),
|
||||
write_access=(
|
||||
(user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
(user.role == 'admin' and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
or user.id == model.user_id
|
||||
or model.id in writable_model_ids
|
||||
),
|
||||
@@ -130,10 +127,8 @@ async def get_models(
|
||||
###########################
|
||||
|
||||
|
||||
@router.get("/base", response_model=list[ModelResponse])
|
||||
async def get_base_models(
|
||||
user=Depends(get_admin_user), db: Session = Depends(get_session)
|
||||
):
|
||||
@router.get('/base', response_model=list[ModelResponse])
|
||||
async def get_base_models(user=Depends(get_admin_user), db: Session = Depends(get_session)):
|
||||
return Models.get_base_models(db=db)
|
||||
|
||||
|
||||
@@ -142,11 +137,9 @@ async def get_base_models(
|
||||
###########################
|
||||
|
||||
|
||||
@router.get("/tags", response_model=list[str])
|
||||
async def get_model_tags(
|
||||
user=Depends(get_verified_user), db: Session = Depends(get_session)
|
||||
):
|
||||
if user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL:
|
||||
@router.get('/tags', response_model=list[str])
|
||||
async def get_model_tags(user=Depends(get_verified_user), db: Session = Depends(get_session)):
|
||||
if user.role == 'admin' and BYPASS_ADMIN_ACCESS_CONTROL:
|
||||
models = Models.get_models(db=db)
|
||||
else:
|
||||
models = Models.get_models_by_user_id(user.id, db=db)
|
||||
@@ -155,8 +148,8 @@ async def get_model_tags(
|
||||
for model in models:
|
||||
if model.meta:
|
||||
meta = model.meta.model_dump()
|
||||
for tag in meta.get("tags", []):
|
||||
tags_set.add((tag.get("name")))
|
||||
for tag in meta.get('tags', []):
|
||||
tags_set.add((tag.get('name')))
|
||||
|
||||
tags = [tag for tag in tags_set]
|
||||
tags.sort()
|
||||
@@ -168,15 +161,15 @@ async def get_model_tags(
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/create", response_model=Optional[ModelModel])
|
||||
@router.post('/create', response_model=Optional[ModelModel])
|
||||
async def create_new_model(
|
||||
request: Request,
|
||||
form_data: ModelForm,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
if user.role != "admin" and not has_permission(
|
||||
user.id, "workspace.models", request.app.state.config.USER_PERMISSIONS, db=db
|
||||
if user.role != 'admin' and not has_permission(
|
||||
user.id, 'workspace.models', request.app.state.config.USER_PERMISSIONS, db=db
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_401_UNAUTHORIZED,
|
||||
@@ -212,15 +205,15 @@ async def create_new_model(
|
||||
############################
|
||||
|
||||
|
||||
@router.get("/export", response_model=list[ModelModel])
|
||||
@router.get('/export', response_model=list[ModelModel])
|
||||
async def export_models(
|
||||
request: Request,
|
||||
user=Depends(get_verified_user),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
if user.role != "admin" and not has_permission(
|
||||
if user.role != 'admin' and not has_permission(
|
||||
user.id,
|
||||
"workspace.models_export",
|
||||
'workspace.models_export',
|
||||
request.app.state.config.USER_PERMISSIONS,
|
||||
db=db,
|
||||
):
|
||||
@@ -229,7 +222,7 @@ async def export_models(
|
||||
detail=ERROR_MESSAGES.UNAUTHORIZED,
|
||||
)
|
||||
|
||||
if user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL:
|
||||
if user.role == 'admin' and BYPASS_ADMIN_ACCESS_CONTROL:
|
||||
return Models.get_models(db=db)
|
||||
else:
|
||||
return Models.get_models_by_user_id(user.id, db=db)
|
||||
@@ -244,16 +237,16 @@ class ModelsImportForm(BaseModel):
|
||||
models: list[dict]
|
||||
|
||||
|
||||
@router.post("/import", response_model=bool)
|
||||
@router.post('/import', response_model=bool)
|
||||
async def import_models(
|
||||
request: Request,
|
||||
user=Depends(get_verified_user),
|
||||
form_data: ModelsImportForm = (...),
|
||||
db: Session = Depends(get_session),
|
||||
):
|
||||
if user.role != "admin" and not has_permission(
|
||||
if user.role != 'admin' and not has_permission(
|
||||
user.id,
|
||||
"workspace.models_import",
|
||||
'workspace.models_import',
|
||||
request.app.state.config.USER_PERMISSIONS,
|
||||
db=db,
|
||||
):
|
||||
@@ -266,43 +259,36 @@ async def import_models(
|
||||
if isinstance(data, list):
|
||||
# Batch-fetch all existing models in one query to avoid N+1
|
||||
model_ids = [
|
||||
model_data.get("id")
|
||||
model_data.get('id')
|
||||
for model_data in data
|
||||
if model_data.get("id") and is_valid_model_id(model_data.get("id"))
|
||||
if model_data.get('id') and is_valid_model_id(model_data.get('id'))
|
||||
]
|
||||
existing_models = {
|
||||
model.id: model
|
||||
for model in (
|
||||
Models.get_models_by_ids(model_ids, db=db) if model_ids else []
|
||||
)
|
||||
model.id: model for model in (Models.get_models_by_ids(model_ids, db=db) if model_ids else [])
|
||||
}
|
||||
|
||||
for model_data in data:
|
||||
# Here, you can add logic to validate model_data if needed
|
||||
model_id = model_data.get("id")
|
||||
model_id = model_data.get('id')
|
||||
|
||||
if model_id and is_valid_model_id(model_id):
|
||||
existing_model = existing_models.get(model_id)
|
||||
if existing_model:
|
||||
# Update existing model
|
||||
model_data["meta"] = model_data.get("meta", {})
|
||||
model_data["params"] = model_data.get("params", {})
|
||||
model_data['meta'] = model_data.get('meta', {})
|
||||
model_data['params'] = model_data.get('params', {})
|
||||
|
||||
updated_model = ModelForm(
|
||||
**{**existing_model.model_dump(), **model_data}
|
||||
)
|
||||
updated_model = ModelForm(**{**existing_model.model_dump(), **model_data})
|
||||
Models.update_model_by_id(model_id, updated_model, db=db)
|
||||
else:
|
||||
# Insert new model
|
||||
model_data["meta"] = model_data.get("meta", {})
|
||||
model_data["params"] = model_data.get("params", {})
|
||||
model_data['meta'] = model_data.get('meta', {})
|
||||
model_data['params'] = model_data.get('params', {})
|
||||
new_model = ModelForm(**model_data)
|
||||
Models.insert_new_model(
|
||||
user_id=user.id, form_data=new_model, db=db
|
||||
)
|
||||
Models.insert_new_model(user_id=user.id, form_data=new_model, db=db)
|
||||
return True
|
||||
else:
|
||||
raise HTTPException(status_code=400, detail="Invalid JSON format")
|
||||
raise HTTPException(status_code=400, detail='Invalid JSON format')
|
||||
except Exception as e:
|
||||
log.exception(e)
|
||||
raise HTTPException(status_code=500, detail=str(e))
|
||||
@@ -317,7 +303,7 @@ class SyncModelsForm(BaseModel):
|
||||
models: list[ModelModel] = []
|
||||
|
||||
|
||||
@router.post("/sync", response_model=list[ModelModel])
|
||||
@router.post('/sync', response_model=list[ModelModel])
|
||||
async def sync_models(
|
||||
request: Request,
|
||||
form_data: SyncModelsForm,
|
||||
@@ -337,33 +323,31 @@ class ModelIdForm(BaseModel):
|
||||
|
||||
|
||||
# Note: We're not using the typical url path param here, but instead using a query parameter to allow '/' in the id
|
||||
@router.get("/model", response_model=Optional[ModelAccessResponse])
|
||||
async def get_model_by_id(
|
||||
id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)
|
||||
):
|
||||
@router.get('/model', response_model=Optional[ModelAccessResponse])
|
||||
async def get_model_by_id(id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)):
|
||||
model = Models.get_model_by_id(id, db=db)
|
||||
if model:
|
||||
if (
|
||||
(user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
(user.role == 'admin' and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
or model.user_id == user.id
|
||||
or AccessGrants.has_access(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_id=model.id,
|
||||
permission="read",
|
||||
permission='read',
|
||||
db=db,
|
||||
)
|
||||
):
|
||||
return ModelAccessResponse(
|
||||
**model.model_dump(),
|
||||
write_access=(
|
||||
(user.role == "admin" and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
(user.role == 'admin' and BYPASS_ADMIN_ACCESS_CONTROL)
|
||||
or user.id == model.user_id
|
||||
or AccessGrants.has_access(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_id=model.id,
|
||||
permission="write",
|
||||
permission='write',
|
||||
db=db,
|
||||
)
|
||||
),
|
||||
@@ -385,7 +369,7 @@ async def get_model_by_id(
|
||||
###########################
|
||||
|
||||
|
||||
@router.get("/model/profile/image")
|
||||
@router.get('/model/profile/image')
|
||||
def get_model_profile_image(id: str, user=Depends(get_verified_user)):
|
||||
model = Models.get_model_by_id(id)
|
||||
|
||||
@@ -393,21 +377,21 @@ def get_model_profile_image(id: str, user=Depends(get_verified_user)):
|
||||
etag = f'"{model.updated_at}"' if model.updated_at else None
|
||||
|
||||
if model.meta.profile_image_url:
|
||||
if model.meta.profile_image_url.startswith("http"):
|
||||
if model.meta.profile_image_url.startswith('http'):
|
||||
return Response(
|
||||
status_code=status.HTTP_302_FOUND,
|
||||
headers={"Location": model.meta.profile_image_url},
|
||||
headers={'Location': model.meta.profile_image_url},
|
||||
)
|
||||
elif model.meta.profile_image_url.startswith("data:image"):
|
||||
elif model.meta.profile_image_url.startswith('data:image'):
|
||||
try:
|
||||
header, base64_data = model.meta.profile_image_url.split(",", 1)
|
||||
header, base64_data = model.meta.profile_image_url.split(',', 1)
|
||||
image_data = base64.b64decode(base64_data)
|
||||
image_buffer = io.BytesIO(image_data)
|
||||
media_type = header.split(";")[0].lstrip("data:")
|
||||
media_type = header.split(';')[0].lstrip('data:')
|
||||
|
||||
headers = {"Content-Disposition": "inline"}
|
||||
headers = {'Content-Disposition': 'inline'}
|
||||
if etag:
|
||||
headers["ETag"] = etag
|
||||
headers['ETag'] = etag
|
||||
|
||||
return StreamingResponse(
|
||||
image_buffer,
|
||||
@@ -417,9 +401,9 @@ def get_model_profile_image(id: str, user=Depends(get_verified_user)):
|
||||
except Exception as e:
|
||||
pass
|
||||
|
||||
return FileResponse(f"{STATIC_DIR}/favicon.png")
|
||||
return FileResponse(f'{STATIC_DIR}/favicon.png')
|
||||
else:
|
||||
return FileResponse(f"{STATIC_DIR}/favicon.png")
|
||||
return FileResponse(f'{STATIC_DIR}/favicon.png')
|
||||
|
||||
|
||||
############################
|
||||
@@ -427,20 +411,18 @@ def get_model_profile_image(id: str, user=Depends(get_verified_user)):
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/model/toggle", response_model=Optional[ModelResponse])
|
||||
async def toggle_model_by_id(
|
||||
id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)
|
||||
):
|
||||
@router.post('/model/toggle', response_model=Optional[ModelResponse])
|
||||
async def toggle_model_by_id(id: str, user=Depends(get_verified_user), db: Session = Depends(get_session)):
|
||||
model = Models.get_model_by_id(id, db=db)
|
||||
if model:
|
||||
if (
|
||||
user.role == "admin"
|
||||
user.role == 'admin'
|
||||
or model.user_id == user.id
|
||||
or AccessGrants.has_access(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_id=model.id,
|
||||
permission="write",
|
||||
permission='write',
|
||||
db=db,
|
||||
)
|
||||
):
|
||||
@@ -451,7 +433,7 @@ async def toggle_model_by_id(
|
||||
else:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error updating function"),
|
||||
detail=ERROR_MESSAGES.DEFAULT('Error updating function'),
|
||||
)
|
||||
else:
|
||||
raise HTTPException(
|
||||
@@ -470,7 +452,7 @@ async def toggle_model_by_id(
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/model/update", response_model=Optional[ModelModel])
|
||||
@router.post('/model/update', response_model=Optional[ModelModel])
|
||||
async def update_model_by_id(
|
||||
form_data: ModelForm,
|
||||
user=Depends(get_verified_user),
|
||||
@@ -487,21 +469,19 @@ async def update_model_by_id(
|
||||
model.user_id != user.id
|
||||
and not AccessGrants.has_access(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_id=model.id,
|
||||
permission="write",
|
||||
permission='write',
|
||||
db=db,
|
||||
)
|
||||
and user.role != "admin"
|
||||
and user.role != 'admin'
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
)
|
||||
|
||||
model = Models.update_model_by_id(
|
||||
form_data.id, ModelForm(**form_data.model_dump()), db=db
|
||||
)
|
||||
model = Models.update_model_by_id(form_data.id, ModelForm(**form_data.model_dump()), db=db)
|
||||
return model
|
||||
|
||||
|
||||
@@ -516,7 +496,7 @@ class ModelAccessGrantsForm(BaseModel):
|
||||
access_grants: list[dict]
|
||||
|
||||
|
||||
@router.post("/model/access/update", response_model=Optional[ModelModel])
|
||||
@router.post('/model/access/update', response_model=Optional[ModelModel])
|
||||
async def update_model_access_by_id(
|
||||
request: Request,
|
||||
form_data: ModelAccessGrantsForm,
|
||||
@@ -528,7 +508,7 @@ async def update_model_access_by_id(
|
||||
# Non-preset models (e.g. direct Ollama/OpenAI models) may not have a DB
|
||||
# entry yet. Create a minimal one so access grants can be stored.
|
||||
if not model:
|
||||
if user.role != "admin":
|
||||
if user.role != 'admin':
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_403_FORBIDDEN,
|
||||
detail=ERROR_MESSAGES.ACCESS_PROHIBITED,
|
||||
@@ -546,19 +526,19 @@ async def update_model_access_by_id(
|
||||
if not model:
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
|
||||
detail=ERROR_MESSAGES.DEFAULT("Error creating model entry"),
|
||||
detail=ERROR_MESSAGES.DEFAULT('Error creating model entry'),
|
||||
)
|
||||
|
||||
if (
|
||||
model.user_id != user.id
|
||||
and not AccessGrants.has_access(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_id=model.id,
|
||||
permission="write",
|
||||
permission='write',
|
||||
db=db,
|
||||
)
|
||||
and user.role != "admin"
|
||||
and user.role != 'admin'
|
||||
):
|
||||
raise HTTPException(
|
||||
status_code=status.HTTP_400_BAD_REQUEST,
|
||||
@@ -570,12 +550,10 @@ async def update_model_access_by_id(
|
||||
user.id,
|
||||
user.role,
|
||||
form_data.access_grants,
|
||||
"sharing.public_models",
|
||||
'sharing.public_models',
|
||||
)
|
||||
|
||||
AccessGrants.set_access_grants(
|
||||
"model", form_data.id, form_data.access_grants, db=db
|
||||
)
|
||||
AccessGrants.set_access_grants('model', form_data.id, form_data.access_grants, db=db)
|
||||
|
||||
return Models.get_model_by_id(form_data.id, db=db)
|
||||
|
||||
@@ -585,7 +563,7 @@ async def update_model_access_by_id(
|
||||
############################
|
||||
|
||||
|
||||
@router.post("/model/delete", response_model=bool)
|
||||
@router.post('/model/delete', response_model=bool)
|
||||
async def delete_model_by_id(
|
||||
form_data: ModelIdForm,
|
||||
user=Depends(get_verified_user),
|
||||
@@ -599,13 +577,13 @@ async def delete_model_by_id(
|
||||
)
|
||||
|
||||
if (
|
||||
user.role != "admin"
|
||||
user.role != 'admin'
|
||||
and model.user_id != user.id
|
||||
and not AccessGrants.has_access(
|
||||
user_id=user.id,
|
||||
resource_type="model",
|
||||
resource_type='model',
|
||||
resource_id=model.id,
|
||||
permission="write",
|
||||
permission='write',
|
||||
db=db,
|
||||
)
|
||||
):
|
||||
@@ -618,9 +596,7 @@ async def delete_model_by_id(
|
||||
return result
|
||||
|
||||
|
||||
@router.delete("/delete/all", response_model=bool)
|
||||
async def delete_all_models(
|
||||
user=Depends(get_admin_user), db: Session = Depends(get_session)
|
||||
):
|
||||
@router.delete('/delete/all', response_model=bool)
|
||||
async def delete_all_models(user=Depends(get_admin_user), db: Session = Depends(get_session)):
|
||||
result = Models.delete_all_models(db=db)
|
||||
return result
|
||||
|
||||
Reference in New Issue
Block a user