【LLM】RedisSearch 向量相似性搜索在 SpringBoot 中的实现

整理不易，请不要吝啬你的赞和收藏。

• 已安装 Redis Stack ，如何安装请参考 docker compose 安装 Redis Stack 。
• 已对 Redis 作为向量库有了解，如不了解请参考 Redis 作为向量库入门指南。
• 项目中引入了 Spring AI Alibaba ，如何引入请参考 Spring AI Alibaba 的简单使用 。

• 项目中引入了 Redis ，如何引入请参考 SpringBoot 引入 redis 。

<dependency>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-data-redis</artifactId>
    <exclusions>
        <exclusion>
            <groupId>io.lettuce</groupId>
            <artifactId>lettuce-core</artifactId>
        </exclusion>
    </exclusions>
</dependency>

<dependency>
    <groupId>redis.clients</groupId>
    <artifactId>jedis</artifactId>
</dependency>

<parent>
    <groupId>org.springframework.boot</groupId>
    <artifactId>spring-boot-starter-parent</artifactId>
    <version>3.2.4</version>
</parent>

<dependency>
    <groupId>com.alibaba.cloud.ai</groupId>
    <artifactId>spring-ai-alibaba-starter</artifactId>
</dependency>

spring:
  data:
    redis:
      database: 0
      host: 127.0.0.1
      port: 6379
      password: gusy1234
      jedis:
        pool:
          # 连接池最大连接数
          max-active: 8
          # 连接池最大空闲连接数
          max-idle: 8
          # 连接池最大阻塞等待时间，负值表示没有限制
          max-wait: 0
          # 连接池最小空闲连接数
          min-idle: 2
      # 连接超时时间（毫秒）
      timeout: 1000

@Bean
public UnifiedJedis unifiedJedis() {
    HostAndPort hostAndPort = new HostAndPort(host, port);
    // 设置 Jedis 客户端配置
    DefaultJedisClientConfig jedisClientConfig = DefaultJedisClientConfig.builder()
            .password(password)
            .database(database)
            .build();

    // 设置连接池参数
    ConnectionPoolConfig poolConfig = new ConnectionPoolConfig();
    poolConfig.setMaxTotal(maxActive);
    poolConfig.setMaxIdle(maxIdle);
    poolConfig.setMinIdle(minIdle);

    PooledConnectionProvider provider = new PooledConnectionProvider(hostAndPort, jedisClientConfig, poolConfig);
    return new UnifiedJedis(provider);
}

@Bean
public JedisPooled jedisPooled() throws URISyntaxException {
    ConnectionPoolConfig poolConfig = new ConnectionPoolConfig();
    // 池中最大活跃连接数，默认 8
    poolConfig.setMaxTotal(maxActive);
    // 池中最大空闲连接数，默认 8
    poolConfig.setMaxIdle(maxIdle);
    // 池中的最小空闲连接数，默认 0
    poolConfig.setMinIdle(minIdle);
    // 启用等待连接变为可用
    poolConfig.setBlockWhenExhausted(true);
    // 等待连接变为可用的最大秒数，jdk17 以上支持，低版本可用 setMaxWaitMillis
    poolConfig.setMaxWait(Duration.ofSeconds(1));
    // 控制检查池中空闲连接的间隔时间，jdk17 以上支持，低版本可用 setTimeBetweenEvictionRunsMillis
    poolConfig.setTimeBetweenEvictionRuns(Duration.ofSeconds(1));
    return new JedisPooled(poolConfig, new URI(redisUri));
}

    public List<Embedding> textEmbedding(List<String> texts) {
        // 调用百炼 Embedding 模型
        EmbeddingResponse embeddingResponse = dashScopeEmbeddingModel.call(
                new EmbeddingRequest(texts,
                        DashScopeEmbeddingOptions.builder()
                                .withModel(DashScopeApi.EmbeddingModel.EMBEDDING_V3.getValue()) // 设置使用的模型
                                .withTextType(DashScopeApi.DEFAULT_EMBEDDING_TEXT_TYPE) // 设置文本类型
                                .withDimensions(1024)   // 设置向量维度，可选768、1024、1536
                                .build()));
        List<Embedding> results = embeddingResponse.getResults();
        // 打印向量
//        if (CollectionUtils.isEmpty(results)) {
//            int tempSize = results.size();
//            for (int i = 0; i < tempSize; i++) {
//                float[] embeddingValue = results.get(i).getOutput();
//                log.info("embeddingValue:{}", embeddingValue.toString());
//            }
//        }
        return results;
    }

@Autowired
private RedisTemplate redisTemplate;

/**
 * 存入向量信息
 *
 * @param key          redis KEY
 * @param vectorField  向量存储字段名
 * @param vectorValue  向量值
 * @param contentField 向量文本内容存储字段名
 * @param content      向量文本内容，增加存储文本内容，方便查看向量的内容
 */
public void addDocument(String key, String vectorField, float[] vectorValue, String contentField, String content) {
    // 将向量值转为二进制
    byte[] vectorByte = CommonUtil.floatArrayToByteArray(vectorValue);
    // 组装字段map
    Map<String, Object> fieldMap = new LinkedMap<>();
    fieldMap.put(contentField, content.getBytes(StandardCharsets.UTF_8));
    fieldMap.put(vectorField, vectorByte);
    // 入库
    redisTemplate.opsForHash().putAll(key, fieldMap);
}

@Autowired
private JedisPooled jedisPooled;

/**
 * 相似度搜索
 *
 * @param queryVector 查询的向量内容
 * @param k           返回 k 个最相似内容
 * @param indexName   索引名称
 * @param vectorField 向量存储字段名
 * @return
 */
public SearchResult similaritySearch(float[] queryVector, int k, String indexName, String vectorField) {
    // 组装查询
    // 同：FT.SEARCH embedding_index "* => [KNN 3 @vector_field $query_vector AS distance]"
    //              PARAMS 2 query_vector "\x12\xa9\xf5\x6c"
    //              SORTBY distance
    //              DIALECT 4
    Query q = new Query("*=>[KNN $K @" + vectorField + " $query_vector AS distance]").
            returnFields("content", "distance").
            addParam("K", k).
            addParam("query_vector", CommonUtil.floatArrayToByteArray(queryVector)).
            setSortBy("distance", true).
            dialect(4);
    // 使用 UnifiedJedis 执行 FT.SEARCH 语句
    try {
        SearchResult result = jedisPooled.ftSearch(indexName, q);
        log.info("redis相似度搜索，result:{}", JSONObject.toJSONString(result));
        return result;
    }
    catch (Exception e) {
        log.warn("redis相似度搜索，异常:", e);
        throw new ErrorCodeException(e.getMessage());
    }
}

/**
 * float[] 转二进制
 *
 * @param floats
 * @return
 */
public static byte[] floatArrayToByteArray(float[] floats) {
    byte[] bytes = new byte[Float.BYTES * floats.length];
    ByteBuffer.wrap(bytes).order(ByteOrder.LITTLE_ENDIAN).asFloatBuffer().put(floats);
    return bytes;
}

<dependency>
    <groupId>com.alibaba.cloud.ai</groupId>
    <artifactId>spring-ai-alibaba-starter</artifactId>
</dependency>

<dependency>
    <groupId>org.springframework.ai</groupId>
    <artifactId>spring-ai-redis-store-spring-boot-starter</artifactId>
</dependency>

spring:
  ai:
    dashscope:
      # 阿里百炼平台申请的 api_key -->
      api-key: your api_key
    chat:
      client:
        enabled: true

spring:
  data:
    redis:
      # uri 为 spring ai 中使用 redis 作为 Vector Store 的配置
      uri: redis://:gusy1234@127.0.0.1:6379/0
  ai:
    # 这里我多包了一层 mine ,因为如果多加一层，spring 会自动装配
    mine:
      # redis 作为向量库的配置    
      vectorstore:
        redis:
          # 是否初始化索引信息
          initialize-schema: true
          # 索引名称，默认 spring-ai-index
          index-name: spring-ai-index
          # 向量字段 key 的前缀，默认 embedding:
          prefix: 'embedding-ai:'
          # 文档批处理计算嵌入的策略。选项为 TOKEN_COUNT 或 FIXED_SIZE
          batching-strategy: TOKEN_COUNT

    @Resource
    private DashScopeEmbeddingModel dashScopeEmbeddingModel;

    @Resource
    private JedisPooled jedisPooled;

    @Bean("vectorStoreWithDashScope")
    public VectorStore vectorStore() {
        return new RedisVectorStore(RedisVectorStore.RedisVectorStoreConfig.builder()
                .withIndexName(indexName)
                .withPrefix(prefix)
//                .withEmbeddingFieldName("embedding")  // 向量字段名，默认 embedding
//                .withContentFieldName("content")  // 文本字段名，默认 content
                .build(), embeddingModel, jedisPooled, initializeSchema); // initializeSchema 为是否初始化索引配置信息
    }

@Autowired
@Qualifier("vectorStoreWithDashScope")
private VectorStore vectorStore;

// 向量入库主要代码
List<String> inputInfos = List.of("文本1","文本2");
List<Document> documentList = new ArrayList<>();
inputInfos.forEach(text -> {
    documentList.add(new Document(text));
});
vectorStore.add(documentList);

// 相似度查询主要代码
List<Document> result = vectorStore.similaritySearch(org.springframework.ai.vectorstore.SearchRequest.defaults()
        .withQuery(inputInfos.get(0))    // 查询的内容
        .withTopK(3)
        // 一个值从 0 到 1 的双精度数，接近1的值表示更高的相似性。默认为为0.75。
        .withSimilarityThreshold(0.6)
);