dataset metadata update

api/core/workflow/nodes/code/code_node.py

@@ -195,7 +195,7 @@ class CodeNode(BaseNode[CodeNodeData]):
             if output_config.type == "object":
                 # check if output is object
                 if not isinstance(result.get(output_name), dict):
-                    if isinstance(result.get(output_name), type(None)):
+                    if result.get(output_name) is None:
                         transformed_result[output_name] = None
                     else:
                         raise OutputValidationError(
@@ -223,7 +223,7 @@ class CodeNode(BaseNode[CodeNodeData]):
             elif output_config.type == "array[number]":
                 # check if array of number available
                 if not isinstance(result[output_name], list):
-                    if isinstance(result[output_name], type(None)):
+                    if result[output_name] is None:
                         transformed_result[output_name] = None
                     else:
                         raise OutputValidationError(
@@ -244,7 +244,7 @@ class CodeNode(BaseNode[CodeNodeData]):
             elif output_config.type == "array[string]":
                 # check if array of string available
                 if not isinstance(result[output_name], list):
-                    if isinstance(result[output_name], type(None)):
+                    if result[output_name] is None:
                         transformed_result[output_name] = None
                     else:
                         raise OutputValidationError(
@@ -265,7 +265,7 @@ class CodeNode(BaseNode[CodeNodeData]):
             elif output_config.type == "array[object]":
                 # check if array of object available
                 if not isinstance(result[output_name], list):
-                    if isinstance(result[output_name], type(None)):
+                    if result[output_name] is None:
                         transformed_result[output_name] = None
                     else:
                         raise OutputValidationError(

api/core/workflow/nodes/knowledge_retrieval/entities.py

@@ -1,8 +1,10 @@
+from collections.abc import Sequence
 from typing import Any, Literal, Optional
 
-from pydantic import BaseModel
+from pydantic import BaseModel, Field
 
 from core.workflow.nodes.base import BaseNodeData
+from core.workflow.nodes.llm.entities import VisionConfig
 
 
 class RerankingModelConfig(BaseModel):
@@ -73,11 +75,44 @@ class SingleRetrievalConfig(BaseModel):
     model: ModelConfig
 
 
+SupportedComparisonOperator = Literal[
+    # for string or array
+    "contains",
+    "not contains",
+    "starts with",
+    "ends with",
+    "is",
+    "is not",
+    "empty",
+    "is not empty",
+    # for number
+    "=",
+    "≠",
+    ">",
+    "<",
+    "≥",
+    "≤",
+    # for time
+    "before",
+    "after",
+]
+
+
+class Condition(BaseModel):
+    """
+    Conditon detail
+    """
+
+    metadata_name: str
+    comparison_operator: SupportedComparisonOperator
+    value: str | Sequence[str] | None = None
+
+
 class MetadataFilteringCondition(BaseModel):
     """
     Metadata Filtering Condition.
     """
-    
+
     logical_operator: Optional[Literal["and", "or"]] = "and"
     conditions: Optional[list[Condition]] = Field(default=None, deprecated=True)
 
@@ -93,5 +128,7 @@ class KnowledgeRetrievalNodeData(BaseNodeData):
     retrieval_mode: Literal["single", "multiple"]
     multiple_retrieval_config: Optional[MultipleRetrievalConfig] = None
     single_retrieval_config: Optional[SingleRetrievalConfig] = None
-    metadata_filtering_mode: Optional[Literal["disabled", "automatic", "manual"]]  = "disabled"
-    metadata_filtering_conditions: Optional[dict[str, Any]] = None
+    metadata_filtering_mode: Optional[Literal["disabled", "automatic", "manual"]] = "disabled"
+    metadata_model_config: Optional[ModelConfig] = None
+    metadata_filtering_conditions: Optional[MetadataFilteringCondition] = None
+    vision: VisionConfig = Field(default_factory=VisionConfig)

api/core/workflow/nodes/knowledge_retrieval/exc.py

@@ -16,3 +16,7 @@ class ModelNotSupportedError(KnowledgeRetrievalNodeError):
 
 class ModelQuotaExceededError(KnowledgeRetrievalNodeError):
     """Raised when the model provider quota is exceeded."""
+
+
+class InvalidModelTypeError(KnowledgeRetrievalNodeError):
+    """Raised when the model is not a Large Language Model."""

api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py

@@ -1,6 +1,8 @@
+import json
 import logging
+from collections import defaultdict
 from collections.abc import Mapping, Sequence
-from typing import Any, cast
+from typing import Any, Optional, cast
 
 from sqlalchemy import func
 
@@ -9,21 +11,38 @@ from core.app.entities.app_invoke_entities import ModelConfigWithCredentialsEnti
 from core.entities.agent_entities import PlanningStrategy
 from core.entities.model_entities import ModelStatus
 from core.model_manager import ModelInstance, ModelManager
-from core.model_runtime.entities.model_entities import ModelFeature, ModelType
+from core.model_runtime.entities.message_entities import PromptMessageRole
+from core.model_runtime.entities.model_entities import ModelFeature, ModelPropertyKey, ModelType
 from core.model_runtime.model_providers.__base.large_language_model import LargeLanguageModel
+from core.prompt.advanced_prompt_transform import AdvancedPromptTransform
+from core.prompt.simple_prompt_transform import ModelMode
 from core.rag.datasource.retrieval_service import RetrievalService
 from core.rag.retrieval.dataset_retrieval import DatasetRetrieval
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from core.variables import StringSegment
 from core.workflow.entities.node_entities import NodeRunResult
-from core.workflow.nodes.base import BaseNode
 from core.workflow.nodes.enums import NodeType
+from core.workflow.nodes.event.event import ModelInvokeCompletedEvent
+from core.workflow.nodes.knowledge_retrieval.template_prompts import (
+    METADATA_FILTER_ASSISTANT_PROMPT_1,
+    METADATA_FILTER_ASSISTANT_PROMPT_2,
+    METADATA_FILTER_COMPLETION_PROMPT,
+    METADATA_FILTER_SYSTEM_PROMPT,
+    METADATA_FILTER_USER_PROMPT_1,
+    METADATA_FILTER_USER_PROMPT_3,
+)
+from core.workflow.nodes.list_operator.exc import InvalidConditionError
+from core.workflow.nodes.llm.entities import LLMNodeChatModelMessage, LLMNodeCompletionModelPromptTemplate
+from core.workflow.nodes.llm.node import LLMNode
+from core.workflow.nodes.question_classifier.template_prompts import QUESTION_CLASSIFIER_USER_PROMPT_2
 from extensions.ext_database import db
-from models.dataset import Dataset, Document
+from libs.json_in_md_parser import parse_and_check_json_markdown
+from models.dataset import Dataset, DatasetMetadata, Document
 from models.workflow import WorkflowNodeExecutionStatus
 
 from .entities import KnowledgeRetrievalNodeData
 from .exc import (
+    InvalidModelTypeError,
     KnowledgeRetrievalNodeError,
     ModelCredentialsNotInitializedError,
     ModelNotExistError,
@@ -42,13 +61,14 @@ default_retrieval_model = {
 }
 
 
-class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]):
+class KnowledgeRetrievalNode(LLMNode):
     _node_data_cls = KnowledgeRetrievalNodeData
     _node_type = NodeType.KNOWLEDGE_RETRIEVAL
 
     def _run(self) -> NodeRunResult:
+        node_data = cast(KnowledgeRetrievalNodeData, self.node_data)
         # extract variables
-        variable = self.graph_runtime_state.variable_pool.get(self.node_data.query_variable_selector)
+        variable = self.graph_runtime_state.variable_pool.get(node_data.query_variable_selector)
         if not isinstance(variable, StringSegment):
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.FAILED,
@@ -63,7 +83,7 @@ class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]):
             )
         # retrieve knowledge
         try:
-            results = self._fetch_dataset_retriever(node_data=self.node_data, query=query)
+            results = self._fetch_dataset_retriever(node_data=node_data, query=query)
             outputs = {"result": results}
             return NodeRunResult(
                 status=WorkflowNodeExecutionStatus.SUCCEEDED, inputs=variables, process_data=None, outputs=outputs
@@ -95,8 +115,8 @@ class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]):
             db.session.query(Document.dataset_id, func.count(Document.id).label("available_document_count"))
             .filter(
                 Document.indexing_status == "completed",
-                Document.enabled == True,  # noqa: E712
-                Document.archived == False,  # noqa: E712
+                Document.enabled == True,
+                Document.archived == False,
                 Document.dataset_id.in_(dataset_ids),
             )
             .group_by(Document.dataset_id)
@@ -117,6 +137,9 @@ class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]):
             if not dataset:
                 continue
             available_datasets.append(dataset)
+        metadata_filter_document_ids = self._get_metadata_filter_condition(
+            [dataset.id for dataset in available_datasets], query, node_data
+        )
         all_documents = []
         dataset_retrieval = DatasetRetrieval()
         if node_data.retrieval_mode == DatasetRetrieveConfigEntity.RetrieveStrategy.SINGLE.value:
@@ -146,6 +169,7 @@ class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]):
                     model_config=model_config,
                     model_instance=model_instance,
                     planning_strategy=planning_strategy,
+                    metadata_filter_document_ids=metadata_filter_document_ids,
                 )
         elif node_data.retrieval_mode == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE.value:
             if node_data.multiple_retrieval_config is None:
@@ -221,8 +245,8 @@ class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]):
                     dataset = Dataset.query.filter_by(id=segment.dataset_id).first()
                     document = Document.query.filter(
                         Document.id == segment.document_id,
-                        Document.enabled == True,  # noqa: E712
-                        Document.archived == False,  # noqa: E712
+                        Document.enabled == True,
+                        Document.archived == False,
                     ).first()
                     if dataset and document:
                         source = {
@@ -258,6 +282,134 @@ class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]):
                 item["metadata"]["position"] = position
         return retrieval_resource_list
 
+    def _get_metadata_filter_condition(
+        self, dataset_ids: list, query: str, node_data: KnowledgeRetrievalNodeData
+    ) -> dict[str, list[str]]:
+        document_query = db.session.query(Document.id).filter(
+            Document.dataset_id.in_(dataset_ids),
+            Document.indexing_status == "completed",
+            Document.enabled == True,
+            Document.archived == False,
+        )
+        if node_data.metadata_filtering_mode == "disabled":
+            return None
+        elif node_data.metadata_filtering_mode == "automatic":
+            automatic_metadata_filters = self._automatic_metadata_filter_func(dataset_ids, query, node_data)
+            if automatic_metadata_filters:
+                for filter in automatic_metadata_filters:
+                    self._process_metadata_filter_func(
+                        filter.get("condition"), filter.get("metadata_name"), filter.get("value"), document_query
+                    )
+        elif node_data.metadata_filtering_mode == "manual":
+            for condition in node_data.metadata_filtering_conditions.conditions:
+                metadata_name = condition.metadata_name
+                expected_value = condition.value
+                if isinstance(expected_value, str):
+                    expected_value = self.graph_runtime_state.variable_pool.convert_template(expected_value).text
+                self._process_metadata_filter_func(
+                    condition.comparison_operator, metadata_name, expected_value, document_query
+                )
+        else:
+            raise ValueError("Invalid metadata filtering mode")
+        documnents = document_query.all()
+        # group by dataset_id
+        metadata_filter_document_ids = defaultdict(list)
+        for document in documnents:
+            metadata_filter_document_ids[document.dataset_id].append(document.id)
+        return metadata_filter_document_ids
+
+    def _automatic_metadata_filter_func(
+        self, dataset_ids: list, query: str, node_data: KnowledgeRetrievalNodeData
+    ) -> list[dict[str, Any]]:
+        # get all metadata field
+        metadata_fields = db.session.query(DatasetMetadata).filter(DatasetMetadata.dataset_id.in_(dataset_ids)).all()
+        all_metadata_fields = [metadata_field.field_name for metadata_field in metadata_fields]
+        # get metadata model config
+        metadata_model_config = node_data.metadata_model_config
+        if metadata_model_config is None:
+            raise ValueError("metadata_model_config is required")
+        # get metadata model instance
+        # fetch model config
+        model_instance, model_config = self._fetch_model_config(node_data.metadata_model_config)
+        # fetch prompt messages
+        prompt_template = self._get_prompt_template(
+            node_data=node_data,
+            query=query or "",
+            metadata_fields=all_metadata_fields,
+        )
+        prompt_messages, stop = self._fetch_prompt_messages(
+            prompt_template=prompt_template,
+            sys_query=query,
+            memory=None,
+            model_config=model_config,
+            sys_files=[],
+            vision_enabled=node_data.vision.enabled,
+            vision_detail=node_data.vision.configs.detail,
+            variable_pool=self.graph_runtime_state.variable_pool,
+            jinja2_variables=[],
+        )
+
+        result_text = ""
+        try:
+            # handle invoke result
+            generator = self._invoke_llm(
+                node_data_model=node_data.model,
+                model_instance=model_instance,
+                prompt_messages=prompt_messages,
+                stop=stop,
+            )
+
+            for event in generator:
+                if isinstance(event, ModelInvokeCompletedEvent):
+                    result_text = event.text
+                    break
+
+            result_text_json = parse_and_check_json_markdown(result_text, [])
+            automatic_metadata_filters = []
+            if "metadata_map" in result_text_json:
+                metadata_map = result_text_json["metadata_map"]
+                for item in metadata_map:
+                    if item.get("metadata_field_name") in all_metadata_fields:
+                        automatic_metadata_filters.append(
+                            {
+                                "metadata_name": item.get("metadata_field_name"),
+                                "value": item.get("metadata_field_value"),
+                                "condition": item.get("comparison_operator"),
+                            }
+                        )
+        except Exception as e:
+            return None
+        return automatic_metadata_filters
+
+    def _process_metadata_filter_func(*, condition: str, metadata_name: str, value: str, query):
+        match condition:
+            case "contains":
+                query = query.filter(Document.doc_metadata[metadata_name].like(f"%{value}%"))
+            case "not contains":
+                query = query.filter(Document.doc_metadata[metadata_name].notlike(f"%{value}%"))
+            case "start with":
+                query = query.filter(Document.doc_metadata[metadata_name].like(f"{value}%"))
+            case "end with":
+                query = query.filter(Document.doc_metadata[metadata_name].like(f"%{value}"))
+            case "is", "=":
+                query = query.filter(Document.doc_metadata[metadata_name] == value)
+            case "is not", "≠":
+                query = query.filter(Document.doc_metadata[metadata_name] != value)
+            case "is empty":
+                query = query.filter(Document.doc_metadata[metadata_name].is_(None))
+            case "is not empty":
+                query = query.filter(Document.doc_metadata[metadata_name].isnot(None))
+            case "before", "<":
+                query = query.filter(Document.doc_metadata[metadata_name] < value)
+            case "after", ">":
+                query = query.filter(Document.doc_metadata[metadata_name] > value)
+            case "≤", ">=":
+                query = query.filter(Document.doc_metadata[metadata_name] <= value)
+            case "≥", ">=":
+                query = query.filter(Document.doc_metadata[metadata_name] >= value)
+            case _:
+                raise InvalidConditionError(f"Invalid condition: {condition}")
+
     @classmethod
     def _extract_variable_selector_to_variable_mapping(
         cls,
@@ -343,3 +495,94 @@ class KnowledgeRetrievalNode(BaseNode[KnowledgeRetrievalNodeData]):
             parameters=completion_params,
             stop=stop,
         )
+
+    def _calculate_rest_token(
+        self,
+        node_data: KnowledgeRetrievalNodeData,
+        query: str,
+        model_config: ModelConfigWithCredentialsEntity,
+        context: Optional[str],
+    ) -> int:
+        prompt_transform = AdvancedPromptTransform(with_variable_tmpl=True)
+        prompt_template = self._get_prompt_template(node_data, query, None, 2000)
+        prompt_messages = prompt_transform.get_prompt(
+            prompt_template=prompt_template,
+            inputs={},
+            query="",
+            files=[],
+            context=context,
+            memory_config=node_data.memory,
+            memory=None,
+            model_config=model_config,
+        )
+        rest_tokens = 2000
+
+        model_context_tokens = model_config.model_schema.model_properties.get(ModelPropertyKey.CONTEXT_SIZE)
+        if model_context_tokens:
+            model_instance = ModelInstance(
+                provider_model_bundle=model_config.provider_model_bundle, model=model_config.model
+            )
+
+            curr_message_tokens = model_instance.get_llm_num_tokens(prompt_messages)
+
+            max_tokens = 0
+            for parameter_rule in model_config.model_schema.parameter_rules:
+                if parameter_rule.name == "max_tokens" or (
+                    parameter_rule.use_template and parameter_rule.use_template == "max_tokens"
+                ):
+                    max_tokens = (
+                        model_config.parameters.get(parameter_rule.name)
+                        or model_config.parameters.get(parameter_rule.use_template or "")
+                    ) or 0
+
+            rest_tokens = model_context_tokens - max_tokens - curr_message_tokens
+            rest_tokens = max(rest_tokens, 0)
+
+        return rest_tokens
+
+    def _get_prompt_template(self, node_data: KnowledgeRetrievalNodeData, metadata_fields: list, query: str):
+        model_mode = ModelMode.value_of(node_data.metadata_model_config.mode)
+        input_text = query
+        memory_str = ""
+
+        prompt_messages: list[LLMNodeChatModelMessage] = []
+        if model_mode == ModelMode.CHAT:
+            system_prompt_messages = LLMNodeChatModelMessage(
+                role=PromptMessageRole.SYSTEM, text=METADATA_FILTER_SYSTEM_PROMPT
+            )
+            prompt_messages.append(system_prompt_messages)
+            user_prompt_message_1 = LLMNodeChatModelMessage(
+                role=PromptMessageRole.USER, text=METADATA_FILTER_USER_PROMPT_1
+            )
+            prompt_messages.append(user_prompt_message_1)
+            assistant_prompt_message_1 = LLMNodeChatModelMessage(
+                role=PromptMessageRole.ASSISTANT, text=METADATA_FILTER_ASSISTANT_PROMPT_1
+            )
+            prompt_messages.append(assistant_prompt_message_1)
+            user_prompt_message_2 = LLMNodeChatModelMessage(
+                role=PromptMessageRole.USER, text=QUESTION_CLASSIFIER_USER_PROMPT_2
+            )
+            prompt_messages.append(user_prompt_message_2)
+            assistant_prompt_message_2 = LLMNodeChatModelMessage(
+                role=PromptMessageRole.ASSISTANT, text=METADATA_FILTER_ASSISTANT_PROMPT_2
+            )
+            prompt_messages.append(assistant_prompt_message_2)
+            user_prompt_message_3 = LLMNodeChatModelMessage(
+                role=PromptMessageRole.USER,
+                text=METADATA_FILTER_USER_PROMPT_3.format(
+                    input_text=input_text,
+                    metadata_fields=json.dumps(metadata_fields, ensure_ascii=False),
+                ),
+            )
+            prompt_messages.append(user_prompt_message_3)
+            return prompt_messages
+        elif model_mode == ModelMode.COMPLETION:
+            return LLMNodeCompletionModelPromptTemplate(
+                text=METADATA_FILTER_COMPLETION_PROMPT.format(
+                    input_text=input_text,
+                    metadata_fields=json.dumps(metadata_fields, ensure_ascii=False),
+                )
+            )
+
+        else:
+            raise InvalidModelTypeError(f"Model mode {model_mode} not support.")