Imported from google/adk-samples (
core/rag-vector-search/AGENTS.md). Install upstream withnpx skills add google/adk-samples --skill rag-vector-search. Copyright stays with the author.
RAG Agent — Vector Search
Intent
A clone-and-study RAG agent grounded on Vertex AI Vector Search 2.0, with a Kubeflow Pipelines (KFP) pipeline that loads, chunks, and ingests documents into a Collection that auto-generates embeddings server-side. The hard/interesting parts are Terraform (Collection + auto-embedding config + IAM) and data handling (the ingestion pipeline); the agent itself is a thin ADK wrapper around a semantic-search tool.
When To Use
- The user wants explicit control over chunking/ingestion (a custom KFP pipeline) instead of a fully-managed connector.
- The user needs semantic/vector search over their own documents with server-side auto-embedding (Vector Search 2.0).
- The user wants a reproducible BigQuery-staged ingestion + Terraform example.
Eval
- Scenarios Path:
tests/eval/datasets/(config:tests/eval/eval_config.yaml) - Minimum Score: not enforced —
custom_response_qualitygraded 1–5 (aim ≥ 4)
End-to-end flow
- Terraform (
make setup-infra) enables APIs, creates a pipeline GCS bucket, a runner service account + IAM, and a VS 2.0 Collection whosetext_embeddingfield is configured to auto-embed withgemini-embedding-001. - Ingestion (
make data-ingestion) runs a KFP pipeline that fetches docs, converts HTML→markdown, chunks them, stages + dedups them in BigQuery, then batch-creates data objects in the Collection. Embeddings are generated by the Collection, not the pipeline (objects are sent with empty vectors). - Agent (
make playground) answers questions; itsretrieve_docstool runs semantic search over the Collection and feeds chunks back to Gemini.
Most interesting files to study (in order)
Data ingestion (the bulk of the logic)
data_ingestion/data_ingestion_pipeline/components/process_data.py— the core transform. A KFP@componentthat: fetches data (a synthetic Q&A CTE in BigQuery — the documented place to swap in your own source), converts HTML→ markdown (markdownify), chunks withRecursiveCharacterTextSplitter(chunk_size=1500,chunk_overlap=20), writes a time-partitioned incremental table, and builds a deduped table (latestcreation_timestampperquestion_id). Note the deterministicchunk_id=question_id__<index>that makes re-ingestion idempotent.data_ingestion/data_ingestion_pipeline/components/ingest_data.py— how data lands in VS 2.0: reads the deduped table and callsbatch_create_data_objectswith emptyvectors: {}(auto-embedding), capped at 250 objects/batch, swallowingAlreadyExistsso re-runs skip existing chunks.data_ingestion/data_ingestion_pipeline/pipeline.py— KFP wiring:process_data→ingest_datawithset_retry. Small but shows the DAG and parameters.data_ingestion/data_ingestion_pipeline/submit_pipeline.py— run modes:--local(KFPSubprocessRunner, whatmake data-ingestionuses) vs remote (aiplatform.PipelineJob.submit), plus--cron-schedule/--schedule-onlyforPipelineJobSchedule. Readparse_argsfor the env-var contract.data_ingestion/README.md— the ingestion runbook (local + remote).
Terraform (provisioning + the embedding config)
infra/terraform/scripts/setup_vector_search_collection.py— the most important infra file. Defines the Collection'sdata_schema(question_id,text_chunk,full_text_md) and thevector_schema.text_embeddingauto-embedding config:dense_vectordimensions=3072,vertex_embedding_configmodel_id=gemini-embedding-001,text_template="{text_chunk}",task_type=RETRIEVAL_DOCUMENT. Idempotent via aget_collectioncheck. (delete_vector_search_collection.pyis the teardown.)infra/terraform/vector_search.tf— the pipeline bucket plus anull_resourcethat runs the create script on apply and the delete script ondestroy(local-exec+triggers). Good example of wrapping an API not yet in the provider.infra/terraform/vector_search_iam.tf+vector_search_variables.tf— runner service account and itspipelines_roles(BigQuery, Storage, Vertex AI,vectorsearch.dataObjectWriter), assigned viasetproduct.infra/terraform/apis.tf— enabled services (notevectorsearch.googleapis.com) and the Vertex AI service identity.locals.tf,variables.tf,datastore_outputs.tf,vars/env.tfvarsare the supporting scaffolding.
Agent (thin layer — read last)
app/retrievers.py—search_collection():SemanticSearchwithsearch_field="text_embedding",task_type="RETRIEVAL_QUERY",top_k, andoutput_fields(question_id,text_chunk,full_text_md). These fields must match the schema set insetup_vector_search_collection.py— that symmetry between write-side and read-side is the key thing to notice.app/agent.py— ADKAgent(Geminigemini-flash-latest) exposing theretrieve_docstool (wrapssearch_collectionwith error handling); readsVECTOR_SEARCH_COLLECTIONfrom env.app/__init__.pyexportsapp.
Data handling
- Chunking: recursive character splitting (1500/20). Tune in
pipeline.pydefaults orprocess_dataargs. - BigQuery staging: an append-only incremental table feeds a replace-mode
deduped table (one row per chunk from the latest ingestion of each
question_id) — the deduped table is the ingestion source of truth. - Embeddings: never computed in the pipeline. Objects ship with empty vectors
and VS 2.0 embeds them server-side using the Collection's
text_embeddingconfig. The read path uses the same model implicitly viaRETRIEVAL_QUERY. - Idempotency: deterministic
chunk_id+AlreadyExistsskip means re-runs don't duplicate. Caveat: if a doc's chunk count shrinks, orphaned chunks from a prior run are not deleted.
Gotchas / things to know
- Order matters:
make setup-infra→make data-ingestion→ query. Querying before ingestion returns "No relevant documents found." google-cloud-vectorsearchis the v1beta preview API (vectorsearch_v1beta).- The integration test (
tests/integration/test_agent.py) makes a live Gemini call; it is skipped without ADC, and the retriever is mocked whenINTEGRATION_TEST=TRUE(set bymake test). app/agent.pyloads.envand resolves the project fromGOOGLE_CLOUD_PROJECT, falling back togoogle.auth.default()at import when it is unset — so credential-less imports fail unless that var is set (the test imports lazily to allow skipping).- The default dataset is synthetic Q&A generated inline in BigQuery; replace
the CTE in
process_data.fetch_sample_datato ingest real documents. tests/eval/eval_config.yamluses theagents-clieval format (LLM-judgecustom_response_quality+ a turn-count function); it is not wired intomake test.
Where to run things
Makefile targets: make setup-infra (provision), make install,
make data-ingestion (run the KFP pipeline locally), make playground
(local ADK web UI), make test, make lint.
Eval lives under tests/eval/ (eval_config.yaml + datasets/).
Reuse (copy as-is)
infra/terraform/(Collection + IAM) is self-contained — copy the whole directory (includingscripts/) and setproject_idinvars/env.tfvars. Requiresuvfor thelocal-execprovisioner scripts.data_ingestion/is an independent, self-contained KFP project (its ownpyproject.toml+uv.lock) — copy the directory and run it standalone. Point it at your Collection via--project/--region/--collection-id, and swap thefetch_sample_dataCTE inprocess_data.pyfor your own source.deployment/cloudbuild.yamlschedules the ingestion pipeline in CI viasubmit_pipeline.py(a recurring Vertex AIPipelineJobSchedule) — copy it and set the_*substitutions.- There is no code coupling to
app/: the agent reaches the Collection purely through theVECTOR_SEARCH_COLLECTIONenv var.