近年来，向量嵌入技术已成为自然语言处理(NLP)和语义搜索的核心。与传统的关键词搜索不同，向量数据库通过比较文本的向量表示(嵌入)来理解文本的语义含义。本示例展示如何结合OpenAI嵌入、Go语言和PostgreSQL数据库(以及pgvector扩展)构建一个语义搜索引擎。
什么是嵌入？
嵌入是文本(或其他数据)在高维空间中的向量表示。语义相似的文本在该空间中具有相近的向量。将嵌入存储在PostgreSQL(使用pgvector扩展)等数据库中，可以实现快速准确的相似性搜索。
为什么PostgreSQL使用pgvector？
pgvector是一个流行的PostgreSQL扩展，它添加了向量数据类型，支持：
- 存储嵌入作为向量列
- 执行近似或精确最近邻搜索
- 使用标准SQL进行查询
应用程序概述
- 调用OpenAI的嵌入API将输入文本转换为向量嵌入。
- 使用pgvector扩展将这些嵌入存储在PostgreSQL数据库中。
- 查询嵌入以查找数据库中最语义相似的条目。
先决条件
- Go语言安装(推荐1.19版本)。
- 本地或云端部署并运行的PostgreSQL数据库。
- 在PostgreSQL中安装pgvector扩展。(安装说明请参考pgvector的GitHub页面。)
- 可访问OpenAI API密钥。
Makefile用于本地测试，包含PostgreSQL/pgvector和Docker相关任务：
pgvector:
@docker run -d \
--name pgvector \
-e postgres_user=admin \
-e postgres_password=admin \
-e postgres_db=vectordb \
-v pgvector_data:/var/lib/postgresql/data \
-p 5432:5432 \
pgvector/pgvector:pg17
psql:
@psql -h localhost -u admin -d vectordb
确保已安装pgvector。然后在PostgreSQL数据库中执行：
create extension if not exists vector;
完整代码
package main
import (
"context"
"fmt"
"log"
"os"
"strings"
"github.com/jackc/pgx/v5/pgxpool"
"github.com/joho/godotenv"
"github.com/sashabaranov/go-openai"
)
func floats32ToString(floats []float32) string {
strVals := make([]string, len(floats))
for i, val := range floats {
strVals[i] = fmt.Sprintf("%f", val)
}
joined := strings.Join(strVals, ", ")
return "[" + joined + "]"
}
func main() {
err := godotenv.Load()
if err != nil {
log.Fatal("Error loading .env file")
}
dbpool, err := pgxpool.New(context.Background(), os.Getenv("DATABASE_URL"))
if err != nil {
fmt.Fprintf(os.Stderr, "Unable to create connection pool: %v\n", err)
os.Exit(1)
}
defer dbpool.Close()
_, err = dbpool.Exec(context.Background(), "CREATE EXTENSION IF NOT EXISTS vector;")
if err != nil {
log.Fatalf("Failed to create extension: %v\n", err)
os.Exit(1)
}
createTableSQL := `
CREATE TABLE IF NOT EXISTS documents (
id SERIAL PRIMARY KEY,
content TEXT,
embedding vector(1536)
);
`
_, err = dbpool.Exec(context.Background(), createTableSQL)
if err != nil {
log.Fatalf("Failed to create table: %v\n", err)
}
createIndexSQL := `
CREATE INDEX IF NOT EXISTS documents_embedding_idx
ON documents USING ivfflat (embedding vector_l2_ops) WITH (lists = 100);
`
_, err = dbpool.Exec(context.Background(), createIndexSQL)
if err != nil {
log.Fatalf("Failed to create index: %v\n", err)
}
apiKey := os.Getenv("OPENAI_API_KEY")
if apiKey == "" {
log.Fatal("OPENAI_API_KEY is not set")
}
openaiClient := openai.NewClient(apiKey)
docs := []string{
"PostgreSQL is an advanced open-source relational database.",
"OpenAI provides GPT-based models to generate text embeddings.",
"pgvector allows storing embeddings in a Postgres database.",
}
for _, doc := range docs {
err = insertDocument(context.Background(), dbpool, openaiClient, doc)
if err != nil {
log.Printf("Failed to insert document '%s': %v\n", doc, err)
}
}
queryText := "How to store embeddings in Postgres?"
similarDocs, err := searchSimilarDocuments(context.Background(), dbpool, openaiClient, queryText, 5)
if err != nil {
log.Fatalf("Search failed: %v\n", err)
}
fmt.Println("=== Most Similar Documents ===")
for _, doc := range similarDocs {
fmt.Printf("- %s\n", doc)
}
}
func insertDocument(ctx context.Context, dbpool *pgxpool.Pool, client *openai.Client, content string) error {
embedResp, err := client.CreateEmbeddings(ctx, openai.EmbeddingRequest{
Model: openai.AdaEmbeddingV2,
Input: []string{content},
})
if err != nil {
return fmt.Errorf("CreateEmbeddings API call failed: %w", err)
}
embedding := embedResp.Data[0].Embedding
embeddingStr := floats32ToString(embedding)
insertSQL := `
INSERT INTO documents (content, embedding)
VALUES ($1, $2::vector)
`
_, err = dbpool.Exec(ctx, insertSQL, content, embeddingStr)
if err != nil {
return fmt.Errorf("failed to insert document: %w", err)
}
return nil
}
func searchSimilarDocuments(ctx context.Context, pool *pgxpool.Pool, client *openai.Client, query string, k int) ([]string, error) {
embedResp, err := client.CreateEmbeddings(ctx, openai.EmbeddingRequest{
Model: openai.AdaEmbeddingV2,
Input: []string{query},
})
if err != nil {
return nil, fmt.Errorf("CreateEmbeddings API call failed: %w", err)
}
queryEmbedding := embedResp.Data[0].Embedding
queryEmbeddingStr := floats32ToString(queryEmbedding)
selectSQL := fmt.Sprintf(`
SELECT content
FROM documents
ORDER BY embedding <-> '%s'::vector
LIMIT %d;
`, queryEmbeddingStr, k)
rows, err := pool.Query(ctx, selectSQL)
if err != nil {
return nil, fmt.Errorf("failed to query documents: %w", err)
}
defer rows.Close()
var contents []string
for rows.Next() {
var content string
if err := rows.Scan(&content); err != nil {
return nil, fmt.Errorf("failed to scan row: %w", err)
}
contents = append(contents, content)
}
if err = rows.Err(); err != nil {
return nil, fmt.Errorf("row iteration error: %w", err)
}
return contents, nil
}
结论
OpenAI嵌入、Go语言和PostgreSQL数据库中的pgvector扩展提供了一种简便的语义搜索引擎构建方案。通过将文本表示为向量并利用数据库索引的优势，我们实现了从传统的基于关键词的搜索到基于语义理解的搜索的转变。
