fix RetrievalMethod StrEnum (#26768)

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Co-authored-by: crazywoola <100913391+crazywoola@users.noreply.github.com>

api/core/datasource/entities/datasource_entities.py

@@ -1,5 +1,5 @@
 import enum
-from enum import Enum
+from enum import StrEnum
 from typing import Any
 
 from pydantic import BaseModel, Field, ValidationInfo, field_validator
@@ -218,7 +218,7 @@ class DatasourceLabel(BaseModel):
     icon: str = Field(..., description="The icon of the tool")
 
 
-class DatasourceInvokeFrom(Enum):
+class DatasourceInvokeFrom(StrEnum):
     """
     Enum class for datasource invoke
     """

api/core/model_runtime/entities/provider_entities.py

@@ -1,5 +1,5 @@
 from collections.abc import Sequence
-from enum import Enum, StrEnum, auto
+from enum import StrEnum, auto
 
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
@@ -7,7 +7,7 @@ from core.model_runtime.entities.common_entities import I18nObject
 from core.model_runtime.entities.model_entities import AIModelEntity, ModelType
 
 
-class ConfigurateMethod(Enum):
+class ConfigurateMethod(StrEnum):
     """
     Enum class for configurate method of provider model.
     """

api/core/rag/datasource/retrieval_service.py

@@ -34,7 +34,7 @@ class RetrievalService:
     @classmethod
     def retrieve(
         cls,
-        retrieval_method: str,
+        retrieval_method: RetrievalMethod,
         dataset_id: str,
         query: str,
         top_k: int,
@@ -56,7 +56,7 @@ class RetrievalService:
         # Optimize multithreading with thread pools
         with ThreadPoolExecutor(max_workers=dify_config.RETRIEVAL_SERVICE_EXECUTORS) as executor:  # type: ignore
             futures = []
-            if retrieval_method == "keyword_search":
+            if retrieval_method == RetrievalMethod.KEYWORD_SEARCH:
                 futures.append(
                     executor.submit(
                         cls.keyword_search,
@@ -220,7 +220,7 @@ class RetrievalService:
         score_threshold: float | None,
         reranking_model: dict | None,
         all_documents: list,
-        retrieval_method: str,
+        retrieval_method: RetrievalMethod,
         exceptions: list,
         document_ids_filter: list[str] | None = None,
     ):

api/core/rag/entities/event.py

@@ -1,11 +1,11 @@
 from collections.abc import Mapping
-from enum import Enum
+from enum import StrEnum
 from typing import Any
 
 from pydantic import BaseModel, Field
 
 
-class DatasourceStreamEvent(Enum):
+class DatasourceStreamEvent(StrEnum):
     """
     Datasource Stream event
     """

api/core/rag/index_processor/index_processor_base.py

@@ -7,6 +7,7 @@ from typing import TYPE_CHECKING, Any, Optional
 from configs import dify_config
 from core.rag.extractor.entity.extract_setting import ExtractSetting
 from core.rag.models.document import Document
+from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from core.rag.splitter.fixed_text_splitter import (
     EnhanceRecursiveCharacterTextSplitter,
     FixedRecursiveCharacterTextSplitter,
@@ -49,7 +50,7 @@ class BaseIndexProcessor(ABC):
     @abstractmethod
     def retrieve(
         self,
-        retrieval_method: str,
+        retrieval_method: RetrievalMethod,
         query: str,
         dataset: Dataset,
         top_k: int,

api/core/rag/index_processor/processor/paragraph_index_processor.py

@@ -14,6 +14,7 @@ from core.rag.extractor.extract_processor import ExtractProcessor
 from core.rag.index_processor.constant.index_type import IndexType
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor
 from core.rag.models.document import Document
+from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from core.tools.utils.text_processing_utils import remove_leading_symbols
 from libs import helper
 from models.dataset import Dataset, DatasetProcessRule
@@ -106,7 +107,7 @@ class ParagraphIndexProcessor(BaseIndexProcessor):
 
     def retrieve(
         self,
-        retrieval_method: str,
+        retrieval_method: RetrievalMethod,
         query: str,
         dataset: Dataset,
         top_k: int,

api/core/rag/index_processor/processor/parent_child_index_processor.py

@@ -16,6 +16,7 @@ from core.rag.extractor.extract_processor import ExtractProcessor
 from core.rag.index_processor.constant.index_type import IndexType
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor
 from core.rag.models.document import ChildDocument, Document, ParentChildStructureChunk
+from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from extensions.ext_database import db
 from libs import helper
 from models.dataset import ChildChunk, Dataset, DatasetProcessRule, DocumentSegment
@@ -161,7 +162,7 @@ class ParentChildIndexProcessor(BaseIndexProcessor):
 
     def retrieve(
         self,
-        retrieval_method: str,
+        retrieval_method: RetrievalMethod,
         query: str,
         dataset: Dataset,
         top_k: int,

api/core/rag/index_processor/processor/qa_index_processor.py

@@ -21,6 +21,7 @@ from core.rag.extractor.extract_processor import ExtractProcessor
 from core.rag.index_processor.constant.index_type import IndexType
 from core.rag.index_processor.index_processor_base import BaseIndexProcessor
 from core.rag.models.document import Document, QAStructureChunk
+from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from core.tools.utils.text_processing_utils import remove_leading_symbols
 from libs import helper
 from models.dataset import Dataset
@@ -141,7 +142,7 @@ class QAIndexProcessor(BaseIndexProcessor):
 
     def retrieve(
         self,
-        retrieval_method: str,
+        retrieval_method: RetrievalMethod,
         query: str,
         dataset: Dataset,
         top_k: int,

api/core/rag/retrieval/dataset_retrieval.py

@@ -364,7 +364,7 @@ class DatasetRetrieval:
                     top_k = retrieval_model_config["top_k"]
                     # get retrieval method
                     if dataset.indexing_technique == "economy":
-                        retrieval_method = "keyword_search"
+                        retrieval_method = RetrievalMethod.KEYWORD_SEARCH
                     else:
                         retrieval_method = retrieval_model_config["search_method"]
                     # get reranking model
@@ -623,7 +623,7 @@ class DatasetRetrieval:
                 if dataset.indexing_technique == "economy":
                     # use keyword table query
                     documents = RetrievalService.retrieve(
-                        retrieval_method="keyword_search",
+                        retrieval_method=RetrievalMethod.KEYWORD_SEARCH,
                         dataset_id=dataset.id,
                         query=query,
                         top_k=top_k,

api/core/rag/retrieval/retrieval_methods.py

@@ -1,7 +1,7 @@
-from enum import Enum
+from enum import StrEnum
 
 
-class RetrievalMethod(Enum):
+class RetrievalMethod(StrEnum):
     SEMANTIC_SEARCH = "semantic_search"
     FULL_TEXT_SEARCH = "full_text_search"
     HYBRID_SEARCH = "hybrid_search"

api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py

@@ -172,7 +172,7 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
             if dataset.indexing_technique == "economy":
                 # use keyword table query
                 documents = RetrievalService.retrieve(
-                    retrieval_method="keyword_search",
+                    retrieval_method=RetrievalMethod.KEYWORD_SEARCH,
                     dataset_id=dataset.id,
                     query=query,
                     top_k=retrieval_model.get("top_k") or 4,

api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py

@@ -130,7 +130,7 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
             if dataset.indexing_technique == "economy":
                 # use keyword table query
                 documents = RetrievalService.retrieve(
-                    retrieval_method="keyword_search",
+                    retrieval_method=RetrievalMethod.KEYWORD_SEARCH,
                     dataset_id=dataset.id,
                     query=query,
                     top_k=self.top_k,

api/core/workflow/enums.py

@@ -1,7 +1,7 @@
-from enum import Enum, StrEnum
+from enum import StrEnum
 
 
-class NodeState(Enum):
+class NodeState(StrEnum):
     """State of a node or edge during workflow execution."""
 
     UNKNOWN = "unknown"

api/core/workflow/graph_engine/layers/execution_limits.py

@@ -10,7 +10,7 @@ When limits are exceeded, the layer automatically aborts execution.
 
 import logging
 import time
-from enum import Enum
+from enum import StrEnum
 from typing import final
 
 from typing_extensions import override
@@ -24,7 +24,7 @@ from core.workflow.graph_events import (
 from core.workflow.graph_events.node import NodeRunFailedEvent, NodeRunSucceededEvent
 
 
-class LimitType(Enum):
+class LimitType(StrEnum):
     """Types of execution limits that can be exceeded."""
 
     STEP_LIMIT = "step_limit"

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

@@ -2,6 +2,7 @@ from typing import Literal, Union
 
 from pydantic import BaseModel
 
+from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from core.workflow.nodes.base import BaseNodeData
 
 
@@ -63,7 +64,7 @@ class RetrievalSetting(BaseModel):
     Retrieval Setting.
     """
 
-    search_method: Literal["semantic_search", "keyword_search", "full_text_search", "hybrid_search"]
+    search_method: RetrievalMethod
     top_k: int
     score_threshold: float | None = 0.5
     score_threshold_enabled: bool = False