使用rds_ai插件实现RAG_云数据库 RDS(RDS)-阿里云帮助中心

RDS PostgreSQL的rds_ai插件集成了阿里云百炼的先进模型，包括通义千问、通用文本向量和通用文本排序等。本文将介绍如何通过rds_ai插件在数据库内部仅使用SQL的方式实现检索增强生成（Retrieval-Augmented Generation，简称RAG）技术。

您可以加入RDS PostgreSQL插件交流钉钉群（103525002795），进行咨询、交流和反馈，获取更多关于插件的信息。

前提条件

例大版本为RDS PostgreSQL 16。
实例内核小版本为20241230及以上。
如需升级内核小版本，请参见升级内核小版本。
已创建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');
```
3. （可选）如果通过私网连接（PrivateLink）建立RDS PostgreSQL专有网络 VPC与阿里云百炼平台之间的网络连接，则需要执行如下命令，更新目标模型的URL地址。
```
SELECT rds_ai.update_model('qwen-plus', 'uri', '<通过私网连接后的URL地址>'); 
```
  通过私网连接（PrivateLink）建立RDS PostgreSQL专有网络 VPC与阿里云百炼平台之间的网络连接后，模型URL地址请参见步骤二：通过终端节点私网访问阿里云百炼平台。例如，当VPC地域为华北2（北京）时，将https://dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation中公网域名dashscope.aliyuncs.com替换为已经获取到的自定义服务域名vpc-cn-beijing.dashscope.aliyuncs.com。
  当需要更新多个模型的URL地址时，SQL示例如下：
```
SELECT rds_ai.update_model(
    model_name,
    'uri',
    'https://vpc-cn-beijing.dashscope.aliyuncs.com/api/v1/services/aigc/text-generation/generation'
) 
FROM rds_ai.model_list 
WHERE model_name IN ('qwen-max', 'qwen-turbo', 'qwen-plus');
```

步骤二：处理数据

将知识库文档数据通过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是用于访问数据库的最常用标准化语言，它是由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'
);