使用rds_ai外掛程式實現RAG - ApsaraDB RDS

RDS PostgreSQL的rds_ai外掛程式整合了阿里雲百鍊的先進模型，包括千問、通用文本向量等。本文將介紹如何通過rds_ai外掛程式在資料庫內部僅使用SQL的方式實現檢索增強產生（Retrieval-Augmented Generation，簡稱RAG）技術。

前提條件

執行個體版本滿足以下要求：

執行個體大版本	核心小版本
PostgreSQL16	20241230及以上
PostgreSQL 14、15、17	20250430或以上
PostgreSQL18	20251130及以上

如需升級核心小版本，請參見升級核心小版本。

已建立RDS PostgreSQL高許可權帳號，如何建立高許可權帳號請參見建立帳號。
本文使用的模型為阿里雲大模型服務平台百鍊提供，請先前往阿里雲百鍊開通服務，並擷取API-KEY。具體操作，請參見擷取API Key。
已完成網路設定，以允許RDS PostgreSQL訪問外部模型。

方案介紹

本方案通過rds_ai外掛程式，利用SQL在資料庫內部簡單實現檢索增強產生（Retrieval-Augmented Generation，RAG）。該解決方案特別適用於相對不複雜的應用情境，您可在無須編寫任何代碼的情況下，輕鬆實現RAG。

您無需擔心各類RAG架構、運行環境及複雜組件的維護與營運。同時，對RAG的深層理解和使用經驗也不再是必要條件。只需通過SQL語句，便可簡便地完成RAG操作，顯著降低了技術門檻，提升了使用效率。更多資訊請參見rds_ai。

本方案的主要流程如下：

文檔內容分割：使用SQL將文檔內容劃分為多個塊（chunk）。
文本向量化：使用rds_ai外掛程式，將文本轉換為向量。
多路召回：使用rds_ai外掛程式，實現多種召回方式，包括向量召回和關鍵詞召回等。
結果重排序：通過rds_ai外掛程式對多路召回的結果進行重排序（rerank），最佳化最終結果的相關性。
調用大語言模型：使用rds_ai外掛程式調用大語言模型（LLM）服務，返回問題的最佳答案。

步驟一：建立外掛程式並完成基本配置

安裝rds_ai外掛程式和pg_jieba外掛程式。pg_jieba外掛程式為中文分詞外掛程式，詳情請參見中文分詞（pg_jieba）。
1. 訪問RDS執行個體列表，在上方選擇地區，然後單擊目標執行個體ID。
2. 在左側導覽列單擊外掛程式管理。
3. 外掛程式市場頁面，單擊rds_ai外掛程式的安裝。
4. 在彈出的視窗中選擇目標資料庫和帳號後，單擊安裝，將外掛程式安裝至目標資料庫。
5. 待執行個體狀態由維護執行個體中變為運行中時，重複上述操作，安裝pg_jieba外掛程式。
  
  說明
  外掛程式安裝大約需要1分鐘，請您重新整理頁面以查看執行個體狀態。
串連PostgreSQL執行個體，在目標資料庫的public schema中完成rds_ai外掛程式的基礎配置。
1. 配置所有預設模型的API-KEY。
```
-- 為rds_ai.model_list中所有的模型設定api key
SELECT rds_ai.update_model(model_name,'token','sk-****') FROM rds_ai.model_list;
```
2. rds_ai對模型的遠程調用是基於pgsql-http外掛程式實現的，進行以下逾時設定，以中斷長時間執行的調用。
  
  重要
  以下逾時設定為會話層級。如果您在使用rds_ai外掛程式時建立新的會話，請在該新會話中重新設定逾時時間。
```
-- 佈建要求的逾時時間，單位毫秒
SET http.timeout_msec TO 200000;

SELECT http.http_set_curlopt('CURLOPT_TIMEOUT', '200000');

-- 設定連線逾時時間
SELECT http.http_set_curlopt('CURLOPT_CONNECTTIMEOUT_MS', '200000');
```

步驟二：處理資料

將知識庫文檔資料通過SQL語句插入至資料庫中。在此過程中，需要特別注意文本的轉義問題，可以採用轉義工具來實現，例如將分行符號轉義為\n。

建立測試表和索引。

CREATE TABLE doc (
    id SERIAL PRIMARY KEY,
    title VARCHAR(255),
    content TEXT
);

CREATE TABLE chunk (
    id SERIAL PRIMARY KEY,
    doc_id INTEGER NOT NULL,
    text TEXT,
    embedding VECTOR(1024),
    ts_vector_extra tsvector
);

-- 建立索引
CREATE INDEX idx_doc_id ON chunk (doc_id);

CREATE INDEX ON chunk USING hnsw (embedding vector_cosine_ops);

CREATE INDEX chunk_text_gin ON chunk USING gin (ts_vector_extra);

CREATE INDEX ON chunk USING hnsw (embedding vector_cosine_ops) WITH (m = 16, ef_construction = 64);

建立split_text函數用於切分長文本資訊，採用固定字元長度進行分割，並設定分割前後兩個塊的重複長度。對於更複雜的切分方式，建議使用專門的文本切分架構進行實現。

CREATE OR REPLACE FUNCTION split_text(
    input_text TEXT,
    chunk_size INT,
    chunk_overlap INT
) RETURNS SETOF TEXT AS $$
DECLARE
    current_idx INT;
    start_idx INT;
    chunk TEXT;
BEGIN
    -- 邊界值處理
    IF chunk_overlap >= chunk_size THEN
        RAISE EXCEPTION 'chunk_overlap must be less than chunk_size';
    END IF;

    current_idx := 1;

    LOOP
        -- 起始位置
        start_idx := current_idx - chunk_overlap;
        
        IF start_idx < 1 THEN
            start_idx := 1;
        END IF;

        chunk := substr(input_text, start_idx, chunk_size);

        IF chunk IS NULL OR length(chunk) = 0 THEN
            EXIT;
        END IF;

        RETURN NEXT chunk;
        current_idx := current_idx + chunk_size - chunk_overlap;
    END LOOP;
END;
$$ LANGUAGE plpgsql;

依次建立觸發器，以確保在向doc表插入資料時，能夠自動向chunk表插入相應的資料。

如果文檔處理和分割塊的情境較為複雜，請參照步驟三：文檔轉向量，實現批量文本轉向量的處理。

-- 建立觸發器，只要給doc表增加和刪除，自動執行split_text函數插入chunk表
CREATE OR REPLACE FUNCTION insert_into_chunk() 
RETURNS TRIGGER AS $$
BEGIN
    INSERT INTO chunk (doc_id, text)
    SELECT NEW.id, result 
    FROM split_text(NEW.content, 300, 50) AS result;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER after_insert_doc
AFTER INSERT ON doc
FOR EACH ROW
EXECUTE FUNCTION insert_into_chunk();

CREATE OR REPLACE FUNCTION delete_from_chunk() 
RETURNS TRIGGER AS $$
BEGIN
    DELETE FROM chunk WHERE doc_id = OLD.id;
    RETURN OLD;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER after_delete_doc
AFTER DELETE ON doc
FOR EACH ROW
EXECUTE FUNCTION delete_from_chunk();

-- 建立觸發器，自動為chunk表的text內容執行rds_ai.embed函數轉為向量
CREATE OR REPLACE FUNCTION update_chunk_embedding()
RETURNS TRIGGER AS $$
BEGIN
    NEW.embedding := rds_ai.embed(NEW.text)::vector;
    RETURN NEW;
END;
$$ LANGUAGE plpgsql;

CREATE TRIGGER set_chunk_embedding
BEFORE INSERT OR UPDATE ON chunk
FOR EACH ROW
EXECUTE FUNCTION update_chunk_embedding();

--建立觸發器，自動為chunk表的text內容轉為tsvector，用於關鍵詞檢索
CREATE TRIGGER embedding_tsvector_update
BEFORE UPDATE OR INSERT ON chunk
FOR EACH ROW
EXECUTE PROCEDURE tsvector_update_trigger(
    'ts_vector_extra',
    'public.jiebacfg',
    'text'
);

插入知識庫文檔資料。例如：

INSERT INTO doc (title, content) VALUES 
  ('PostgreSQL簡史', '現在被稱為PostgreSQL的對象-關係型資料庫管理系統是從加州大學伯克利分校寫的POSTGRES軟體包發展而來的。經過二十多年的發展，PostgreSQL是世界上可以獲得的最先進的開來源資料庫。\n\n2.1. 伯克利的POSTGRES專案\n由Michael Stonebraker教授領導的POSTGRES專案是由防務進階研究專案局（DARPA）、陸軍研究辦公室（ARO）、國家科學基金（NSF） 以及 ESL, Inc 共同贊助的。 POSTGRES的實現始於 1986 年。該系統最初的概念詳見[ston86] 。 最初的資料模型定義見[rowe87]。當時的規則系統設計在[ston87a]裡描述。儲存管理器的理論基礎和體繫結構在[ston87b]裡有詳細描述。\n\n從那以後，POSTGRES經歷了幾次主要的版本更新。 第一個“示範性”系統在 1987 年便可使用了， 並且在 1988 年的ACM-SIGMOD大會上展出。在 1989 年6月發布了版本 1（見[ston90a]）給一些外部的使用者使用。 為了回應使用者對第一個規則系統（[ston89]）的批評，規則系統被重新設計了（[ston90b]），在1990年6月發布了使用新規則系統的版本 2。 版本 3 在1991年出現，增加了多儲存管理器的支援， 並且改進了查詢執行器、重寫了規則系統。直到Postgres95發布前（見下文）的後續版本大多把工作都集中在移植性和可靠性上。\n\nPOSTGRES已經被用於實現很多不同的研究和生產應用。這些應用程式套件括： 一個財務資料分析系統、一個噴氣引擎效能監視軟體包、一個小行星追蹤資料庫、一個醫學資訊資料庫和一些地理資訊系統。POSTGRES還被許多大學用於教學用途。最後，Illustra Information Technologies（後來併入Informix， 而Informix現在被IBM所擁有） 拿到代碼並使之商業化。在 1992 年末POSTGRES成為Sequoia 2000科學計算專案的主要資料管理器。\n\n在 1993 年間，外部使用者社區的數量幾乎翻番。隨著使用者的增加， 用於原始碼維護的時間日益增加並佔用了太多本應該用於資料庫研究的時間，為了減少支援的負擔，伯克利的POSTGRES專案在版本 4.2 時正式終止。\n\n2.2. Postgres95\n在 1994 年，Andrew Yu 和 Jolly Chen 向POSTGRES中增加了 SQL 語言的解譯器。並隨後用新名字Postgres95將原始碼發布到互連網上供大家使用， 成為最初POSTGRES伯克利代碼的開源繼承者。\n\nPostgres95的原始碼都是完全的 ANSI C，而且代碼量減少了25%。許多內部修改提高了效能和可維護性。Postgres95的1.0.x版本在進行 Wisconsin Benchmark 測試時大概比POSTGRES的版本4.2 快 30-50%。除了修正了一些錯誤，下面的是一些主要提升：\n\n原來的查詢語言 PostQUEL 被SQL取代（在伺服器端實現）。介面庫libpq被按照PostQUEL命名。在PostgreSQL之前還不支援子查詢（見下文），但它們可以在Postgres95中由使用者定義的SQL函數類比。聚集合函式被重新實現。同時還增加了對GROUP BY 查詢子句的支援。\n\n新增加了一個利用GNU的Readline進行互動 SQL 查詢的程式（psql）。這個程式很大程度上取代了老的monitor程式。\n\n增加了新的前端庫（libpgtcl）， 用以支援基於Tcl的用戶端。一個樣本 shell（pgtclsh），提供了新的 Tcl 命令用於Tcl程式和Postgres95伺服器之間的互動。\n\n徹底重寫了大對象的介面。保留了將大對象倒轉（Inversion ）作為儲存大對象的唯一機制（去掉了倒轉（Inversion ）檔案系統）。\n\n去掉了執行個體級的規則系統。但規則仍然以重寫規則的形式存在。\n\n在發布的源碼中增加了一個介紹SQL和Postgres95特性的簡短教程。\n\n用GNU的make（取代了BSD的make）來編譯。Postgres95可以使用不打補丁的GCC編譯（修正了雙精確度資料對齊問題）。\n\n\n2.3. PostgreSQL\n到了 1996 年， 很明顯“Postgres95”這個名字已經跟不上時代了。於是我們選擇了一個新名字PostgreSQL來反映與最初的POSTGRES和最新的具有SQL能力的版本之間的關係。同時版本號碼也從 6.0 開始， 將版本號碼放回到最初由伯克利POSTGRES專案開始的序列中。\n\n很多人會因為傳統或者更容易發音而繼續用“Postgres”來指代PostgreSQL（現在很少用全大寫字母）。這種用法也被廣泛接受為一種暱稱或別名。\n\nPostgres95的開發重點放在標識和理解後端代碼的現有問題上。PostgreSQL的開發重點則轉到了一些有爭議的特性和功能上面，當然各個方面的工作同時都在進行。\n\n自那以來，PostgreSQL發生的變化可以在附錄 E中找到。'),
  ('MySQL資料庫管理系統概述', 'MySQL是最流行的開放源碼SQL資料庫管理系統，它是由MySQL AB公司開發、發布並支援的。MySQL AB是由多名MySQL開發人創辦的一家商業公司。它是一家第二代開放源碼公司，結合了開放源碼價值取向、方法和成功的商業模型。\n\n在MySQL的網站（http://www.mysql.com/）上，給出了關於MySQL和MySQL的最新資訊。\n\n·         MySQL是一種資料庫管理系統。\n\n資料庫是資料的結構化集合。它可以是任何東西，從簡單的購物清單到畫展，或商業網路中的海量資訊。要想將資料添加到資料庫，或訪問、處理電腦資料庫中儲存的資料，需要使用資料庫管理系統，如MySQL伺服器。電腦是處理大量資料的理想工具，因此，資料庫管理系統在計算方面扮演著關鍵的中心角色，或是作為獨立的工具 + 生產力，或是作為其他應用程式的組成部分。\n\nMySQL是一種關聯資料庫管理系統。\n\n關聯資料庫將資料儲存在不同的表中，而不是將所有資料放在一個大的倉庫內。這樣就增加了速度並提高了靈活性。MySQL的SQL指得是“結構化查詢語言 (SQL)”。SQL是用於訪問資料庫的最常用標準化語言，它是由ANSI/ISO SQL標準定義的。SQL標準自1986年以來不斷演化發展，有數種版本。在本手冊中，“SQL-92”指得是1992年發布的標準，“SQL:1999”指得是1999年發布的標準，“SQL:2003”指得是標準的目前的版本。我們採用術語“SQL標準”標示SQL標準的目前的版本。\n\nMySQL軟體是一種開放源碼軟體。\n\n“開放源碼”意味著任何人都能使用和改變軟體。任何人都能從Internet下載MySQL軟體，而無需支付任何費用。如果願意，你可以研究源碼並進行恰當的更改，以滿足你自己的需求。MySQL軟體採用了GPL（GNU通用公用許可證），http://www.mysql.com/company/legal/licensing/）。\n\nMySQL資料庫伺服器具有快速、可靠和便於使用的特點。\n\n如果它正是你所尋找的，不妨一試。MySQL伺服器還有一套實用的特性集合，這些特性是通過與我們使用者的密切合作而開發的。在我們的基準測試首頁上，給出了MySQL伺服器和其他資料庫管理員的比較結果。請參見7.1.4 “MySQL基準套件”。\n\nMySQL伺服器最初是為處理大型資料庫而開發的，與已有的解決方案相比，它的速度更快，多年以來，它已成功用於眾多要求很高的生產環境。儘管MySQL始終在不斷髮展，但目前MySQL伺服器已能提供豐富和有用的功能。它具有良好的連通性、速度和安全性，這使的MySQL十分適合於訪問Internet上的資料庫。\n\nMySQL伺服器工作在用戶端/伺服器模式下，或嵌入式系統中。\n\nMySQL資料庫軟體是一種用戶端/伺服器系統，由支援不同後端的1個多線程SQL伺服器，數種不同的用戶端程式和庫，眾多管理工具和廣泛的應用編程介面API組成。\n\n我們還能以嵌入式多線程庫的形式提供MySQL伺服器，你可以將其連結到你的應用程式，從而獲得更小、更快、和更易管理的產品。\n\n有大量可用的共用MySQL軟體。\n\n你所喜歡的應用程式和語言均支援MySQL資料庫伺服器，這種情況十分可能。\n\n“MySQL”的正式發音是“My Ess Que Ell”（而不是“my sequel”）,但我們並不介意你的發音方式是“my sequel”或其他當地方式。'),
  ('什麼是 SQL Server？', 'Microsoft SQL Server 是一種關聯式資料庫管理系統 (RDBMS)。 應用程式和工具串連到 SQL Server 執行個體或資料庫，並使用 Transact-SQL (T-SQL) 進行通訊。\n\n部署選項\n可在 Windows 或 Linux 上安裝SQL Server，將其部署在 Linux 容器中，或者部署在 Azure 虛擬機器或其他虛擬機器平台上。 你之前可能將它稱為“裝箱產品”。\n\n支援哪些 SQL Server 版本取決於你的許可協議，但就本文檔而言，我們指的是 SQL Server 2016 (13.x) 及更高版本。 有關 SQL Server 2014 (12.x) 和以前版本的文檔，請參閱 SQL Server 先前版本對應的文檔。 若要瞭解當前支援哪些 SQL Server 版本，請參閱 SQL Server 終止支援選項。\n\n以下產品和服務也使用基礎 SQL Server 資料庫引擎：\n\nAzure SQL 資料庫\nAzure SQL 受管理的執行個體\nMicrosoft Analytics Platform System (PDW)\nAzure Synapse Analytics\nAzure SQL Edge\n有關 Windows 上 SQL Server 各版本支援的功能列表，請參閱：\n\nSQL Server 2022 各個版本及其支援的功能\nSQL Server 2019 各個版本及其支援的功能\n版本和 SQL Server 2017 支援的功能\n版本和 SQL Server 2016 支援的功能\nSQL Server 元件和技術\n本部分介紹 SQL Server 中提供的一些關鍵技術。\n\n組件\t說明\n資料庫引擎\t資料庫引擎是用於儲存、處理和保護資料的核心服務。 資料庫引擎提供了受控訪問和交易處理，以滿足企業內最苛刻的資料消費應用程式的要求。 資料庫引擎還通過商務持續性和資料庫恢複 (SQL Server) 為保持商務持續性提供全面的支援。\n機器學習服務 (MLS)\tSQL Server 機器學習服務支援使用常用的 R 和 Python 語言，將機器學習整合到企業工作流程中。\n\n機器學習服務（資料庫內）將 R 和 Python 與 SQL Server 整合，方便使用者通過調用預存程序，輕鬆產生、重新定型模型，並對模型評分。 機器學習伺服器對 R 和 Python 提供企業級支援，使用者無需使用 SQL Server。\nIntegration Services (SSIS)\tSQL Server Integration Services 是一個可用於構建高效能資料整合解決方案的平台，其中包括為資料倉儲提供擷取、轉換和下載 (ETL) 處理的包。\nAnalysis Services (SSAS)\tSQL Server Analysis Services 是一個用於個人、團隊和公司商業智慧的分析資料平台和工具集。 伺服器和用戶端設計器通過使用 Power Pivot、Excel 和 SharePoint Server 環境，支援傳統的 OLAP 解決方案、新的表格建模解決方案以及自助式分析和協作。 Analysis Services 還包括資料採礦，以便您可以發現隱藏在大量資料中的模式和關係。\nReporting Services (SSRS)\tSQL Server Reporting Services 提供支援 Web 的企業級報表功能。 從而使使用者可以建立從多個資料來源提取內容的報表，發布各種格式的報表，以及集中管理安全性和訂閱。\n複製\tSQL Server 複製是一組技術，用於將資料和資料庫物件從一個資料庫複寫並分發到另一個資料庫，然後在資料庫之間進行同步以保持一致性。 使用複製時，可以通過區域網路和廣域網路、撥號連線、無線串連和 Internet，將資料分發到不同位置以及分發給遠端使用者或移動使用者。\nData Quality Services (DQS)\tData Quality Services 向您提供知識驅動型資料清理解決方案。 DQS 使您可以產生知識庫，然後使用此知識庫，同時採用電腦輔助方法和互動方法，執行資料更正和消除重複的資料。 您可以使用雲端式的引用資料服務，並可以產生一個資料管理解決方案將 DQS 與 SQL Server Integration Services 和 Master Data Services 相整合。\nMaster Data Services (MDS)\tMaster Data Services 是 SQL Server 的主要資料管理解決方案。 基於 Master Data Services 產生的解決方案可協助確保報表和分析均基於適當的資訊。 使用 Master Data Services，可以為主要資料建立中央存放庫，並隨著主要資料隨時間變化而維護一個可審核的安全性實體記錄。\n')
;

查看分割後的chunk表，文本已被分割至text欄位，並自動轉換為向量和tsvector類型。
```
SELECT * FROM chunk;
```

步驟三：文檔轉向量

如果已經建立了觸發器，在插入doc表時自動更新chunk表，則可以跳過這部分內容。

如果文檔處理和分割塊的情境較為複雜，建議使用LangChain等RAG架構進行編程處理。可以在資料庫中對已存在的塊表執行嵌入操作。本步驟將介紹如何逐次調用rds_ai.embed函數以及如何調用自訂函數實現批量文本轉向量的處理。

逐次調用

使用rds_ai.embed函數，將 chunk 表中每條記錄的文本資訊轉換為向量表示，並將這些向量儲存到embedding列。
```
--全表embedding，逐次調用
UPDATE chunk
SET embedding = rds_ai.embed(text)::vector;
```

批量調用

增加批量調用自訂模型調用text-embedding-v3-batch。

SELECT rds_ai.add_model(
    'text-embedding-v3-batch',          -- 模型名
    'POST',                             -- 請求方式 
    ARRAY[('Authorization', 'Bearer %s')]::http.http_header[],  -- 要求標頭
    'https://dashscope.aliyuncs.com/compatible-mode/v1/embeddings', -- 請求的URL
    'application/json',                 -- 請求內容格式
    '{
        "model": "text-embedding-v3",
        "input": %s
    }',                               -- 請求體
    'SELECT %L::jsonb->''data''' -- 處理響應的SQL
);

配置新增模型的API-KEY。

SELECT rds_ai.update_model(
    'text-embedding-v3-batch',
    'token',
    'sk-xxxxxx'
);

測試模型。

SELECT * 
FROM rds_ai.invoke_model(
    'text-embedding-v3-batch',
    ARRAY['["01","02","03"]']
);

批量調用速度是逐次調用的若干倍，然而建議採用逐次調用的方式，具體選擇應根據實際測試結果進行確認。

DO $$
DECLARE
    batch_size INT := 20;         -- Batch size for embedding calls, can be adjusted based on the model
    pointer INT := 0;             -- Pointer to traverse the table
    record_count INT;             -- Total number of records in the table
    input_text TEXT;              -- Concatenated string of fields to be embedded in ["item1", "item2",..., "item n"] format
    json_results JSON;            -- JSON response from embedding API call
    json_item JSON;               -- Single element from the JSON array response containing index and embedding fields
    idx INT;                      -- Index from the JSON response
BEGIN
    -- Get the total number of records
    SELECT COUNT(*) INTO record_count FROM chunk;

    -- Loop through each batch
    WHILE pointer < record_count LOOP
        -- Construct input
        SELECT json_agg(text::TEXT) INTO input_text
        FROM (
            SELECT text
            FROM chunk
            ORDER BY id
            LIMIT batch_size OFFSET pointer
        ) AS subquery;

        -- Call the embedding model
        json_results := rds_ai.invoke_model('text-embedding-v3-batch', ARRAY[input_text]);

        -- Loop through the model results and update corresponding rows
        FOR idx IN 0..json_array_length(json_results) - 1 LOOP
            json_item := json_results->idx;
            UPDATE chunk
            SET embedding = (SELECT (json_item->>'embedding')::VECTOR(1024))
            WHERE id = (SELECT id FROM chunk ORDER BY id LIMIT 1 OFFSET pointer + idx);
        END LOOP;

        -- Update the pointer for the next batch
        pointer := pointer + batch_size;
        RAISE NOTICE '%/% has done.', pointer, record_count;  
    END LOOP;
END;
$$ LANGUAGE plpgsql;

步驟四：多路召回

以下以關鍵詞召回與向量召回的雙路召回為例，若僅使用向量召回，建議直接調用rds_ai.retrieve函數。

CREATE OR REPLACE FUNCTION multi_retrieve(query TEXT)
RETURNS TABLE(ret_chunk text, score numeric, method text, rank int) AS $$
DECLARE
    rec RECORD;
BEGIN
    -- Retrieve by keyword
    FOR rec IN
        SELECT
            text,
            subquery.score,
            'retrive_by_key_word' AS method,
            RANK() OVER (ORDER BY subquery.score DESC) AS rank_id
        FROM (
            SELECT
                text,
                ts_rank(
                    ts_vector_extra,
                    to_tsquery(replace(
                        text(plainto_tsquery('jiebacfg', query)), '&', '|'
                    ))
                ) AS score
            FROM chunk
            WHERE ts_vector_extra @@ to_tsquery(replace(
                text(plainto_tsquery('jiebacfg', query)), '&', '|'
            ))
        ) AS subquery
        ORDER BY subquery.score DESC
        LIMIT 5
    LOOP
        ret_chunk := rec.text;
        score := rec.score;
        method := rec.method;
        rank := rec.rank_id;
        RETURN NEXT;
    END LOOP;

    -- Retrieve by vector
    FOR rec IN
        SELECT
            *,
            'retrive_by_vector' AS method,
            RANK() OVER (ORDER BY distance) AS rank_id
        FROM rds_ai.retrieve(
            query, 
            'public', 
            'chunk', 
            'text', 
            'embedding', 
            distance_type => 'cosine', 
            topn => 5
        )
    LOOP
        ret_chunk := rec.chunk;
        score := rec.distance;
        method := rec.method;
        rank := rec.rank_id;
        RETURN NEXT;
    END LOOP;
END;
$$ LANGUAGE plpgsql;

多路召回測試。

SELECT * 
FROM multi_retrieve('介紹一下PostgreSQL資料庫的來源') 
ORDER BY METHOD, RANK;

步驟五：結果重排序

多路召回

本步驟介紹了兩種結果重排序的方法：使用rds_ai.rank函數和使用RFF函數。

使用rds_ai.rank函數

DROP FUNCTION get_reranked_results(text, integer);
CREATE OR REPLACE FUNCTION get_reranked_results(query TEXT, top_n INT DEFAULT 1)
RETURNS TABLE(score FLOAT, value TEXT) AS $$
DECLARE
    result_array TEXT[];
BEGIN
    SELECT array_agg(ret_chunk)
    INTO result_array
    FROM multi_retrieve(query);

    RETURN QUERY
    SELECT *
    FROM rds_ai.rank(query, result_array)
    ORDER BY score_value DESC
    LIMIT top_n;
END $$ LANGUAGE plpgsql;

SELECT * FROM get_reranked_results('介紹一下PostgreSQL資料庫的來源');

使用RFF函數

-- 常見rff函數對應的彙總函式
CREATE TYPE score_agg_state AS (
    vector_score numeric,
    keyword_score numeric
);

CREATE OR REPLACE FUNCTION score_agg_transfn(state score_agg_state, rank numeric, method text)
RETURNS score_agg_state AS $$
BEGIN
    CASE method
        WHEN 'retrive_by_vector' THEN
            state.vector_score := COALESCE(1 / (60+rank), 0);
        WHEN 'retrive_by_key_word' THEN
            state.keyword_score := COALESCE(1 / (60+rank), 0);
    END CASE;
    RETURN state;
END;
$$ LANGUAGE plpgsql;

CREATE OR REPLACE FUNCTION score_agg_finalfn(state score_agg_state)
RETURNS numeric AS $$
BEGIN
    RETURN COALESCE(state.vector_score, 0) + COALESCE(state.keyword_score, 0);
END;
$$ LANGUAGE plpgsql;

CREATE AGGREGATE rff_function(numeric, text)(
    SFUNC = score_agg_transfn,
    STYPE = score_agg_state,
    FINALFUNC = score_agg_finalfn,
    INITCOND = '(0, 0)'
);

調用函數對結果進行重新排序。

SELECT
    ret_chunk,
    rff_function(rank, method) AS score
FROM
    multi_retrieve('介紹一下PostgreSQL資料庫的來源')
GROUP BY
    ret_chunk
ORDER BY
    score DESC;

向量召回

使用向量召回，並希望通過rank函數進行一次精細排序，可以執行如下SQL：

DROP FUNCTION get_reranked_results(text, integer);
CREATE OR REPLACE FUNCTION get_reranked_results(query TEXT, top_n INT DEFAULT 1)
RETURNS TABLE(score FLOAT, value TEXT) AS $$
DECLARE
    result_array TEXT[];
BEGIN
    SELECT array_agg(chunk)
    INTO result_array
    FROM rds_ai.retrieve(query, 'public', 'chunk', 'text', 'embedding');
    RETURN QUERY
    SELECT *
    FROM rds_ai.rank(query, result_array)
    ORDER BY score_value DESC
    LIMIT top_n;
END $$ LANGUAGE plpgsql;

-- 測試
SELECT * FROM get_reranked_results('介紹一下PostgreSQL資料庫的來源');

步驟六：調用大模型返回結果

多路召回

建立轉義函數escape_for_json。這是因為調用rds_ai.prompt函數最終需要將參數組裝為JSON格式，並通過HTTP發送至阿里雲百鍊服務端，因此在資料庫中對換行等特殊字元進行轉義是必要的。
```
CREATE OR REPLACE FUNCTION escape_for_json(input TEXT)
RETURNS TEXT AS $$
BEGIN
    RETURN replace(replace(input, '"', '\\"'), E'\n', '\\n');
END;
$$ LANGUAGE plpgsql;
```

建立RAG函數。本文以使用rds_ai.rank函數進行多路召回為例。

CREATE OR REPLACE FUNCTION rag(query TEXT) RETURNS TEXT AS $$
DECLARE
  prompt_content TEXT;
  result TEXT;
BEGIN
  -- 使用 get_reranked_results 函數來擷取排序後的結果並彙總為一個數組
  WITH rank_result AS (
    SELECT string_agg(value, ',') AS prompt_content_pre 
    FROM get_reranked_results(query)
  )
  -- 從 rank_result 查詢產生提示並擷取 AI 模型的回答
  SELECT 
    '基於下面的內容' || prompt_content_pre || '回答我的問題，' || query
  INTO prompt_content
  FROM rank_result;

  -- 使用 rds_ai.prompt 函數擷取 AI 模型的回答
  SELECT 
    rds_ai.prompt(escape_for_json(prompt_content))
  INTO result;

  -- 返回結果
  RETURN result;
END;
$$ LANGUAGE plpgsql;

調用RAG函數返回結果。

SELECT rag('介紹一下PostgreSQL資料庫的來源');

向量召回

使用向量檢索，可以直接調用 rds_ai.rag 函數。

SELECT *
FROM rds_ai.rag (
    '介紹一下PostgreSQL資料庫的來源', 
    'public', 
    'chunk', 
    'text', 
    'embedding'
);