feat(knowledge): Support ark rerank (#852)

backend/application/base/appinfra/app_infra.go

@@ -59,6 +59,7 @@ import (
 	"github.com/coze-dev/coze-studio/backend/infra/impl/document/parser/builtin"
 	"github.com/coze-dev/coze-studio/backend/infra/impl/document/parser/ppstructure"
 	"github.com/coze-dev/coze-studio/backend/infra/impl/document/rerank/rrf"
+	vikingReranker "github.com/coze-dev/coze-studio/backend/infra/impl/document/rerank/vikingdb"
 	"github.com/coze-dev/coze-studio/backend/infra/impl/document/searchstore/elasticsearch"
 	"github.com/coze-dev/coze-studio/backend/infra/impl/document/searchstore/milvus"
 	"github.com/coze-dev/coze-studio/backend/infra/impl/document/searchstore/vikingdb"
@@ -145,7 +146,7 @@ func Init(ctx context.Context) (*AppDependencies, error) {
 		return nil, fmt.Errorf("init knowledge event bus producer failed, err=%w", err)
 	}
 
-	deps.Reranker = rrf.NewRRFReranker(0)
+	deps.Reranker = initReranker()
 
 	deps.Rewriter, err = initRewriter(ctx)
 	if err != nil {
@@ -207,6 +208,26 @@ func initSearchStoreManagers(ctx context.Context, es es.Client) ([]searchstore.M
 	return []searchstore.Manager{esSearchstoreManager, mgr}, nil
 }
 
+func initReranker() rerank.Reranker {
+	rerankerType := os.Getenv("RERANK_TYPE")
+	switch rerankerType {
+	case "vikingdb":
+		return vikingReranker.NewReranker(getVikingRerankerConfig())
+	case "rrf":
+		return rrf.NewRRFReranker(0)
+	default:
+		return rrf.NewRRFReranker(0)
+	}
+}
+func getVikingRerankerConfig() *vikingReranker.Config {
+	return &vikingReranker.Config{
+		AK:     os.Getenv("VIKINGDB_RERANK_AK"),
+		SK:     os.Getenv("VIKINGDB_RERANK_SK"),
+		Domain: os.Getenv("VIKINGDB_RERANK_HOST"),
+		Region: os.Getenv("VIKINGDB_RERANK_REGION"),
+		Model:  os.Getenv("VIKINGDB_RERANK_MODEL"),
+	}
+}
 func initRewriter(ctx context.Context) (messages2query.MessagesToQuery, error) {
 	rewriterChatModel, _, err := getBuiltinChatModel(ctx, "M2Q_")
 	if err != nil {

backend/domain/knowledge/service/retrieve.go

@@ -390,6 +390,10 @@ func (k *knowledgeSVC) nl2SqlExec(ctx context.Context, doc *model.KnowledgeDocum
 		}
 		replaceMap[doc.Name].ColumnMap[doc.TableInfo.Columns[i].Name] = convert.ColumnIDToRDBField(doc.TableInfo.Columns[i].ID)
 	}
+	virtualColumnMap := map[string]*entity.TableColumn{}
+	for i := range doc.TableInfo.Columns {
+		virtualColumnMap[convert.ColumnIDToRDBField(doc.TableInfo.Columns[i].ID)] = doc.TableInfo.Columns[i]
+	}
 	parsedSQL, err := sqlparser.NewSQLParser().ParseAndModifySQL(sql, replaceMap)
 	if err != nil {
 		logs.CtxErrorf(ctx, "parse sql failed: %v", err)
@@ -423,6 +427,32 @@ func (k *knowledgeSVC) nl2SqlExec(ctx context.Context, doc *model.KnowledgeDocum
 			prefix := "sql:" + sql + ";result:"
 			d.Content = prefix + string(byteData)
 		} else {
+			transferMap := map[string]string{}
+			for cName, val := range resp.ResultSet.Rows[i] {
+				column, found := virtualColumnMap[cName]
+				if !found {
+					logs.CtxInfof(ctx, "column not found, name: %s", cName)
+					continue
+				}
+				columnData, err := convert.ParseAnyData(column, val)
+				if err != nil {
+					logs.CtxErrorf(ctx, "parse any data failed: %v", err)
+					return nil, errorx.New(errno.ErrKnowledgeColumnParseFailCode, errorx.KV("msg", err.Error()))
+				}
+				if columnData.Type == document.TableColumnTypeString {
+					columnData.ValString = ptr.Of(k.formatSliceContent(ctx, columnData.GetStringValue()))
+				}
+				if columnData.Type == document.TableColumnTypeImage {
+					columnData.ValImage = ptr.Of(k.formatSliceContent(ctx, columnData.GetStringValue()))
+				}
+				transferMap[column.Name] = columnData.GetNullableStringValue()
+			}
+			byteData, err := sonic.Marshal(transferMap)
+			if err != nil {
+				logs.CtxErrorf(ctx, "marshal sql resp failed: %v", err)
+				return nil, err
+			}
+			d.Content = string(byteData)
 			d.ID = strconv.FormatInt(id, 10)
 		}
 		d.WithScore(1)

backend/infra/impl/document/rerank/vikingdb/vikingdb.go

@@ -35,7 +35,8 @@ import (
 type Config struct {
 	AK     string
 	SK     string
-
+	Domain string
+	Model  string
 	Region string // default cn-north-1
 }
 
@@ -43,6 +44,12 @@ func NewReranker(config *Config) rerank.Reranker {
 	if config.Region == "" {
 		config.Region = "cn-north-1"
 	}
+	if config.Domain == "" {
+		config.Domain = domain
+	}
+	if config.Model == "" {
+		config.Model = defaultModel
+	}
 	return &reranker{config: config}
 }
 
@@ -78,12 +85,32 @@ type rerankResp struct {
 func (r *reranker) Rerank(ctx context.Context, req *rerank.Request) (*rerank.Response, error) {
 	rReq := &rerankReq{
 		Datas:       make([]rerankData, 0, len(req.Data)),
-		RerankModel: defaultModel,
+		RerankModel: r.config.Model,
 	}
-
+	sorted := make([]*rerank.Data, 0)
 	var flat []*rerank.Data
+	visited := map[string]bool{}
 	for _, channel := range req.Data {
-		flat = append(flat, channel...)
+		if len(channel) == 0 {
+			continue
+		}
+		for _, item := range channel {
+			if item == nil || item.Document == nil {
+				continue
+			}
+			if item.Document.ID == "" {
+				sorted = append(sorted, &rerank.Data{
+					Document: item.Document,
+					Score:    1,
+				})
+				continue
+			}
+			if visited[item.Document.ID] {
+				continue
+			}
+			visited[item.Document.ID] = true
+			flat = append(flat, item)
+		}
 	}
 
 	for _, item := range flat {
@@ -117,7 +144,6 @@ func (r *reranker) Rerank(ctx context.Context, req *rerank.Request) (*rerank.Res
 		return nil, fmt.Errorf("[Rerank] failed, code=%d, msg=%v", rResp.Code, rResp.Message)
 	}
 
-	sorted := make([]*rerank.Data, 0, len(rResp.Data.Scores))
 	for i, score := range rResp.Data.Scores {
 		sorted = append(sorted, &rerank.Data{
 			Document: flat[i].Document,
@@ -143,7 +169,7 @@ func (r *reranker) Rerank(ctx context.Context, req *rerank.Request) (*rerank.Res
 func (r *reranker) prepareRequest(body []byte) *http.Request {
 	u := url.URL{
 		Scheme: "https",
-		Host:   domain,
+		Host:   r.config.Domain,
 		Path:   "/api/knowledge/service/rerank",
 	}
 	req, _ := http.NewRequest(http.MethodPost, u.String(), bytes.NewReader(body))

backend/infra/impl/document/rerank/vikingdb/vikingdb_test.go

@@ -1,48 +0,0 @@
-/*
- * Copyright 2025 coze-dev Authors
- *
- * Licensed under the Apache License, Version 2.0 (the "License");
- * you may not use this file except in compliance with the License.
- * You may obtain a copy of the License at
- *
- *     http://www.apache.org/licenses/LICENSE-2.0
- *
- * Unless required by applicable law or agreed to in writing, software
- * distributed under the License is distributed on an "AS IS" BASIS,
- * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- * See the License for the specific language governing permissions and
- * limitations under the License.
- */
-
-package vikingdb
-
-//func TestRun(t *testing.T) {
-//	AK := os.Getenv("test_ak")
-//	SK := os.Getenv("test_sk")
-//
-//	r := NewReranker(&Config{
-//		AK: AK,
-//		SK: SK,
-//	})
-//	resp, err := r.Rerank(context.Background(), &rerank.Request{
-//		Data: [][]*knowledge.RetrieveSlice{
-//			{
-//				{Slice: & entity. Slice {PlainText: "According to the Guinness World Records website, the blue whale is currently the largest animal known in the world, with a body length of up to 30 meters, which is equivalent to the length of a Boeing 737 aircraft"}},
-//				{Slice: & entity. Slice {PlainText: "An adult female bowhead whale can grow to 22 meters long, while a male whale can grow to 18 meters long"}},
-//			},
-//		},
-//		Query: "What is the largest whale in the world?"
-//		TopN:  nil,
-//	})
-//	assert.NoError(t, err)
-//
-//	for _, item := range resp.Sorted {
-//		fmt.Println(item.Slice.PlainText, item.Score)
-//	}
-//	According to the Guinness World Records website, the blue whale is the largest known animal in the world, with a body length of up to 30 meters, which is equivalent to the length of a Boeing 737 aircraft 6209664529733573
-//	//An adult female bowhead whale can grow up to 22 meters long, while a male whale can grow up to 18 meters 4269785303456468
-//
-//	fmt.Println(resp.TokenUsage)
-//	// 95
-//
-//}

docker/.env.debug.example

@@ -139,6 +139,16 @@ export GEMINI_EMBEDDING_LOCATION=""                  # (string, optional) Gemini
 export HTTP_EMBEDDING_ADDR=""   # (string, required) http embedding address
 export HTTP_EMBEDDING_DIMS=1024 # (string, required) http embedding dimensions
 
+# Settings for Rerank
+# If you want to use the rerank-related functions in the knowledge base feature，You need to set up the rerank configuration.
+export RERANK_TYPE="" # current support `vikingdb`,`rrf`,default:rrf
+# vikingdb rerank
+export VIKINGDB_RERANK_HOST="" # optional,default:api-knowledgebase.mlp.cn-beijing.volces.com
+export VIKINGDB_RERANK_REGION="" # optional,default:cn-north-1
+export VIKINGDB_RERANK_AK="" # required
+export VIKINGDB_RERANK_SK="" # required
+export VIKINGDB_RERANK_MODEL="" # optional,default:base-multilingual-rerank,also support m3-v2-rerank
+
 # Settings for OCR
 # If you want to use the OCR-related functions in the knowledge base feature，You need to set up the OCR configuration.
 # Currently, Coze Studio has built-in Volcano OCR.

docker/.env.example

@@ -137,6 +137,16 @@ export GEMINI_EMBEDDING_LOCATION=""                  # (string, optional) Gemini
 export HTTP_EMBEDDING_ADDR=""   # (string, required) http embedding address
 export HTTP_EMBEDDING_DIMS=1024 # (string, required) http embedding dimensions
 
+# Settings for Rerank
+# If you want to use the rerank-related functions in the knowledge base feature，You need to set up the rerank configuration.
+export RERANK_TYPE="" # current support `vikingdb`,`rrf`,default:rrf
+# vikingdb rerank
+export VIKINGDB_RERANK_HOST="" # optional,default:api-knowledgebase.mlp.cn-beijing.volces.com
+export VIKINGDB_RERANK_REGION="" # optional,default:cn-north-1
+export VIKINGDB_RERANK_AK="" # required
+export VIKINGDB_RERANK_SK="" # required
+export VIKINGDB_RERANK_MODEL="" # optional,default:base-multilingual-rerank,also support m3-v2-rerank
+
 # Settings for OCR
 # If you want to use the OCR-related functions in the knowledge base feature，You need to set up the OCR configuration.
 # Currently, Coze Studio has built-in Volcano OCR.