estel/dify
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

Fix/reranking mode is null (#7012)

Browse Source
This commit is contained in:
Jyong
2024-08-06 19:12:04 +08:00
committed by GitHub
parent c110888aee
commit 28d4e5b045
5 changed files with 20 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
api/core/tools/tool/dataset_retriever/dataset_multi_retriever_tool.py
View File

@@ -177,10 +177,12 @@ class DatasetMultiRetrieverTool(DatasetRetrieverBaseTool):
dataset_id=dataset.id,
query=query,
top_k=self.top_k,
score_threshold=retrieval_model['score_threshold']
score_threshold=retrieval_model.get('score_threshold', .0)
if retrieval_model['score_threshold_enabled'] else None,
reranking_model=retrieval_model['reranking_model']
reranking_model=retrieval_model.get('reranking_model', None)
if retrieval_model['reranking_enable'] else None,
reranking_mode=retrieval_model.get('reranking_mode')
if retrieval_model.get('reranking_mode') else 'reranking_model',
weights=retrieval_model.get('weights', None),
)

10
api/core/tools/tool/dataset_retriever/dataset_retriever_tool.py
View File

@@ -14,6 +14,7 @@ default_retrieval_model = {
'reranking_provider_name': '',
'reranking_model_name': ''
},
'reranking_mode': 'reranking_model',
'top_k': 2,
'score_threshold_enabled': False
}
@@ -71,14 +72,15 @@ class DatasetRetrieverTool(DatasetRetrieverBaseTool):
else:
if self.top_k > 0:
# retrieval source
documents = RetrievalService.retrieve(retrival_method=retrieval_model['search_method'],
documents = RetrievalService.retrieve(retrival_method=retrieval_model.get('search_method', 'semantic_search'),
dataset_id=dataset.id,
query=query,
top_k=self.top_k,
score_threshold=retrieval_model['score_threshold']
score_threshold=retrieval_model.get('score_threshold', .0)
if retrieval_model['score_threshold_enabled'] else None,
reranking_model=retrieval_model['reranking_model']
if retrieval_model['reranking_enable'] else None,
reranking_model=retrieval_model.get('reranking_model', None),
reranking_mode=retrieval_model.get('reranking_mode')
if retrieval_model.get('reranking_mode') else 'reranking_model',
weights=retrieval_model.get('weights', None),
)
else: