multi-retrival use dataset's top-k (#8416)

2024-09-14 21:48:44 +08:00 · 2024-09-14 21:48:44 +08:00 · 7e611ffbf3
commit 7e611ffbf3
parent 65162a87b6
2 changed files with 6 additions and 3 deletions
--- a/api/core/rag/retrieval/dataset_retrieval.py
+++ b/api/core/rag/retrieval/dataset_retrieval.py
@ -426,7 +426,7 @@ class DatasetRetrieval:
                        retrieval_method=retrieval_model["search_method"],
                        dataset_id=dataset.id,
                        query=query,
-                        top_k=top_k,
+                        top_k=retrieval_model.get("top_k") or 2,
                        score_threshold=retrieval_model.get("score_threshold", 0.0)
                        if retrieval_model["score_threshold_enabled"]
                        else 0.0,