From c19a822e1b4982e611c1c9e5fcd9bf7158b2bcc4 Mon Sep 17 00:00:00 2001
From: Jake Armstrong <65635253+jakearmstrong59@users.noreply.github.com>
Date: Thu, 9 Apr 2026 06:13:04 +0200
Subject: [PATCH] refactor: deduplicate DefaultRetrievalModelDict TypedDict
 into retrieval_service.py (#34758)

---
 .../dataset_multi_retriever_tool.py             |  3 +--
 .../dataset_retriever/dataset_retriever_tool.py | 17 ++---------------
 2 files changed, 3 insertions(+), 17 deletions(-)

diff --git a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py
index c72bdf02ed..03e3c5918d 100644
--- a/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py
+++ b/api/core/tools/utils/dataset_retriever/dataset_multi_retriever_tool.py
@@ -7,14 +7,13 @@ from sqlalchemy import select
 
 from core.callback_handler.index_tool_callback_handler import DatasetIndexToolCallbackHandler
 from core.model_manager import ModelManager
-from core.rag.datasource.retrieval_service import RetrievalService
+from core.rag.datasource.retrieval_service import DefaultRetrievalModelDict, RetrievalService
 from core.rag.entities import RetrievalSourceMetadata
 from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from core.rag.models.document import Document as RagDocument
 from core.rag.rerank.rerank_model import RerankModelRunner
 from core.rag.retrieval.retrieval_methods import RetrievalMethod
 from core.tools.utils.dataset_retriever.dataset_retriever_base_tool import DatasetRetrieverBaseTool
-from core.tools.utils.dataset_retriever.dataset_retriever_tool import DefaultRetrievalModelDict
 from extensions.ext_database import db
 from models.dataset import Dataset, Document, DocumentSegment
 
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 a346eb53c4..6a189fa6aa 100644
--- a/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py
+++ b/api/core/tools/utils/dataset_retriever/dataset_retriever_tool.py
@@ -1,11 +1,10 @@
-from typing import NotRequired, TypedDict, cast
+from typing import cast
 
 from pydantic import BaseModel, Field
 from sqlalchemy import select
 
 from core.app.app_config.entities import DatasetRetrieveConfigEntity, ModelConfig
-from core.rag.data_post_processor.data_post_processor import RerankingModelDict, WeightsDict
-from core.rag.datasource.retrieval_service import RetrievalService
+from core.rag.datasource.retrieval_service import DefaultRetrievalModelDict, RetrievalService
 from core.rag.entities import DocumentContext, RetrievalSourceMetadata
 from core.rag.index_processor.constant.index_type import IndexTechniqueType
 from core.rag.models.document import Document as RetrievalDocument
@@ -17,18 +16,6 @@ from models.dataset import Dataset
 from models.dataset import Document as DatasetDocument
 from services.external_knowledge_service import ExternalDatasetService
 
-
-class DefaultRetrievalModelDict(TypedDict):
-    search_method: RetrievalMethod
-    reranking_enable: bool
-    reranking_model: RerankingModelDict
-    reranking_mode: NotRequired[str]
-    weights: NotRequired[WeightsDict | None]
-    score_threshold: NotRequired[float]
-    top_k: int
-    score_threshold_enabled: bool
-
-
 default_retrieval_model: DefaultRetrievalModelDict = {
     "search_method": RetrievalMethod.SEMANTIC_SEARCH,
     "reranking_enable": False,