From 1aa6a2126ff2c4972ef98226dbc427f00dd0f493 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=87=91=E9=B9=8F=E7=A8=8B?=
 <jinpengcheng01@corp.netease.com>
Date: Mon, 28 Apr 2025 19:41:49 +0800
Subject: [PATCH] fix: metadata filter not work

---
 api/core/agent/cot_agent_runner.py                            | 4 ++++
 api/core/agent/fc_agent_runner.py                             | 4 ++++
 .../tools/utils/dataset_retriever/dataset_retriever_tool.py   | 4 ++++
 3 files changed, 12 insertions(+)

diff --git a/api/core/agent/cot_agent_runner.py b/api/core/agent/cot_agent_runner.py
index feb8abf6ef..eee96fccd5 100644
--- a/api/core/agent/cot_agent_runner.py
+++ b/api/core/agent/cot_agent_runner.py
@@ -69,6 +69,10 @@ class CotAgentRunner(BaseAgentRunner, ABC):
         tool_instances, prompt_messages_tools = self._init_prompt_tools()
         self._prompt_messages_tools = prompt_messages_tools
 
+        metadata_filtering_conditions = app_config.dataset.retrieve_config.metadata_filtering_conditions
+        for key, dataset_retriever_tool in tool_instances.items():
+            dataset_retriever_tool.retrieval_tool.metadata_filtering_conditions = metadata_filtering_conditions
+
         function_call_state = True
         llm_usage: dict[str, Optional[LLMUsage]] = {"usage": None}
         final_answer = ""
diff --git a/api/core/agent/fc_agent_runner.py b/api/core/agent/fc_agent_runner.py
index a1110e7709..2f0f534479 100644
--- a/api/core/agent/fc_agent_runner.py
+++ b/api/core/agent/fc_agent_runner.py
@@ -45,6 +45,10 @@ class FunctionCallAgentRunner(BaseAgentRunner):
         # convert tools into ModelRuntime Tool format
         tool_instances, prompt_messages_tools = self._init_prompt_tools()
 
+        metadata_filtering_conditions = app_config.dataset.retrieve_config.metadata_filtering_conditions
+        for key, dataset_retriever_tool in tool_instances.items():
+            dataset_retriever_tool.retrieval_tool.metadata_filtering_conditions = metadata_filtering_conditions
+
         assert app_config.agent
 
         iteration_step = 1
diff --git a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py
index f5838c3b76..031ca517b2 100644
--- a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py
+++ b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py
@@ -6,6 +6,7 @@ from core.rag.datasource.retrieval_service import RetrievalService
 from core.rag.entities.context_entities import DocumentContext
 from core.rag.models.document import Document as RetrievalDocument
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
+from core.rag.entities.metadata_entities import MetadataCondition
 from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool
 from extensions.ext_database import db
 from models.dataset import Dataset
@@ -33,6 +34,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
     args_schema: type[BaseModel] = DatasetRetrieverToolInput
     description: str = "use this to retrieve a dataset. "
     dataset_id: str
+    metadata_filtering_conditions: MetadataCondition
 
     @classmethod
     def from_dataset(cls, dataset: Dataset, **kwargs):
@@ -46,6 +48,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
             tenant_id=dataset.tenant_id,
             dataset_id=dataset.id,
             description=description,
+            metadata_filtering_conditions= MetadataCondition(),
             **kwargs,
         )
 
@@ -65,6 +68,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
                 dataset_id=dataset.id,
                 query=query,
                 external_retrieval_parameters=dataset.retrieval_model,
+                metadata_condition=self.metadata_filtering_conditions
             )
             for external_document in external_documents:
                 document = RetrievalDocument(