Build RAG retrieval with permissions

Applies to VAST Cluster 5.5.

Retrieval-augmented generation (RAG) on VAST: chunk documents, embed them, load the embeddings into a vector-indexed table, and retrieve the nearest chunks for each question, returning only chunks from documents the user may read. The worked example shows the permissions part end to end. The pipeline section after it says where each production step is documented.

Top: the pipeline sketch. Bottom: the worked example: the same search returns fewer chunks for the user whose policy has the row filter.

Worked example: keep restricted documents out of retrieval

A RAG application must not hand a user's question context from documents that user isn't allowed to read. With the row filter by S3 path, you tag each chunk with the S3 path of the document it came from. The cluster then returns a chunk only to users who may read that document. Your search SQL doesn't change. (Row and column security)

What your platform team sets up (Admin task: send to your platform team.) Two things:

an identity or bucket policy that denies s3:GetObject on the documents a user mustn't see;
for that user, the row and column security statement:

{
  "Effect": "RowColumnSecurity",
  "Resource": ["my-bucket/*"],
  "RowFilterByS3PathAuth": "True"
}

This example uses two users with the same access: both may read open/ but are denied restricted/. Only the second user's policy has the row filter. Their keys are in RAG_ALL_ACCESS_KEY_ID / RAG_ALL_SECRET_ACCESS_KEY (no row filter) and RAG_FILTERED_ACCESS_KEY_ID / RAG_FILTERED_SECRET_ACCESS_KEY (row filter). The data owner's keys are in the usual AWS_* variables, and the Query Engine endpoint is in QE_ENDPOINT.

1. Set up and store the source documents. The documents live in the bucket: a public report under open/ and salary data under restricted/:

import os

import adbc_driver_manager.dbapi as dbapi
import adbc_driver_vastdb
import boto3
import pyarrow as pa
import vastdb
from vastdb import VectorIndexSpec

BUCKET = os.environ["VASTDB_BUCKET"]
SCHEMA = "vs_rag_security"

s3 = boto3.client("s3", endpoint_url=os.environ["AWS_S3_ENDPOINT_URL"])
s3.put_object(Bucket=BUCKET, Key="open/report.txt", Body=b"Public annual report.")
s3.put_object(Bucket=BUCKET, Key="restricted/salary.csv", Body=b"role,salary\nCEO,1\nCFO,2\n")
print(sorted(o["Key"] for o in s3.list_objects_v2(Bucket=BUCKET)["Contents"]
             if o["Key"].startswith(("open/report", "restricted/salary"))))

['open/report.txt', 'restricted/salary.csv']

Decide both at creation. A vector index and the vastdb_s3_path_auth column can only be defined when you create the table (limits, Row and column security). Plan them together; see design a vector table.

2. Create the table. Add a vastdb_s3_path_auth column next to the chunk text and its embedding:

columns = pa.schema([
    ("id", pa.int64()),
    ("chunk", pa.utf8()),
    ("vastdb_s3_path_auth", pa.utf8()),
    ("embedding", pa.list_(pa.field("item", pa.float32(), nullable=False), 4)),
])

session = vastdb.connect()
with session.transaction() as tx:
    schema = tx.bucket(BUCKET).create_schema(SCHEMA, fail_if_exists=False)
    schema.create_table("chunks", columns, fail_if_exists=False,
                        vector_index=VectorIndexSpec(column="embedding", distance_metric="cosine"))
    print("Created", SCHEMA + ".chunks")

Created vs_rag_security.chunks

3. Tag each chunk with its document's path. Each value is a quoted "bucket/object" string. A chunk drawn from several documents lists them all, separated by commas. In a real pipeline the embeddings come from your model. Here, small fixed vectors stand in for them:

def path(key):
    return f'"{BUCKET}/{key}"'

chunks = pa.table({
    "id": [1, 2, 3, 4, 5],
    "chunk": ["Report: revenue grew", "Report: new offices", "Salary: CEO", "Salary: CFO", "Summary of both files"],
    "vastdb_s3_path_auth": [path("open/report.txt"), path("open/report.txt"),
                            path("restricted/salary.csv"), path("restricted/salary.csv"),
                            path("open/report.txt") + "," + path("restricted/salary.csv")],
    "embedding": [[1, 0, 0, 0], [0.9, 0.1, 0, 0], [0, 1, 0, 0], [0, 0.9, 0.1, 0], [0.5, 0.5, 0, 0]],
}, schema=columns)

with session.transaction() as tx:
    tx.bucket(BUCKET).schema(SCHEMA).table("chunks").insert(chunks)
print("Inserted", chunks.num_rows, "chunks")

Inserted 5 chunks

4. Search as each user. Run the same search, for a question about salaries, as both users:

def search(user, question_embedding, k=5):
    sql = f"""SELECT id, chunk FROM "{BUCKET}/{SCHEMA}".chunks
              ORDER BY array_cosine_distance(embedding::FLOAT[4], {question_embedding}::FLOAT[4])
              LIMIT {k}"""
    with dbapi.connect(driver=adbc_driver_vastdb.get_driver_path(), autocommit=True, db_kwargs={
            "vast.db.endpoint": os.environ["QE_ENDPOINT"],
            "vast.db.access_key": os.environ[f"{user}_ACCESS_KEY_ID"],
            "vast.db.secret_key": os.environ[f"{user}_SECRET_ACCESS_KEY"]}) as conn:
        with conn.cursor() as cur:
            cur.execute(sql)
            return sorted(cur.fetchall())

question = [0.0, 1.0, 0.0, 0.0]   # "What is the CEO's salary?"
for user in ("RAG_ALL", "RAG_FILTERED"):
    print(user)
    for row in search(user, question):
        print("  ", row)

RAG_ALL
   (1, 'Report: revenue grew')
   (2, 'Report: new offices')
   (3, 'Salary: CEO')
   (4, 'Salary: CFO')
   (5, 'Summary of both files')
RAG_FILTERED
   (1, 'Report: revenue grew')
   (2, 'Report: new offices')

The user without the row filter gets the salary chunks, even though their policy denies reading restricted/salary.csv. A GetObject deny alone doesn't hide rows. The user with the row filter gets only chunks from the open report. Chunk 5, which lists both files, is hidden too, because a row is returned only if the user may read every path it lists. The same filter applies to the Python SDK's select() for that user.

Quote every path. An unquoted path is rejected when you insert it:

try:
    bad = chunks.slice(0, 1).set_column(2, "vastdb_s3_path_auth", pa.array([f"{BUCKET}/open/report.txt"]))
    with session.transaction() as tx:
        tx.bucket(BUCKET).schema(SCHEMA).table("chunks").insert(bad)
except vastdb.errors.BadRequest as e:
    print(e.status, e.code, e.message)

400 InvalidArgument Invalid Argument

5. Generate the answer. Pass the retrieved chunks to your language model as context. A user without access to the salary file gets an answer built only from what they may read, so the model can't leak it. This step depends on your model provider. Here's an example with an OpenAI-compatible endpoint:

from sentence_transformers import SentenceTransformer
from openai import OpenAI

model = SentenceTransformer("nomic-ai/nomic-embed-text-v1.5", trust_remote_code=True)
question_text = "What is the salary of the CEO?"
question = model.encode([question_text])[0].tolist()   # same model as the stored chunks

context = "\n".join(chunk for _, chunk in search("RAG_FILTERED", question, k=5))
client = OpenAI(base_url="<your LLM endpoint>", api_key=os.environ["LLM_API_KEY"])
answer = client.chat.completions.create(
    model="<model>",
    messages=[{"role": "system", "content": "Answer from the context. If it isn't there, say you don't know."},
              {"role": "user", "content": f"Context: {context}\n\nQuestion: {question_text}"}])
print(answer.choices[0].message.content)

In a real pipeline the embedding dimension comes from the model (VECTOR(FLOAT, 768) for nomic-embed-text v1.5), and the chunks come from your document loader: for example, one chunk per section of a PDF and one per row of a CSV.

6. Clean up.

with session.transaction() as tx:
    schema = tx.bucket(BUCKET).schema(SCHEMA)
    for table in schema.tables():
        table.drop()
    schema.drop()
s3.delete_object(Bucket=BUCKET, Key="open/report.txt")
s3.delete_object(Bucket=BUCKET, Key="restricted/salary.csv")
print("Dropped", SCHEMA, "and the two source files")

Dropped vs_rag_security and the two source files

The pipeline in production

A sketch of the full pipeline. Each step links to the page with the details.

Chunk and embed. Split your documents into chunks and run each through an embedding model, with Spark, a VAST DataEngine pipeline or your own code. Write Parquet files with one row per chunk: an ID, the text, the embedding (a fixed-length list of 32-bit floats with the model's dimension), and any columns you'll filter on, such as language.
Land the files. Upload the Parquet files to an S3 bucket on the cluster, through the ordinary (protocols) endpoints.
Create the table. Create a table with a vector index on the embedding column, using the metric the model was trained for, for example cosine. Add the vastdb_s3_path_auth column now if you'll filter by document permissions (create an index):
```
CREATE TABLE "kb/rag".chunks (
    id BIGINT,
    language VARCHAR,
    text VARCHAR,
    embedding VECTOR(FLOAT, 768)
)
VECTOR CLUSTER BY embedding (metric = 'cosine');
```
Import. Load the files with the Python SDK's import_files() (ingest). If you insert Arrow tables with the SDK's insert() instead, keep each batch to a few thousand rows when the table has vastdb_s3_path_auth: with SDK 2.1.0 a batch larger than the SDK's request size fails with trying to insert to a table without values for column=vastdb_s3_path_auth.
Wait for clustering. The chunks can be searched straight away, but searches are fastest once the table's clustering grade is A or A+ (check the index first).
Retrieve. For each question, embed it with the same model and search through the Query Engine endpoint with ADBC: array_cosine_distance in ORDER BY, your filters in WHERE, and LIMIT k (query design). In this fragment, replace <question_embedding> with the question's vector:
```
SELECT id, text,
       array_cosine_distance(embedding::FLOAT[768], <question_embedding>::FLOAT[768]) AS dist
FROM "kb/rag".chunks
WHERE language = 'en'
ORDER BY dist
LIMIT 5;
```
Generate. Build the prompt for the language model from the returned IDs and text.

For a complete application built this way, see the Research assistant agent: DataEngine ingests uploaded documents, VAST DataBase holds the vectors and the conversation history, and a chat agent answers on top (reference application on GitHub, not verified against these docs).

Don't run the search through Trino or Spark: their connectors don't use the vector index (limits).

To return only chunks the user may read, tag each chunk as in the worked example and see row filter by S3 path. If your application searches with one service account on behalf of many users, pass each user's name in the vast.db.end_user connection option: the search then returns only the chunks that user may read.

Build RAG retrieval with permissions

Worked example: keep restricted documents out of retrieval

The pipeline in production

See also