使用 OpenAI、Go 和 PostgreSQL (pgvector) 建立語意搜尋引擎

近年來，向量嵌入已成為現代自然語言處理 (NLP) 和語義搜尋的基礎。向量資料庫不再依賴關鍵字搜索，而是透過數值表示（嵌入）來比較文字的「含義」。本範例示範如何利用 OpenAI 嵌入、Go 和具有 pgvector 擴充功能的 PostgreSQL 來建立一個語意搜尋引擎。

什麼是嵌入？

嵌入是文字（或其他資料）在高維空間中的向量表示。如果兩段文字在語意上相似，則它們的向量在該空間中會彼此靠近。透過將嵌入儲存在像 PostgreSQL（帶有 pgvector 擴充功能）這樣的資料庫中，我們可以快速且準確地執行相似性搜尋。

為什麼選擇 PostgreSQL 和 pgvector？

pgvector 是一個流行的擴展，它將向量資料類型加入 PostgreSQL 中。它使您能夠：

• 將嵌入儲存為向量列
• 執行近似或精確的最近鄰搜尋
• 使用標準 SQL 執行查詢

應用程式概述

1. 呼叫 OpenAI 的嵌入 API 將輸入文字轉換為向量嵌入。
2. 使用 pgvector 擴充功能將這些嵌入儲存在 PostgreSQL 中。
3. 查詢嵌入以尋找資料庫中最語意相似的條目。

先決條件

• 已安裝 Go（建議 1.19 ）。
• 已安裝並執行 PostgreSQL（本機或代管）。
• 在 PostgreSQL 中安裝 pgvector 擴充。 （有關安裝說明，請參閱 pgvector 的 GitHub 頁面。）
• 具有嵌入存取權限的 OpenAI API 金鑰。

用於本機測試的包含與 postgres/pgvector 和 Docker 相關的任務的 Makefile。

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, ", ")

    // pgvector 需要方括号表示法才能输入向量，例如 [0.1, 0.2, 0.3]
    return "[" + joined + "]"
}

func main() {
    // 加载环境变量
    err := godotenv.Load()
    if err != nil {
        log.Fatal("加载 .env 文件出错")
    }

    // 创建连接池
    dbpool, err := pgxpool.New(context.Background(), os.Getenv("DATABASE_URL"))
    if err != nil {
        fmt.Fprintf(os.Stderr, "无法创建连接池：%v\n", err)
        os.Exit(1)
    }
    defer dbpool.Close()

    // 1. 确保已启用 pgvector 扩展
    _, err = dbpool.Exec(context.Background(), "CREATE EXTENSION IF NOT EXISTS vector;")
    if err != nil {
        log.Fatalf("创建扩展失败：%v\n", err)
        os.Exit(1)
    }

    // 2. 创建表（如果不存在）
    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("创建表失败：%v\n", err)
    }

    // 3. 创建索引（如果不存在）
    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("创建索引失败：%v\n", err)
    }

    // 4. 初始化 OpenAI 客户端
    apiKey := os.Getenv("OPENAI_API_KEY")
    if apiKey == "" {
        log.Fatal("未设置 OPENAI_API_KEY")
    }
    openaiClient := openai.NewClient(apiKey)

    // 5. 插入示例文档
    docs := []string{
        "PostgreSQL 是一个先进的开源关系数据库。",
        "OpenAI 提供基于 GPT 的模型来生成文本嵌入。",
        "pgvector 允许将嵌入存储在 Postgres 数据库中。",
    }

    for _, doc := range docs {
        err = insertDocument(context.Background(), dbpool, openaiClient, doc)
        if err != nil {
            log.Printf("插入文档“%s”失败：%v\n", doc, err)
        }
    }

    // 6. 查询相似性
    queryText := "如何在 Postgres 中存储嵌入？"
    similarDocs, err := searchSimilarDocuments(context.Background(), dbpool, openaiClient, queryText, 5)
    if err != nil {
        log.Fatalf("搜索失败：%v\n", err)
    }

    fmt.Println("=== 最相似的文档 ===")
    for _, doc := range similarDocs {
        fmt.Printf("- %s\n", doc)
    }
}

// insertDocument 使用 OpenAI API 为 `content` 生成嵌入，并将其插入 documents 表中。
func insertDocument(ctx context.Context, dbpool *pgxpool.Pool, client *openai.Client, content string) error {
    // 1) 从 OpenAI 获取嵌入
    embedResp, err := client.CreateEmbeddings(ctx, openai.EmbeddingRequest{
        Model: openai.AdaEmbeddingV2, // "text-embedding-ada-002"
        Input: []string{content},
    })
    if err != nil {
        return fmt.Errorf("CreateEmbeddings API 调用失败：%w", err)
    }

    // 2) 将嵌入转换为 pgvector 的方括号字符串
    embedding := embedResp.Data[0].Embedding // []float32
    embeddingStr := floats32ToString(embedding)

    // 3) 插入 PostgreSQL
    insertSQL := `
        INSERT INTO documents (content, embedding)
        VALUES (, ::vector)
    `
    _, err = dbpool.Exec(ctx, insertSQL, content, embeddingStr)
    if err != nil {
        return fmt.Errorf("插入文档失败：%w", err)
    }

    return nil
}

// searchSimilarDocuments 获取用户查询的嵌入，并根据向量相似性返回前 k 个相似的文档。
func searchSimilarDocuments(ctx context.Context, pool *pgxpool.Pool, client *openai.Client, query string, k int) ([]string, error) {
    // 1) 通过 OpenAI 获取用户查询的嵌入
    embedResp, err := client.CreateEmbeddings(ctx, openai.EmbeddingRequest{
        Model: openai.AdaEmbeddingV2, // "text-embedding-ada-002"
        Input: []string{query},
    })
    if err != nil {
        return nil, fmt.Errorf("CreateEmbeddings API 调用失败：%w", err)
    }

    // 2) 将 OpenAI 嵌入转换为 pgvector 的方括号字符串格式
    queryEmbedding := embedResp.Data[0].Embedding // []float32
    queryEmbeddingStr := floats32ToString(queryEmbedding)
    // 例如 "[0.123456, 0.789012, ...]"

    // 3) 构建按向量相似性排序的 SELECT 语句
    selectSQL := fmt.Sprintf(`
        SELECT content
        FROM documents
        ORDER BY embedding <-> '%s'::vector
        LIMIT %d;
    `, queryEmbeddingStr, k)

    // 4) 运行查询
    rows, err := pool.Query(ctx, selectSQL)
    if err != nil {
        return nil, fmt.Errorf("查询文档失败：%w", err)
    }
    defer rows.Close()

    // 5) 读取匹配的文档
    var contents []string
    for rows.Next() {
        var content string
        if err := rows.Scan(&content); err != nil {
            return nil, fmt.Errorf("扫描行失败：%w", err)
        }
        contents = append(contents, content)
    }
    if err = rows.Err(); err != nil {
        return nil, fmt.Errorf("行迭代错误：%w", err)
    }

    return contents, nil
}

結論

PostgreSQL 中的 OpenAI 嵌入、Go 和 pgvector 為建立語意搜尋應用程式提供了一種直接的解決方案。透過將文字表示為向量並利用資料庫索引的功能，我們從傳統的基於關鍵字的搜尋轉向按上下文和含義進行搜尋。

This revised output maintains the original language style, rephrases sentences for originality, and keeps the image in the same format and location. The code is also slightly improved for larity and ability. 帶 clarity comments.

以上是使用 OpenAI、Go 和 PostgreSQL (pgvector) 建立語意搜尋引擎的詳細內容。更多資訊請關注PHP中文網其他相關文章！