Add the ability to store the KnowledgeGraphWriter results in vector indexes via Neo4jEmbeddingStore (#217)

vga91 · web-flow · commit 065c581f4c43 · 2025-06-03T22:17:53.000+08:00
diff --git a/content-retrievers/langchain4j-community-neo4j-retriever/src/main/java/dev/langchain4j/community/rag/content/retriever/neo4j/KnowledgeGraphWriter.java b/content-retrievers/langchain4j-community-neo4j-retriever/src/main/java/dev/langchain4j/community/rag/content/retriever/neo4j/KnowledgeGraphWriter.java
@@ -7,7 +7,14 @@
 
 import dev.langchain4j.Experimental;
 import dev.langchain4j.community.data.document.graph.GraphDocument;
+import dev.langchain4j.community.data.document.graph.GraphNode;
+import dev.langchain4j.community.store.embedding.neo4j.Neo4jEmbeddingStore;
 import dev.langchain4j.data.document.Document;
+import dev.langchain4j.data.document.Metadata;
+import dev.langchain4j.data.embedding.Embedding;
+import dev.langchain4j.data.segment.TextSegment;
+import dev.langchain4j.model.embedding.EmbeddingModel;
+import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -34,23 +41,38 @@ public class KnowledgeGraphWriter {
     final String sanitizedTextProperty;
 
     private final Neo4jGraph graph;
+    private final Neo4jEmbeddingStore embeddingStore;
+    private EmbeddingModel embeddingModel = null;
 
     public KnowledgeGraphWriter(
             Neo4jGraph graph,
             String idProperty,
             String label,
             String textProperty,
             String relType,
-            String constraintName) {
+            String constraintName,
+            Neo4jEmbeddingStore embeddingStore,
+            EmbeddingModel embeddingModel) {
         this.graph = ensureNotNull(graph, "graph");
+
+        this.embeddingStore = embeddingStore;
+        final boolean storeIsNull = this.embeddingStore == null;
+        if (!storeIsNull) {
+            this.embeddingModel = ensureNotNull(embeddingModel, "embeddingModel");
+        }
+
         this.label = getOrDefault(label, DEFAULT_LABEL);
         this.relType = getOrDefault(relType, DEFAULT_REL_TYPE);
+
         this.idProperty = getOrDefault(idProperty, DEFAULT_ID_PROP);
         this.textProperty = getOrDefault(textProperty, DEFAULT_TEXT_PROP);
         this.constraintName = getOrDefault(constraintName, DEFAULT_CONS_NAME);
 
         /* sanitize labels and property names, to prevent from Cypher Injections */
-        this.sanitizedLabel = sanitizeOrThrows(this.label, "label");
+
+        // if embeddingStore then label is taken from there getSanitizedLabel()
+        this.sanitizedLabel =
+                storeIsNull ? sanitizeOrThrows(this.label, "label") : this.embeddingStore.getSanitizedLabel();
         this.sanitizedRelType = sanitizeOrThrows(this.relType, "relType");
         this.sanitizedIdProperty = sanitizeOrThrows(this.idProperty, "idProperty");
         this.sanitizedTextProperty = sanitizeOrThrows(this.textProperty, "textProperty");
@@ -78,9 +100,10 @@ public void addGraphDocuments(List<GraphDocument> graphDocuments, boolean includ
 
             // Import nodes
             Map<String, Object> nodeParams = new HashMap<>();
-            nodeParams.put(
-                    "data", graphDoc.nodes().stream().map(Neo4jUtils::toMap).toList());
-
+            if (embeddingStore == null) {
+                nodeParams.put(
+                        "rows", graphDoc.nodes().stream().map(Neo4jUtils::toMap).toList());
+            }
             if (includeSource) {
                 // create a copyOf metadata, not to update existing graphDoc,
                 // subsequent tests could potentially fail
@@ -93,8 +116,7 @@ public void addGraphDocuments(List<GraphDocument> graphDocuments, boolean includ
                 nodeParams.put("document", document);
             }
 
-            String nodeImportQuery = getNodeImportQuery(includeSource);
-            graph.executeWrite(nodeImportQuery, nodeParams);
+            insertNodes(includeSource, graphDoc, nodeParams);
 
             // Import relationships
             List<Map<String, String>> relData = graphDoc.relationships().stream()
@@ -105,23 +127,60 @@ public void addGraphDocuments(List<GraphDocument> graphDocuments, boolean includ
                             "target_label", rel.targetNode().type(),
                             "type", rel.type().replace(" ", "_").toUpperCase()))
                     .toList();
-
             String relImportQuery = getRelImportQuery();
             graph.executeWrite(relImportQuery, Map.of("data", relData));
         }
     }
 
+    private void insertNodes(boolean includeSource, GraphDocument graphDoc, Map<String, Object> nodeParams) {
+        if (embeddingStore == null) {
+            String nodeImportQuery = getNodeImportQuery(includeSource);
+            graph.executeWrite(nodeImportQuery, nodeParams);
+            return;
+        }
+
+        if (includeSource) {
+            final String creationQuery = mergeSourceWithDocs(true)
+                    + """
+                            SET source += row.%3$s
+                            WITH row, source
+                            CALL db.create.setNodeVectorProperty(source, $embeddingProperty, row.%4$s)
+                            RETURN count(*)""";
+            embeddingStore.setEntityCreationQuery(creationQuery);
+            embeddingStore.setAdditionalParams(nodeParams);
+        }
+
+        // we save the ids, otherwise it create UUID properties and the merge with import relationships doesn't work
+        List<String> ids = new ArrayList<>();
+        List<TextSegment> segments = new ArrayList<>();
+        for (GraphNode node : graphDoc.nodes()) {
+            final Map<String, String> properties = new HashMap<>(node.properties());
+            properties.put("type", node.type());
+            final String id = node.id();
+            final TextSegment segment = TextSegment.from(id, Metadata.from(properties));
+            ids.add(id);
+            segments.add(segment);
+        }
+
+        final List<Embedding> embeddings = embeddingModel.embedAll(segments).content();
+
+        this.embeddingStore.addAll(ids, embeddings, segments);
+    }
+
     private String getNodeImportQuery(boolean includeSource) {
 
+        return mergeSourceWithDocs(includeSource) + "WITH source, row "
+                + "SET source:$(row.type) "
+                + "RETURN count(*) as total";
+    }
+
+    private String mergeSourceWithDocs(boolean includeSource) {
         String includeDocsQuery = getIncludeDocsQuery(includeSource);
         final String withDocsRel = includeSource ? String.format("MERGE (d)-[:%s]->(source) ", relType) : "";
 
-        return includeDocsQuery + "UNWIND $data AS row "
-                + String.format("MERGE (source:%1$s {%2$s: row.id}) ", sanitizedLabel, sanitizedIdProperty)
-                + withDocsRel
-                + "WITH source, row "
-                + "SET source:$(row.type) "
-                + "RETURN count(*) as total";
+        return includeDocsQuery + "UNWIND $rows AS row \n"
+                + String.format("MERGE (source:%1$s {%2$s: row.id}) \n", sanitizedLabel, sanitizedIdProperty)
+                + withDocsRel;
     }
 
     private String getIncludeDocsQuery(boolean includeSource) {
@@ -160,6 +219,8 @@ public static class Builder {
         private String relType;
         private String constraintName;
         private Neo4jGraph graph;
+        private Neo4jEmbeddingStore embeddingStore;
+        private EmbeddingModel embeddingModel;
 
         /**
          * @param graph the {@link Neo4jGraph} (required)
@@ -212,8 +273,30 @@ public Builder constraintName(String constraintName) {
             return this;
         }
 
+        /**
+         * Sets the optional embedding store used to store texts as vectors via
+         * {@link Neo4jEmbeddingStore#add(dev.langchain4j.data.embedding.Embedding)}.
+         *
+         * @param embeddingStore the {@link Neo4jEmbeddingStore} instance to store vector embeddings (optional)
+         */
+        public Builder embeddingStore(Neo4jEmbeddingStore embeddingStore) {
+            this.embeddingStore = embeddingStore;
+            return this;
+        }
+
+        /**
+         * Sets the embedding model to be used for embedding text, if {@code embeddingStore} is provided.
+         *
+         * @param embeddingModel the {@link EmbeddingModel} used to generate embeddings
+         */
+        public Builder embeddingModel(EmbeddingModel embeddingModel) {
+            this.embeddingModel = embeddingModel;
+            return this;
+        }
+
         public KnowledgeGraphWriter build() {
-            return new KnowledgeGraphWriter(graph, idProperty, label, textProperty, relType, constraintName);
+            return new KnowledgeGraphWriter(
+                    graph, idProperty, label, textProperty, relType, constraintName, embeddingStore, embeddingModel);
         }
     }
 }
diff --git a/content-retrievers/langchain4j-community-neo4j-retriever/src/test/java/dev/langchain4j/community/rag/content/retriever/neo4j/Neo4jKnowledgeGraphWriterBaseTest.java b/content-retrievers/langchain4j-community-neo4j-retriever/src/test/java/dev/langchain4j/community/rag/content/retriever/neo4j/Neo4jKnowledgeGraphWriterBaseTest.java
@@ -87,7 +87,7 @@ abstract class Neo4jKnowledgeGraphWriterBaseTest {
     public static final String VALUE_KEANU = "value3";
     public static String USERNAME = "neo4j";
     public static String ADMIN_PASSWORD = "adminPass";
-    private static final String NEO4J_VERSION = System.getProperty("neo4jVersion", "2025.01.0-enterprise");
+    private static final String NEO4J_VERSION = System.getProperty("neo4jVersion", "2025.04.0-enterprise");
 
     public static String CAT_ON_THE_TABLE = "Sylvester the cat is on the table";
     public static String KEANU_REEVES_ACTED = "Keanu Reeves acted in Matrix";
diff --git a/content-retrievers/langchain4j-community-neo4j-retriever/src/test/java/dev/langchain4j/community/rag/content/retriever/neo4j/Neo4jKnowledgeGraphWriterTest.java b/content-retrievers/langchain4j-community-neo4j-retriever/src/test/java/dev/langchain4j/community/rag/content/retriever/neo4j/Neo4jKnowledgeGraphWriterTest.java
@@ -5,10 +5,17 @@
 import static org.mockito.ArgumentMatchers.argThat;
 import static org.mockito.Mockito.when;
 
+import dev.langchain4j.community.store.embedding.neo4j.Neo4jEmbeddingStore;
+import dev.langchain4j.data.embedding.Embedding;
 import dev.langchain4j.data.message.AiMessage;
 import dev.langchain4j.data.message.ChatMessage;
+import dev.langchain4j.data.segment.TextSegment;
 import dev.langchain4j.model.chat.ChatModel;
 import dev.langchain4j.model.chat.response.ChatResponse;
+import dev.langchain4j.model.embedding.EmbeddingModel;
+import dev.langchain4j.model.embedding.onnx.allminilml6v2q.AllMiniLmL6V2QuantizedEmbeddingModel;
+import dev.langchain4j.store.embedding.EmbeddingMatch;
+import dev.langchain4j.store.embedding.EmbeddingSearchRequest;
 import java.util.List;
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.extension.ExtendWith;
@@ -56,4 +63,96 @@ void testWrongConstraintName() {
             assertThat(e.getMessage()).contains("Error executing query: CREATE CONSTRAINT 111");
         }
     }
+
+    private static final String LABEL_CUSTOM = "Custom";
+    private static final EmbeddingModel EMBEDDING_MODEL = new AllMiniLmL6V2QuantizedEmbeddingModel();
+
+    @Test
+    void testKnowledgeGraphWithEmbeddingStoreAndNullEmbeddingModel() {
+
+        final Neo4jEmbeddingStore embeddingStore = Neo4jEmbeddingStore.builder()
+                .withBasicAuth(neo4jContainer.getBoltUrl(), USERNAME, ADMIN_PASSWORD)
+                .dimension(384)
+                .label(LABEL_CUSTOM)
+                .build();
+
+        try {
+            knowledgeGraphWriter = KnowledgeGraphWriter.builder()
+                    .graph(neo4jGraph)
+                    .embeddingStore(embeddingStore)
+                    .build();
+            fail("Should fail due to null embeddingModel");
+        } catch (Exception e) {
+            assertThat(e.getMessage()).contains("embeddingModel cannot be null");
+        }
+    }
+
+    @Test
+    void testKnowledgeGraphWithEmbeddingStore() {
+
+        final Neo4jEmbeddingStore embeddingStore = Neo4jEmbeddingStore.builder()
+                .withBasicAuth(neo4jContainer.getBoltUrl(), USERNAME, ADMIN_PASSWORD)
+                .dimension(384)
+                .label(LABEL_CUSTOM)
+                .build();
+
+        testKnowledgeGraphWithEmbeddingStoreCommon(embeddingStore, false);
+    }
+
+    @Test
+    void testKnowledgeGraphWithEmbeddingStoreAndIncludeSource() {
+        final Neo4jEmbeddingStore embeddingStore = Neo4jEmbeddingStore.builder()
+                .withBasicAuth(neo4jContainer.getBoltUrl(), USERNAME, ADMIN_PASSWORD)
+                .dimension(384)
+                .label(LABEL_CUSTOM)
+                .build();
+
+        testKnowledgeGraphWithEmbeddingStoreCommon(embeddingStore, true);
+    }
+
+    @Test
+    void testKnowledgeGraphWithEmbeddingStoreRetrievalQueryAndIncludeSource() {
+
+        final Neo4jEmbeddingStore embeddingStore = Neo4jEmbeddingStore.builder()
+                .withBasicAuth(neo4jContainer.getBoltUrl(), USERNAME, ADMIN_PASSWORD)
+                .dimension(384)
+                .label(LABEL_CUSTOM)
+                .retrievalQuery(
+                        """
+                        MATCH (node)<-[r:HAS_ENTITY]-(d:Document)
+                        WITH d, collect(DISTINCT {chunk: node, score: score}) AS chunks, avg(score) as avg_score
+                        RETURN d.text AS text, avg_score AS score, properties(d) AS metadata
+                        ORDER BY score DESC
+                        LIMIT $maxResults
+                        """)
+                .build();
+
+        testKnowledgeGraphWithEmbeddingStoreCommon(embeddingStore, true);
+    }
+
+    private static void testKnowledgeGraphWithEmbeddingStoreCommon(
+            Neo4jEmbeddingStore embeddingStore, boolean includeSource) {
+        final String text = "keanu reeves";
+        final Embedding queryEmbedding = EMBEDDING_MODEL.embed(text).content();
+        final EmbeddingSearchRequest request = EmbeddingSearchRequest.builder()
+                .queryEmbedding(queryEmbedding)
+                .minScore(0.9)
+                .build();
+        final List<EmbeddingMatch<TextSegment>> matchesBefore =
+                embeddingStore.search(request).matches();
+        assertThat(matchesBefore).isEmpty();
+
+        knowledgeGraphWriter = KnowledgeGraphWriter.builder()
+                .graph(neo4jGraph)
+                .embeddingStore(embeddingStore)
+                .embeddingModel(EMBEDDING_MODEL)
+                .build();
+
+        knowledgeGraphWriter.addGraphDocuments(graphDocs, includeSource);
+
+        final List<EmbeddingMatch<TextSegment>> matches =
+                embeddingStore.search(request).matches();
+        assertThat(matches).hasSize(1);
+        assertThat(matches.get(0).embedded().text()).containsIgnoringCase(text);
+    }
 }
diff --git a/embedding-stores/langchain4j-community-neo4j/src/main/java/dev/langchain4j/community/store/embedding/neo4j/Neo4jEmbeddingStore.java b/embedding-stores/langchain4j-community-neo4j/src/main/java/dev/langchain4j/community/store/embedding/neo4j/Neo4jEmbeddingStore.java
@@ -131,7 +131,7 @@ public class Neo4jEmbeddingStore implements EmbeddingStore<TextSegment> {
     private final String sanitizedLabel;
     private final String textProperty;
     private final String retrievalQuery;
-    private final String entityCreationQuery;
+    private String entityCreationQuery;
     private final Set<String> notMetaKeys;
     private Map<String, Object> additionalParams;
 
@@ -282,6 +282,10 @@ public void setAdditionalParams(final Map<String, Object> additionalParams) {
         this.additionalParams = additionalParams;
     }
 
+    public void setEntityCreationQuery(final String entityCreationQuery) {
+        this.entityCreationQuery = entityCreationQuery;
+    }
+
     /*
     Methods with `@Override`
     */
