apoc.vectordb.WeaviateHandler Maven / Gradle / Ivy

Go to download
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.ml.RestAPIConfig.BODY_KEY;
import static apoc.ml.RestAPIConfig.METHOD_KEY;
import static apoc.util.MapUtil.map;
import static apoc.vectordb.VectorEmbeddingConfig.FIELDS_KEY;
import static apoc.vectordb.VectorEmbeddingConfig.METADATA_KEY;
import static apoc.vectordb.VectorEmbeddingConfig.VECTOR_KEY;

public class WeaviateHandler implements VectorDbHandler {

    @Override
    public String getUrl(String hostOrKey) {
        String url = new UrlResolver("http", "localhost", 8000)
                .getUrl("weaviate", hostOrKey);
        return url + "/v1";
    }

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

    @Override
    public String getLabel() {
        return "Weaviate";
    }

    // -- embedding handler
    static class WeaviateEmbeddingHandler implements VectorEmbeddingHandler {

        @Override
        public  VectorEmbeddingConfig fromGet(Map config, ProcedureCallContext procedureCallContext, List ids, String collection) {
            config.putIfAbsent(BODY_KEY, null);
            return populateApiBodyRequest(getVectorEmbeddingConfig(config), Map.of());
        }

        @Override
        public VectorEmbeddingConfig fromQuery(Map config, ProcedureCallContext procedureCallContext, List vector, Object filter, long limit, String collection) {
            List fields = procedureCallContext.outputFields().toList();
            config.putIfAbsent(METHOD_KEY, "POST");
            VectorEmbeddingConfig vectorEmbeddingConfig = getVectorEmbeddingConfig(config);

            List list = (List) config.get(FIELDS_KEY);
            if (list == null) {
                throw new RuntimeException("You have to define `field` list of parameter to be returned");
            }
            Object fieldList = String.join("\n", list);

            filter = filter == null
                    ? ""
                    : ", where: " + filter;

            String includeVector = (fields.contains("vector") && vectorEmbeddingConfig.isAllResults()) ? ",vector" : "";
            String additional = "_additional {id, distance " + includeVector  + "}";
            String query = """
                  {
                      Get {
                        %s(limit: %s, nearVector: {vector: %s } %s) {%s  %s}
                      }
                  }
                    """.formatted(
                    collection, limit, vector, filter, fieldList, additional
            );

            Map additionalBodies = map("query", query);

            return populateApiBodyRequest(vectorEmbeddingConfig, additionalBodies);
        }

        private static VectorEmbeddingConfig getVectorEmbeddingConfig(Map config) {
            config.putIfAbsent(VECTOR_KEY, "vector");
            config.putIfAbsent(METADATA_KEY, "properties");

            return new VectorEmbeddingConfig(config);
        }
    }

}