Skip to content

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

Open
pgentile wants to merge 1 commit into
google:mainfrom
pgentile:rag_retrieval_similarity_top_k
Open

feat: configure similarity top-K in the VertexAiRagRetrieval tool#1442
pgentile wants to merge 1 commit into
google:mainfrom
pgentile:rag_retrieval_similarity_top_k

Conversation

@pgentile

@pgentile pgentile commented Aug 19, 2026

Copy link
Copy Markdown

Link to Issue or Description of Change

Problem:

Add support for the similarity top-K parameter in the VertexAiRagRetrieval tool

Solution:

New similarityTopK constructor parameter in the constructor of the VertexAiRagRetrieval class.

The similarity top-K is then added to the LLM request, when the model is a Gemini one and when the GOOGLE_GENAI_USE_VERTEXAI env var is enabled. Otherwise, the similarity top-K parameter is transmitted to the retrieveContexts method.

Testing Plan

Updates of the VertexAiRagRetrievalTest unit test.

Unit Tests:

  • I have added or updated unit tests for my change.
  • All unit tests pass locally.
[INFO] +-- com.google.adk.tools.retrieval.VertexAiRagRetrievalTest
[INFO] | +-- [OK] processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig - 1.298 ss
[INFO] | +-- [OK] processLlmRequest_gemini2Model_addVertexRagStoreToConfig - 0.006 ss

...

[INFO] +-- com.google.adk.tools.retrieval.VertexAiRagRetrievalTest
[INFO] | +-- [OK] processLlmRequest_otherModel_doNotAddVertexRagStoreToConfig - 0.084 ss
[INFO] | +-- [OK] processLlmRequest_gemini2Model_addVertexRagStoreToConfig - 0.001 ss
[INFO] | +-- [OK] runAsync_withResults_returnsContexts - 0.190 ss
[INFO] | +-- [OK] runAsync_noResults_returnsNoResultFoundMessage - 0.003 ss
[INFO] | +-- [OK] processLlmRequest_gemini2Model_withoutSimilarityTopK_addVertexRagStoreToConfig - 0 ss

Manual End-to-End (E2E) Tests:

  • Use the tool new parameter in a production agent.
  • Check both code branches based on the GOOGLE_GENAI_USE_VERTEXAI env var.
  • Count retrieved chunks from the RAG corpus. The number of chunks must match the configured similarity top-K

Checklist

  • I have read the CONTRIBUTING.md document.
  • My pull request contains a single commit.
  • I have performed a self-review of my own code.
  • I have commented my code, particularly in hard-to-understand areas.
  • I have added tests that prove my fix is effective or that my feature works.
  • New and existing unit tests pass locally with my changes.
  • I have manually tested my changes end-to-end.
  • Any dependent changes have been merged and published in downstream modules.

Additional context

N/A

@pgentile
pgentile marked this pull request as ready for review August 19, 2026 15:07
@pgentile
pgentile force-pushed the rag_retrieval_similarity_top_k branch from 71cd444 to 2493cd1 Compare August 19, 2026 15:07
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant