fix reranking model field error (#6654)

langgenius · Jul 25, 2024 · 91fd852 · 91fd852
1 parent 4ec9a87
commit 91fd852
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/api/core/workflow/nodes/knowledge_retrieval/entities.py b/api/core/workflow/nodes/knowledge_retrieval/entities.py
@@ -46,8 +46,8 @@ class MultipleRetrievalConfig(BaseModel):
     score_threshold: Optional[float] = None
     reranking_mode: str = 'reranking_model'
     reranking_enable: bool = True
-    reranking_model: RerankingModelConfig
-    weights: WeightedScoreConfig
+    reranking_model: Optional[RerankingModelConfig] = None
+    weights: Optional[WeightedScoreConfig] = None
 
 
 class ModelConfig(BaseModel):

diff --git a/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py b/api/core/workflow/nodes/knowledge_retrieval/knowledge_retrieval_node.py
@@ -140,8 +140,8 @@ def _fetch_dataset_retriever(self, node_data: KnowledgeRetrievalNodeData, query:
         elif node_data.retrieval_mode == DatasetRetrieveConfigEntity.RetrieveStrategy.MULTIPLE.value:
             if node_data.multiple_retrieval_config.reranking_mode == 'reranking_model':
                 reranking_model = {
-                    'reranking_provider_name': node_data.multiple_retrieval_config.reranking_model['provider'],
-                    'reranking_model_name': node_data.multiple_retrieval_config.reranking_model['name']
+                    'reranking_provider_name': node_data.multiple_retrieval_config.reranking_model.provider,
+                    'reranking_model_name': node_data.multiple_retrieval_config.reranking_model.model
                 }
                 weights = None
             elif node_data.multiple_retrieval_config.reranking_mode == 'weighted_score':