com.agentsflex.store.redis.RedisVectorStore Maven / Gradle / Ivy

Go to download
/*
 *  Copyright (c) 2023-2025, Agents-Flex ([email protected]).
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  

 *  http://www.apache.org/licenses/LICENSE-2.0
 *  

 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */
package com.agentsflex.store.redis;

import com.agentsflex.core.document.Document;
import com.agentsflex.core.store.DocumentStore;
import com.agentsflex.core.store.SearchWrapper;
import com.agentsflex.core.store.StoreOptions;
import com.agentsflex.core.store.StoreResult;
import com.alibaba.fastjson.JSON;
import kotlin.collections.ArrayDeque;
import redis.clients.jedis.*;
import redis.clients.jedis.json.Path2;
import redis.clients.jedis.search.*;
import redis.clients.jedis.search.schemafields.SchemaField;
import redis.clients.jedis.search.schemafields.TextField;
import redis.clients.jedis.search.schemafields.VectorField;

import java.net.URI;
import java.nio.ByteBuffer;
import java.nio.ByteOrder;
import java.nio.FloatBuffer;
import java.text.MessageFormat;
import java.util.*;

public class RedisVectorStore extends DocumentStore {

    private final RedisVectorStoreConfig config;
    private final JedisPooled jedis;
    private final Set redisIndexesCache = new HashSet<>();

    public RedisVectorStore(RedisVectorStoreConfig config) {
        this.config = config;
        this.jedis = new JedisPooled(
            URI.create(config.getUri())
        );
    }


    private void createSchemaIfNecessary(String indexName) {
        if (redisIndexesCache.contains(indexName)) {
            return;
        }

        // 检查 indexName 是否存在
        Set existIndexes = this.jedis.ftList();
        if (existIndexes != null && existIndexes.contains(indexName)) {
            redisIndexesCache.add(indexName);
            return;
        }

        FTCreateParams ftCreateParams = FTCreateParams.createParams()
            .on(IndexDataType.JSON)
            .addPrefix(this.config.getStorePrefix());

        jedis.ftCreate(indexName, ftCreateParams, schemaFields());
        redisIndexesCache.add(indexName);
    }


    private Iterable schemaFields() {
        Map vectorAttrs = new HashMap<>();
        //支持  COSINE: 余弦距离 , IP: 内积距离, L2: 欧几里得距离
        vectorAttrs.put("DISTANCE_METRIC", "COSINE");
        vectorAttrs.put("TYPE", "FLOAT32");
        vectorAttrs.put("DIM", this.getEmbeddingModel().dimensions());

        List fields = new ArrayList<>();
        fields.add(TextField.of(jsonPath("text")).as("text").weight(1.0));

        fields.add(VectorField.builder()
            .fieldName(jsonPath("vector"))
            .algorithm(VectorField.VectorAlgorithm.HNSW)
            .attributes(vectorAttrs)
            .as("vector")
            .build());

        return fields;
    }

    private String jsonPath(String field) {
        return "$." + field;
    }


    @Override
    public StoreResult storeInternal(List documents, StoreOptions options) {
        String indexName = createIndexName(options);
        createSchemaIfNecessary(indexName);

        try (Pipeline pipeline = jedis.pipelined();) {
            for (Document document : documents) {
                java.util.Map fields = new HashMap<>();
                fields.put("text", document.getContent());
                fields.put("vector", document.getVector());

                String key = config.getStorePrefix() + document.getId();
                pipeline.jsonSetWithEscape(key, Path2.of("$"), fields);
            }

            List