>백엔드 개발 >Golang >OpenAI, Go 및 PostgreSQL을 사용하여 의미 체계 검색 엔진 구축(pgVector)

OpenAI, Go 및 PostgreSQL을 사용하여 의미 체계 검색 엔진 구축(pgVector)

Linda Hamilton
Linda Hamilton원래의
2025-01-15 11:09:44764검색

Building a Semantic Search Engine with OpenAI, Go, and PostgreSQL (pgvector)

최근 몇 년 동안 벡터 임베딩은 현대 자연어 처리(NLP) 및 의미 검색의 기초가 되었습니다. 벡터 데이터베이스는 키워드 검색에 의존하는 대신 숫자 표현(임베딩)을 통해 텍스트의 "의미"를 비교합니다. 이 예에서는 OpenAI 임베딩, Go 및 PostgreSQL과 pgVector 확장을 사용하여 의미 체계 검색 엔진을 만드는 방법을 보여줍니다.

임베딩이란 무엇인가요?

임베딩은 고차원 공간에서 텍스트(또는 기타 데이터)를 벡터로 표현하는 것입니다. 두 텍스트 조각이 의미상 유사하다면 해당 벡터는 이 공간에서 서로 가까워집니다. PostgreSQL(pgVector 확장명 포함)과 같은 데이터베이스에 임베딩을 저장함으로써 유사성 검색을 빠르고 정확하게 수행할 수 있습니다.

PostgreSQL과 pgVector를 선택하는 이유는 무엇입니까?

pgVector는 PostgreSQL에 벡터 데이터 유형을 추가하는 널리 사용되는 확장입니다. 이를 통해 다음을 수행할 수 있습니다.

  • 임베딩을 벡터 열로 저장
  • 대략적이거나 정확한 최근접 이웃 검색 수행
  • 표준 SQL을 사용하여 쿼리 실행

앱 개요

  1. OpenAI의 임베딩 API를 호출하여 입력 텍스트를 벡터 임베딩으로 변환합니다.
  2. PgVector 확장을 사용하여 PostgreSQL에 이러한 임베딩을 저장하세요.
  3. 임베딩을 쿼리하여 데이터베이스에서 의미상 가장 유사한 항목을 찾습니다.

전제조건

  • 설치하세요(1.19 권장).
  • PostgreSQL이 설치되어 실행 중입니다(로컬 또는 호스트).
  • PostgreSQL에 pgVector 확장을 설치합니다. (설치 지침은 pgVector의 GitHub 페이지를 참조하세요.)
  • 액세스가 내장된 OpenAI API 키입니다.

로컬 테스트를 위한 postgres/pgVector 및 Docker 관련 작업이 포함된 Makefile입니다.

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

pgVector가 설치되어 있는지 확인하세요. 그런 다음 PostgreSQL 데이터베이스에서:

<code class="language-sql">CREATE EXTENSION IF NOT EXISTS vector;</code>

전체 코드

<code class="language-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, ", ")

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

결론

PostgreSQL, Go 및 pgVector의 OpenAI 임베딩은 의미 체계 검색 애플리케이션 구축을 위한 간단한 솔루션을 제공합니다. 텍스트를 벡터로 표현하고 데이터베이스 색인의 힘을 활용함으로써 우리는 전통적인 키워드 기반 검색에서 문맥과 의미에 따른 검색으로 전환합니다.

이 수정된 출력은 원래의 언어 스타일을 유지하고, 독창성을 위해 문장을 바꾸며, 이미지를 동일한 형식과 위치로 유지합니다. 코드도 약간 개선되어 명확성과 가독성이 향상되었습니다.

위 내용은 OpenAI, Go 및 PostgreSQL을 사용하여 의미 체계 검색 엔진 구축(pgVector)의 상세 내용입니다. 자세한 내용은 PHP 중국어 웹사이트의 기타 관련 기사를 참조하세요!

성명:
본 글의 내용은 네티즌들의 자발적인 기여로 작성되었으며, 저작권은 원저작자에게 있습니다. 본 사이트는 이에 상응하는 법적 책임을 지지 않습니다. 표절이나 침해가 의심되는 콘텐츠를 발견한 경우 admin@php.cn으로 문의하세요.
이전 기사:DSA 브레이슬릿다음 기사:DSA 브레이슬릿