go-redis 使用RedisSearch进行向量的检索与查询

Redis 允许您在哈希JSON对象中索引向量字段(更多信息请参阅向量参考页面)。向量字段可以存储文本embedding等内容,文本embedding是 AI 生成的向量表示,用于表示文本片段中的语义信息。两个embedding之间的向量距离表明它们在语义上的相似程度。通过比较从查询文本生成的embedding与存储在哈希或 JSON 字段中的embedding的相似性,Redis 可以检索与查询的含义密切相关的文档。

创建索引


func doCreateIndex() (err error) {
    var (
        RedisKeyPrefix = "doc:"
        IndexName      = "vactor_test"
        Dimension      = 2560 //| 实际字节数 | 你传进来的 blob 长度 | 10240 字节 | → 10240 ÷ 4 = 2560 维
    )
    cli := NewRedisStackClient("localhost:6379", "", 0)
    ctx := context.Background()
    // 确保在错误时关闭连接
    defer func() {
        if err != nil {
            cli.Client.Close()
        }
    }()
    if err = cli.Client.Ping(ctx).Err(); err != nil {
        return fmt.Errorf("failed to connect to Redis: %w", err)
    }
    indexName := fmt.Sprintf("%s%s", RedisKeyPrefix, IndexName)
    // 检查是否存在索引
    exists, err := cli.Client.Do(ctx, "FT.INFO", indexName).Result()
    if err != nil {
        if !strings.Contains(err.Error(), "Unknown index name") {
            return fmt.Errorf("failed to check if index exists: %w", err)
        }
        err = nil
    } else if exists != nil {
        return nil
    }
    // Create new index
    createIndexArgs := []interface{}{
        "FT.CREATE", indexName,
        "ON", "HASH", //-- 数据载体:JSON 或 HASH
        "PREFIX", "1", RedisKeyPrefix, // -- 只扫描以 doc: 开头的键
        "SCHEMA",
        "content", "TEXT", //-- 业务字段示例content, 文本类型
        "genre", "TAG", //-- 业务字段示例metadata, TAG类型
        "embedding", "VECTOR", "FLAT", //-- 业务字段示例vector, 向量类型
        "6",               //-- 6 = 接下来 6 个参数
        "TYPE", "FLOAT32", //向量元素类型
        "DIM", Dimension, //向量维度,与模型一致
        "DISTANCE_METRIC", "COSINE", //距离算法:COSINE / L2 / IP
    }
    if err = cli.Client.Do(ctx, createIndexArgs...).Err(); err != nil {
        return fmt.Errorf("failed to create index: %w", err)
    }
    // 验证索引是否创建成功
    if _, err = cli.Client.Do(ctx, "FT.INFO", indexName).Result(); err != nil {
        return fmt.Errorf("failed to verify index creation: %w", err)
    }
    return nil
}

索引:doc:vactor_test只扫描以 doc: 开头的键 , 包含三个字段:

  • content: 内容, TEXT文本类型
  • genre: 类型, TAG标签类型
  • embedding:向量, VECTOR 类型,FLAT标识向量检索的方式(HNSW 适合在线、高并发近似搜索;若数据量很小可用 FLAT(暴力线性扫描)), FLAT的原理: 不做任何近似或聚类(与 HNSW / IVF 不同)。每次 FT.SEARCH … KNN 都把查询向量与索引中的 每一条向量 计算距离,再排序取 Top-k。因此 内存占用高(需要完整存储原始向量),但 召回率 100 %,也无额外调参。适用场景为数据量 ≤ 几万条或延迟要求不高。

DIM , 标识向量维度, 必须与模型一直, 这里使用的是字节ARK的模型,维度为 2560, 在 给 RediSearch 建索引时,你需把 DIM 设成实际要用的那个值, 在使用时,因一开始不确定维度数,设置的值384, 系统报以下错误:

Could not add vector with blob size 10240 (expected size 1536)"

期望字节数: DIM × 4, 即 384 × 4 = 1536 字节, 但实际ARK返回的向量维度为10240 * 4 = 2560(除4是因为float32 占4字节), 故删除索引后重建索引, 修改DIM值为2560, 上述问题解决, 删除索引:

func dropIndex() error {
    cli := NewRedisStackClient("localhost:6379", "", 0)
    ctx := context.Background()
    return cli.Client.FTDropIndex(ctx, "doc:vactor_test").Err()
}

DISTANCE_METRIC 用来告诉 RediSearch 向量之间的距离如何计算。目前支持三种,选其一即可:

取值 全称 公式(简化) 适用场景
COSINE Cosine Similarity 1 − cos(θ) 文本 / 语义向量,长度已归一化
L2 Euclidean Distance √Σ(xi − yi)² 通用数值向量,维度量纲一致
IP Inner Product − Σ(xi·yi) 已归一化且需要最大化内积

基于字节ARK的embedding 模型对文本向量化,并存在在redis中


func addEmbeddins() (err error) {
    cli := NewRedisStackClient("localhost:6379", "", 0)
    ctx := context.Background()
    // 确保在错误时关闭连接
    defer func() {
        if err != nil {
            cli.Client.Close()
        }
    }()
    sentences := []string{
        "That is a very happy person",
        "That is a happy dog",
        "Today is a sunny day",
    }
    tags := []string{
        "persons", "pets", "weather",
    }
    config := &ark.EmbeddingConfig{
        Model:  os.Getenv("ARK_EMBEDDING_MODEL"),
        APIKey: os.Getenv("ARK_API_KEY"),
    }
    eb, err := ark.NewEmbedder(ctx, config)
    if err != nil {
        return err
    }
    embeddings, err := eb.EmbedStrings(ctx, sentences)
    if err != nil {
        return err
    }
    for i, emb := range embeddings {
        buffer := utils.Vector2Bytes(emb)
        if err != nil {
            return err
        }
        count, err := cli.Client.HSet(ctx,
            fmt.Sprintf("doc:%v", i),
            map[string]any{
                "content":   sentences[i],
                "genre":     tags[i],
                "embedding": buffer,
            },
        ).Result()
        if err != nil {
            return err
        }
    }
    return nil
}

字节开源的大模型应用框架eino-ext 已封装Ark了对应的Embedder组件, 用于文档的向量化,需引入依赖: github.com/cloudwego/eino-ext/components/embedding/ark
Embedder组件的EmbedStrings方法,入参为string类型的切片, 出参为一个float64类型的二维数组, 对应每个字符串向量值, 存储时需要将float64转为float32 后再转bytes 切片存储:

func Vector2Bytes(vector []float64) []byte {
    float32Arr := make([]float32, len(vector))
    for i, v := range vector {
        float32Arr[i] = float32(v)
    }
    bytes := make([]byte, len(float32Arr)*4)
    for i, v := range float32Arr {
        binary.LittleEndian.PutUint32(bytes[i*4:], math.Float32bits(v))
    }
    return bytes
}

存入Redis 中的数据如下图所示:

image.png

向量查询

func doVectorQuery() (err error) {
    var (
        RedisKeyPrefix = "doc:"
        IndexName      = "vactor_test"
        queryStr = []string{"That is a happy person"}
    )
    cli := NewRedisStackClient("localhost:6379", "", 0)
    ctx := context.Background()
    config := &ark.EmbeddingConfig{
        Model:  os.Getenv("ARK_EMBEDDING_MODEL"),
        APIKey: os.Getenv("ARK_API_KEY"),
    }
    eb, err := ark.NewEmbedder(ctx, config)
    if err != nil {
        return err
    }
    embeddings, err := eb.EmbedStrings(ctx, queryStr)
    if err != nil {
        panic(err)
    }
    buffer := utils.FloatsToBytes(utils.ToFloat32(embeddings[0]))
    if err != nil {
        panic(err)
    }
    indexName := fmt.Sprintf("%s%s", RedisKeyPrefix, IndexName)

    //*=>[ ... ]
    //RediSearch 的 “向量查询子句” 语法糖,* 代表“所有文档”,=> 后面放向量运算。
    //KNN 3
    //取 3 个最近邻(k-nearest-neighbors)。
    //@embedding
    //指定 要比较的向量字段(索引里必须事先声明为 VECTOR 类型)。
    //$vec
    //用户传入的查询向量,需在 PARAMS 里绑定,例如 PARAMS 2 vec <base64或float数组>。
    //AS vector_distance
    //把计算出的距离(或相似度)作为 返回字段,后续可在 RETURN / SORTBY / DIALECT 里引用。
    results, err := cli.Client.FTSearchWithArgs(ctx,
        indexName,
        "*=>[KNN 3 @embedding $vec AS vector_distance]",
        &redis.FTSearchOptions{
            Return: []redis.FTSearchReturn{
                {FieldName: "vector_distance"},
                {FieldName: "content"},
            },
            DialectVersion: 2,
            Params: map[string]any{
                "vec": buffer,
            },
        },
    ).Result()

    if err != nil {
        panic(err)
    }

    for _, doc := range results.Docs {
        fmt.Printf(
            "ID: %v, Distance:%v, Content:'%v'\n",
            doc.ID, doc.Fields["vector_distance"], doc.Fields["content"],
        )
    }
    return nil
}

向量查询语句That is a happy person, 同样需要使用的Ark的Embedder组件进行向量化, 将Floate64类型转为float32类型后转bytes切片, 进行匹配查询:

    results, err := cli.Client.FTSearchWithArgs(ctx,
        indexName,
        "*=>[KNN 3 @embedding $vec AS vector_distance]",
        &redis.FTSearchOptions{
            Return: []redis.FTSearchReturn{
                {FieldName: "vector_distance"},
                {FieldName: "content"},
            },
            DialectVersion: 2,
            Params: map[string]any{
                "vec": buffer,
            },
        },
    ).Result()
  • =>[ ... ]
    RediSearch 的 “向量查询子句” 语法糖,
    代表“所有文档”,=> 后面放向量运算。
  • KNN 3
    取 3 个最近邻(k-nearest-neighbors)。
  • @embedding
    指定 要比较的向量字段(索引里必须事先声明为 VECTOR 类型)。
  • $vec
    用户传入的查询向量,需在 Params里绑定, 统一放入到map里。
  • AS vector_distance
    把计算出的距离(或相似度)作为 返回字段,后续可在 RETURN / SORTBY / DIALECT 里引用。

上述查询语句返回除ID之外声明的vector_distance和content字段, DialectVersion(在 Redis/RediSearch 中常写作 DIALECT)是 查询语法版本号,用来告诉 RediSearch 用哪一套解析器去解释你的 FT.SEARCH、FT.AGGREGATE 等命令。向量检索、参数绑定、JSON 多值字段 必须 DIALECT ≥ 2,否则会报错 “syntax error”。

That is a happy person 语句的查询返回输出:


image.png
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容