使用 OpenAI、Go 和 PostgreSQL (pgvector) 构建语义搜索引擎

下载万磁搜索绿色版

一分耕耘，一分收获！既然都打开这篇《使用 OpenAI、Go 和 PostgreSQL (pgvector) 构建语义搜索引擎》，就坚持看下去，学下去吧！本文主要会给大家讲到等等知识点，如果大家对本文有好的建议或者看到有不足之处，非常欢迎大家积极提出！在后续文章我会继续更新Golang相关的内容，希望对大家都有所帮助！

近年来，向量嵌入技术已成为自然语言处理(NLP)和语义搜索的核心。与传统的关键词搜索不同，向量数据库通过比较文本的向量表示(嵌入)来理解文本的语义含义。本示例展示如何结合OpenAI嵌入、Go语言和PostgreSQL数据库(以及pgvector扩展)构建一个语义搜索引擎。

什么是嵌入？

嵌入是文本(或其他数据)在高维空间中的向量表示。语义相似的文本在该空间中具有相近的向量。将嵌入存储在PostgreSQL(使用pgvector扩展)等数据库中，可以实现快速准确的相似性搜索。

为什么PostgreSQL使用pgvector？

pgvector是一个流行的PostgreSQL扩展，它添加了向量数据类型，支持：

• 存储嵌入作为向量列
• 执行近似或精确最近邻搜索
• 使用标准SQL进行查询

应用程序概述

1. 调用OpenAI的嵌入API将输入文本转换为向量嵌入。
2. 使用pgvector扩展将这些嵌入存储在PostgreSQL数据库中。
3. 查询嵌入以查找数据库中最语义相似的条目。

先决条件

• Go语言安装(推荐1.19版本)。
• 本地或云端部署并运行的PostgreSQL数据库。
• 在PostgreSQL中安装pgvector扩展。(安装说明请参考pgvector的GitHub页面。)
• 可访问OpenAI API密钥。

Makefile用于本地测试，包含PostgreSQL/pgvector和Docker相关任务：

<code>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</code>

确保已安装pgvector。然后在PostgreSQL数据库中执行：

<code>create extension if not exists vector;</code>

完整代码

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扩展提供了一种简便的语义搜索引擎构建方案。通过将文本表示为向量并利用数据库索引的优势，我们实现了从传统的基于关键词的搜索到基于语义理解的搜索的转变。

理论要掌握，实操不能落！以上关于《使用 OpenAI、Go 和 PostgreSQL (pgvector) 构建语义搜索引擎》的详细介绍，大家都掌握了吧！如果想要继续提升自己的能力，那么就来关注golang学习网公众号吧！