Skip to content

Commit f8862cb

Browse files
ByteDanceLiuYangByteDanceLiuYang
committed
Merge 'rerank-api' into 'main'
rerank api See merge request: !16
2 parents ace0981 + 50a7c3d commit f8862cb

9 files changed

Lines changed: 227 additions & 2 deletions

File tree

‎README.md‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ To use the SDK in your Maven project, add the following dependency to your `pom.
2525
<dependency>
2626
<groupId>com.volcengine</groupId>
2727
<artifactId>vikingdb-java-sdk</artifactId>
28-
<version>0.0.4</version> <!-- Please check for the latest version -->
28+
<version>0.1.0</version> <!-- Please check for the latest version -->
2929
</dependency>
3030
```
3131

@@ -114,6 +114,10 @@ VectorService service = new VectorService("your-vikingdb-host", auth);
114114

115115
- `embedding()`
116116

117+
### Rerank Operations
118+
119+
- `rerank()`
120+
117121
### Error Handling
118122

119123
The SDK throws exceptions for network issues or API errors. It's recommended to wrap SDK calls in a try-catch block.

‎pom.xml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@
66

77
<groupId>com.volcengine</groupId>
88
<artifactId>vikingdb-java-sdk</artifactId>
9-
<version>0.0.4</version>
9+
<version>0.1.0</version>
1010
<name>vikingdb-java-sdk</name>
1111
<url>https://open.volcengineapi.com</url>
1212
<description>VikingDB Java SDK</description>

‎src/main/java/com/volcengine/vikingdb/runtime/vector/api/ApiInfo.java‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,4 +33,5 @@ public class ApiInfo {
3333
public static final ApiInfo SORT = new ApiInfo("Sort", Method.POST, "/api/vikingdb/data/sort");
3434

3535
public static final ApiInfo Embedding = new ApiInfo("Embedding", Method.POST, "/api/vikingdb/embedding");
36+
public static final ApiInfo Rerank = new ApiInfo("Rerank", Method.POST, "/api/vikingdb/rerank");
3637
}

‎src/main/java/com/volcengine/vikingdb/runtime/vector/model/request/RerankRequest.java‎

Lines changed: 39 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,39 @@
1+
// Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package com.volcengine.vikingdb.runtime.vector.model.request;
5+
6+
import com.fasterxml.jackson.annotation.JsonProperty;
7+
import lombok.AllArgsConstructor;
8+
import lombok.Data;
9+
import lombok.EqualsAndHashCode;
10+
import lombok.NoArgsConstructor;
11+
import lombok.experimental.SuperBuilder;
12+
13+
import java.util.List;
14+
15+
@NoArgsConstructor
16+
@AllArgsConstructor
17+
@SuperBuilder(toBuilder = true)
18+
@Data
19+
@EqualsAndHashCode()
20+
public class RerankRequest {
21+
@JsonProperty("model_name")
22+
private String modelName;
23+
24+
@JsonProperty("model_version")
25+
private String modelVersion;
26+
27+
@JsonProperty("data")
28+
private List<List<FullModalData>> data;
29+
30+
@JsonProperty("query")
31+
private List<FullModalData> query;
32+
33+
@JsonProperty("instruction")
34+
private String instruction;
35+
36+
@JsonProperty("return_origin_data")
37+
private Boolean returnOriginData;
38+
39+
}

‎src/main/java/com/volcengine/vikingdb/runtime/vector/model/response/FullModalData.java‎

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
// Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package com.volcengine.vikingdb.runtime.vector.model.response;
5+
6+
7+
import com.fasterxml.jackson.annotation.JsonProperty;
8+
import lombok.AllArgsConstructor;
9+
import lombok.Data;
10+
import lombok.EqualsAndHashCode;
11+
import lombok.NoArgsConstructor;
12+
import lombok.experimental.SuperBuilder;
13+
14+
@NoArgsConstructor
15+
@AllArgsConstructor
16+
@SuperBuilder(toBuilder = true)
17+
@Data
18+
@EqualsAndHashCode()
19+
public class FullModalData {
20+
@JsonProperty("text")
21+
private String text;
22+
23+
@JsonProperty("image")
24+
private Object image;
25+
26+
@JsonProperty("video")
27+
private Object video;
28+
29+
}

‎src/main/java/com/volcengine/vikingdb/runtime/vector/model/response/RerankItem.java‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
// Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package com.volcengine.vikingdb.runtime.vector.model.response;
5+
6+
import com.fasterxml.jackson.annotation.JsonProperty;
7+
import lombok.AllArgsConstructor;
8+
import lombok.Data;
9+
import lombok.NoArgsConstructor;
10+
11+
import java.util.List;
12+
13+
@NoArgsConstructor
14+
@AllArgsConstructor
15+
@Data
16+
public class RerankItem {
17+
@JsonProperty("id")
18+
private Integer id;
19+
20+
@JsonProperty("score")
21+
private Float score;
22+
23+
@JsonProperty("origin_data")
24+
private List<FullModalData> originData;
25+
}

‎src/main/java/com/volcengine/vikingdb/runtime/vector/model/response/RerankResult.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,23 @@
1+
// Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package com.volcengine.vikingdb.runtime.vector.model.response;
5+
6+
import com.fasterxml.jackson.annotation.JsonProperty;
7+
import lombok.AllArgsConstructor;
8+
import lombok.Data;
9+
import lombok.NoArgsConstructor;
10+
11+
import java.util.List;
12+
13+
@NoArgsConstructor
14+
@AllArgsConstructor
15+
@Data
16+
public class RerankResult {
17+
@JsonProperty("data")
18+
private List<RerankItem> data;
19+
20+
@JsonProperty("token_usage")
21+
private Object tokenUsage;
22+
23+
}

‎src/main/java/com/volcengine/vikingdb/runtime/vector/service/VectorService.java‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -379,4 +379,24 @@ public DataApiResponse<EmbeddingResult> embedding(EmbeddingRequest request, Requ
379379
).callApi(apiClient, addition);
380380
}
381381

382+
/**
383+
* Rerank documents for the given data inputs.
384+
*
385+
* @param request The request object containing the data to be reranked.
386+
* @return An API response containing the reranked results.
387+
*/
388+
public DataApiResponse<RerankResult> rerank(RerankRequest request)
389+
throws ApiClientException, VectorApiException {
390+
return rerank(request, null);
391+
}
392+
393+
public DataApiResponse<RerankResult> rerank(RerankRequest request, RequestAddition addition)
394+
throws ApiClientException, VectorApiException {
395+
return new ApiRequestBinding<>(
396+
ApiInfo.Rerank,
397+
request,
398+
new TypeReference<DataApiResponse<RerankResult>>() {}
399+
).callApi(apiClient, addition);
400+
}
401+
382402
}

‎src/test/java/com/volcengine/vikingdb/runtime/vector/rerank/RerankTest.java‎

Lines changed: 84 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,84 @@
1+
// Copyright (c) 2025 Beijing Volcano Engine Technology Co., Ltd.
2+
// SPDX-License-Identifier: Apache-2.0
3+
4+
package com.volcengine.vikingdb.runtime.vector.rerank;
5+
6+
import com.volcengine.vikingdb.runtime.core.ApiClientTest;
7+
import com.volcengine.vikingdb.runtime.core.auth.AuthWithAkSk;
8+
import com.volcengine.vikingdb.runtime.enums.Scheme;
9+
import com.volcengine.vikingdb.runtime.exception.ApiClientException;
10+
import com.volcengine.vikingdb.runtime.exception.VectorApiException;
11+
import com.volcengine.vikingdb.runtime.vector.TestVar;
12+
import com.volcengine.vikingdb.runtime.vector.model.request.EmbeddingDataItem;
13+
import com.volcengine.vikingdb.runtime.vector.model.request.EmbeddingRequest;
14+
import com.volcengine.vikingdb.runtime.vector.model.request.FullModalData;
15+
import com.volcengine.vikingdb.runtime.vector.model.request.RerankRequest;
16+
import com.volcengine.vikingdb.runtime.vector.model.response.DataApiResponse;
17+
import com.volcengine.vikingdb.runtime.vector.model.response.EmbeddingResult;
18+
import com.volcengine.vikingdb.runtime.vector.model.response.RerankResult;
19+
import com.volcengine.vikingdb.runtime.vector.service.VectorService;
20+
import org.junit.jupiter.api.BeforeAll;
21+
import org.junit.jupiter.api.Test;
22+
23+
import java.io.IOException;
24+
import java.io.InputStream;
25+
import java.util.Arrays;
26+
import java.util.Collections;
27+
import java.util.Properties;
28+
29+
public class RerankTest {
30+
private static VectorService service;
31+
32+
@BeforeAll
33+
public static void setUp() throws Exception {
34+
Properties properties = new Properties();
35+
try (InputStream input = ApiClientTest.class.getClassLoader().getResourceAsStream("test-config.properties")) {
36+
if (input == null) {
37+
throw new IOException("Unable to find test-config.properties");
38+
}
39+
properties.load(input);
40+
}
41+
42+
service = new VectorService(
43+
Scheme.fromString(properties.getProperty("scheme")),
44+
properties.getProperty("host"),
45+
new AuthWithAkSk(properties.getProperty("ak"), properties.getProperty("sk"))
46+
);
47+
}
48+
49+
@Test
50+
public void testRerankForText() throws Exception {
51+
RerankRequest request = RerankRequest.builder()
52+
.modelName("doubao-seed-rerank")
53+
.modelVersion("251028")
54+
.query(Collections.singletonList(FullModalData.builder().text("apple").build()))
55+
.instruction("Whether the Document answers the Query or matches the content retrieval intent")
56+
.returnOriginData(true)
57+
.data(Arrays.asList(
58+
Collections.singletonList(FullModalData.builder().text("This is an apple.").build()),
59+
Collections.singletonList(FullModalData.builder().text("My name is John.").build()),
60+
Collections.singletonList(FullModalData.builder().text("How many fruits are there?").build())
61+
))
62+
.build();
63+
DataApiResponse<RerankResult> response = service.rerank(request);
64+
System.out.println(response);
65+
}
66+
67+
@Test
68+
public void testEmbeddingForMultiModal() throws Exception {
69+
RerankRequest request = RerankRequest.builder()
70+
.modelName("doubao-seed-rerank")
71+
.modelVersion("251028")
72+
.query(Collections.singletonList(FullModalData.builder().text("iceberg").build()))
73+
.instruction("Whether the Document answers the Query or matches the content retrieval intent")
74+
.returnOriginData(true)
75+
.data(Arrays.asList(
76+
Collections.singletonList(FullModalData.builder().text("Here is a mountain and it is very beautiful.").build()),
77+
Collections.singletonList(FullModalData.builder().image("https://ark-project.tos-cn-beijing.volces.com/images/view.jpeg").build()),
78+
Collections.singletonList(FullModalData.builder().video("https://ark-project.tos-cn-beijing.volces.com/doc_video/ark_vlm_video_input.mp4").build())
79+
)).build();
80+
DataApiResponse<RerankResult> response = service.rerank(request);
81+
System.out.println(response);
82+
}
83+
84+
}

0 commit comments

Comments
 (0)