setup.sql

-- Build the hello_dist_rag knowledge base from ysqlsh instead of load.py.
-- Everything load.py does is these few SQL statements; run them with:
--
--   ./bin/ysqlsh -U yugabyte -f setup.sql
--
-- A RAG worker must be running (see README.md) or the queued work sits idle.

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

-- 1. Register the document source and capture its UUID.
SELECT dist_rag.create_source(
  r_source_uri := 's3://yugabytedb-sample-data/hello-dist-rag/'
) AS source_id \gset

-- 2. Create the vector index. This creates the backing table
--    public.hello_dist_rag and its ybhnsw vector index automatically.
--    Chunking config must be explicit (the worker rejects empty
--    chunk_params); "args" is a JSON string of LangChain splitter kwargs.
SELECT dist_rag.init_vector_index(
  r_index_name := 'hello_dist_rag',
  r_sources := ARRAY[:'source_id']::UUID[],
  r_ai_provider := 'OPENAI',
  r_embedding_model_params := '{"model": "text-embedding-3-small", "dimensions": 1536}'::jsonb,
  r_chunk_params := '{"splitter": "recursive_character", "args": "{\"chunk_size\": 1000, \"chunk_overlap\": 100}"}'::jsonb
);

-- Wait until a worker finishes crawling the source (status = COMPLETED),
-- re-running this until it flips:
SELECT status FROM dist_rag.sources WHERE id = :'source_id';

-- 3. Queue every discovered document for chunking and embedding.
SELECT dist_rag.build_index(r_index_name := 'hello_dist_rag');

-- Monitor progress until every document shows COMPLETED:
SELECT document_name, document_status, chunks_processed,
       embeddings_persisted, last_error_message
FROM dist_rag.vector_index_pipeline_details
WHERE index_name = 'hello_dist_rag';

-- The finished knowledge base is a plain table; query it with pgvector:
-- SELECT chunk_text FROM public.hello_dist_rag LIMIT 5;