From 3bfe802c4857e6cf3b75ef96c0dbdb34e31911b9 Mon Sep 17 00:00:00 2001
From: heheer <heheer@sealos.io>
Date: Mon, 24 Feb 2025 17:43:31 +0800
Subject: [PATCH] fix collection folder tags filter (#3853)

* fix collection folder tags filter

* add comment

* fix
---
 .../service/core/dataset/search/controller.ts | 80 +++++++++++++++++--
 1 file changed, 73 insertions(+), 7 deletions(-)

diff --git a/packages/service/core/dataset/search/controller.ts b/packages/service/core/dataset/search/controller.ts
index 762ae687d..de8ea8fdc 100644
--- a/packages/service/core/dataset/search/controller.ts
+++ b/packages/service/core/dataset/search/controller.ts
@@ -200,6 +200,62 @@ export async function searchDatasetData(
       forbidCollectionIdList: collections.map((item) => String(item._id))
     };
   };
+
+  async function getAllCollectionIds({
+    teamId,
+    datasetIds,
+    parentCollectionIds
+  }: {
+    teamId: string;
+    datasetIds: string[];
+    parentCollectionIds: string[];
+  }): Promise<string[]> {
+    if (!parentCollectionIds.length) {
+      return [];
+    }
+    const collections = await MongoDatasetCollection.find(
+      {
+        teamId,
+        datasetId: { $in: datasetIds },
+        _id: { $in: parentCollectionIds }
+      },
+      '_id type',
+      {
+        ...readFromSecondary
+      }
+    ).lean();
+
+    const resultIds = new Set(collections.map((item) => String(item._id)));
+
+    const folderIds = collections
+      .filter((item) => item.type === 'folder')
+      .map((item) => String(item._id));
+
+    // Get all child collection ids
+    if (folderIds.length) {
+      const childCollections = await MongoDatasetCollection.find(
+        {
+          teamId,
+          datasetId: { $in: datasetIds },
+          parentId: { $in: folderIds }
+        },
+        '_id',
+        {
+          ...readFromSecondary
+        }
+      ).lean();
+
+      const childIds = await getAllCollectionIds({
+        teamId,
+        datasetIds,
+        parentCollectionIds: childCollections.map((item) => String(item._id))
+      });
+
+      childIds.forEach((id) => resultIds.add(id));
+    }
+
+    return Array.from(resultIds);
+  }
   /* 
     Collection metadata filter
     标签过滤：
@@ -326,13 +382,23 @@ export async function searchDatasetData(
       }
 
       // Concat tag and time
-      if (tagCollectionIdList && createTimeCollectionIdList) {
-        return tagCollectionIdList.filter((id) => createTimeCollectionIdList!.includes(id));
-      } else if (tagCollectionIdList) {
-        return tagCollectionIdList;
-      } else if (createTimeCollectionIdList) {
-        return createTimeCollectionIdList;
-      }
+      const finalIds = (() => {
+        if (tagCollectionIdList && createTimeCollectionIdList) {
+          return tagCollectionIdList.filter((id) =>
+            (createTimeCollectionIdList as string[]).includes(id)
+          );
+        }
+
+        return tagCollectionIdList || createTimeCollectionIdList;
+      })();
+
+      return finalIds
+        ? await getAllCollectionIds({
+            teamId,
+            datasetIds,
+            parentCollectionIds: finalIds
+          })
+        : undefined;
     } catch (error) {}
   };
   const embeddingRecall = async ({