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

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

@@ -33,9 +33,10 @@ import (
 )
 
 type Config struct {
-	AK string
-	SK string
-
+	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
-//
-//}