apoc.vectordb.MilvusHandler Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of apoc-extended Show documentation
Extended package for Neo4j Procedures
There is a newer version: 5.26.0
package apoc.vectordb;

import apoc.util.UrlResolver;
import org.neo4j.internal.kernel.api.procs.ProcedureCallContext;

import java.util.List;
import java.util.Map;

import static apoc.util.MapUtil.map;
import static apoc.vectordb.VectorEmbeddingConfig.FIELDS_KEY;
import static apoc.vectordb.VectorEmbeddingConfig.META_AS_SUBKEY_KEY;
import static apoc.vectordb.VectorEmbeddingConfig.SCORE_KEY;

public class MilvusHandler implements VectorDbHandler {

    @Override
    public String getUrl(String hostOrKey) {
        String url = new UrlResolver("http", "localhost", 19530).getUrl("milvus", hostOrKey);
        return url + "/v2/vectordb";
    }

    @Override
    public VectorEmbeddingHandler getEmbedding() {
        return new MilvusEmbeddingHandler();
    }

    @Override
    public String getLabel() {
        return "Milvus";
    }
    
    // -- embedding handler
    static class MilvusEmbeddingHandler implements VectorEmbeddingHandler {

        @Override
        public  VectorEmbeddingConfig fromGet(Map config, ProcedureCallContext procedureCallContext, List ids, String collection) {
            List fields = procedureCallContext.outputFields().toList();
            Map additionalBodies = map("id", ids);

            return getVectorEmbeddingConfig(config, fields, collection, additionalBodies);
        }

        @Override
        public VectorEmbeddingConfig fromQuery(Map config, ProcedureCallContext procedureCallContext, List vector, Object filter, long limit, String collection) {
            config.putIfAbsent(SCORE_KEY, "distance");

            List fields = procedureCallContext.outputFields().toList();
            Map additionalBodies = map("data", List.of(vector),
                    "limit", limit);
            if (filter != null) {
                additionalBodies.put("filter", filter);
            }

            return getVectorEmbeddingConfig(config, fields, collection, additionalBodies);
        }

        private VectorEmbeddingConfig getVectorEmbeddingConfig(Map config, List procFields, String collection, Map additionalBodies) {
            config.putIfAbsent(META_AS_SUBKEY_KEY, false);

            List listFields = (List) config.get(FIELDS_KEY);
            if (listFields == null) {
                throw new RuntimeException("You have to define `field` list of parameter to be returned");
            }
            if (procFields.contains("vector") && !listFields.contains("vector")) {
                listFields.add("vector");
            }
            VectorEmbeddingConfig conf = new VectorEmbeddingConfig(config);
            additionalBodies.put("collectionName", collection);
            additionalBodies.put("outputFields", listFields);

            return populateApiBodyRequest(conf, additionalBodies);
        }
    }
}