Build a RAG system

Prerequisites

Before you begin, ensure that your environment meets the following requirements:

• Engine: PostgreSQL 16 (kernel minor version 2.0.16.10.11.0 or later).

• Node: Add an AI node with GPU specifications and set a database account for the AI node. For more information, see Add and manage AI nodes.

  Note

  • If you added an AI node when you created the cluster, you can directly set the database account for the AI node.

  • The database account for the AI node must have read and write permissions to access the target database.

• Endpoint: Use the cluster endpoint to connect to the PolarDB cluster.

How it works

The core RAG process in PolarDB involves processing and storing external documents, then generating answers by combining vector search with a large language model (LLM).

1. Data preparation and loading: Load your source documents into a database table.

2. Document chunking: Call the polar_ai.ai_rag_chunking function to split long documents into smaller, semantically complete text chunks. This improves retrieval accuracy.

3. Embedding: Call the polar_ai.ai_rag_text_embedding function to convert each text chunk into a high-dimensional vector, which captures its semantic information, and store it with the original text.

4. Create a vector index: Create an HNSW index on the vector column to significantly accelerate similarity searches.

5. Retrieval: When a user asks a question, first embed the question itself. Then, perform a vector similarity search in the database to quickly recall the text chunks that are most relevant to the question.

6. rerank: To further improve relevance, call the polar_ai.ai_rag_rerank function to perform a secondary ranking on the initially recalled text chunks and select the most critical information.

7. Generate answer: Use the top reranked text chunks as context. Submit this context, along with the user's original question, to the large language model (LLM). The model then generates a final, precise, natural-language answer based on the provided context.

Prepare the environment

This section covers the prerequisites for using the RAG feature.

1. Log on to the cluster with a privileged account.

2. Get the node token.

   Go to the PolarDB console. On the details page of your target cluster, in the Database Nodes section, find the AI node, click View Node Token, and then record the Node Token.

3. Create the polar_ai extension and configure the node token.

   Run the following commands in your database to create the extension and set the key to access large model services.

   Important

   Ensure you have set a database account for the AI node.

   -- Create the extension.
   CREATE EXTENSION polar_ai;
   
   -- Set your key. Replace sk-xxx with your key for the large model service.
   SELECT polar_ai._ai_nl2sql_alter_token('sk-xxx');

4. Create the vector extension.

   -- Provides vector data types, vector indexes, and similarity calculation capabilities.
   CREATE EXTENSION vector;

Deploy the RAG service

After you set the database account for the AI node, you can use any account to deploy the RAG service. Run the following command to load and start the models and services required for RAG on the AI node:

SELECT polar_ai.ai_nl2sql_deployModel('RAG');

Prepare data

This section uses a PDF document as an example to show how to import, chunk, and embed it in the database.

1. Log on to the cluster.

2. Create data tables to store the source document and the processed data. Run the following SQL statements to create two tables.

   • file_content: Stores the binary content of the uploaded source document.

   • file_chunk: Stores the text content of the document chunks and their corresponding vectors.

   -- Stores the source document
   CREATE TABLE file_content (
     id SERIAL PRIMARY KEY, -- Unique identifier for the document
     name TEXT,             -- Document name
     content BYTEA          -- Binary content of the document
   );
   
   -- Stores chunks and vectors
   CREATE TABLE file_chunk (
     file_id INTEGER,       -- Associated document ID
     chunk_id INTEGER,      -- Chunk ID within the document
     chunk_content TEXT,    -- Text content of the chunk
     embedding VECTOR(1024) -- Chunk vector. The dimension must match the model output.
   );

3. Load local PDF documents into the file_content table. We recommend using one of the following two methods to import the documents.

   • Method 1 (Recommended): Use the polar_ai_util extension to load the binary content of a file from OSS.

     Note

     You must use a Privileged Account to install the extension.

     CREATE EXTENSION polar_ai_util; -- Used to load documents from OSS
     
     INSERT INTO file_content(name, content)
     SELECT 'PolarAI.pdf', polar_ai.AI_LOADFILE('oss://<access-key-id>:<access-key-secret>@oss-cn-beijing-internal.aliyuncs.com/your-bucket/path/to/PolarAI.pdf');

   • Method 2: Use a client program to write the data. You can write a script in any programming language that supports PostgreSQL, such as Python with the psycopg2 library, to read the local file content and INSERT it into the file_content table in bytea format.

     Python sample script

     // You must first install the psycopg2-binary extension library.
     
     import psycopg2
     import os
     
     # --- Modify your configuration here ---
     # Database connection information
     DB_CONFIG = {
         "host": "your_database_host",  # Example: "localhost" or an IP address
         "dbname": "your_database_name",
         "user": "your_username",
         "password": "your_password",
         "port": "5432"  # The default is 5432
     }
     
     # Local path to the file you want to upload
     LOCAL_FILE_PATH = "/path/to/your/PolarAI.pdf"
     # --- End of configuration ---
     
     
     def upload_file_to_db(file_path, db_params):
         """
         Connects to PostgreSQL, reads a local file, and inserts it into the database as BYTEA.
         """
         # Check if the file exists
         if not os.path.exists(file_path):
             print(f"Error: File '{file_path}' not found.")
             return
     
         conn = None
         try:
             # 1. Connect to the PostgreSQL database
             print("Connecting to the PostgreSQL database...")
             conn = psycopg2.connect(**db_params)
             cursor = conn.cursor()
             print("Connection successful!")
     
             # 2. Read the file content in binary mode ('rb')
             file_name = os.path.basename(file_path)
             print(f"Reading binary content of file '{file_name}'...")
             with open(file_path, 'rb') as f:
                 file_content = f.read()
             print(f"File read complete. Size: {len(file_content)} bytes.")
     
             # 3. Prepare the SQL INSERT statement
             # Use parameterized queries (%s) to prevent SQL injection
             sql_insert = "INSERT INTO file_content (name, content) VALUES (%s, %s) RETURNING id;"
     
             # 4. Execute the insert operation
             print("Inserting file content into the 'file_content' table...")
             cursor.execute(sql_insert, (file_name, file_content))
     
             # Get the ID of the newly inserted record
             inserted_id = cursor.fetchone()[0]
     
             # 5. Commit the transaction
             conn.commit()
             print("---" * 10)
             print(f"Success! File '{file_name}' has been uploaded to the database.")
             print(f"The new record ID is: {inserted_id}")
             print("---" * 10)
     
         except psycopg2.Error as e:
             print(f"Database error: {e}")
             if conn:
                 conn.rollback()  # If an error occurs, roll back the transaction
         except Exception as e:
             print(f"An unexpected error occurred: {e}")
         finally:
             # 6. Close the database connection
             if conn:
                 cursor.close()
                 conn.close()
                 print("Database connection closed.")
     
     
     # --- Script entry point ---
     if __name__ == "__main__":
         upload_file_to_db(LOCAL_FILE_PATH, DB_CONFIG)

4. Chunk the imported document and store the results in the file_chunk table.

   INSERT INTO file_chunk(file_id, chunk_id, chunk_content)
   SELECT id,  -- Document ID
     (polar_ai.ai_rag_chunking(content, 'pdf')).* -- Chunking function
   FROM file_content;

5. Embed the content to facilitate retrieval.

   UPDATE file_chunk SET embedding = polar_ai.ai_rag_text_embedding(chunk_content);

6. Create a high-performance HNSW index on the embedding column for faster vector retrieval. HNSW is an efficient approximate nearest neighbor search algorithm suitable for high-dimensional vector data.

   -- Use an HNSW index and L2 (Euclidean) distance
   CREATE INDEX ON file_chunk USING HNSW (embedding vector_l2_ops);

Build a RAG query

Integrate vector retrieval, reranking, and LLM calls to implement a complete RAG Q&A process.

RAG execution steps

1. Convert the question into a vector.

2. Perform a vector search or hybrid search to find the required text chunks.

3. Rerank the retrieved text chunks to select the most relevant ones.

4. Submit the most relevant text chunks to the large language model to generate an answer.

Procedure

1. Create the custom function ai_summarize: Create a custom SQL function ai_summarize to encapsulate the logic for calling a large language model to generate the final answer. This function combines the question and the retrieved context into a structured prompt and calls the underlying text generation function.

   CREATE OR REPLACE FUNCTION public.ai_summarize(question text, contents text[])
   RETURNS text
       LANGUAGE plpgsql
       AS $function$
       DECLARE
          body text;
       BEGIN
          body = format('Answer the question based on the provided materials: %I. The provided materials are: %I', question, contents::text);
          RETURN polar_ai.ai_text_generation(body, '_polar4ai/_polar4ai_nl2sql_tongyi');
       END;
       $function$;

2. Execute an end-to-end RAG query. The following SQL statement demonstrates the entire RAG process, from vectorizing the question to generating a final answer:

   -- RAG
   SELECT public.ai_summarize('What are the advantages of PolarAI?', array_agg(chunks))
   FROM (
     -- rerank
     SELECT unnest(polar_ai.ai_rag_rerank('What are the advantages of PolarAI?', array_agg(chunk_content))) as chunks
     FROM (
         -- vector search
         SELECT chunk_content
         FROM 
           file_chunk
         ORDER BY 
           -- Vectorize the question
           polar_ai.ai_rag_text_embedding('What are the advantages of PolarAI?')::vector(1024) <-> embedding ASC
         LIMIT 10
     ) t
     LIMIT 5 
   );

Related SQL