In recent years, vector embeddings have become the foundation of modern natural language processing (NLP) and semantic search. Instead of relying on keyword searches, vector databases compare the "meaning" of text through numerical representations (embeddings). This example demonstrates how to create a semantic search engine using OpenAI embedding, Go, and PostgreSQL with the pgvector extension.

What is embedding?

Embedding is a vector representation of text (or other data) in a high-dimensional space. If two pieces of text are semantically similar, their vectors will be close to each other in this space. By storing embeddings in a database like PostgreSQL (with the pgvector extension), we can perform similarity searches quickly and accurately.

Why choose PostgreSQL and pgvector?

pgvector is a popular extension that adds vector data types to PostgreSQL. It enables you to:

• Store embeddings as vector columns
• Perform an approximate or exact nearest neighbor search
• Run queries using standard SQL

App Overview

1. Call OpenAI’s embedding API to convert input text into vector embeddings.
2. Use the pgvector extension to store these embeddings in PostgreSQL.
3. Query embeddings to find the most semantically similar entries in the database.

Prerequisites

• Go installed (1.19 recommended).
• PostgreSQL installed and running (local or hosted).
• Install the pgvector extension in PostgreSQL. (See pgvector’s GitHub page for installation instructions.)
• OpenAI API key with embedded access.

Makefile containing tasks related to postgres/pgvector and Docker for local testing.

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

Make sure pgvector is installed. Then, in your PostgreSQL database:

CREATE EXTENSION IF NOT EXISTS vector;

Full 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 {
        // 將每個(gè)浮點(diǎn)數(shù)格式化為字符串
        strVals[i] = fmt.Sprintf("%f", val)
    }

    // 使用逗號(hào) + 空格連接它們
    joined := strings.Join(strVals, ", ")

    // pgvector 需要方括號(hào)表示法才能輸入向量，例如 [0.1, 0.2, 0.3]
    return "[" + joined + "]"
}

func main() {
    // 加載環(huán)境變量
    err := godotenv.Load()
    if err != nil {
        log.Fatal("加載 .env 文件出錯(cuò)")
    }

    // 創(chuàng)建連接池
    dbpool, err := pgxpool.New(context.Background(), os.Getenv("DATABASE_URL"))
    if err != nil {
        fmt.Fprintf(os.Stderr, "無(wú)法創(chuàng)建連接池：%v\n", err)
        os.Exit(1)
    }
    defer dbpool.Close()

    // 1. 確保已啟用 pgvector 擴(kuò)展
    _, err = dbpool.Exec(context.Background(), "CREATE EXTENSION IF NOT EXISTS vector;")
    if err != nil {
        log.Fatalf("創(chuàng)建擴(kuò)展失?。?v\n", err)
        os.Exit(1)
    }

    // 2. 創(chuàng)建表（如果不存在）
    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("創(chuàng)建表失?。?v\n", err)
    }

    // 3. 創(chuàng)建索引（如果不存在）
    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("創(chuàng)建索引失?。?v\n", err)
    }

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

    // 5. 插入示例文檔
    docs := []string{
        "PostgreSQL 是一個(gè)先進(jìn)的開(kāi)源關(guān)系數(shù)據(jù)庫(kù)。",
        "OpenAI 提供基于 GPT 的模型來(lái)生成文本嵌入。",
        "pgvector 允許將嵌入存儲(chǔ)在 Postgres 數(shù)據(jù)庫(kù)中。",
    }

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

    // 6. 查詢相似性
    queryText := "如何在 Postgres 中存儲(chǔ)嵌入？"
    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 調(diào)用失?。?w", err)
    }

    // 2) 將嵌入轉(zhuǎn)換為 pgvector 的方括號(hào)字符串
    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 獲取用戶查詢的嵌入，并根據(jù)向量相似性返回前 k 個(gè)相似的文檔。
func searchSimilarDocuments(ctx context.Context, pool *pgxpool.Pool, client *openai.Client, query string, k int) ([]string, error) {
    // 1) 通過(guò) 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 調(diào)用失?。?w", err)
    }

    // 2) 將 OpenAI 嵌入轉(zhuǎn)換為 pgvector 的方括號(hào)字符串格式
    queryEmbedding := embedResp.Data[0].Embedding // []float32
    queryEmbeddingStr := floats32ToString(queryEmbedding)
    // 例如 "[0.123456, 0.789012, ...]"

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

    // 4) 運(yùn)行查詢
    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("行迭代錯(cuò)誤：%w", err)
    }

    return contents, nil
}

Conclusion

OpenAI embeddings in PostgreSQL, Go and pgvector provide a straightforward solution for building semantic search applications. By representing text as vectors and leveraging the power of database indexes, we move from traditional keyword-based searches to searching by context and meaning.

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 clarity and readability. The key changes include more descriptive variable names and comments.

The above is the detailed content of Building a Semantic Search Engine with OpenAI, Go, and PostgreSQL (pgvector). For more information, please follow other related articles on the PHP Chinese website!