Skip to content

Commit b0bdd98

Browse files
copybara-githubcopybara-github
committed
Merge pull request #1442 from pgentile:rag_retrieval_similarity_top_k
PiperOrigin-RevId: 1005486394
2 parents 7096019 + 71bd1cd commit b0bdd98

3 files changed

Lines changed: 229 additions & 18 deletions

File tree

‎core/pom.xml‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,11 @@
274274
<GOOGLE_GENAI_USE_VERTEXAI>true</GOOGLE_GENAI_USE_VERTEXAI>
275275
</environmentVariables>
276276
<!-- run a second variant of the following tests -->
277-
<test>VertexAiRagRetrievalTest#processLlmRequest_gemini2Model_addVertexRagStoreToConfig, VertexAiRagRetrievalTest#processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig</test>
277+
<test>
278+
VertexAiRagRetrievalTest#processLlmRequest_gemini2Model_addVertexRagStoreToConfig,
279+
VertexAiRagRetrievalTest#processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig,
280+
VertexAiRagRetrievalTest#processLlmRequest_gemini2Model_withoutSimilarityTopK_addVertexRagStoreToConfig
281+
</test>
278282
</configuration>
279283
</execution>
280284
<execution>

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

Lines changed: 96 additions & 8 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;
56-
private final List<RagResource> ragResources;
58+
@Nullable private final List<RagResource> ragResources;
5759
private final Double vectorDistanceThreshold;
60+
@Nullable private final Integer similarityTopK;
5861
private final VertexRagStore vertexRagStore;
5962
private final RetrieveContextsRequest.VertexRagStore apiVertexRagStore;
6063

@@ -65,11 +68,24 @@ public VertexAiRagRetrieval(
6568
String parent,
6669
@Nullable List<RagResource> ragResources,
6770
@Nullable Double vectorDistanceThreshold) {
68-
super(name, description);
69-
this.vertexRagServiceClient = vertexRagServiceClient;
70-
this.parent = parent;
71-
this.ragResources = ragResources;
72-
this.vectorDistanceThreshold = vectorDistanceThreshold;
71+
this(
72+
builder()
73+
.name(name)
74+
.description(description)
75+
.vertexRagServiceClient(vertexRagServiceClient)
76+
.parent(parent)
77+
.ragResources(ragResources)
78+
.vectorDistanceThreshold(vectorDistanceThreshold));
79+
}
80+
81+
private VertexAiRagRetrieval(Builder builder) {
82+
super(builder.name, builder.description);
83+
84+
this.vertexRagServiceClient = builder.vertexRagServiceClient;
85+
this.parent = builder.parent;
86+
this.ragResources = builder.ragResources;
87+
this.vectorDistanceThreshold = builder.vectorDistanceThreshold;
88+
this.similarityTopK = builder.similarityTopK;
7389

7490
// For Gemini 2
7591
VertexRagStore.Builder vertexRagStoreBuilder = VertexRagStore.builder();
@@ -86,6 +102,9 @@ public VertexAiRagRetrieval(
86102
if (this.vectorDistanceThreshold != null) {
87103
vertexRagStoreBuilder.vectorDistanceThreshold(this.vectorDistanceThreshold);
88104
}
105+
if (this.similarityTopK != null) {
106+
vertexRagStoreBuilder.similarityTopK(this.similarityTopK);
107+
}
89108
this.vertexRagStore = vertexRagStoreBuilder.build();
90109

91110
// For runAsync
@@ -100,6 +119,67 @@ public VertexAiRagRetrieval(
100119
this.apiVertexRagStore = apiVertexRagStoreBuilder.build();
101120
}
102121

122+
public static Builder builder() {
123+
return new Builder();
124+
}
125+
126+
/** Builder for {@link VertexAiRagRetrieval}. */
127+
public static class Builder {
128+
private String name;
129+
private String description;
130+
private VertexRagServiceClient vertexRagServiceClient;
131+
private String parent;
132+
@Nullable private List<RagResource> ragResources;
133+
@Nullable private Double vectorDistanceThreshold;
134+
@Nullable private Integer similarityTopK;
135+
136+
@CanIgnoreReturnValue
137+
public Builder name(String name) {
138+
this.name = name;
139+
return this;
140+
}
141+
142+
@CanIgnoreReturnValue
143+
public Builder description(String description) {
144+
this.description = description;
145+
return this;
146+
}
147+
148+
@CanIgnoreReturnValue
149+
public Builder vertexRagServiceClient(VertexRagServiceClient vertexRagServiceClient) {
150+
this.vertexRagServiceClient = vertexRagServiceClient;
151+
return this;
152+
}
153+
154+
@CanIgnoreReturnValue
155+
public Builder parent(String parent) {
156+
this.parent = parent;
157+
return this;
158+
}
159+
160+
@CanIgnoreReturnValue
161+
public Builder ragResources(@Nullable List<RagResource> ragResources) {
162+
this.ragResources = ragResources;
163+
return this;
164+
}
165+
166+
@CanIgnoreReturnValue
167+
public Builder vectorDistanceThreshold(@Nullable Double vectorDistanceThreshold) {
168+
this.vectorDistanceThreshold = vectorDistanceThreshold;
169+
return this;
170+
}
171+
172+
@CanIgnoreReturnValue
173+
public Builder similarityTopK(@Nullable Integer similarityTopK) {
174+
this.similarityTopK = similarityTopK;
175+
return this;
176+
}
177+
178+
public VertexAiRagRetrieval build() {
179+
return new VertexAiRagRetrieval(this);
180+
}
181+
}
182+
103183
@Override
104184
@CanIgnoreReturnValue
105185
public Completable processLlmRequest(
@@ -135,10 +215,18 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> args, ToolContex
135215
return Single.fromCallable(
136216
() -> {
137217
logger.info("Retrieving context for query: {}", query);
218+
219+
RagQuery.Builder queryBuilder = RagQuery.newBuilder().setText(query);
220+
221+
if (similarityTopK != null) {
222+
queryBuilder.setRagRetrievalConfig(
223+
RagRetrievalConfig.newBuilder().setTopK(similarityTopK).build());
224+
}
225+
138226
RetrieveContextsRequest retrieveContextsRequest =
139227
RetrieveContextsRequest.newBuilder()
140228
.setParent(this.parent)
141-
.setQuery(RagQuery.newBuilder().setText(query))
229+
.setQuery(queryBuilder.build())
142230
.setVertexRagStore(this.apiVertexRagStore)
143231
.build();
144232
logger.info("Request to VertexRagService: {}", retrieveContextsRequest);

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

Lines changed: 128 additions & 9 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;
@@ -107,6 +108,53 @@ public void runAsync_withResults_returnsContexts() throws Exception {
107108
verify(vertexRagServiceClient).retrieveContexts(eq(expectedRequest));
108109
}
109110

111+
@Test
112+
public void runAsync_withSimilarityTopK_setsRagRetrievalConfig() throws Exception {
113+
ImmutableList<RagResource> ragResources =
114+
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
115+
Double vectorDistanceThreshold = 0.5;
116+
Integer similarityTopK = 10;
117+
VertexAiRagRetrieval tool =
118+
VertexAiRagRetrieval.builder()
119+
.name("testTool")
120+
.description("test description")
121+
.vertexRagServiceClient(vertexRagServiceClient)
122+
.parent("projects/test-project/locations/us-central1")
123+
.ragResources(ragResources)
124+
.vectorDistanceThreshold(vectorDistanceThreshold)
125+
.similarityTopK(similarityTopK)
126+
.build();
127+
String query = "test query";
128+
ToolContext toolContext = buildToolContext();
129+
RetrieveContextsRequest expectedRequest =
130+
RetrieveContextsRequest.newBuilder()
131+
.setParent("projects/test-project/locations/us-central1")
132+
.setQuery(
133+
RagQuery.newBuilder()
134+
.setText(query)
135+
.setRagRetrievalConfig(
136+
RagRetrievalConfig.newBuilder().setTopK(similarityTopK).build()))
137+
.setVertexRagStore(
138+
com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.newBuilder()
139+
.addAllRagResources(ragResources)
140+
.setVectorDistanceThreshold(vectorDistanceThreshold))
141+
.build();
142+
when(vertexRagServiceClient.retrieveContexts(eq(expectedRequest)))
143+
.thenReturn(
144+
RetrieveContextsResponse.newBuilder()
145+
.setContexts(
146+
RagContexts.newBuilder()
147+
.addContexts(RagContexts.Context.newBuilder().setText("context1"))
148+
.addContexts(RagContexts.Context.newBuilder().setText("context2")))
149+
.build());
150+
151+
Map<String, Object> result =
152+
tool.runAsync(ImmutableMap.of("query", query), toolContext).blockingGet();
153+
154+
assertThat(result).containsExactly("response", ImmutableList.of("context1", "context2"));
155+
verify(vertexRagServiceClient).retrieveContexts(eq(expectedRequest));
156+
}
157+
110158
@Test
111159
public void runAsync_noResults_returnsNoResultFoundMessage() throws Exception {
112160
ImmutableList<RagResource> ragResources =
@@ -154,14 +202,17 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
154202
ImmutableList<RagResource> ragResources =
155203
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
156204
Double vectorDistanceThreshold = 0.5;
205+
Integer similarityTopK = 10;
157206
VertexAiRagRetrieval tool =
158-
new VertexAiRagRetrieval(
159-
"testTool",
160-
"test description",
161-
vertexRagServiceClient,
162-
"projects/test-project/locations/us-central1",
163-
ragResources,
164-
vectorDistanceThreshold);
207+
VertexAiRagRetrieval.builder()
208+
.name("testTool")
209+
.description("test description")
210+
.vertexRagServiceClient(vertexRagServiceClient)
211+
.parent("projects/test-project/locations/us-central1")
212+
.ragResources(ragResources)
213+
.vectorDistanceThreshold(vectorDistanceThreshold)
214+
.similarityTopK(similarityTopK)
215+
.build();
165216
LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("gemini-2-pro");
166217
ToolContext toolContext = buildToolContext();
167218

@@ -181,7 +232,8 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
181232
VertexRagStoreRagResource.builder()
182233
.ragCorpus("corpus1")
183234
.build()))
184-
.vectorDistanceThreshold(0.5)
235+
.vectorDistanceThreshold(vectorDistanceThreshold)
236+
.similarityTopK(similarityTopK)
185237
.build())
186238
.build())
187239
.build());
@@ -212,7 +264,9 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
212264
}
213265

214266
@Test
215-
public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
267+
public void processLlmRequest_gemini2Model_withoutSimilarityTopK_addVertexRagStoreToConfig() {
268+
// This test's behavior depends on the GOOGLE_GENAI_USE_VERTEXAI environment variable
269+
boolean useVertexAi = Boolean.parseBoolean(System.getenv("GOOGLE_GENAI_USE_VERTEXAI"));
216270
ImmutableList<RagResource> ragResources =
217271
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
218272
Double vectorDistanceThreshold = 0.5;
@@ -224,6 +278,71 @@ public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
224278
"projects/test-project/locations/us-central1",
225279
ragResources,
226280
vectorDistanceThreshold);
281+
LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("gemini-2-pro");
282+
ToolContext toolContext = buildToolContext();
283+
284+
tool.processLlmRequest(llmRequestBuilder, toolContext).blockingAwait();
285+
286+
if (useVertexAi) {
287+
// Assert that VertexRagStore is added to the config without similarityTopK
288+
assertThat(llmRequestBuilder.build().config().get().tools().get())
289+
.containsExactly(
290+
Tool.builder()
291+
.retrieval(
292+
Retrieval.builder()
293+
.vertexRagStore(
294+
VertexRagStore.builder()
295+
.ragResources(
296+
ImmutableList.of(
297+
VertexRagStoreRagResource.builder()
298+
.ragCorpus("corpus1")
299+
.build()))
300+
.vectorDistanceThreshold(vectorDistanceThreshold)
301+
.build())
302+
.build())
303+
.build());
304+
} else {
305+
// Assert that the function declaration is added instead
306+
assertThat(llmRequestBuilder.build().config().get().tools().get())
307+
.containsExactly(
308+
Tool.builder()
309+
.functionDeclarations(
310+
ImmutableList.of(
311+
FunctionDeclaration.builder()
312+
.name("testTool")
313+
.description("test description")
314+
.parameters(
315+
Schema.builder()
316+
.properties(
317+
ImmutableMap.of(
318+
"query",
319+
Schema.builder()
320+
.description("The query to retrieve.")
321+
.type("STRING")
322+
.build()))
323+
.type("OBJECT")
324+
.build())
325+
.build()))
326+
.build());
327+
}
328+
}
329+
330+
@Test
331+
public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
332+
ImmutableList<RagResource> ragResources =
333+
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
334+
Double vectorDistanceThreshold = 0.5;
335+
Integer similarityTopK = 10;
336+
VertexAiRagRetrieval tool =
337+
VertexAiRagRetrieval.builder()
338+
.name("testTool")
339+
.description("test description")
340+
.vertexRagServiceClient(vertexRagServiceClient)
341+
.parent("projects/test-project/locations/us-central1")
342+
.ragResources(ragResources)
343+
.vectorDistanceThreshold(vectorDistanceThreshold)
344+
.similarityTopK(similarityTopK)
345+
.build();
227346
LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("other-model");
228347
ToolContext toolContext = buildToolContext();
229348
GenerateContentConfig initialConfig = GenerateContentConfig.builder().build();

0 commit comments

Comments
 (0)