Skip to content

Commit 71cd444

Browse files
pgentilepgentile
committed
feat: configure similarity top-K in the VertexAiRagRetrieval tool
1 parent 527ff4e commit 71cd444

2 files changed

Lines changed: 124 additions & 9 deletions

File tree

‎core/src/main/java/com/google/adk/tools/retrieval/VertexAiRagRetrieval.java‎

Lines changed: 35 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,6 +23,7 @@
2323
import com.google.adk.utils.ModelNameUtils;
2424
import com.google.cloud.aiplatform.v1.RagContexts;
2525
import com.google.cloud.aiplatform.v1.RagQuery;
26+
import com.google.cloud.aiplatform.v1.RagRetrievalConfig;
2627
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest;
2728
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.RagResource;
2829
import com.google.cloud.aiplatform.v1.RetrieveContextsResponse;
@@ -47,14 +48,16 @@
4748
* A retrieval tool that fetches context from Vertex AI RAG.
4849
*
4950
* <p>This tool allows to retrieve relevant information based on a query using Vertex AI RAG
50-
* service. It supports configuration of rag resources and a vector distance threshold.
51+
* service. It supports configuration of rag resources, a vector distance threshold, and similarity
52+
* top-k.
5153
*/
5254
public class VertexAiRagRetrieval extends BaseRetrievalTool {
5355
private static final Logger logger = LoggerFactory.getLogger(VertexAiRagRetrieval.class);
5456
private final VertexRagServiceClient vertexRagServiceClient;
5557
private final String parent;
5658
private final List<RagResource> ragResources;
5759
private final Double vectorDistanceThreshold;
60+
private final Integer similarityTopK;
5861
private final VertexRagStore vertexRagStore;
5962
private final RetrieveContextsRequest.VertexRagStore apiVertexRagStore;
6063

@@ -65,11 +68,30 @@ public VertexAiRagRetrieval(
6568
String parent,
6669
@Nullable List<RagResource> ragResources,
6770
@Nullable Double vectorDistanceThreshold) {
71+
this(
72+
name,
73+
description,
74+
vertexRagServiceClient,
75+
parent,
76+
ragResources,
77+
vectorDistanceThreshold,
78+
null);
79+
}
80+
81+
public VertexAiRagRetrieval(
82+
String name,
83+
String description,
84+
VertexRagServiceClient vertexRagServiceClient,
85+
String parent,
86+
@Nullable List<RagResource> ragResources,
87+
@Nullable Double vectorDistanceThreshold,
88+
@Nullable Integer similarityTopK) {
6889
super(name, description);
6990
this.vertexRagServiceClient = vertexRagServiceClient;
7091
this.parent = parent;
7192
this.ragResources = ragResources;
7293
this.vectorDistanceThreshold = vectorDistanceThreshold;
94+
this.similarityTopK = similarityTopK;
7395

7496
// For Gemini 2
7597
VertexRagStore.Builder vertexRagStoreBuilder = VertexRagStore.builder();
@@ -86,6 +108,9 @@ public VertexAiRagRetrieval(
86108
if (this.vectorDistanceThreshold != null) {
87109
vertexRagStoreBuilder.vectorDistanceThreshold(this.vectorDistanceThreshold);
88110
}
111+
if (this.similarityTopK != null) {
112+
vertexRagStoreBuilder.similarityTopK(this.similarityTopK);
113+
}
89114
this.vertexRagStore = vertexRagStoreBuilder.build();
90115

91116
// For runAsync
@@ -135,10 +160,18 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
135160
return Single.fromCallable(
136161
() -> {
137162
logger.info("Retrieving context for query: {}", query);
163+
164+
RagQuery.Builder queryBuilder = RagQuery.newBuilder().setText(query);
165+
166+
if (similarityTopK != null) {
167+
queryBuilder.setRagRetrievalConfig(
168+
RagRetrievalConfig.newBuilder().setTopK(similarityTopK).build());
169+
}
170+
138171
RetrieveContextsRequest retrieveContextsRequest =
139172
RetrieveContextsRequest.newBuilder()
140173
.setParent(this.parent)
141-
.setQuery(RagQuery.newBuilder().setText(query))
174+
.setQuery(queryBuilder.build())
142175
.setVertexRagStore(this.apiVertexRagStore)
143176
.build();
144177
logger.info("Request to VertexRagService: {}", retrieveContextsRequest);

‎core/src/test/java/com/google/adk/tools/retrieval/VertexAiRagRetrievalTest.java‎

Lines changed: 89 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,7 @@
2929
import com.google.adk.tools.ToolContext;
3030
import com.google.cloud.aiplatform.v1.RagContexts;
3131
import com.google.cloud.aiplatform.v1.RagQuery;
32+
import com.google.cloud.aiplatform.v1.RagRetrievalConfig;
3233
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest;
3334
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.RagResource;
3435
import com.google.cloud.aiplatform.v1.RetrieveContextsResponse;
@@ -72,20 +73,26 @@ public void runAsync_withResults_returnsContexts() throws Exception {
7273
ImmutableList<RagResource> ragResources =
7374
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
7475
Double vectorDistanceThreshold = 0.5;
76+
Integer similarityTopK = 10;
7577
VertexAiRagRetrieval tool =
7678
new VertexAiRagRetrieval(
7779
"testTool",
7880
"test description",
7981
vertexRagServiceClient,
8082
"projects/test-project/locations/us-central1",
8183
ragResources,
82-
vectorDistanceThreshold);
84+
vectorDistanceThreshold,
85+
similarityTopK);
8386
String query = "test query";
8487
ToolContext toolContext = buildToolContext();
8588
RetrieveContextsRequest expectedRequest =
8689
RetrieveContextsRequest.newBuilder()
8790
.setParent("projects/test-project/locations/us-central1")
88-
.setQuery(RagQuery.newBuilder().setText(query))
91+
.setQuery(
92+
RagQuery.newBuilder()
93+
.setText(query)
94+
.setRagRetrievalConfig(
95+
RagRetrievalConfig.newBuilder().setTopK(similarityTopK).build()))
8996
.setVertexRagStore(
9097
com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.newBuilder()
9198
.addAllRagResources(ragResources)
@@ -112,20 +119,26 @@ public void runAsync_noResults_returnsNoResultFoundMessage() throws Exception {
112119
ImmutableList<RagResource> ragResources =
113120
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
114121
Double vectorDistanceThreshold = 0.5;
122+
Integer similarityTopK = 10;
115123
VertexAiRagRetrieval tool =
116124
new VertexAiRagRetrieval(
117125
"testTool",
118126
"test description",
119127
vertexRagServiceClient,
120128
"projects/test-project/locations/us-central1",
121129
ragResources,
122-
vectorDistanceThreshold);
130+
vectorDistanceThreshold,
131+
similarityTopK);
123132
String query = "test query";
124133
ToolContext toolContext = buildToolContext();
125134
RetrieveContextsRequest expectedRequest =
126135
RetrieveContextsRequest.newBuilder()
127136
.setParent("projects/test-project/locations/us-central1")
128-
.setQuery(RagQuery.newBuilder().setText(query))
137+
.setQuery(
138+
RagQuery.newBuilder()
139+
.setText(query)
140+
.setRagRetrievalConfig(
141+
RagRetrievalConfig.newBuilder().setTopK(similarityTopK).build()))
129142
.setVertexRagStore(
130143
com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.newBuilder()
131144
.addAllRagResources(ragResources)
@@ -154,14 +167,16 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
154167
ImmutableList<RagResource> ragResources =
155168
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
156169
Double vectorDistanceThreshold = 0.5;
170+
Integer similarityTopK = 10;
157171
VertexAiRagRetrieval tool =
158172
new VertexAiRagRetrieval(
159173
"testTool",
160174
"test description",
161175
vertexRagServiceClient,
162176
"projects/test-project/locations/us-central1",
163177
ragResources,
164-
vectorDistanceThreshold);
178+
vectorDistanceThreshold,
179+
similarityTopK);
165180
LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("gemini-2-pro");
166181
ToolContext toolContext = buildToolContext();
167182

@@ -181,7 +196,8 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
181196
VertexRagStoreRagResource.builder()
182197
.ragCorpus("corpus1")
183198
.build()))
184-
.vectorDistanceThreshold(0.5)
199+
.vectorDistanceThreshold(vectorDistanceThreshold)
200+
.similarityTopK(similarityTopK)
185201
.build())
186202
.build())
187203
.build());
@@ -212,7 +228,9 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
212228
}
213229

214230
@Test
215-
public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
231+
public void processLlmRequest_gemini2Model_withoutSimilarityTopK_addVertexRagStoreToConfig() {
232+
// This test's behavior depends on the GOOGLE_GENAI_USE_VERTEXAI environment variable
233+
boolean useVertexAi = Boolean.parseBoolean(System.getenv("GOOGLE_GENAI_USE_VERTEXAI"));
216234
ImmutableList<RagResource> ragResources =
217235
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
218236
Double vectorDistanceThreshold = 0.5;
@@ -224,6 +242,70 @@ public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
224242
"projects/test-project/locations/us-central1",
225243
ragResources,
226244
vectorDistanceThreshold);
245+
LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("gemini-2-pro");
246+
ToolContext toolContext = buildToolContext();
247+
248+
tool.processLlmRequest(llmRequestBuilder, toolContext).blockingAwait();
249+
250+
if (useVertexAi) {
251+
// Assert that VertexRagStore is added to the config without similarityTopK
252+
assertThat(llmRequestBuilder.build().config().get().tools().get())
253+
.containsExactly(
254+
Tool.builder()
255+
.retrieval(
256+
Retrieval.builder()
257+
.vertexRagStore(
258+
VertexRagStore.builder()
259+
.ragResources(
260+
ImmutableList.of(
261+
VertexRagStoreRagResource.builder()
262+
.ragCorpus("corpus1")
263+
.build()))
264+
.vectorDistanceThreshold(vectorDistanceThreshold)
265+
.build())
266+
.build())
267+
.build());
268+
} else {
269+
// Assert that the function declaration is added instead
270+
assertThat(llmRequestBuilder.build().config().get().tools().get())
271+
.containsExactly(
272+
Tool.builder()
273+
.functionDeclarations(
274+
ImmutableList.of(
275+
FunctionDeclaration.builder()
276+
.name("testTool")
277+
.description("test description")
278+
.parameters(
279+
Schema.builder()
280+
.properties(
281+
ImmutableMap.of(
282+
"query",
283+
Schema.builder()
284+
.description("The query to retrieve.")
285+
.type("STRING")
286+
.build()))
287+
.type("OBJECT")
288+
.build())
289+
.build()))
290+
.build());
291+
}
292+
}
293+
294+
@Test
295+
public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
296+
ImmutableList<RagResource> ragResources =
297+
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
298+
Double vectorDistanceThreshold = 0.5;
299+
Integer similarityTopK = 10;
300+
VertexAiRagRetrieval tool =
301+
new VertexAiRagRetrieval(
302+
"testTool",
303+
"test description",
304+
vertexRagServiceClient,
305+
"projects/test-project/locations/us-central1",
306+
ragResources,
307+
vectorDistanceThreshold,
308+
similarityTopK);
227309
LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("other-model");
228310
ToolContext toolContext = buildToolContext();
229311
GenerateContentConfig initialConfig = GenerateContentConfig.builder().build();

0 commit comments

Comments
 (0)