refac

Co-Authored-By: Classic298 <27028174+Classic298@users.noreply.github.com>

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)

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)