From 02b06c420e244a429acb1e3183b6fc46fb482d07 Mon Sep 17 00:00:00 2001
From: jyong <718720800@qq.com>
Date: Tue, 24 Sep 2024 23:52:01 +0800
Subject: [PATCH] add external_retrieval_model

---
 api/fields/dataset_fields.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/api/fields/dataset_fields.py b/api/fields/dataset_fields.py
index 7aee087d78..b32423f10c 100644
--- a/api/fields/dataset_fields.py
+++ b/api/fields/dataset_fields.py
@@ -38,6 +38,10 @@ dataset_retrieval_model_fields = {
     "score_threshold_enabled": fields.Boolean,
     "score_threshold": fields.Float,
 }
+external_retrieval_model_fields = {
+    "top_k": fields.Integer,
+    "score_threshold": fields.Float,
+}
 
 tag_fields = {"id": fields.String, "name": fields.String, "type": fields.String}
 
@@ -69,6 +73,7 @@ dataset_detail_fields = {
     "retrieval_model_dict": fields.Nested(dataset_retrieval_model_fields),
     "tags": fields.List(fields.Nested(tag_fields)),
     "external_knowledge_info": fields.Nested(external_knowledge_info_fields),
+    "external_retrieval_model": fields.Nested(external_retrieval_model_fields, allow_null=True),
 }
 
 dataset_query_detail_fields = {