diff --git a/core/src/main/java/com/google/adk/tools/retrieval/VertexAiRagRetrieval.java b/core/src/main/java/com/google/adk/tools/retrieval/VertexAiRagRetrieval.java
index a2720aae5..789086113 100644
--- a/core/src/main/java/com/google/adk/tools/retrieval/VertexAiRagRetrieval.java
+++ b/core/src/main/java/com/google/adk/tools/retrieval/VertexAiRagRetrieval.java
@@ -23,6 +23,7 @@
import com.google.adk.utils.ModelNameUtils;
import com.google.cloud.aiplatform.v1.RagContexts;
import com.google.cloud.aiplatform.v1.RagQuery;
+import com.google.cloud.aiplatform.v1.RagRetrievalConfig;
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest;
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.RagResource;
import com.google.cloud.aiplatform.v1.RetrieveContextsResponse;
@@ -47,14 +48,16 @@
* A retrieval tool that fetches context from Vertex AI RAG.
*
* This tool allows to retrieve relevant information based on a query using Vertex AI RAG
- * service. It supports configuration of rag resources and a vector distance threshold.
+ * service. It supports configuration of rag resources, a vector distance threshold, and similarity
+ * top-k.
*/
public class VertexAiRagRetrieval extends BaseRetrievalTool {
private static final Logger logger = LoggerFactory.getLogger(VertexAiRagRetrieval.class);
private final VertexRagServiceClient vertexRagServiceClient;
private final String parent;
- private final List ragResources;
+ @Nullable private final List ragResources;
private final Double vectorDistanceThreshold;
+ @Nullable private final Integer similarityTopK;
private final VertexRagStore vertexRagStore;
private final RetrieveContextsRequest.VertexRagStore apiVertexRagStore;
@@ -65,11 +68,24 @@ public VertexAiRagRetrieval(
String parent,
@Nullable List ragResources,
@Nullable Double vectorDistanceThreshold) {
- super(name, description);
- this.vertexRagServiceClient = vertexRagServiceClient;
- this.parent = parent;
- this.ragResources = ragResources;
- this.vectorDistanceThreshold = vectorDistanceThreshold;
+ this(
+ builder()
+ .name(name)
+ .description(description)
+ .vertexRagServiceClient(vertexRagServiceClient)
+ .parent(parent)
+ .ragResources(ragResources)
+ .vectorDistanceThreshold(vectorDistanceThreshold));
+ }
+
+ private VertexAiRagRetrieval(Builder builder) {
+ super(builder.name, builder.description);
+
+ this.vertexRagServiceClient = builder.vertexRagServiceClient;
+ this.parent = builder.parent;
+ this.ragResources = builder.ragResources;
+ this.vectorDistanceThreshold = builder.vectorDistanceThreshold;
+ this.similarityTopK = builder.similarityTopK;
// For Gemini 2
VertexRagStore.Builder vertexRagStoreBuilder = VertexRagStore.builder();
@@ -86,6 +102,9 @@ public VertexAiRagRetrieval(
if (this.vectorDistanceThreshold != null) {
vertexRagStoreBuilder.vectorDistanceThreshold(this.vectorDistanceThreshold);
}
+ if (this.similarityTopK != null) {
+ vertexRagStoreBuilder.similarityTopK(this.similarityTopK);
+ }
this.vertexRagStore = vertexRagStoreBuilder.build();
// For runAsync
@@ -100,6 +119,67 @@ public VertexAiRagRetrieval(
this.apiVertexRagStore = apiVertexRagStoreBuilder.build();
}
+ public static Builder builder() {
+ return new Builder();
+ }
+
+ /** Builder for {@link VertexAiRagRetrieval}. */
+ public static class Builder {
+ private String name;
+ private String description;
+ private VertexRagServiceClient vertexRagServiceClient;
+ private String parent;
+ @Nullable private List ragResources;
+ @Nullable private Double vectorDistanceThreshold;
+ @Nullable private Integer similarityTopK;
+
+ @CanIgnoreReturnValue
+ public Builder name(String name) {
+ this.name = name;
+ return this;
+ }
+
+ @CanIgnoreReturnValue
+ public Builder description(String description) {
+ this.description = description;
+ return this;
+ }
+
+ @CanIgnoreReturnValue
+ public Builder vertexRagServiceClient(VertexRagServiceClient vertexRagServiceClient) {
+ this.vertexRagServiceClient = vertexRagServiceClient;
+ return this;
+ }
+
+ @CanIgnoreReturnValue
+ public Builder parent(String parent) {
+ this.parent = parent;
+ return this;
+ }
+
+ @CanIgnoreReturnValue
+ public Builder ragResources(@Nullable List ragResources) {
+ this.ragResources = ragResources;
+ return this;
+ }
+
+ @CanIgnoreReturnValue
+ public Builder vectorDistanceThreshold(@Nullable Double vectorDistanceThreshold) {
+ this.vectorDistanceThreshold = vectorDistanceThreshold;
+ return this;
+ }
+
+ @CanIgnoreReturnValue
+ public Builder similarityTopK(@Nullable Integer similarityTopK) {
+ this.similarityTopK = similarityTopK;
+ return this;
+ }
+
+ public VertexAiRagRetrieval build() {
+ return new VertexAiRagRetrieval(this);
+ }
+ }
+
@Override
@CanIgnoreReturnValue
public Completable processLlmRequest(
@@ -135,10 +215,18 @@ public Single> runAsync(Map args, ToolContex
return Single.fromCallable(
() -> {
logger.info("Retrieving context for query: {}", query);
+
+ RagQuery.Builder queryBuilder = RagQuery.newBuilder().setText(query);
+
+ if (similarityTopK != null) {
+ queryBuilder.setRagRetrievalConfig(
+ RagRetrievalConfig.newBuilder().setTopK(similarityTopK).build());
+ }
+
RetrieveContextsRequest retrieveContextsRequest =
RetrieveContextsRequest.newBuilder()
.setParent(this.parent)
- .setQuery(RagQuery.newBuilder().setText(query))
+ .setQuery(queryBuilder.build())
.setVertexRagStore(this.apiVertexRagStore)
.build();
logger.info("Request to VertexRagService: {}", retrieveContextsRequest);
diff --git a/core/src/test/java/com/google/adk/tools/retrieval/VertexAiRagRetrievalTest.java b/core/src/test/java/com/google/adk/tools/retrieval/VertexAiRagRetrievalTest.java
index 1b8cbf66a..6f8226669 100644
--- a/core/src/test/java/com/google/adk/tools/retrieval/VertexAiRagRetrievalTest.java
+++ b/core/src/test/java/com/google/adk/tools/retrieval/VertexAiRagRetrievalTest.java
@@ -29,6 +29,7 @@
import com.google.adk.tools.ToolContext;
import com.google.cloud.aiplatform.v1.RagContexts;
import com.google.cloud.aiplatform.v1.RagQuery;
+import com.google.cloud.aiplatform.v1.RagRetrievalConfig;
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest;
import com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.RagResource;
import com.google.cloud.aiplatform.v1.RetrieveContextsResponse;
@@ -107,6 +108,53 @@ public void runAsync_withResults_returnsContexts() throws Exception {
verify(vertexRagServiceClient).retrieveContexts(eq(expectedRequest));
}
+ @Test
+ public void runAsync_withSimilarityTopK_setsRagRetrievalConfig() throws Exception {
+ ImmutableList ragResources =
+ ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
+ Double vectorDistanceThreshold = 0.5;
+ Integer similarityTopK = 10;
+ VertexAiRagRetrieval tool =
+ VertexAiRagRetrieval.builder()
+ .name("testTool")
+ .description("test description")
+ .vertexRagServiceClient(vertexRagServiceClient)
+ .parent("projects/test-project/locations/us-central1")
+ .ragResources(ragResources)
+ .vectorDistanceThreshold(vectorDistanceThreshold)
+ .similarityTopK(similarityTopK)
+ .build();
+ String query = "test query";
+ ToolContext toolContext = buildToolContext();
+ RetrieveContextsRequest expectedRequest =
+ RetrieveContextsRequest.newBuilder()
+ .setParent("projects/test-project/locations/us-central1")
+ .setQuery(
+ RagQuery.newBuilder()
+ .setText(query)
+ .setRagRetrievalConfig(
+ RagRetrievalConfig.newBuilder().setTopK(similarityTopK).build()))
+ .setVertexRagStore(
+ com.google.cloud.aiplatform.v1.RetrieveContextsRequest.VertexRagStore.newBuilder()
+ .addAllRagResources(ragResources)
+ .setVectorDistanceThreshold(vectorDistanceThreshold))
+ .build();
+ when(vertexRagServiceClient.retrieveContexts(eq(expectedRequest)))
+ .thenReturn(
+ RetrieveContextsResponse.newBuilder()
+ .setContexts(
+ RagContexts.newBuilder()
+ .addContexts(RagContexts.Context.newBuilder().setText("context1"))
+ .addContexts(RagContexts.Context.newBuilder().setText("context2")))
+ .build());
+
+ Map result =
+ tool.runAsync(ImmutableMap.of("query", query), toolContext).blockingGet();
+
+ assertThat(result).containsExactly("response", ImmutableList.of("context1", "context2"));
+ verify(vertexRagServiceClient).retrieveContexts(eq(expectedRequest));
+ }
+
@Test
public void runAsync_noResults_returnsNoResultFoundMessage() throws Exception {
ImmutableList ragResources =
@@ -154,14 +202,17 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
ImmutableList ragResources =
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
Double vectorDistanceThreshold = 0.5;
+ Integer similarityTopK = 10;
VertexAiRagRetrieval tool =
- new VertexAiRagRetrieval(
- "testTool",
- "test description",
- vertexRagServiceClient,
- "projects/test-project/locations/us-central1",
- ragResources,
- vectorDistanceThreshold);
+ VertexAiRagRetrieval.builder()
+ .name("testTool")
+ .description("test description")
+ .vertexRagServiceClient(vertexRagServiceClient)
+ .parent("projects/test-project/locations/us-central1")
+ .ragResources(ragResources)
+ .vectorDistanceThreshold(vectorDistanceThreshold)
+ .similarityTopK(similarityTopK)
+ .build();
LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("gemini-2-pro");
ToolContext toolContext = buildToolContext();
@@ -181,7 +232,8 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
VertexRagStoreRagResource.builder()
.ragCorpus("corpus1")
.build()))
- .vectorDistanceThreshold(0.5)
+ .vectorDistanceThreshold(vectorDistanceThreshold)
+ .similarityTopK(similarityTopK)
.build())
.build())
.build());
@@ -212,7 +264,9 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
}
@Test
- public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
+ public void processLlmRequest_gemini2Model_withoutSimilarityTopK_addVertexRagStoreToConfig() {
+ // This test's behavior depends on the GOOGLE_GENAI_USE_VERTEXAI environment variable
+ boolean useVertexAi = Boolean.parseBoolean(System.getenv("GOOGLE_GENAI_USE_VERTEXAI"));
ImmutableList ragResources =
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
Double vectorDistanceThreshold = 0.5;
@@ -224,6 +278,71 @@ public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
"projects/test-project/locations/us-central1",
ragResources,
vectorDistanceThreshold);
+ LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("gemini-2-pro");
+ ToolContext toolContext = buildToolContext();
+
+ tool.processLlmRequest(llmRequestBuilder, toolContext).blockingAwait();
+
+ if (useVertexAi) {
+ // Assert that VertexRagStore is added to the config without similarityTopK
+ assertThat(llmRequestBuilder.build().config().get().tools().get())
+ .containsExactly(
+ Tool.builder()
+ .retrieval(
+ Retrieval.builder()
+ .vertexRagStore(
+ VertexRagStore.builder()
+ .ragResources(
+ ImmutableList.of(
+ VertexRagStoreRagResource.builder()
+ .ragCorpus("corpus1")
+ .build()))
+ .vectorDistanceThreshold(vectorDistanceThreshold)
+ .build())
+ .build())
+ .build());
+ } else {
+ // Assert that the function declaration is added instead
+ assertThat(llmRequestBuilder.build().config().get().tools().get())
+ .containsExactly(
+ Tool.builder()
+ .functionDeclarations(
+ ImmutableList.of(
+ FunctionDeclaration.builder()
+ .name("testTool")
+ .description("test description")
+ .parameters(
+ Schema.builder()
+ .properties(
+ ImmutableMap.of(
+ "query",
+ Schema.builder()
+ .description("The query to retrieve.")
+ .type("STRING")
+ .build()))
+ .type("OBJECT")
+ .build())
+ .build()))
+ .build());
+ }
+ }
+
+ @Test
+ public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
+ ImmutableList ragResources =
+ ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
+ Double vectorDistanceThreshold = 0.5;
+ Integer similarityTopK = 10;
+ VertexAiRagRetrieval tool =
+ VertexAiRagRetrieval.builder()
+ .name("testTool")
+ .description("test description")
+ .vertexRagServiceClient(vertexRagServiceClient)
+ .parent("projects/test-project/locations/us-central1")
+ .ragResources(ragResources)
+ .vectorDistanceThreshold(vectorDistanceThreshold)
+ .similarityTopK(similarityTopK)
+ .build();
LlmRequest.Builder llmRequestBuilder = LlmRequest.builder().model("other-model");
ToolContext toolContext = buildToolContext();
GenerateContentConfig initialConfig = GenerateContentConfig.builder().build();