diff --git a/api/apps/restful_apis/dataset_api.py b/api/apps/restful_apis/dataset_api.py index eaaf0d44d3..d540958b63 100644 --- a/api/apps/restful_apis/dataset_api.py +++ b/api/apps/restful_apis/dataset_api.py @@ -321,6 +321,14 @@ def list_datasets(tenant_id): type: string required: false description: Dataset ID to filter. + - in: query + name: ids + type: array + required: false + items: + type: string + collectionFormat: multi + description: Dataset IDs to filter. - in: query name: name type: string diff --git a/api/apps/services/dataset_api_service.py b/api/apps/services/dataset_api_service.py index 264f6dccfb..63f1a0eaac 100644 --- a/api/apps/services/dataset_api_service.py +++ b/api/apps/services/dataset_api_service.py @@ -403,6 +403,7 @@ def list_datasets(tenant_id: str, args: dict): :return: (success, result) or (success, error_message) """ kb_id = args.get("id") + kb_ids = args.get("ids") name = args.get("name") page = int(args.get("page", 1)) page_size = int(args.get("page_size", 30)) @@ -419,10 +420,13 @@ def list_datasets(tenant_id: str, args: dict): # unknown type, default to True desc = True + if kb_id and kb_ids: + return False, f"Should not provide both 'id':{kb_id} and 'ids'{kb_ids}" if kb_id: kbs = KnowledgebaseService.get_kb_by_id(kb_id, tenant_id) if not kbs: return False, f"User '{tenant_id}' lacks permission for dataset '{kb_id}'" + if name: kbs = KnowledgebaseService.get_kb_by_name(name, tenant_id) if not kbs: @@ -438,7 +442,12 @@ def list_datasets(tenant_id: str, args: dict): tenants = TenantService.get_joined_tenants_by_user_id(tenant_id) tenant_ids = [m["tenant_id"] for m in tenants] query_user_id = tenant_id - kbs, total = KnowledgebaseService.get_list(tenant_ids, query_user_id, page, page_size, orderby, desc, kb_id, name, keywords, parser_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)}'""" + 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/db/services/knowledgebase_service.py b/api/db/services/knowledgebase_service.py index 0cf2500360..dfa668f18d 100644 --- a/api/db/services/knowledgebase_service.py +++ b/api/db/services/knowledgebase_service.py @@ -463,7 +463,7 @@ class KnowledgebaseService(CommonService): @classmethod @DB.connection_context() - def get_list(cls, joined_tenant_ids, user_id, page_number, items_per_page, orderby, desc, id, name, keywords, parser_id=None): + def get_list(cls, joined_tenant_ids, user_id, page_number, items_per_page, orderby, desc, id, name, keywords, parser_id=None, ids=None): # Get list of knowledge bases with filtering and pagination # Args: # joined_tenant_ids: List of tenant IDs @@ -482,6 +482,8 @@ class KnowledgebaseService(CommonService): kbs = cls.model.select() if id: kbs = kbs.where(cls.model.id == id) + if ids: + kbs = kbs.where(cls.model.id.in_(ids)) if name: kbs = kbs.where(cls.model.name == name) if keywords: @@ -501,6 +503,12 @@ class KnowledgebaseService(CommonService): return list(kbs.dicts()), total + @classmethod + @DB.connection_context() + def get_accessible_ids(cls, joined_tenant_ids, user_id, ids): + kbs = cls.model.select(cls.model.id).where(cls.model.id.in_(ids), cls._visibility_and_status_filter(joined_tenant_ids, user_id)) + return {kb.id for kb in kbs} + @classmethod @DB.connection_context() def get_owner_filter(cls, joined_tenant_ids, user_id): diff --git a/api/utils/validation_utils.py b/api/utils/validation_utils.py index 20847b6f4f..450ac95045 100644 --- a/api/utils/validation_utils.py +++ b/api/utils/validation_utils.py @@ -19,7 +19,8 @@ import pathlib import re from collections import Counter import string -from typing import Annotated, Any, Literal +from types import UnionType +from typing import Annotated, Any, Literal, Union, get_args, get_origin from uuid import UUID from quart import Request @@ -33,6 +34,15 @@ from api.utils.pagination_utils import validate_rest_api_page_size from common.constants import RetCode +def _is_list_annotation(annotation: Any) -> bool: + origin = get_origin(annotation) + if origin is list: + return True + if origin in (Union, UnionType): + return any(_is_list_annotation(arg) for arg in get_args(annotation)) + return False + + async def validate_and_parse_json_request( request: Request, validator: type[BaseModel], *, extras: dict[str, Any] | None = None, exclude_unset: bool = False ) -> tuple[dict[str, Any] | None, str | None]: @@ -161,6 +171,10 @@ def validate_and_parse_request_args(request: Request, validator: type[BaseModel] - Preserves type conversion from Pydantic validation """ args = request.args.to_dict(flat=True) + for field_name, field_info in validator.model_fields.items(): + query_name = field_info.alias or field_name + if query_name in request.args and _is_list_annotation(field_info.annotation): + args[query_name] = [value for item in request.args.getlist(query_name) for value in item.split(",") if value] # Handle ext parameter: parse JSON string to dict if it's a string if "ext" in args and isinstance(args["ext"], str): @@ -1002,9 +1016,26 @@ class BaseListReq(BaseModel): class ListDatasetReq(BaseListReq): """Request model for listing datasets.""" + ids: Annotated[list[str] | None, Field(default=None)] include_parsing_status: Annotated[bool, Field(default=False)] ext: Annotated[dict, Field(default={})] + @field_validator("ids", mode="after") + @classmethod + def validate_ids(cls, v_list: list[str] | None) -> list[str] | None: + if v_list is None: + return None + + ids_list = [] + for v in v_list: + ids_list.append(validate_uuid1_hex(v)) + + duplicates = [item for item, count in Counter(ids_list).items() if count > 1] + if duplicates: + raise PydanticCustomError("duplicate_uuids", "Duplicate ids: '{duplicate_ids}'", {"duplicate_ids": ", ".join(duplicates)}) + + return ids_list + # ---- File Management Request Models ---- diff --git a/sdk/python/ragflow_sdk/ragflow.py b/sdk/python/ragflow_sdk/ragflow.py index a339fe1d3f..4932e02d4d 100644 --- a/sdk/python/ragflow_sdk/ragflow.py +++ b/sdk/python/ragflow_sdk/ragflow.py @@ -95,7 +95,9 @@ class RAGFlow: return _list[0] raise Exception("Dataset %s not found" % name) - def list_datasets(self, page: int = 1, page_size: int = 30, orderby: str = "create_time", desc: bool = True, id: str | None = None, name: str | None = None) -> list[DataSet]: + def list_datasets( + self, page: int = 1, page_size: int = 30, orderby: str = "create_time", desc: bool = True, id: str | None = None, ids: list[str] | None = None, name: str | None = None + ) -> list[DataSet]: res = self.get( "/datasets", { @@ -104,6 +106,7 @@ class RAGFlow: "orderby": orderby, "desc": desc, "id": id, + "ids": ids, "name": name, }, ) diff --git a/test/unit_test/api/apps/services/test_dataset_api_service_list_datasets.py b/test/unit_test/api/apps/services/test_dataset_api_service_list_datasets.py index 89bf7d217b..f7bedd44d5 100644 --- a/test/unit_test/api/apps/services/test_dataset_api_service_list_datasets.py +++ b/test/unit_test/api/apps/services/test_dataset_api_service_list_datasets.py @@ -77,6 +77,7 @@ def _identity_remap(source_data, key_aliases=None): def _load_list_datasets_module(monkeypatch, *, kbs, parsing_status_by_kb): parsing_status_mock = MagicMock(return_value=parsing_status_by_kb) get_list_mock = MagicMock(return_value=(list(kbs), len(kbs))) + get_accessible_ids_mock = MagicMock(return_value={kb["id"] for kb in kbs}) _stub( monkeypatch, @@ -127,6 +128,7 @@ def _load_list_datasets_module(monkeypatch, *, kbs, parsing_status_by_kb): "api.db.services.knowledgebase_service", KnowledgebaseService=SimpleNamespace( get_list=get_list_mock, + get_accessible_ids=get_accessible_ids_mock, ), validate_dataset_embedding_models=lambda kbs: None, ) @@ -168,6 +170,11 @@ def _load_list_datasets_module(monkeypatch, *, kbs, parsing_status_by_kb): "common.misc_utils", thread_pool_exec=MagicMock(), ) + _stub( + monkeypatch, + "rag.advanced_rag.knowlege_compile.wiki", + WIKI_PAGE_COMPILE_KWD="wiki", + ) repo_root = Path(__file__).resolve().parents[5] module_path = repo_root / "api" / "apps" / "services" / "dataset_api_service.py" @@ -235,6 +242,7 @@ def test_list_datasets_with_include_parsing_status_true_attaches_counts(monkeypa assert ok is True parsing_status_mock.assert_called_once_with(["kb-a", "kb-b"]) + assert "parsing_status" in payload["data"][0] by_id = {r["id"]: r for r in payload["data"]} assert by_id["kb-a"]["parsing_status"] == status_by_kb["kb-a"] assert by_id["kb-b"]["parsing_status"] == status_by_kb["kb-b"] @@ -258,7 +266,25 @@ def test_list_datasets_with_include_parsing_status_string_true(monkeypatch): assert ok is True parsing_status_mock.assert_called_once_with(["kb-a", "kb-b"]) - assert "parsing_status" in payload["data"][0] + + +def test_list_datasets_with_ids_filters_query_once(monkeypatch): + """ids filter is checked once and pushed into the list query.""" + module, _, get_list_mock = _load_list_datasets_module( + monkeypatch, + kbs=_stub_kbs(), + parsing_status_by_kb={}, + ) + + ok, payload = module.list_datasets( + "tenant-1", + {"page": 1, "page_size": 30, "ids": ["kb-a", "kb-b"]}, + ) + + assert ok is True + assert payload["total"] == 2 + module.KnowledgebaseService.get_accessible_ids.assert_called_once_with(["tenant-1"], "tenant-1", ["kb-a", "kb-b"]) + get_list_mock.assert_called_once_with(["tenant-1"], "tenant-1", 1, 30, "create_time", True, None, None, "", None, ["kb-a", "kb-b"]) def test_list_datasets_with_include_parsing_status_false_skips_helper(monkeypatch):