由 John Doe 十二月 5, 2024
摘要：在本文中，我们将学习如何使用 PostgreSQL、pgvector、ollama、Llama3 和 Go 构建检索增强生成系统（大模型 RAG）。
目录
借助检索增强生成系统（大模型 RAG），您可以创建一个 AI 助手，该助手可以根据您现有的内部知识库（如 Wiki、手册、培训和参考资料）中包含的信息，来回答问题。
请继续阅读，了解如何使用 PostgreSQL、pgvector、ollama 和不到 200 行的 Go 代码，来构建自己的 RAG 系统。
概述
我们将使用一个故事中的几段，来作为我们的 “文档语料库”。对于每个文档，我们将使用 Meta 开源的 LLM Llama3 生成文档的嵌入词，并使用 ollama 在本地托管。然后，我们将文档和它的嵌入词存储到一个 PostgreSQL 表中。嵌入词将使用 pgvector 扩展进行存储和访问。
对于查询，我们将从与用户查询最相关的表中检索单个文档，并再次使用 llama3 生成响应。
ollama 提供了类似 OpenAI 的 HTTP API，我们将使用这些 API 来生成嵌入词和聊天响应。
对于 Go 代码，我们将使用 jackc/pgx 和 pgvector-go 与 PostgreSQL 通信，并使用 ollama 客户端 API 的包进行 HTTP API 的调用。
使用 Ollama 运行模型
Ollama 是最近发布的一个工具，它允许您在本地机器上运行流行的开源模型（见列表）。它提供了一个与模型无关的 OpenAI 风格的 REST API，可以在客户端应用程序中使用。请参阅有关安装和运行 ollama 的文档。
在将 ollama 安装到机器后，我们就可以简单地在 ollama 中运行 llama3 模型：
$ ollama pull llama3
pulling manifest
pulling 00e1317cbf74... 100% | 4.7 GB
pulling 4fa551d4f938... 100% | 12 KB
pulling 8ab4849b038c... 100% | 254 B
pulling 577073ffcc6c... 100% | 110 B
pulling ad1518640c43... 100% | 483 B
verifying sha256 digest
writing manifest
removing any unused layers
success
请注意，ollama 的 HTTP 服务器默认在 127.0.0.1:11434 上进行监听。如果您在一台机器上运行 ollama，并从另一台机器访问它，请参阅文档。
使用 curl 进行测试：
$ curl -X POST http://192.168.0.107:11434/api/embeddings -d '{
"model": "llama3",
"prompt":"Researchers have found that cats display behavioral patterns and social traits that were previously thought to be unique to humans and primates."
}'
{"embedding":[-4.317752838134766,-1.9526829719543457,-2.889270305633545,0.6311468482017517,1.636886477470398,1.2741944789886475,-3.402531147003174,-1.3292253017425537,
[...snip...]
安装 pgvector
pgvector 是可用于 PostgreSQL 版本 12 到 16 的一个 PostgreSQL 扩展。如果您使用的是 pgdg APT 仓库，则可以使用下面的命令进行安装：
sudo apt install postgresql-16-pgvector
不需要重新启动 PostgreSQL 服务器。pgvector 也可以在其他系统上使用（docker、yum、许多托管服务环境等）。有关更多信息，请参阅安装说明。
安装后，我们可以给一个数据库设置扩展，比如 ragdemo ：
CREATE EXTENSION vector;
我们再创建一个表，来存储文档和它的嵌入词：
CREATE TABLE items (id serial PRIMARY KEY, doc text, embedding vector(4096));
类型 vector 是由 pgvector 扩展提供的。(4096) 表示我们要存储维度为 4096 的向量。类型 vector (4096) 的每个值基本上是一个包含 4096 个 4 字节浮点数的数组。
要存储的文档
我们将使用 Arthur Conan Doyle 的夏洛克·福尔摩斯故事《博斯库姆谷之谜》中的 4 段。你可以在这里阅读完整的故事。各个段落作为单独的文件，与演示代码 git 仓库中的代码一起存档。
代码
该代码位于 GitHub 上的 git 仓库中，在 MIT 许可证下发布。要自己编译和运行它，首先确保你有一个正常工作的 Go 运行环境，然后：
git clone https://github.com/rapidloop/ragdemo.git
cd ragdemo
go build
连接到 PostgreSQL
该代码使用流行的 jackc/pgx 驱动程序，连接到使用环境变量 DATABASE_URL 指定的 PostgreSQL 数据库。格式与 psql 可以使用的格式相同。其他标准的 PostgreSQL 环境变量，如 PGHOST、PGDATABASE 等也可被使用。一些 DATABASE_URL 的示例如下：
export DATABASE_URL=postgres:///ragdemo?host=/var/run/postgresql
export DATABASE_URL=postgres://myuser@myhost:6432/ragdemo
export PGPASSWORD=mysecretpass
它会执行下面的 SQL，来插入文档和它的嵌入词：
INSERT INTO items (doc, embedding) VALUES ($1, $2)
并执行下面的 SQL，来获取最相关的文档：
SELECT doc FROM items ORDER BY embedding <-> $1 LIMIT 1
在两个 vector(1024) 类型的值之间，可以使用由 pgvector 扩展提供的 <-> 运算符。该运算符会返回两个向量之间的 L2 距离。有关其他运算符和更多信息，请参阅 pgvector 文档。
与 Ollama API 对话
Ollama 提供了简单的 HTTP REST API，类似于 OpenAI 的 HTTP API。Ollama 本身是用 Go 编写的，并提供了一个简单包装的包，来提供对这些 API 的类型化访问。
使用下面代码创建客户端：
cli, err := api.ClientFromEnvironment()
在服务器不在同一台计算机上运行的情况下，可以在客户端使用环境变量 OLLAMA_HOST，以提供服务器的 URL：
export OLLAMA_HOST=192.168.0.107:11434
以下是使用 llama3 模型生成嵌入词的请求：
req := &api.EmbeddingRequest{
Model: "llama3",
Prompt: doc,
}
resp, err := cli.Embeddings(context.Background(), req)
// resp.Embedding is a []float64 with 4096 entries
嵌入词会从 float64[] 转换为 float32[] ，然后转换为 pgvector，再使用 pgx 的 API 将其存储在数据库中。
要创建聊天响应，而不进行流式处理，会使用下面的代码：
stream := false
req2 := &api.ChatRequest{
Model: "llama3",
Stream: &stream,
Messages: []api.Message{
{
Role: "user",
Content: fmt.Sprintf(`Using the given reference text, succinctly answer the question that follows:
reference text is:
%s
end of reference text. The question is:
%s
`, doc, prompt),
},
},
}
var response string
resp2fn := func(resp2 api.ChatResponse) error {
response = strings.TrimSpace(resp2.Message.Content)
return nil
}
err = cli.Chat(context.Background(), req2, resp2fn)
实际的聊天响应位于变量 response 中。该代码还显示了用于生成步骤的提示词。
命令行接口
该代码使用了几个标志，来允许用户：
- 将存储为文件的文档插入到数据库中，或者
- 使用给定的提示词，来检索和生成响应
这是命令行的用法：
$ ./ragdemo
Usage:
./ragdemo -insert {path-to-doc-file}
./ragdemo -query {query-text}
Environment variables:
DATABASE_URL url of database, like postgres://host/dbname
OLLAMA_HOST url or host:port of OLLAMA server
PG* standard postgres env. vars are understood
将它们整合在一起
存储文档
当使用 -insert 和一个文档的路径调用时， ragdemo 会读取文件的内容，通过 Ollama 的 HTTP API 与 Llama3 模型通信，并检索文档的嵌入词。然后，它使用 pgx 的 API，将文档和它的嵌入词存储到 items 表中。
请注意，成功时不会输出任何内容：
$ ./ragdemo -insert doc1
$ ./ragdemo -insert doc2
$ ./ragdemo -insert doc3
$ ./ragdemo -insert doc4
查询 RAG 系统
当使用 -query 后跟用户提示（通常是一个问题或几个关键字）调用时， ragdemo 首先也会为该提示生成嵌入词。然后，它会将该嵌入词与 items 表中的所有其他嵌入词进行比较，并找到“最接近”的嵌入词。
“最接近”的那个，是使用 pgvector 提供的 “L2 距离运算符” <-> 来确定的。距离值最小的就是 “最接近”的，因此，回答用户提示的也是最相关的。 ragdemo 仅检索单个最相关的文档，并将其用于生成。
在最后一步生成的过程中，ragdemo 会使用检索到的文档，并使用上述代码中所示的提示再次调用 Llama3，以创建聊天响应。提示包含了原始提示，并使用检索到的文档进行增强。然后会输出响应，最终会看到像魔术一样的结果：
$ ./ragdemo -query "What was the verdict of McCarthy's inquest?"
The verdict of McCarthy's inquest was "wilful murder".
$ ./ragdemo -query "What did the gamekeeper see?"
According to the reference text, William Crowder, the game-keeper, lost sight of Mr. McCarthy and his son after a certain time. He didn't actually see anything significant in this case.
$ ./ragdemo -query "Did anyone see the McCarthys quarrelling?"
Yes, Patience Moran, a 14-year-old girl who was picking flowers in the woods, saw Mr. McCarthy and his son having a violent quarrel at the border of the wood and close to the lake. She heard Mr. McCarthy using strong language and saw the son raise his hand as if to strike his father.
$ ./ragdemo -query "So who killed McCarthy?"
According to the reference text, young Mr. McCarthy was arrested and a verdict of "wilful murder" was returned at the inquest. He was brought before the magistrates and the case was referred to the next Assizes. Therefore, based on the information provided, it appears that young Mr. McCarthy killed his father, Mr. McCarthy.
后续步骤
这只是一个基本的 RAG 演示，它使用了 PostgreSQL、pgvector、ollama、Llama3 和 Go。以下是一些进一步的说明：
- 有一些模型重点设计用于生成嵌入词，它们在该特定工作中比 llama3 会更有效。可以查看该排行榜作为参考。
- 如果您的文档不是英文的，可能会有更适合的模型。
- L2 距离可能并不是您所选模型的最佳距离计算方法。pgvector 也支持其他方法。
- 扫描整个表并对每个嵌入词计算 L2 距离，性能不可随数据量而扩展。可查看 pgvector 索引支持和其他技术。
- 在生成过程中检索和使用更多的文档，以及包括使用关键字匹配或其他搜索技术获取的文档，也可以提高生成输出的相关性和质量。
- 调整生成提示词，并尝试具有不同能力的大模型执行生成的步骤，甚至向用户显示多个大模型生成的输出应该也会很有用。