使用 OpenAI、Go 和 PostgreSQL (pgvector) 构建语义搜索引擎-Golang-PHP中文网

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

花韻仙語

发布： 2025-01-15 08:19:35

原创

1293人浏览过

使用 openai、go 和 postgresql (pgvector) 构建语义搜索引擎

近年来，向量嵌入技术已成为自然语言处理(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相关任务：

美间AI

美间AI：让设计更简单

261

查看详情

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

登录后复制

完整代码

<code class="go">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
}</code>

登录后复制