Skip to content

feat: configure similarity top-K in the VertexAiRagRetrieval tool #1442

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
pgentile wants to merge 1 commit into
base: main
Choose a base branch
from
+229 −18
Open

feat: configure similarity top-K in the VertexAiRagRetrieval tool #1442

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 5 additions & 1 deletion core/pom.xml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -274,7 +274,11 @@
<GOOGLE_GENAI_USE_VERTEXAI>true</GOOGLE_GENAI_USE_VERTEXAI>
</environmentVariables>
<!-- run a second variant of the following tests -->
<test>VertexAiRagRetrievalTest#processLlmRequest_gemini2Model_addVertexRagStoreToConfig, VertexAiRagRetrievalTest#processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig</test>
<test>
VertexAiRagRetrievalTest#processLlmRequest_gemini2Model_addVertexRagStoreToConfig,
VertexAiRagRetrievalTest#processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig,
VertexAiRagRetrievalTest#processLlmRequest_gemini2Model_withoutSimilarityTopK_addVertexRagStoreToConfig
</test>
</configuration>
</execution>
<execution>
Expand Down
104 changes: 96 additions & 8 deletions core/src/main/java/com/google/adk/tools/retrieval/VertexAiRagRetrieval.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand All @@ -47,14 +48,16 @@
* A retrieval tool that fetches context from Vertex AI RAG.
*
* <p>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<RagResource> ragResources;
@Nullable private final List<RagResource> ragResources;
private final Double vectorDistanceThreshold;
@Nullable private final Integer similarityTopK;
private final VertexRagStore vertexRagStore;
private final RetrieveContextsRequest.VertexRagStore apiVertexRagStore;

Expand All @@ -65,11 +68,24 @@ public VertexAiRagRetrieval(
String parent,
@Nullable List<RagResource> 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();
Expand All @@ -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
Expand All @@ -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<RagResource> 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<RagResource> 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(
Expand Down Expand Up @@ -135,10 +215,18 @@ public Single<Map<String, Object>> runAsync(Map<String, Object> 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);
Expand Down
137 changes: 128 additions & 9 deletions core/src/test/java/com/google/adk/tools/retrieval/VertexAiRagRetrievalTest.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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<RagResource> 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<String, Object> 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<RagResource> ragResources =
Expand Down Expand Up @@ -154,14 +202,17 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
ImmutableList<RagResource> 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();

Expand All @@ -181,7 +232,8 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
VertexRagStoreRagResource.builder()
.ragCorpus("corpus1")
.build()))
.vectorDistanceThreshold(0.5)
.vectorDistanceThreshold(vectorDistanceThreshold)
.similarityTopK(similarityTopK)
.build())
.build())
.build());
Expand Down Expand Up @@ -212,7 +264,9 @@ public void processLlmRequest_gemini2Model_addVertexRagStoreToConfig() {
}

@Test
public void processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig() {
public void processLlmRequest_gemini2Model_withoutSimilarityTopK_addVertexRagStoreToConfig() {
Comment thread
pgentile marked this conversation as resolved.
// This test's behavior depends on the GOOGLE_GENAI_USE_VERTEXAI environment variable
boolean useVertexAi = Boolean.parseBoolean(System.getenv("GOOGLE_GENAI_USE_VERTEXAI"));
ImmutableList<RagResource> ragResources =
ImmutableList.of(RagResource.newBuilder().setRagCorpus("corpus1").build());
Double vectorDistanceThreshold = 0.5;
Expand All @@ -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<RagResource> 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();
Expand Down
Loading