Distributed RAG setup TECH PREVIEW

Distributed RAG consists of two components:

• pg_dist_rag PostgreSQL extension

  Manages Retrieval-Augmented Generation (RAG) pipelines from SQL. It registers document sources (currently accessed via S3 buckets), coordinates distributed preprocessing and embedding generation, and stores vectors in pgvector indexes backed by YugabyteDB.

  With pg_dist_rag, you can do the following:

  • Point a vector index at a document source instead of building custom ETL pipelines.
  • Chunk documents and generate embeddings using a configured AI provider.
  • Monitor pipeline progress and retry failed documents from SQL views.
  • Query generated embeddings with standard pgvector similarity search.

• RAG workers

  Long-running Python services that do the heavy lifting: crawling sources, parsing, chunking, and generating embeddings. Workers run outside the database, on dedicated nodes or VMs, so pipeline failures and resource spikes never affect the database itself. AI workloads and database workloads scale independently.

Set up YugabyteDB

Prerequisites

• YugabyteDB v2026.1.1 or later
• pgvector extension (vector type support)
• An OpenAI API key
• Cloud credentials (AWS S3) when reading documents from object storage

Enable the extension

Install pgvector first, then create the pg_dist_rag extension:

CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS pg_dist_rag;

This creates the dist_rag schema with tables, types, functions, and views for managing sources, indexes, documents, pipelines, and the work queue.

Build a vector index

The typical workflow has four steps: create sources, initialize a vector index, optionally add more sources, and build the index, all managed using the following four functions.

1. Create a source

Register a document source URI. Currently, only S3 is supported. This queues a CREATE_SOURCE task in the work queue.

-- Minimal: just a URI
SELECT dist_rag.create_source(
  r_source_uri := 's3://my-bucket/documents/'
);

Returns a UUID source ID.

2. Initialize a vector index

Create a named vector index, optionally associating it with existing sources. This creates a backing table in the target schema and an HNSW index on the embeddings column.

-- Create an index with sources attached
SELECT dist_rag.init_vector_index(
  r_index_name := 'my_knowledge_base',
  r_sources := ARRAY['<source_uuid_1>', '<source_uuid_2>']::UUID[],
  r_embedding_model_params := '{"dimensions": 1536}'::jsonb,
  r_ai_provider := 'OPENAI'
);

-- Create an empty index (add sources later)
SELECT dist_rag.init_vector_index(
  r_index_name := 'my_empty_index',
  r_embedding_model_params := '{"dimensions": 1536}'::jsonb
);

Returns a UUID vector index ID.

3. Add a source to an existing index

Use add_source_to_index() to attach additional sources to an already-created vector index, optionally with custom chunking parameters.

SELECT dist_rag.add_source_to_index(
  r_index_id := '<index_uuid>',
  r_source_id := '<source_uuid>',
  r_chunk_params := '{"chunk_size": 512, "chunk_overlap": 50, "strategy": "recursive"}'::jsonb
);

4. Build the index

Kick off preprocessing for all documents across all sources in an index. Each document gets a PREPROCESS task queued in the work queue. RAG agent workers claim tasks from the queue and process documents in parallel.

Provide exactly one of r_index_id or r_index_name:

SELECT dist_rag.build_index(r_index_id := '<index_uuid>');

SELECT dist_rag.build_index(r_index_name := 'my_knowledge_base');

Schema reference

Tables

Enum types

Views

Deploy RAG workers

Prerequisites

• Network connectivity to the YugabyteDB database (port 5433), S3, and the embedding provider API
• Python 3.11
• For PDF workers, the OCR system packages. For example, on Ubuntu: tesseract-ocr, poppler-utils, and libpq-dev.

Setup

Create the pg_dist_rag extension before starting workers. Workers poll the dist_rag.work_queue table, which is created by CREATE EXTENSION pg_dist_rag.

To deploy a worker:

1. Change to the install directory:

   cd <yugabytedb-install-dir>/python/ai/rag_agent
   

2. Create a virtual environment and install dependencies:

   uv venv --python=3.11
   source .venv/bin/activate
   uv pip install -r requirements.txt
   

3. Configure the worker:

   export YUGABYTEDB_CONNECTION_STRING="postgresql://yugabyte:<password>@<node-address>:5433/yugabyte"
   export OPENAI_API_KEY=<your-openai-api-key>
   export AWS_REGION=us-east-1
   export AWS_ACCESS_KEY_ID=<key-id>           # omit for public buckets
   export AWS_SECRET_ACCESS_KEY=<secret-key>   # omit for public buckets
   

4. Start the worker:

   python __main__.py
   

   To start a PDF worker, additionally set the worker type before starting:

   export WORKER_DOCUMENT_TYPE=PDF
   

Environment variables

Workers are configured entirely using the following environment variables.

Scale workers

Ingestion throughput scales horizontally with the number of workers: each worker claims a different document from the work queue, so adding workers increases parallelism with no configuration changes.

Because workers are decoupled from the database, you can:

• Run zero workers when no ingestion is happening; queued work resumes when a worker starts.
• Add TEXT and PDF workers independently, sized to the mix of documents in your sources.
• Place PDF workers on GPU-optimized instances without changing your database node shapes.

Monitor pipelines

Two views are available for observing pipeline progress and statistics.

Detailed per-document pipeline status:

SELECT index_name, document_name, pipeline_status, chunks_processed,
       embeddings_persisted, current_step, last_error_message
FROM dist_rag.vector_index_pipeline_details
WHERE index_name = 'my_knowledge_base';

Aggregated stats per document:

SELECT index_name, document_name, calls, total_chunks_processed,
       total_embeddings_persisted, completion_rate_percent
FROM dist_rag.pipeline_stats
WHERE index_name = 'my_knowledge_base';

Query embeddings

After a pipeline completes, embeddings are stored in the backing table created for the vector index (for example, public.my_knowledge_base). Query it using standard pgvector operators:

SELECT id, chunk_text, metadata_filters,
       embeddings <=> '[0.1, 0.2, 0.3]'::vector AS distance
FROM public.my_knowledge_base
ORDER BY embeddings <=> '[0.1, 0.2, 0.3]'::vector
LIMIT 10;

Each row includes:

Combine metadata filters with vector similarity search:

SELECT chunk_text, embeddings <=> $1 AS distance
FROM public.my_knowledge_base
WHERE metadata_filters @> '{"type": "documentation"}'::jsonb
ORDER BY distance
LIMIT 10;

Complete example

The following shows an example end-to-end:

CREATE EXTENSION IF NOT EXISTS vector;
CREATE EXTENSION IF NOT EXISTS pg_dist_rag;

-- Create two document sources, capturing the returned IDs
SELECT dist_rag.create_source(
  r_source_uri := 's3://company-docs/engineering/',
  r_metadata := '{"team": "engineering", "access": "internal"}'::jsonb
) AS eng_source_id \gset

SELECT dist_rag.create_source(
  r_source_uri := 's3://company-docs/product/',
  r_metadata := '{"team": "product", "access": "internal"}'::jsonb
) AS product_source_id \gset

-- Initialize a vector index with both sources
SELECT dist_rag.init_vector_index(
  r_index_name := 'engineering_kb',
  r_sources := ARRAY[:'eng_source_id', :'product_source_id']::UUID[],
  r_ai_provider := 'OPENAI',
  r_embedding_model_params := '{"dimensions": 1536, "model": "text-embedding-ada-002"}'::jsonb
);

-- Build the index (queues all documents for preprocessing)
SELECT dist_rag.build_index(r_index_name := 'engineering_kb');

-- Monitor progress
SELECT index_name, document_name, pipeline_status, chunks_processed, current_step
FROM dist_rag.vector_index_pipeline_details
WHERE index_name = 'engineering_kb';

SELECT document_name, calls, total_chunks_processed, completion_rate_percent
FROM dist_rag.pipeline_stats
WHERE index_name = 'engineering_kb';

Learn more

• pg_dist_rag extension
• RAG agent source code