diff --git a/backend/open_webui/routers/ollama.py b/backend/open_webui/routers/ollama.py index 4f43c41d3..737ed7153 100644 --- a/backend/open_webui/routers/ollama.py +++ b/backend/open_webui/routers/ollama.py @@ -46,6 +46,7 @@ from open_webui.internal.db import get_session from open_webui.models.models import Models from open_webui.models.access_grants import AccessGrants +from open_webui.models.groups import Groups from open_webui.utils.misc import ( calculate_sha256, ) @@ -425,15 +426,24 @@ async def get_all_models(request: Request, user: UserModel = None): async def get_filtered_models(models, user, db=None): # Filter models based on user access control + model_ids = [model["model"] for model in models.get("models", [])] + model_infos = { + m.id: m for m in Models.get_models_by_ids(model_ids, db=db) + } + user_group_ids = { + g.id for g in Groups.get_groups_by_member_id(user.id, db=db) + } + filtered_models = [] for model in models.get("models", []): - model_info = Models.get_model_by_id(model["model"], db=db) + model_info = model_infos.get(model["model"]) if model_info: if user.id == model_info.user_id or AccessGrants.has_access( user_id=user.id, resource_type="model", resource_id=model_info.id, permission="read", + user_group_ids=user_group_ids, db=db, ): filtered_models.append(model) @@ -1643,15 +1653,24 @@ async def get_openai_models( if user.role == "user" and not BYPASS_MODEL_ACCESS_CONTROL: # Filter models based on user access control + model_ids = [model["id"] for model in models] + model_infos = { + m.id: m for m in Models.get_models_by_ids(model_ids, db=db) + } + user_group_ids = { + g.id for g in Groups.get_groups_by_member_id(user.id, db=db) + } + filtered_models = [] for model in models: - model_info = Models.get_model_by_id(model["id"], db=db) + model_info = model_infos.get(model["id"]) if model_info: if user.id == model_info.user_id or AccessGrants.has_access( user_id=user.id, resource_type="model", resource_id=model_info.id, permission="read", + user_group_ids=user_group_ids, db=db, ): filtered_models.append(model) diff --git a/backend/open_webui/routers/openai.py b/backend/open_webui/routers/openai.py index ef1eaa2e5..38fd7feed 100644 --- a/backend/open_webui/routers/openai.py +++ b/backend/open_webui/routers/openai.py @@ -25,6 +25,7 @@ from open_webui.internal.db import get_session from open_webui.models.models import Models from open_webui.models.access_grants import AccessGrants +from open_webui.models.groups import Groups from open_webui.config import ( CACHE_DIR, ) @@ -466,15 +467,24 @@ async def get_all_models_responses(request: Request, user: UserModel) -> list: async def get_filtered_models(models, user, db=None): # Filter models based on user access control + model_ids = [model["id"] for model in models.get("data", [])] + model_infos = { + m.id: m for m in Models.get_models_by_ids(model_ids, db=db) + } + user_group_ids = { + g.id for g in Groups.get_groups_by_member_id(user.id, db=db) + } + filtered_models = [] for model in models.get("data", []): - model_info = Models.get_model_by_id(model["id"], db=db) + model_info = model_infos.get(model["id"]) if model_info: if user.id == model_info.user_id or AccessGrants.has_access( user_id=user.id, resource_type="model", resource_id=model_info.id, permission="read", + user_group_ids=user_group_ids, db=db, ): filtered_models.append(model)