Support knowledge metadata filter (#15982)

api/controllers/console/__init__.py

@@ -81,6 +81,7 @@ from .datasets import (
     datasets_segments,
     external,
     hit_testing,
+    metadata,
     website,
 )
 

api/controllers/console/datasets/datasets_document.py

@@ -621,7 +621,7 @@ class DocumentDetailApi(DocumentResource):
             raise InvalidMetadataError(f"Invalid metadata value: {metadata}")
 
         if metadata == "only":
-            response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata}
+            response = {"id": document.id, "doc_type": document.doc_type, "doc_metadata": document.doc_metadata_details}
         elif metadata == "without":
             dataset_process_rules = DatasetService.get_process_rules(dataset_id)
             document_process_rules = document.dataset_process_rule.to_dict()
@@ -682,7 +682,7 @@ class DocumentDetailApi(DocumentResource):
                 "disabled_by": document.disabled_by,
                 "archived": document.archived,
                 "doc_type": document.doc_type,
-                "doc_metadata": document.doc_metadata,
+                "doc_metadata": document.doc_metadata_details,
                 "segment_count": document.segment_count,
                 "average_segment_length": document.average_segment_length,
                 "hit_count": document.hit_count,

api/controllers/console/datasets/metadata.py

@@ -0,0 +1,155 @@
+from flask_login import current_user  # type: ignore  # type: ignore
+from flask_restful import Resource, marshal_with, reqparse  # type: ignore
+from werkzeug.exceptions import NotFound
+
+from controllers.console import api
+from controllers.console.wraps import account_initialization_required, enterprise_license_required, setup_required
+from fields.dataset_fields import dataset_metadata_fields
+from libs.login import login_required
+from services.dataset_service import DatasetService
+from services.entities.knowledge_entities.knowledge_entities import (
+    MetadataArgs,
+    MetadataOperationData,
+)
+from services.metadata_service import MetadataService
+
+
+def _validate_name(name):
+    if not name or len(name) < 1 or len(name) > 40:
+        raise ValueError("Name must be between 1 to 40 characters.")
+    return name
+
+
+def _validate_description_length(description):
+    if len(description) > 400:
+        raise ValueError("Description cannot exceed 400 characters.")
+    return description
+
+
+class DatasetMetadataCreateApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @enterprise_license_required
+    @marshal_with(dataset_metadata_fields)
+    def post(self, dataset_id):
+        parser = reqparse.RequestParser()
+        parser.add_argument("type", type=str, required=True, nullable=True, location="json")
+        parser.add_argument("name", type=str, required=True, nullable=True, location="json")
+        args = parser.parse_args()
+        metadata_args = MetadataArgs(**args)
+
+        dataset_id_str = str(dataset_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        DatasetService.check_dataset_permission(dataset, current_user)
+
+        metadata = MetadataService.create_metadata(dataset_id_str, metadata_args)
+        return metadata, 201
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @enterprise_license_required
+    def get(self, dataset_id):
+        dataset_id_str = str(dataset_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        return MetadataService.get_dataset_metadatas(dataset), 200
+
+
+class DatasetMetadataApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @enterprise_license_required
+    @marshal_with(dataset_metadata_fields)
+    def patch(self, dataset_id, metadata_id):
+        parser = reqparse.RequestParser()
+        parser.add_argument("name", type=str, required=True, nullable=True, location="json")
+        args = parser.parse_args()
+
+        dataset_id_str = str(dataset_id)
+        metadata_id_str = str(metadata_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        DatasetService.check_dataset_permission(dataset, current_user)
+
+        metadata = MetadataService.update_metadata_name(dataset_id_str, metadata_id_str, args.get("name"))
+        return metadata, 200
+
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @enterprise_license_required
+    def delete(self, dataset_id, metadata_id):
+        dataset_id_str = str(dataset_id)
+        metadata_id_str = str(metadata_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        DatasetService.check_dataset_permission(dataset, current_user)
+
+        MetadataService.delete_metadata(dataset_id_str, metadata_id_str)
+        return 200
+
+
+class DatasetMetadataBuiltInFieldApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @enterprise_license_required
+    def get(self):
+        built_in_fields = MetadataService.get_built_in_fields()
+        return {"fields": built_in_fields}, 200
+
+
+class DatasetMetadataBuiltInFieldActionApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @enterprise_license_required
+    def post(self, dataset_id, action):
+        dataset_id_str = str(dataset_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        DatasetService.check_dataset_permission(dataset, current_user)
+
+        if action == "enable":
+            MetadataService.enable_built_in_field(dataset)
+        elif action == "disable":
+            MetadataService.disable_built_in_field(dataset)
+        return 200
+
+
+class DocumentMetadataEditApi(Resource):
+    @setup_required
+    @login_required
+    @account_initialization_required
+    @enterprise_license_required
+    def post(self, dataset_id):
+        dataset_id_str = str(dataset_id)
+        dataset = DatasetService.get_dataset(dataset_id_str)
+        if dataset is None:
+            raise NotFound("Dataset not found.")
+        DatasetService.check_dataset_permission(dataset, current_user)
+
+        parser = reqparse.RequestParser()
+        parser.add_argument("operation_data", type=list, required=True, nullable=True, location="json")
+        args = parser.parse_args()
+        metadata_args = MetadataOperationData(**args)
+
+        MetadataService.update_documents_metadata(dataset, metadata_args)
+
+        return 200
+
+
+api.add_resource(DatasetMetadataCreateApi, "/datasets/<uuid:dataset_id>/metadata")
+api.add_resource(DatasetMetadataApi, "/datasets/<uuid:dataset_id>/metadata/<uuid:metadata_id>")
+api.add_resource(DatasetMetadataBuiltInFieldApi, "/datasets/metadata/built-in")
+api.add_resource(DatasetMetadataBuiltInFieldActionApi, "/datasets/<uuid:dataset_id>/metadata/built-in/<string:action>")
+api.add_resource(DocumentMetadataEditApi, "/datasets/<uuid:dataset_id>/documents/metadata")

api/core/app/app_config/easy_ui_based_app/dataset/manager.py

@@ -1,7 +1,12 @@
 import uuid
 from typing import Optional
 
-from core.app.app_config.entities import DatasetEntity, DatasetRetrieveConfigEntity
+from core.app.app_config.entities import (
+    DatasetEntity,
+    DatasetRetrieveConfigEntity,
+    MetadataFilteringCondition,
+    ModelConfig,
+)
 from core.entities.agent_entities import PlanningStrategy
 from models.model import AppMode
 from services.dataset_service import DatasetService
@@ -78,6 +83,15 @@ class DatasetConfigManager:
                     retrieve_strategy=DatasetRetrieveConfigEntity.RetrieveStrategy.value_of(
                         dataset_configs["retrieval_model"]
                     ),
+                    metadata_filtering_mode=dataset_configs.get("metadata_filtering_mode", "disabled"),
+                    metadata_model_config=ModelConfig(**dataset_configs.get("metadata_model_config"))
+                    if dataset_configs.get("metadata_model_config")
+                    else None,
+                    metadata_filtering_conditions=MetadataFilteringCondition(
+                        **dataset_configs.get("metadata_filtering_conditions", {})
+                    )
+                    if dataset_configs.get("metadata_filtering_conditions")
+                    else None,
                 ),
             )
         else:
@@ -96,6 +110,15 @@ class DatasetConfigManager:
                     weights=dataset_configs.get("weights"),
                     reranking_enabled=dataset_configs.get("reranking_enabled", True),
                     rerank_mode=dataset_configs.get("reranking_mode", "reranking_model"),
+                    metadata_filtering_mode=dataset_configs.get("metadata_filtering_mode", "disabled"),
+                    metadata_model_config=ModelConfig(**dataset_configs.get("metadata_model_config"))
+                    if dataset_configs.get("metadata_model_config")
+                    else None,
+                    metadata_filtering_conditions=MetadataFilteringCondition(
+                        **dataset_configs.get("metadata_filtering_conditions", {})
+                    )
+                    if dataset_configs.get("metadata_filtering_conditions")
+                    else None,
                 ),
             )
 

api/core/app/app_config/entities.py

@@ -1,10 +1,11 @@
 from collections.abc import Sequence
 from enum import Enum, StrEnum
-from typing import Any, Optional
+from typing import Any, Literal, Optional
 
 from pydantic import BaseModel, Field, field_validator
 
 from core.file import FileTransferMethod, FileType, FileUploadConfig
+from core.model_runtime.entities.llm_entities import LLMMode
 from core.model_runtime.entities.message_entities import PromptMessageRole
 from models.model import AppMode
 
@@ -135,6 +136,55 @@ class ExternalDataVariableEntity(BaseModel):
     config: dict[str, Any] = Field(default_factory=dict)
 
 
+SupportedComparisonOperator = Literal[
+    # for string or array
+    "contains",
+    "not contains",
+    "start with",
+    "end with",
+    "is",
+    "is not",
+    "empty",
+    "not empty",
+    # for number
+    "=",
+    "≠",
+    ">",
+    "<",
+    "≥",
+    "≤",
+    # for time
+    "before",
+    "after",
+]
+
+
+class ModelConfig(BaseModel):
+    provider: str
+    name: str
+    mode: LLMMode
+    completion_params: dict[str, Any] = {}
+
+
+class Condition(BaseModel):
+    """
+    Conditon detail
+    """
+
+    name: str
+    comparison_operator: SupportedComparisonOperator
+    value: str | Sequence[str] | None | int | float = None
+
+
+class MetadataFilteringCondition(BaseModel):
+    """
+    Metadata Filtering Condition.
+    """
+
+    logical_operator: Optional[Literal["and", "or"]] = "and"
+    conditions: Optional[list[Condition]] = Field(default=None, deprecated=True)
+
+
 class DatasetRetrieveConfigEntity(BaseModel):
     """
     Dataset Retrieve Config Entity.
@@ -171,6 +221,9 @@ class DatasetRetrieveConfigEntity(BaseModel):
     reranking_model: Optional[dict] = None
     weights: Optional[dict] = None
     reranking_enabled: Optional[bool] = True
+    metadata_filtering_mode: Optional[Literal["disabled", "automatic", "manual"]] = "disabled"
+    metadata_model_config: Optional[ModelConfig] = None
+    metadata_filtering_conditions: Optional[MetadataFilteringCondition] = None
 
 
 class DatasetEntity(BaseModel):

api/core/app/apps/chat/app_runner.py

@@ -180,6 +180,7 @@ class ChatAppRunner(AppRunner):
                 hit_callback=hit_callback,
                 memory=memory,
                 message_id=message.id,
+                inputs=inputs,
             )
 
         # reorganize all inputs and template to prompt messages

api/core/app/apps/completion/app_runner.py

@@ -139,6 +139,7 @@ class CompletionAppRunner(AppRunner):
                 show_retrieve_source=app_config.additional_features.show_retrieve_source,
                 hit_callback=hit_callback,
                 message_id=message.id,
+                inputs=inputs,
             )
 
         # reorganize all inputs and template to prompt messages

api/core/rag/datasource/keyword/jieba/jieba.py

@@ -88,16 +88,17 @@ class Jieba(BaseKeyword):
         keyword_table = self._get_dataset_keyword_table()
 
         k = kwargs.get("top_k", 4)
-
+        document_ids_filter = kwargs.get("document_ids_filter")
         sorted_chunk_indices = self._retrieve_ids_by_query(keyword_table or {}, query, k)
 
         documents = []
         for chunk_index in sorted_chunk_indices:
-            segment = (
-                db.session.query(DocumentSegment)
-                .filter(DocumentSegment.dataset_id == self.dataset.id, DocumentSegment.index_node_id == chunk_index)
-                .first()
+            segment_query = db.session.query(DocumentSegment).filter(
+                DocumentSegment.dataset_id == self.dataset.id, DocumentSegment.index_node_id == chunk_index
             )
+            if document_ids_filter:
+                segment_query = segment_query.filter(DocumentSegment.document_id.in_(document_ids_filter))
+            segment = segment_query.first()
 
             if segment:
                 documents.append(

api/core/rag/datasource/retrieval_service.py

@@ -41,6 +41,7 @@ class RetrievalService:
         reranking_model: Optional[dict] = None,
         reranking_mode: str = "reranking_model",
         weights: Optional[dict] = None,
+        document_ids_filter: Optional[list[str]] = None,
     ):
         if not query:
             return []
@@ -64,6 +65,7 @@ class RetrievalService:
                         top_k=top_k,
                         all_documents=all_documents,
                         exceptions=exceptions,
+                        document_ids_filter=document_ids_filter,
                     )
                 )
             if RetrievalMethod.is_support_semantic_search(retrieval_method):
@@ -79,6 +81,7 @@ class RetrievalService:
                         all_documents=all_documents,
                         retrieval_method=retrieval_method,
                         exceptions=exceptions,
+                        document_ids_filter=document_ids_filter,
                     )
                 )
             if RetrievalMethod.is_support_fulltext_search(retrieval_method):
@@ -130,7 +133,14 @@ class RetrievalService:
 
     @classmethod
     def keyword_search(
-        cls, flask_app: Flask, dataset_id: str, query: str, top_k: int, all_documents: list, exceptions: list
+        cls,
+        flask_app: Flask,
+        dataset_id: str,
+        query: str,
+        top_k: int,
+        all_documents: list,
+        exceptions: list,
+        document_ids_filter: Optional[list[str]] = None,
     ):
         with flask_app.app_context():
             try:
@@ -139,7 +149,10 @@ class RetrievalService:
                     raise ValueError("dataset not found")
 
                 keyword = Keyword(dataset=dataset)
-                documents = keyword.search(cls.escape_query_for_search(query), top_k=top_k)
+
+                documents = keyword.search(
+                    cls.escape_query_for_search(query), top_k=top_k, document_ids_filter=document_ids_filter
+                )
                 all_documents.extend(documents)
             except Exception as e:
                 exceptions.append(str(e))
@@ -156,6 +169,7 @@ class RetrievalService:
         all_documents: list,
         retrieval_method: str,
         exceptions: list,
+        document_ids_filter: Optional[list[str]] = None,
     ):
         with flask_app.app_context():
             try:
@@ -170,6 +184,7 @@ class RetrievalService:
                     top_k=top_k,
                     score_threshold=score_threshold,
                     filter={"group_id": [dataset.id]},
+                    document_ids_filter=document_ids_filter,
                 )
 
                 if documents:

api/core/rag/datasource/vdb/analyticdb/analyticdb_vector.py

@@ -53,7 +53,7 @@ class AnalyticdbVector(BaseVector):
         self.analyticdb_vector.delete_by_metadata_field(key, value)
 
     def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
-        return self.analyticdb_vector.search_by_vector(query_vector)
+        return self.analyticdb_vector.search_by_vector(query_vector, **kwargs)
 
     def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
         return self.analyticdb_vector.search_by_full_text(query, **kwargs)

api/core/rag/datasource/vdb/analyticdb/analyticdb_vector_sql.py

@@ -196,6 +196,11 @@ class AnalyticdbVectorBySql:
         top_k = kwargs.get("top_k", 4)
         if not isinstance(top_k, int) or top_k <= 0:
             raise ValueError("top_k must be a positive integer")
+        document_ids_filter = kwargs.get("document_ids_filter")
+        where_clause = "WHERE 1=1"
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            where_clause += f"AND metadata_->>'document_id' IN ({document_ids})"
         score_threshold = float(kwargs.get("score_threshold") or 0.0)
         with self._get_cursor() as cur:
             query_vector_str = json.dumps(query_vector)
@@ -204,7 +209,7 @@ class AnalyticdbVectorBySql:
                 f"SELECT t.id AS id, t.vector AS vector, (1.0 - t.score) AS score, "
                 f"t.page_content as page_content, t.metadata_ AS metadata_ "
                 f"FROM (SELECT id, vector, page_content, metadata_, vector <=> %s AS score "
-                f"FROM {self.table_name} ORDER BY score LIMIT {top_k} ) t",
+                f"FROM {self.table_name} {where_clause} ORDER BY score LIMIT {top_k} ) t",
                 (query_vector_str,),
             )
             documents = []
@@ -224,12 +229,17 @@ class AnalyticdbVectorBySql:
         top_k = kwargs.get("top_k", 4)
         if not isinstance(top_k, int) or top_k <= 0:
             raise ValueError("top_k must be a positive integer")
+        document_ids_filter = kwargs.get("document_ids_filter")
+        where_clause = ""
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            where_clause += f"AND metadata_->>'document_id' IN ({document_ids})"
         with self._get_cursor() as cur:
             cur.execute(
                 f"""SELECT id, vector, page_content, metadata_, 
                 ts_rank(to_tsvector, to_tsquery_from_text(%s, 'zh_cn'), 32) AS score
                 FROM {self.table_name}
-                WHERE to_tsvector@@to_tsquery_from_text(%s, 'zh_cn')
+                WHERE to_tsvector@@to_tsquery_from_text(%s, 'zh_cn') {where_clause}
                 ORDER BY score DESC
                 LIMIT {top_k}""",
                 (f"'{query}'", f"'{query}'"),

api/core/rag/datasource/vdb/baidu/baidu_vector.py

@@ -123,11 +123,21 @@ class BaiduVector(BaseVector):
 
     def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
         query_vector = [float(val) if isinstance(val, np.float64) else val for val in query_vector]
-        anns = AnnSearch(
-            vector_field=self.field_vector,
-            vector_floats=query_vector,
-            params=HNSWSearchParams(ef=kwargs.get("ef", 10), limit=kwargs.get("top_k", 4)),
-        )
+        document_ids_filter = kwargs.get("document_ids_filter")
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            anns = AnnSearch(
+                vector_field=self.field_vector,
+                vector_floats=query_vector,
+                params=HNSWSearchParams(ef=kwargs.get("ef", 10), limit=kwargs.get("top_k", 4)),
+                filter=f"document_id IN ({document_ids})",
+            )
+        else:
+            anns = AnnSearch(
+                vector_field=self.field_vector,
+                vector_floats=query_vector,
+                params=HNSWSearchParams(ef=kwargs.get("ef", 10), limit=kwargs.get("top_k", 4)),
+            )
         res = self._db.table(self._collection_name).search(
             anns=anns,
             projections=[self.field_id, self.field_text, self.field_metadata],

api/core/rag/datasource/vdb/chroma/chroma_vector.py

@@ -95,7 +95,15 @@ class ChromaVector(BaseVector):
 
     def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
         collection = self._client.get_or_create_collection(self._collection_name)
-        results: QueryResult = collection.query(query_embeddings=query_vector, n_results=kwargs.get("top_k", 4))
+        document_ids_filter = kwargs.get("document_ids_filter")
+        if document_ids_filter:
+            results: QueryResult = collection.query(
+                query_embeddings=query_vector,
+                n_results=kwargs.get("top_k", 4),
+                where={"document_id": {"$in": document_ids_filter}},  # type: ignore
+            )
+        else:
+            results: QueryResult = collection.query(query_embeddings=query_vector, n_results=kwargs.get("top_k", 4))  # type: ignore
         score_threshold = float(kwargs.get("score_threshold") or 0.0)
 
         # Check if results contain data

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

@@ -117,6 +117,9 @@ class ElasticSearchVector(BaseVector):
         top_k = kwargs.get("top_k", 4)
         num_candidates = math.ceil(top_k * 1.5)
         knn = {"field": Field.VECTOR.value, "query_vector": query_vector, "k": top_k, "num_candidates": num_candidates}
+        document_ids_filter = kwargs.get("document_ids_filter")
+        if document_ids_filter:
+            knn["filter"] = {"terms": {"metadata.document_id": document_ids_filter}}
 
         results = self._client.search(index=self._collection_name, knn=knn, size=top_k)
 
@@ -145,6 +148,9 @@ class ElasticSearchVector(BaseVector):
 
     def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
         query_str = {"match": {Field.CONTENT_KEY.value: query}}
+        document_ids_filter = kwargs.get("document_ids_filter")
+        if document_ids_filter:
+            query_str["filter"] = {"terms": {"metadata.document_id": document_ids_filter}}  # type: ignore
         results = self._client.search(index=self._collection_name, query=query_str, size=kwargs.get("top_k", 4))
         docs = []
         for hit in results["hits"]["hits"]:

api/core/rag/datasource/vdb/lindorm/lindorm_vector.py

@@ -168,7 +168,12 @@ class LindormVectorStore(BaseVector):
             raise ValueError("All elements in query_vector should be floats")
 
         top_k = kwargs.get("top_k", 10)
-        query = default_vector_search_query(query_vector=query_vector, k=top_k, **kwargs)
+        document_ids_filter = kwargs.get("document_ids_filter")
+        filters = []
+        if document_ids_filter:
+            filters.append({"terms": {"metadata.document_id": document_ids_filter}})
+        query = default_vector_search_query(query_vector=query_vector, k=top_k, filters=filters, **kwargs)
+
         try:
             params = {}
             if self._using_ugc:
@@ -206,7 +211,10 @@ class LindormVectorStore(BaseVector):
         should = kwargs.get("should")
         minimum_should_match = kwargs.get("minimum_should_match", 0)
         top_k = kwargs.get("top_k", 10)
-        filters = kwargs.get("filter")
+        filters = kwargs.get("filter", [])
+        document_ids_filter = kwargs.get("document_ids_filter")
+        if document_ids_filter:
+            filters.append({"terms": {"metadata.document_id": document_ids_filter}})
         routing = self._routing
         full_text_query = default_text_search_query(
             query_text=query,

api/core/rag/datasource/vdb/milvus/milvus_vector.py

@@ -228,12 +228,18 @@ class MilvusVector(BaseVector):
         """
         Search for documents by vector similarity.
         """
+        document_ids_filter = kwargs.get("document_ids_filter")
+        filter = ""
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            filter = f'metadata["document_id"] in ({document_ids})'
         results = self._client.search(
             collection_name=self._collection_name,
             data=[query_vector],
             anns_field=Field.VECTOR.value,
             limit=kwargs.get("top_k", 4),
             output_fields=[Field.CONTENT_KEY.value, Field.METADATA_KEY.value],
+            filter=filter,
         )
 
         return self._process_search_results(
@@ -249,6 +255,11 @@ class MilvusVector(BaseVector):
         if not self._hybrid_search_enabled or not self.field_exists(Field.SPARSE_VECTOR.value):
             logger.warning("Full-text search is not supported in current Milvus version (requires >= 2.5.0)")
             return []
+        document_ids_filter = kwargs.get("document_ids_filter")
+        filter = ""
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            filter = f'metadata["document_id"] in ({document_ids})'
 
         results = self._client.search(
             collection_name=self._collection_name,
@@ -256,6 +267,7 @@ class MilvusVector(BaseVector):
             anns_field=Field.SPARSE_VECTOR.value,
             limit=kwargs.get("top_k", 4),
             output_fields=[Field.CONTENT_KEY.value, Field.METADATA_KEY.value],
+            filter=filter,
         )
 
         return self._process_search_results(

api/core/rag/datasource/vdb/myscale/myscale_vector.py

@@ -133,6 +133,10 @@ class MyScaleVector(BaseVector):
             if self._metric.upper() == "COSINE" and order == SortOrder.ASC and score_threshold > 0.0
             else ""
         )
+        document_ids_filter = kwargs.get("document_ids_filter")
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            where_str = f"{where_str} AND metadata['document_id'] in ({document_ids})"
         sql = f"""
             SELECT text, vector, metadata, {dist} as dist FROM {self._config.database}.{self._collection_name}
             {where_str} ORDER BY dist {order.value} LIMIT {top_k}

api/core/rag/datasource/vdb/oceanbase/oceanbase_vector.py

@@ -154,6 +154,11 @@ class OceanBaseVector(BaseVector):
         return []
 
     def search_by_vector(self, query_vector: list[float], **kwargs: Any) -> list[Document]:
+        document_ids_filter = kwargs.get("document_ids_filter")
+        where_clause = None
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            where_clause = f"metadata->>'$.document_id' in ({document_ids})"
         ef_search = kwargs.get("ef_search", self._hnsw_ef_search)
         if ef_search != self._hnsw_ef_search:
             self._client.set_ob_hnsw_ef_search(ef_search)
@@ -167,6 +172,7 @@ class OceanBaseVector(BaseVector):
             distance_func=func.l2_distance,
             output_column_names=["text", "metadata"],
             with_dist=True,
+            where_clause=where_clause,
         )
         docs = []
         for text, metadata, distance in cur:

api/core/rag/datasource/vdb/opensearch/opensearch_vector.py

@@ -154,6 +154,9 @@ class OpenSearchVector(BaseVector):
             "size": kwargs.get("top_k", 4),
             "query": {"knn": {Field.VECTOR.value: {Field.VECTOR.value: query_vector, "k": kwargs.get("top_k", 4)}}},
         }
+        document_ids_filter = kwargs.get("document_ids_filter")
+        if document_ids_filter:
+            query["query"] = {"terms": {"metadata.document_id": document_ids_filter}}
 
         try:
             response = self._client.search(index=self._collection_name.lower(), body=query)
@@ -179,6 +182,9 @@ class OpenSearchVector(BaseVector):
 
     def search_by_full_text(self, query: str, **kwargs: Any) -> list[Document]:
         full_text_query = {"query": {"match": {Field.CONTENT_KEY.value: query}}}
+        document_ids_filter = kwargs.get("document_ids_filter")
+        if document_ids_filter:
+            full_text_query["query"]["terms"] = {"metadata.document_id": document_ids_filter}
 
         response = self._client.search(index=self._collection_name.lower(), body=full_text_query)
 

api/core/rag/datasource/vdb/oracle/oraclevector.py

@@ -201,10 +201,15 @@ class OracleVector(BaseVector):
         :return: List of Documents that are nearest to the query vector.
         """
         top_k = kwargs.get("top_k", 4)
+        document_ids_filter = kwargs.get("document_ids_filter")
+        where_clause = ""
+        if document_ids_filter:
+            document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+            where_clause = f"WHERE metadata->>'document_id' in ({document_ids})"
         with self._get_cursor() as cur:
             cur.execute(
                 f"SELECT meta, text, vector_distance(embedding,:1) AS distance FROM {self.table_name}"
-                f" ORDER BY distance fetch first {top_k} rows only",
+                f" {where_clause} ORDER BY distance fetch first {top_k} rows only",
                 [numpy.array(query_vector)],
             )
             docs = []
@@ -257,9 +262,15 @@ class OracleVector(BaseVector):
                     if token not in stop_words:
                         entities.append(token)
             with self._get_cursor() as cur:
+                document_ids_filter = kwargs.get("document_ids_filter")
+                where_clause = ""
+                if document_ids_filter:
+                    document_ids = ", ".join(f"'{id}'" for id in document_ids_filter)
+                    where_clause = f" AND metadata->>'document_id' in ({document_ids}) "
                 cur.execute(
                     f"select meta, text, embedding FROM {self.table_name}"
-                    f" WHERE CONTAINS(text, :1, 1) > 0 order by score(1) desc fetch first {top_k} rows only",
+                    f"WHERE CONTAINS(text, :1, 1) > 0 {where_clause} "
+                    f"order by score(1) desc fetch first {top_k} rows only",
                     [" ACCUM ".join(entities)],
                 )
                 docs = []