feat(tools): port the adk-python v0.1.0 retrieval tools - #854
Open
AmaadMartin wants to merge 6 commits into
Open
feat(tools): port the adk-python v0.1.0 retrieval tools#854AmaadMartin wants to merge 6 commits into
AmaadMartin wants to merge 6 commits into
Conversation
added 4 commits
September 3, 2026 11:35
adk-js covers Vertex AI RAG with VertexRagRetrievalTool but has no
general retrieval-tool surface, so a tool that answers a natural-language
query out of a local index has to reimplement the declaration and the
argument contract from scratch. adk-python has carried that surface since
v0.1.0.
Adds the three modules adk-js lacks, ported from adk-python at ref
v0.1.0:
- BaseRetrievalTool, which fixes the one-argument `query` declaration
every retrieval tool advertises.
- LlamaIndexRetrieval, answering with the top node's text. The
retriever is typed structurally, so llamaindex stays optional
exactly as the TYPE_CHECKING import keeps it optional in Python.
- FilesRetrieval, which indexes a directory. Reading and indexing are
async in LlamaIndexTS where they are sync in Python, so they move
from the constructor to a static create().
Nothing existing changes. VertexRagRetrievalTool is untouched: its
server-side-only behaviour, its fixed name and its BuiltInTool
registration all stand. The v0.1.0 tests for VertexAiRagRetrieval are
ported alongside as a parity set that records, per assertion, where
v0.1.0's model-gated client-side path differs from what adk-js does
today — the Gemini 1.x function-declaration path, the shared tool entry,
and the toolsDict registration that #789 deliberately added.
Documents `BaseRetrievalTool`, `LlamaIndexRetrieval` and `FilesRetrieval` at `docs/guides/tools/retrieval/index.md`, mirroring `core/src/tools/retrieval/`. The three files are one guide rather than three, because a reader choosing between the implementations has to see them side by side. The section that earns the guide is the comparison with the pre-existing `VertexRagRetrievalTool`. The two look interchangeable from the outside and are not: one extends `BaseTool` and retrieves in your process behind a function declaration, the other extends `BuiltInTool` and adds a `retrieval.vertexRagStore` entry so the model retrieves server-side, returning grounding metadata instead of a function response. The guide gives the reader the trade-offs — inspectable results and any backend against no extra round trip — rather than a recommendation. Every snippet was extracted to a scratch file under `samples/` and compiled with `tsc -p samples` against the published types, so the signatures, the `VertexRagStore` fields and the `await FilesRetrieval.create` shape are real. `docs/guides/README.md` is the index the guide skill requires; it is new, since adk-js had no `docs/guides/` tree before this change.
Adds `samples/tools/retrieval/agent.ts`, an `LlmAgent` grounded on a `LlamaIndexRetrieval` tool, plus `samples/tools/README.md` for the new category. The sample ships its corpus in the file behind a hand-written `LlamaIndexRetriever`, which is the honest demonstration of the design: the retriever is a structural interface and nothing in `LlamaIndexRetrieval` imports `llamaindex`, so the tool works without the optional dependency. The `FilesRetrieval` half is real but gated on `ADK_SAMPLE_DOCS_DIR`, because it needs `llamaindex`, `@llamaindex/readers`, a configured embedding model and a directory of documents. The header says so where the workflow samples say `Run (offline, no API key)`, and the category README repeats it with the install command. The category README also records what CI does and does not do here: `tests/integration/docs_samples` resolves `SAMPLES_ROOT` to `samples/workflows`, so lint, Prettier, the license check and `ts:check:samples` reach this sample and nothing executes it. The sample links its guide by relative path rather than an adk.dev URL. Every adk.dev link in the repository is a `graphs/` section a sample ports, and this sample ports none.
Two reviewer corrections. The combined page documented three classes at once. A reader arrives at a guide from search, not at the set, so a page covering all three cannot tell them in its first paragraph whether they are on the right one. Each class now has its own guide, mirroring its source file, and each opens by naming its siblings with the reason to pick each. The client-side versus server-side comparison lives in the BaseRetrievalTool guide, because that is where a reader makes that choice; the other two link to it rather than repeating it. No content was dropped in the split -- the sections were redistributed. Samples now carry their own README, following the revised template: overview, sample inputs, a Mermaid graph, the techniques worth naming, how to run, and the guides to read next. The category README's single retrieval entry became three. The sample's header pointed at the combined guide, which no longer exists; it now points at the LlamaIndexRetrieval guide, the class it is built on. Every relative link in the new files was resolved against the tree, and both anchor targets in the FilesRetrieval guide were confirmed to exist. Prettier, the licence check and `tsc -p samples` are clean.
added 2 commits
September 4, 2026 14:44
…ured The sample asked for ADK_SAMPLE_DOCS_DIR and a LlamaIndex install before its FilesRetrieval half did anything, so a reader who cloned it got one tool and a paragraph explaining the other. Ship three documents beside the agent, default inputDir to them, and catch the optional-peer-dependency failure so the sample runs to completion either way. The guides' "Related samples" links now point at the sample's README rather than straight at agent.ts, and the README gains a Configuration section covering both knobs and how to swap in your own corpus.
adk-python has a matching sample, contributing/samples/integrations/ files_retrieval_agent, which indexes a data/ folder of two documents through a tool called search_documents. Reusing its corpus and its tool name costs nothing and makes the two samples comparable, which is the point of a parity port. The two documents state that ADK is a Python framework and that function tools are Python functions. Both are corrected here rather than carried over; the retrieval section also gains the two classes this port adds.
Collaborator
Author
|
Final look, and one CI note.
It is a harness failure, not an assertion: the test spawns a CLI subprocess and the subprocess dies on Windows. Nothing in the diff touches A2A, streaming, or process spawning. The run's own numbers say the same thing — 4112 tests passed, 0 failed, one suite failed to start. Every other check is green, including State of the PR. 18 files, +1607, −0.
Nothing outstanding from me. Merge when you are happy. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Ports the
tools/retrievalmodules from adk-pythonv0.1.0, additively:nothing already in adk-js is removed, downgraded or duplicated.
18 files, +1607, −0.
core/src/tools/vertex_rag_retrieval_tool.tsisuntouched and not in the diff.
What was missing, and what this adds
BaseRetrievalTool— the shared one-argumentquerydeclarationcore/src/tools/retrieval/base_retrieval_tool.tsLlamaIndexRetrieval— answers from any object withretrievecore/src/tools/retrieval/llama_index_retrieval.tsFilesRetrieval— indexes a directorycore/src/tools/retrieval/files_retrieval.tsVertexRagRetrievalToolDeliberately not added
VertexAiRagRetrievalas a class.VertexRagRetrievalToolalready ownsthat job. Adding a second class under another name would be duplication.
rag_corpora/similarity_top_k/vector_distance_thresholdparameters. adk-js takes the wholeVertexRagStore, which also carriesragRetrievalConfigandstoreContext.Matching v0.1.0 here would remove capability.
Where v0.1.0 and adk-js deliberately disagree
Recorded as explicit DIVERGENCE assertions in the parity test, so the test pins
what adk-js actually does and a future change to it is caught:
toolsDictregistration. v0.1.0's test asserts'rag_retrieval' not in tools_dict. adk-js registers every in-model tool byname on purpose (A hallucinated tool call aborts the whole invocation; onToolErrorCallback is unreachable for it #789) so an explicit call naming it can be routed.
name/descriptionper instance; adk-jshardcodes
vertex_rag_retrieval.{retrieval: {vertexRagStore}}for every model.Not included: v0.1.0's client-side retrieval path, which gates on
model.startswith('gemini-2')and falls back to a function declaration. Addingit means an optional constructor argument on
VertexRagRetrievalTool; happy tofollow up if wanted.
Interface, side by side with Python
Every public and protected member the port adds, against its
v0.1.0original.BaseRetrievalToolv0.1.0class BaseRetrievalTool(BaseTool)abstract class BaseRetrievalTool extends BaseTool_get_declaration() -> types.FunctionDeclaration_getDeclaration(): FunctionDeclarationabstractis explicit here because the class supplies norunAsync; Pythonleaves that implicit. The declaration body is the same — one required
querystring, described identically.
LlamaIndexRetrievalv0.1.0__init__(*, name, description, retriever: BaseRetriever)constructor(params: LlamaIndexRetrievalParams)self.retrieverreadonly retriever: LlamaIndexRetrieverasync run_async(*, args, tool_context) -> Anyoverride async runAsync({args}: RunAsyncToolRequest): Promise<unknown>Keyword-only arguments become one params object, which is how adk-js tools are
already constructed.
retrieverisreadonly. Its type is a structuralLlamaIndexRetrieverrather than llama_index'sBaseRetriever, sollamaindexstays an optional dependency — Python gets the same effect from
TYPE_CHECKING.FilesRetrievalv0.1.0__init__(*, name, description, input_dir), reading the directory and building the index inlinestatic async create(params: FilesRetrievalParams): Promise<FilesRetrieval>__init__constructor(params: FilesRetrievalConstructorParams), taking an already-builtretrieverself.input_dirreadonly inputDir: stringThe one interface change worth a close look.
SimpleDirectoryReader.loadDataand
VectorStoreIndex.fromDocumentsare both async in LlamaIndexTS and sync inPython, so that work cannot run in a constructor. Moving it to a static factory
also lets a caller supply a retriever built some other way.
New exported types, with no Python counterpart
Python refers to llama_index's types under
TYPE_CHECKING. TypeScript needsthem declared at exported positions, so they are exported rather than inlined:
LlamaIndexRetriever,LlamaIndexNode,LlamaIndexNodeWithScore,LlamaIndexRetrievalParams,FilesRetrievalParams,FilesRetrievalConstructorParams.Behaviour differing inside a matching signature
v0.1.0print(f'Loading data from {input_dir}')logger.info(...)retrieve(...)[0]raises a bareIndexErrorErrornaming the toolImportErrorfrom a module-level importErrornaming the package and the install commandTests
All 3 adk-python v0.1.0 retrieval tests are ported, each keeping the Python
function name verbatim so the original is greppable.
Source:
src/google/adk/tests/unittests/tools/retrieval/test_vertex_ai_rag_retrieval.py@v0.1.0Port:
core/test/tools/retrieval/vertex_ai_rag_retrieval_python_parity_test.tstest_vertex_rag_retrieval_for_gemini_1_xtest_vertex_rag_retrieval_for_gemini_1_x_with_another_function_tooltest_vertex_rag_retrieval_for_gemini_2_xv0.1.0 has 3 retrieval tests in total. The other three retrieval test files on
adk-python
maindid not exist at that tag. Plus 14 new tests for the addedmodules, in separate files.
Also included
docs/guides/tools/retrieval/, withdocs/guides/README.mdas the index. This createsdocs/guides/, followingadk-python's layout.
samples/tools/retrieval/, with a README and thedata/folder it indexes. That folder holds the same two documents adk-python's
contributing/samples/integrations/files_retrieval_agentindexes, read by atool of the same name,
search_documents. Its two Python-specific statementsare corrected rather than carried over.
Note
tests/integration/docs_samples/docs_samples_test.tsresolves its root tosamples/workflows, so the newsamples/tools/category is type-checked,linted and licence-checked but never executed. The category README says so.
Verification
ts:checkts:check:samples/lint/format:check/check_license.shdocs:checkdocs_samplesintegrationThe 4
ts:checkerrors are pre-existingCannot find module '@google/adk-devtools'in integration fixtures. Every TypeScript snippet in the guides was compiled
against the published types before being included.
The sample, in
adk web. Runningweb samples/toolswith nothingconfigured:
/list-appsreturns["retrieval"], the agent module loads, andcreating a session succeeds. The
FilesRetrievalhalf reports that it is offand names the command that turns it on, and the agent still starts with
employee_handbook.adk runon the same agent reaches its prompt the sameway. I did not drive a live model turn — this environment has no API key — so
the conversational path rests on the unit tests rather than on a manual check.