From e7572066a4d6c17097010744f22c822e39dc0da4 Mon Sep 17 00:00:00 2001
From: Jyong <76649700+JohnJyong@users.noreply.github.com>
Date: Wed, 19 Mar 2025 15:10:00 +0800
Subject: [PATCH]  fix Weight rerank mode info missed when create dataset
 (#16190)

---
 .../knowledge_entities/knowledge_entities.py    | 17 +++++++++++++++++
 1 file changed, 17 insertions(+)

diff --git a/api/services/entities/knowledge_entities/knowledge_entities.py b/api/services/entities/knowledge_entities/knowledge_entities.py
index 6b0ecd7e33..37c0fb49e5 100644
--- a/api/services/entities/knowledge_entities/knowledge_entities.py
+++ b/api/services/entities/knowledge_entities/knowledge_entities.py
@@ -84,6 +84,22 @@ class RerankingModel(BaseModel):
     reranking_model_name: Optional[str] = None
 
 
+class WeightVectorSetting(BaseModel):
+    vector_weight: float
+    embedding_provider_name: str
+    embedding_model_name: str
+
+
+class WeightKeywordSetting(BaseModel):
+    keyword_weight: float
+
+
+class WeightModel(BaseModel):
+    weight_type: str
+    vector_setting: Optional[WeightVectorSetting] = None
+    keyword_setting: Optional[WeightKeywordSetting] = None
+
+
 class RetrievalModel(BaseModel):
     search_method: Literal["hybrid_search", "semantic_search", "full_text_search"]
     reranking_enable: bool
@@ -92,6 +108,7 @@ class RetrievalModel(BaseModel):
     top_k: int
     score_threshold_enabled: bool
     score_threshold: Optional[float] = None
+    weights: Optional[WeightModel] = None
 
 
 class MetaDataConfig(BaseModel):