# 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.
