diff --git a/api/apps/restful_apis/memory_api.py b/api/apps/restful_apis/memory_api.py index dc59845572..5958924882 100644 --- a/api/apps/restful_apis/memory_api.py +++ b/api/apps/restful_apis/memory_api.py @@ -139,7 +139,7 @@ async def delete_memory(memory_id): @manager.route("/memories", methods=["GET"]) # noqa: F821 @login_required async def list_memory(): - filter_params = {k: request.args.get(k) for k in ["memory_type", "tenant_id", "owner_ids", "storage_type"] if k in request.args} + filter_params = {k: request.args.get(k) for k in ["memory_type", "tenant_id", "owner_ids", "ids", "storage_type"] if k in request.args} keywords = request.args.get("keywords") page = validate_rest_api_page(request.args.get("page", DEFAULT_PAGE)) page_size = validate_rest_api_page_size(request.args.get("page_size", DEFAULT_PAGE_SIZE)) diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 9c9f9961f2..ffee50c209 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -452,9 +452,12 @@ def list_datasets(tenant_id: str, args: dict): query_user_id = tenant_id if kb_ids: accessible_ids = KnowledgebaseService.get_accessible_ids([m["tenant_id"] for m in tenants], tenant_id, kb_ids) - if len(accessible_ids) != len(kb_ids): - denied_ids = [kb_id for kb_id in kb_ids if kb_id not in accessible_ids] - return False, f"""User '{tenant_id}' lacks permission for datasets: '{", ".join(denied_ids)}'""" + denied_ids = [kb_id for kb_id in kb_ids if kb_id not in accessible_ids] + if denied_ids: + logging.warning("User '%s' lacks permission for datasets: '%s'", tenant_id, ", ".join(denied_ids)) + kb_ids = [kb_id for kb_id in kb_ids if kb_id in accessible_ids] + if not kb_ids: + return True, {"data": [], "total": 0} kbs, total = KnowledgebaseService.get_list(tenant_ids, query_user_id, page, page_size, orderby, desc, kb_id, name, keywords, parser_id, kb_ids) users = UserService.get_by_ids([m["tenant_id"] for m in kbs]) user_map = {m.id: m.to_dict() for m in users} diff --git a/api/apps/services/memory_api_service.py b/api/apps/services/memory_api_service.py index 581715341d..78e8f116e9 100644 --- a/api/apps/services/memory_api_service.py +++ b/api/apps/services/memory_api_service.py @@ -13,6 +13,8 @@ # See the License for the specific language governing permissions and # limitations under the License. # +import logging + from api.apps import current_user from api.db import TenantPermission from api.db.services.memory_service import MemoryService @@ -251,6 +253,7 @@ async def list_memory(filter_params: dict, keywords: str, page: int = 1, page_si :param filter_params: { "memory_type": list[str], "tenant_id": list[str], + "ids": list[str], "storage_type": str } :param keywords: str @@ -259,6 +262,18 @@ async def list_memory(filter_params: dict, keywords: str, page: int = 1, page_si """ filter_dict: dict = {"storage_type": filter_params.get("storage_type"), "accessible_user_id": current_user.id} allowed_tenant_ids = _joined_tenant_ids(current_user.id) + + memory_ids = _split_filter_values(filter_params.get("ids")) + if memory_ids: + accessible_memories = _filter_accessible_memories(memory_ids) + accessible_memory_ids = [m.id for m in accessible_memories] + denied_ids = [mid for mid in memory_ids if mid not in accessible_memory_ids] + if denied_ids: + logging.warning("User '%s' lacks permission for memories: '%s'", current_user.id, ", ".join(denied_ids)) + filter_dict["ids"] = accessible_memory_ids + if not accessible_memory_ids: + return {"memory_list": [], "total_count": 0} + tenant_ids = _split_filter_values(filter_params.get("tenant_id") or filter_params.get("owner_ids")) if tenant_ids: filter_dict["tenant_id"] = [tenant_id for tenant_id in tenant_ids if tenant_id in allowed_tenant_ids] diff --git a/api/db/services/memory_service.py b/api/db/services/memory_service.py index 933ea613c2..dac72c60ba 100644 --- a/api/db/services/memory_service.py +++ b/api/db/services/memory_service.py @@ -90,6 +90,8 @@ class MemoryService(CommonService): cls.model.create_date, ] memories = cls.model.select(*fields).join(User, on=(cls.model.tenant_id == User.id)) + if filter_dict.get("ids"): + memories = memories.where(cls.model.id.in_(filter_dict["ids"])) if filter_dict.get("tenant_id"): memories = memories.where(cls.model.tenant_id.in_(filter_dict["tenant_id"])) if filter_dict.get("accessible_user_id"):