add vdb document id index (#16244)

Co-authored-by: crazywoola <[email protected]>

Author: JohnJyong

api/core/rag/datasource/vdb/elasticsearch/elasticsearch_vector.py

@@ -196,7 +196,8 @@ def create_collection(
                         Field.METADATA_KEY.value: {
                             "type": "object",
                             "properties": {
-                                "doc_id": {"type": "keyword"}  # Map doc_id to keyword type
+                                "doc_id": {"type": "keyword"},  # Map doc_id to keyword type
+                                "document_id": {"type": "keyword"},  # Map doc_id to keyword type
                             },
                         },
                     }

api/core/rag/datasource/vdb/field.py

@@ -11,3 +11,4 @@ class Field(Enum):
     TEXT_KEY = "text"
     PRIMARY_KEY = "id"
     DOC_ID = "metadata.doc_id"
+    DOCUMENT_ID = "metadata.document_id"

api/core/rag/datasource/vdb/qdrant/qdrant_vector.py

@@ -134,6 +134,10 @@ def create_collection(self, collection_name: str, vector_size: int):
                 self._client.create_payload_index(
                     collection_name, Field.DOC_ID.value, field_schema=PayloadSchemaType.KEYWORD
                 )
+                # create document_id payload index
+                self._client.create_payload_index(
+                    collection_name, Field.DOCUMENT_ID.value, field_schema=PayloadSchemaType.KEYWORD
+                )
                 # create full text index
                 text_index_params = TextIndexParams(
                     type=TextIndexType.TEXT,

api/core/rag/datasource/vdb/tidb_on_qdrant/tidb_on_qdrant_vector.py

@@ -144,6 +144,10 @@ def create_collection(self, collection_name: str, vector_size: int):
                 self._client.create_payload_index(
                     collection_name, Field.DOC_ID.value, field_schema=PayloadSchemaType.KEYWORD
                 )
+                # create document_id payload index
+                self._client.create_payload_index(
+                    collection_name, Field.DOCUMENT_ID.value, field_schema=PayloadSchemaType.KEYWORD
+                )
                 # create full text index
                 text_index_params = TextIndexParams(
                     type=TextIndexType.TEXT,
@@ -318,23 +322,17 @@ def text_exists(self, id: str) -> bool:
     def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
         from qdrant_client.http import models
 
-        filter = models.Filter(
-            must=[
-                models.FieldCondition(
-                    key="group_id",
-                    match=models.MatchValue(value=self._group_id),
-                ),
-            ],
-        )
+        filter = None
         document_ids_filter = kwargs.get("document_ids_filter")
         if document_ids_filter:
-            if filter.must:
-                filter.must.append(
+            filter = models.Filter(
+                must=[
                     models.FieldCondition(
                         key="metadata.document_id",
                         match=models.MatchAny(any=document_ids_filter),
                     )
-                )
+                ],
+            )
         results = self._client.search(
             collection_name=self._collection_name,
             query_vector=query_vector,
@@ -369,23 +367,17 @@ def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
         """
         from qdrant_client.http import models
 
-        scroll_filter = models.Filter(
-            must=[
-                models.FieldCondition(
-                    key="page_content",
-                    match=models.MatchText(text=query),
-                )
-            ]
-        )
+        scroll_filter = None
         document_ids_filter = kwargs.get("document_ids_filter")
         if document_ids_filter:
-            if scroll_filter.must:
-                scroll_filter.must.append(
+            scroll_filter = models.Filter(
+                must=[
                     models.FieldCondition(
                         key="metadata.document_id",
                         match=models.MatchAny(any=document_ids_filter),
                     )
-                )
+                ]
+            )
         response = self._client.scroll(
             collection_name=self._collection_name,
             scroll_filter=scroll_filter,

api/core/rag/datasource/vdb/tidb_vector/tidb_vector.py

@@ -105,10 +105,12 @@ def _create_collection(self, dimension: int):
                         text TEXT NOT NULL,
                         meta JSON NOT NULL,
                         doc_id VARCHAR(64) AS (JSON_UNQUOTE(JSON_EXTRACT(meta, '$.doc_id'))) STORED,
+                        document_id VARCHAR(64) AS (JSON_UNQUOTE(JSON_EXTRACT(meta, '$.document_id'))) STORED,
                         vector VECTOR<FLOAT>({dimension}) NOT NULL,
                         create_time DATETIME DEFAULT CURRENT_TIMESTAMP,
                         update_time DATETIME DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
                         KEY (doc_id),
+                        KEY (document_id),
                         VECTOR INDEX idx_vector (({tidb_dist_func}(vector))) USING HNSW
                     );
                 """)