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_quality graded 1–5 (aim ≥ 4)

End-to-end flow

1. Terraform (make setup-infra) enables APIs, creates a pipeline GCS bucket, a runner service account + IAM, and a VS 2.0 Collection whose text_embedding field is configured to auto-embed with gemini-embedding-001.
2. 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).
3. Agent (make playground) answers questions; its retrieve_docs tool 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)

1. data_ingestion/data_ingestion_pipeline/components/process_data.py — the core transform. A KFP @component that: fetches data (a synthetic Q&A CTE in BigQuery — the documented place to swap in your own source), converts HTML→ markdown (markdownify), chunks with RecursiveCharacterTextSplitter (chunk_size=1500, chunk_overlap=20), writes a time-partitioned incremental table, and builds a deduped table (latest creation_timestamp per question_id). Note the deterministic chunk_id = question_id__<index> that makes re-ingestion idempotent.
2. data_ingestion/data_ingestion_pipeline/components/ingest_data.py — how data lands in VS 2.0: reads the deduped table and calls batch_create_data_objects with empty vectors: {} (auto-embedding), capped at 250 objects/batch, swallowing AlreadyExists so re-runs skip existing chunks.
3. data_ingestion/data_ingestion_pipeline/pipeline.py — KFP wiring: process_data → ingest_data with set_retry. Small but shows the DAG and parameters.
4. data_ingestion/data_ingestion_pipeline/submit_pipeline.py — run modes: --local (KFP SubprocessRunner, what make data-ingestion uses) vs remote (aiplatform.PipelineJob.submit), plus --cron-schedule/--schedule-only for PipelineJobSchedule. Read parse_args for the env-var contract.
5. data_ingestion/README.md — the ingestion runbook (local + remote).

Terraform (provisioning + the embedding config)

6. infra/terraform/scripts/setup_vector_search_collection.py — the most important infra file. Defines the Collection's data_schema (question_id, text_chunk, full_text_md) and the vector_schema.text_embedding auto-embedding config: dense_vector dimensions=3072, vertex_embedding_config model_id=gemini-embedding-001, text_template="{text_chunk}", task_type=RETRIEVAL_DOCUMENT. Idempotent via a get_collection check. (delete_vector_search_collection.py is the teardown.)
7. infra/terraform/vector_search.tf — the pipeline bucket plus a null_resource that runs the create script on apply and the delete script on destroy (local-exec + triggers). Good example of wrapping an API not yet in the provider.
8. infra/terraform/vector_search_iam.tf + vector_search_variables.tf — runner service account and its pipelines_roles (BigQuery, Storage, Vertex AI, vectorsearch.dataObjectWriter), assigned via setproduct.
9. infra/terraform/apis.tf — enabled services (note vectorsearch.googleapis.com) and the Vertex AI service identity. locals.tf, variables.tf, datastore_outputs.tf, vars/env.tfvars are the supporting scaffolding.

Agent (thin layer — read last)

10. app/retrievers.py — search_collection(): SemanticSearch with search_field="text_embedding", task_type="RETRIEVAL_QUERY", top_k, and output_fields (question_id, text_chunk, full_text_md). These fields must match the schema set in setup_vector_search_collection.py — that symmetry between write-side and read-side is the key thing to notice.
11. app/agent.py — ADK Agent (Gemini gemini-flash-latest) exposing the retrieve_docs tool (wraps search_collection with error handling); reads VECTOR_SEARCH_COLLECTION from env. app/__init__.py exports app.

Data handling

• Chunking: recursive character splitting (1500/20). Tune in pipeline.py defaults or process_data args.
• 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_embedding config. The read path uses the same model implicitly via RETRIEVAL_QUERY.
• Idempotency: deterministic chunk_id + AlreadyExists skip 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-vectorsearch is 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 when INTEGRATION_TEST=TRUE (set by make test).
• app/agent.py loads .env and resolves the project from GOOGLE_CLOUD_PROJECT, falling back to google.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_data to ingest real documents.
• tests/eval/eval_config.yaml uses the agents-cli eval format (LLM-judge custom_response_quality + a turn-count function); it is not wired into make 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 (including scripts/) and set project_id in vars/env.tfvars. Requires uv for the local-exec provisioner scripts.
• data_ingestion/ is an independent, self-contained KFP project (its own pyproject.toml + uv.lock) — copy the directory and run it standalone. Point it at your Collection via --project / --region / --collection-id, and swap the fetch_sample_data CTE in process_data.py for your own source.
• deployment/cloudbuild.yaml schedules the ingestion pipeline in CI via submit_pipeline.py (a recurring Vertex AI PipelineJobSchedule) — copy it and set the _* substitutions.
• There is no code coupling to app/: the agent reaches the Collection purely through the VECTOR_SEARCH_COLLECTION env var.