hello-dist-rag

The hello_rag sample, rebuilt on the pg_dist_rag YugabyteDB extension.

In hello_rag, the application loads documents, chunks them, calls the OpenAI embeddings API, and inserts vectors into a hand-made table — about 100 lines of Python plus llama-index. In this sample, all of that is replaced by three SQL calls:

SELECT dist_rag.create_source(r_source_uri := 's3://yugabytedb-sample-data/hello-dist-rag/');
SELECT dist_rag.init_vector_index(r_index_name := 'hello_dist_rag', ...);
SELECT dist_rag.build_index(r_index_name := 'hello_dist_rag');

A RAG worker shipped with YugabyteDB crawls the S3 bucket, chunks and embeds each document, and writes the vectors into a table that pg_dist_rag creates for you. The question-answering side (ask.py) works exactly like hello_rag's question.py.

What's in the box

File	Purpose
load.py	Builds the knowledge base: the three pg_dist_rag SQL calls, plus progress reporting from the dist_rag monitoring views.
ask.py	Interactive Q&A: embeds your question, runs a pgvector similarity search, and answers with an OpenAI chat model.
setup.sql	The same knowledge-base build as load.py, as plain SQL for ysqlsh.
data/	The Paul Graham essay used by the sample. Upload it to your own S3 bucket if you don't use the default public one.

Note what's not here: no document readers, no chunker, no embedding calls during ingestion, no CREATE TABLE, and no llama-index dependency.

Prerequisites

• YugabyteDB 2026.1 or later with the pg_dist_rag extension (build from yugabyte/yugabyte-db master while the extension is in Tech Preview).
• Python 3.11 (for the RAG worker) and Python 3 for this app.
• An OpenAI API key.
• AWS credentials in the worker's environment. The worker uses them to list the source bucket; any valid credentials work for the public sample bucket. (Downloading the documents themselves falls back to public HTTPS URLs.)

1. Start YugabyteDB with pg_dist_rag enabled

./bin/yugabyted start --advertise_address=127.0.0.1
./bin/ysqlsh -U yugabyte -c "CREATE EXTENSION IF NOT EXISTS vector;" \
                         -c "CREATE EXTENSION IF NOT EXISTS pg_dist_rag;"

No table or index DDL — pg_dist_rag creates the vector table and its ybhnsw index when you initialize the vector index.

2. Start a RAG worker

The worker ships with YugabyteDB under python/ai/rag_agent. It polls the dist_rag.work_queue table and does the crawling, chunking, and embedding:

export YUGABYTEDB_CONNECTION_STRING="postgresql://yugabyte:yugabyte@127.0.0.1:5433/yugabyte"
export OPENAI_API_KEY='your OpenAI key'
export AWS_REGION=us-east-1
export AWS_ACCESS_KEY_ID='your AWS key id'
export AWS_SECRET_ACCESS_KEY='your AWS secret'

cd <yugabytedb-install>/python/ai/rag_agent
python3.11 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt

python __main__.py

Leave it running. Start more workers (on any machine that can reach the database) to process documents in parallel.

3. Run the app

git clone https://github.com/YugabyteDB-Samples/hello-dist-rag.git
cd hello-dist-rag
python3 -m venv .venv && source .venv/bin/activate
pip install -r requirements.txt
export OPENAI_API_KEY='your OpenAI key'

python load.py   # build the knowledge base (three SQL calls + monitoring)
python ask.py    # ask questions against it

load.py and ask.py accept these environment variables (all optional):

Variable	Default	Purpose
YB_CONNECTION_STRING	postgresql://yugabyte:yugabyte@127.0.0.1:5433/yugabyte	Database connection.
SOURCE_URI	s3://yugabytedb-sample-data/hello-dist-rag/	S3 prefix to ingest. Point it at your own bucket to index your own documents.
INDEX_NAME	hello_dist_rag	Vector index (and backing table) name.
EMBEDDING_MODEL	text-embedding-3-small	OpenAI embedding model, used for documents and questions.
EMBEDDING_DIMENSIONS	1536	Embedding width; sizes the vector(N) column.
CHAT_MODEL (ask.py)	gpt-4o	Chat model that generates the answers.

To use your own documents, upload text/Markdown/HTML/CSV/JSON/XML/PDF files to an S3 prefix and set SOURCE_URI. (PDFs need a worker started with WORKER_DOCUMENT_TYPE=PDF.)

How it works

load.py ──SQL──▶ dist_rag.create_source ─┐
                 dist_rag.init_vector_index ─┤   dist_rag.work_queue
                 dist_rag.build_index ───────┘        │ poll
                                                      ▼
                                              RAG worker(s)
                                     crawl S3 → chunk → embed (OpenAI)
                                                      │
                                                      ▼
ask.py ──pgvector search──▶ public.hello_dist_rag (chunk_text, embeddings)

load.py watches dist_rag.sources, dist_rag.documents, and the dist_rag.vector_index_pipeline_details view to report progress; the same views are queryable from ysqlsh at any time.

Learn more

• pg_dist_rag extension documentation
• Hello RAG tutorial — the do-it-yourself version this sample replaces
• pgvector extension