Skip to content

Commit 8a0aeba

Browse files
vga91vga91
committed
[NOID] Fixes #4091: Update RAG docs with vector db examples (#4116)
1 parent d4f43fa commit 8a0aeba

File tree

7 files changed

+228
-5
lines changed

7 files changed

+228
-5
lines changed

‎docs/asciidoc/modules/ROOT/pages/database-integration/vectordb/chroma.adoc‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -148,6 +148,22 @@ CALL apoc.vectordb.chroma.query($host, '<collection_id>',
148148
----
149149

150150

151+
[NOTE]
152+
====
153+
To optimize performances, we can choose what to `YIELD` with the apoc.vectordb.chroma.query and the `apoc.vectordb.chroma.get` procedures.
154+
For example, by executing a `CALL apoc.vectordb.chroma.query(...) YIELD metadata, score, id`, the RestAPI request will have an {"include": ["metadatas", "documents", "distances"]},
155+
so that we do not return the other values that we do not need.
156+
====
157+
158+
It is possible to execute vector db procedures together with the xref::ml/rag.adoc[apoc.ml.rag] as follow:
159+
160+
[source,cypher]
161+
----
162+
CALL apoc.vectordb.chroma.getAndUpdate($host, $collection, [<id1>, <id2>], $conf) YIELD node, metadata, id, vector
163+
WITH collect(node) as paths
164+
CALL apoc.ml.rag(paths, $attributes, $question, $confPrompt) YIELD value
165+
RETURN value
166+
----
151167

152168
.Delete vectors (it leverages https://docs.trychroma.com/usage-guide#deleting-data-from-a-collection[this API])
153169
[source,cypher]

‎docs/asciidoc/modules/ROOT/pages/database-integration/vectordb/qdrant.adoc‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
21
== Qdrant
32

43
Here is a list of all available Qdrant procedures,
@@ -200,7 +199,15 @@ For example, by executing a `CALL apoc.vectordb.qdrant.query(...) YIELD metadata
200199
so that we do not return the other values that we do not need.
201200
====
202201

202+
It is possible to execute vector db procedures together with the xref::ml/rag.adoc[apoc.ml.rag] as follow:
203203

204+
[source,cypher]
205+
----
206+
CALL apoc.vectordb.qdrant.getAndUpdate($host, $collection, [<id1>, <id2>], $conf) YIELD node, metadata, id, vector
207+
WITH collect(node) as paths
208+
CALL apoc.ml.rag(paths, $attributes, $question, $confPrompt) YIELD value
209+
RETURN value
210+
----
204211

205212
.Delete vectors (it leverages https://qdrant.github.io/qdrant/redoc/index.html#tag/points/operation/delete_vectors[this API])
206213
[source,cypher]

‎docs/asciidoc/modules/ROOT/pages/database-integration/vectordb/weaviate.adoc‎

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,3 @@
1-
21
== Weaviate
32

43
Here is a list of all available Weaviate procedures,
@@ -213,7 +212,15 @@ For example, by executing a `CALL apoc.vectordb.weaviate.query(...) YIELD metada
213212
so that we do not return the other values that we do not need.
214213
====
215214

215+
It is possible to execute vector db procedures together with the xref::ml/rag.adoc[apoc.ml.rag] as follow:
216216

217+
[source,cypher]
218+
----
219+
CALL apoc.vectordb.weaviate.getAndUpdate($host, $collection, [<id1>, <id2>], $conf) YIELD score, node, metadata, id, vector
220+
WITH collect(node) as paths
221+
CALL apoc.ml.rag(paths, $attributes, $question, $confPrompt) YIELD value
222+
RETURN value
223+
----
217224

218225
.Delete vectors (it leverages https://weaviate.io/developers/weaviate/api/rest#tag/objects/delete/objects/\{className\}/\{id\}[this API])
219226
[source,cypher]

‎full-it/src/test/java/apoc/full/it/vectordb/ChromaDbTest.java‎

Lines changed: 53 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,23 @@
11
package apoc.full.it.vectordb;
22

3+
import apoc.ml.Prompt;
4+
import apoc.util.TestUtil;
5+
import org.junit.AfterClass;
6+
import org.junit.Before;
7+
import org.junit.BeforeClass;
8+
import org.junit.ClassRule;
9+
import org.junit.Test;
10+
import org.junit.rules.TemporaryFolder;
11+
import org.neo4j.dbms.api.DatabaseManagementService;
12+
import org.neo4j.graphdb.GraphDatabaseService;
13+
import org.neo4j.test.TestDatabaseManagementServiceBuilder;
14+
import org.testcontainers.chromadb.ChromaDBContainer;
15+
16+
import java.util.List;
17+
import java.util.Map;
18+
import java.util.concurrent.atomic.AtomicReference;
19+
20+
import static apoc.ml.RestAPIConfig.HEADERS_KEY;
321
import static apoc.util.MapUtil.map;
422
import static apoc.util.TestUtil.testCall;
523
import static apoc.util.TestUtil.testResult;
@@ -10,7 +28,10 @@
1028
import static apoc.vectordb.VectorDbTestUtil.assertNodesCreated;
1129
import static apoc.vectordb.VectorDbTestUtil.assertRelsCreated;
1230
import static apoc.vectordb.VectorDbTestUtil.dropAndDeleteAll;
31+
import static apoc.vectordb.VectorDbTestUtil.ragSetup;
1332
import static apoc.vectordb.VectorDbUtil.ERROR_READONLY_MAPPING;
33+
import static apoc.vectordb.VectorDbTestUtil.getAuthHeader;
34+
import static apoc.vectordb.VectorDbTestUtil.ragSetup;
1435
import static apoc.vectordb.VectorEmbeddingConfig.ALL_RESULTS_KEY;
1536
import static apoc.vectordb.VectorEmbeddingConfig.MAPPING_KEY;
1637
import static apoc.vectordb.VectorMappingConfig.*;
@@ -27,6 +48,8 @@
2748
import java.util.List;
2849
import java.util.Map;
2950
import java.util.concurrent.atomic.AtomicReference;
51+
52+
import apoc.vectordb.VectorDbTestUtil;
3053
import org.assertj.core.api.Assertions;
3154
import org.junit.AfterClass;
3255
import org.junit.Before;
@@ -62,7 +85,7 @@ public static void setUp() throws Exception {
6285
CHROMA_CONTAINER.start();
6386

6487
HOST = "localhost:" + CHROMA_CONTAINER.getMappedPort(8000);
65-
TestUtil.registerProcedure(db, ChromaDb.class, VectorDb.class);
88+
TestUtil.registerProcedure(db, ChromaDb.class, VectorDb.class, Prompt.class));
6689

6790
testCall(
6891
db,
@@ -452,4 +475,33 @@ public void queryVectorsWithSystemDbStorage() {
452475

453476
assertNodesCreated(db);
454477
}
478+
479+
@Test
480+
public void queryVectorsWithRag() {
481+
String openAIKey = ragSetup(db);
482+
483+
Map<String, Object> conf = map(ALL_RESULTS_KEY, true,
484+
HEADERS_KEY, READONLY_AUTHORIZATION,
485+
MAPPING_KEY, map(NODE_LABEL, "Rag",
486+
ENTITY_KEY, "readID",
487+
METADATA_KEY, "foo")
488+
);
489+
490+
testResult(db,
491+
"""
492+
CALL apoc.vectordb.chroma.getAndUpdate($host, $collection, ['1', '2'], $conf) YIELD node, metadata, id, vector
493+
WITH collect(node) as paths
494+
CALL apoc.ml.rag(paths, $attributes, "Which city has foo equals to one?", $confPrompt) YIELD value
495+
RETURN value
496+
"""
497+
,
498+
map(
499+
"host", HOST,
500+
"conf", conf,
501+
"collection", COLL_ID.get(),
502+
"confPrompt", map(API_KEY_CONF, openAIKey),
503+
"attributes", List.of("city", "foo")
504+
),
505+
VectorDbTestUtil::assertRagWithVectors);
506+
}
455507
}

‎full-it/src/test/java/apoc/full/it/vectordb/QdrantTest.java‎

Lines changed: 41 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -22,17 +22,22 @@
2222
import static org.junit.Assert.assertEquals;
2323
import static org.junit.Assert.assertNotNull;
2424
import static org.junit.Assert.assertNull;
25+
import static org.junit.Assert.assertTrue;
2526
import static org.junit.Assert.fail;
2627
import static org.neo4j.configuration.GraphDatabaseSettings.*;
2728

29+
import apoc.ml.Prompt;
2830
import apoc.util.TestUtil;
2931
import apoc.util.Util;
3032
import apoc.vectordb.Qdrant;
3133
import apoc.vectordb.VectorDb;
3234
import apoc.vectordb.VectorDbTestUtil;
35+
36+
import java.util.List;
3337
import java.util.Map;
3438
import org.assertj.core.api.Assertions;
3539
import org.junit.AfterClass;
40+
import org.junit.Assume;
3641
import org.junit.Before;
3742
import org.junit.BeforeClass;
3843
import org.junit.ClassRule;
@@ -74,7 +79,7 @@ public static void setUp() throws Exception {
7479
QDRANT_CONTAINER.start();
7580

7681
HOST = "localhost:" + QDRANT_CONTAINER.getMappedPort(6333);
77-
TestUtil.registerProcedure(db, Qdrant.class, VectorDb.class);
82+
TestUtil.registerProcedure(db, Qdrant.class, VectorDb.class, Prompt.class);
7883

7984
testCall(
8085
db,
@@ -203,6 +208,41 @@ public void queryVectors() {
203208
});
204209
}
205210

211+
@Test
212+
public void queryVectorsWithRag() {
213+
String openAIKey = System.getenv("OPENAI_KEY");;
214+
Assume.assumeNotNull("No OPENAI_KEY environment configured", openAIKey);
215+
216+
db.executeTransactionally("CREATE (:Rag {readID: 'one'}), (:Rag {readID: 'two'})");
217+
218+
Map<String, Object> conf = map(ALL_RESULTS_KEY, true,
219+
HEADERS_KEY, READONLY_AUTHORIZATION,
220+
MAPPING_KEY, map(NODE_LABEL, "Rag",
221+
ENTITY_KEY, "readID",
222+
METADATA_KEY, "foo")
223+
);
224+
225+
testResult(db,
226+
"""
227+
CALL apoc.vectordb.qdrant.getAndUpdate($host, 'test_collection', [1, 2], $conf) YIELD node, metadata, id, vector
228+
WITH collect(node) as paths
229+
CALL apoc.ml.rag(paths, $attributes, "Which city has foo equals to one?", $confPrompt) YIELD value
230+
RETURN value
231+
"""
232+
,
233+
map(
234+
"host", HOST,
235+
"conf", conf,
236+
"confPrompt", map(API_KEY_CONF, openAIKey),
237+
"attributes", List.of("city", "foo")
238+
),
239+
r -> {
240+
Map<String, Object> row = r.next();
241+
Object value = row.get("value");
242+
assertTrue("The actual value is: " + value, value.toString().contains("Berlin"));
243+
});
244+
}
245+
206246
@Test
207247
public void queryVectorsWithoutVectorResult() {
208248
testResult(

‎full-it/src/test/java/apoc/full/it/vectordb/WeaviateTest.java‎

Lines changed: 64 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,25 @@
11
package apoc.full.it.vectordb;
22

3+
import apoc.ml.Prompt;
4+
import apoc.util.MapUtil;
5+
import apoc.util.TestUtil;
6+
import org.junit.AfterClass;
7+
import org.junit.Assume;
8+
import org.junit.Before;
9+
import org.junit.BeforeClass;
10+
import org.junit.ClassRule;
11+
import org.junit.Test;
12+
import org.junit.rules.TemporaryFolder;
13+
import org.neo4j.dbms.api.DatabaseManagementService;
14+
import org.neo4j.graphdb.Entity;
15+
import org.neo4j.graphdb.GraphDatabaseService;
16+
import org.neo4j.graphdb.ResourceIterator;
17+
import org.neo4j.test.TestDatabaseManagementServiceBuilder;
18+
import org.testcontainers.weaviate.WeaviateContainer;
19+
20+
import java.util.List;
21+
import java.util.Map;
22+
323
import static apoc.ml.RestAPIConfig.HEADERS_KEY;
424
import static apoc.util.TestUtil.testCall;
525
import static apoc.util.TestUtil.testCallEmpty;
@@ -9,7 +29,19 @@
929
import static apoc.vectordb.VectorDbTestUtil.*;
1030
import static apoc.vectordb.VectorDbTestUtil.EntityType.*;
1131
import static apoc.vectordb.VectorDbUtil.ERROR_READONLY_MAPPING;
32+
import static apoc.vectordb.VectorDbTestUtil.EntityType.FALSE;
33+
import static apoc.vectordb.VectorDbTestUtil.EntityType.NODE;
34+
import static apoc.vectordb.VectorDbTestUtil.EntityType.REL;
35+
import static apoc.vectordb.VectorDbTestUtil.assertBerlinResult;
36+
import static apoc.vectordb.VectorDbTestUtil.assertLondonResult;
37+
import static apoc.vectordb.VectorDbTestUtil.assertNodesCreated;
38+
import static apoc.vectordb.VectorDbTestUtil.assertReadOnlyProcWithMappingResults;
39+
import static apoc.vectordb.VectorDbTestUtil.assertRelsCreated;
40+
import static apoc.vectordb.VectorDbTestUtil.dropAndDeleteAll;
41+
import static apoc.vectordb.VectorDbTestUtil.getAuthHeader;
42+
import static apoc.vectordb.VectorDbTestUtil.ragSetup;
1243
import static apoc.vectordb.VectorEmbeddingConfig.ALL_RESULTS_KEY;
44+
import static apoc.vectordb.VectorEmbeddingConfig.FIELDS_KEY;
1345
import static apoc.vectordb.VectorEmbeddingConfig.MAPPING_KEY;
1446
import static apoc.vectordb.VectorMappingConfig.*;
1547
import static org.assertj.core.api.Assertions.assertThat;
@@ -21,6 +53,7 @@
2153
import static org.neo4j.configuration.GraphDatabaseSettings.DEFAULT_DATABASE_NAME;
2254
import static org.neo4j.configuration.GraphDatabaseSettings.SYSTEM_DATABASE_NAME;
2355

56+
import apoc.ml.Prompt;
2457
import apoc.util.MapUtil;
2558
import apoc.util.TestUtil;
2659
import apoc.vectordb.VectorDb;
@@ -82,7 +115,7 @@ public static void setUp() throws Exception {
82115
WEAVIATE_CONTAINER.start();
83116
HOST = WEAVIATE_CONTAINER.getHttpHostAddress();
84117

85-
TestUtil.registerProcedure(db, Weaviate.class, VectorDb.class);
118+
TestUtil.registerProcedure(db, Weaviate.class, VectorDb.class, Prompt.class);
86119

87120
testCall(
88121
db,
@@ -552,4 +585,34 @@ public void queryVectorsWithSystemDbStorage() {
552585

553586
assertNodesCreated(db);
554587
}
588+
589+
@Test
590+
public void queryVectorsWithRag() {
591+
String openAIKey = ragSetup(db);
592+
593+
Map<String, Object> conf = MapUtil.map(
594+
FIELDS_KEY, FIELDS,
595+
ALL_RESULTS_KEY, true,
596+
HEADERS_KEY, READONLY_AUTHORIZATION,
597+
MAPPING_KEY, MapUtil.map(EMBEDDING_KEY, "vect",
598+
NODE_LABEL, "Rag",
599+
ENTITY_KEY, "readID",
600+
METADATA_KEY, "foo")
601+
);
602+
603+
testResult(db,
604+
"CALL apoc.vectordb.weaviate.getAndUpdate($host, 'TestCollection', [$id1], $conf) YIELD score, node, metadata, id, vector\n" +
605+
"WITH collect(node) as paths\n" +
606+
"CALL apoc.ml.rag(paths, $attributes, \"Which city has foo equals to one?\", $confPrompt) YIELD value\n" +
607+
"RETURN value"
608+
,
609+
MapUtil.map(
610+
"host", HOST,
611+
"id1", ID_1,
612+
"conf", conf,
613+
"confPrompt", MapUtil.map(API_KEY_CONF, openAIKey),
614+
"attributes", List.of("city", "foo")
615+
),
616+
VectorDbTestUtil::assertRagWithVectors);
617+
}
555618
}

‎full/src/test/java/apoc/vectordb/VectorDbTestUtil.java‎

Lines changed: 38 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,14 @@
11
package apoc.vectordb;
22

3+
import apoc.util.MapUtil;
4+
import org.junit.Assume;
5+
import org.neo4j.graphdb.Entity;
6+
import org.neo4j.graphdb.GraphDatabaseService;
7+
import org.neo4j.graphdb.ResourceIterator;
8+
import org.neo4j.graphdb.Result;
9+
10+
import java.util.Map;
11+
312
import static apoc.util.TestUtil.testResult;
413
import static apoc.util.Util.map;
514
import static org.junit.Assert.assertEquals;
@@ -89,4 +98,33 @@ private static void assertBerlinProperties(Map props) {
8998
public static Map<String, String> getAuthHeader(String key) {
9099
return map("Authorization", "Bearer " + key);
91100
}
101+
102+
public static void assertReadOnlyProcWithMappingResults(Result r, String node) {
103+
Map<String, Object> row = r.next();
104+
Map<String, Object> props = ((Entity) row.get(node)).getAllProperties();
105+
assertEquals(MapUtil.map("readID", "one"), props);
106+
assertNotNull(row.get("vector"));
107+
assertNotNull(row.get("id"));
108+
109+
row = r.next();
110+
props = ((Entity) row.get(node)).getAllProperties();
111+
assertEquals(MapUtil.map("readID", "two"), props);
112+
assertNotNull(row.get("vector"));
113+
assertNotNull(row.get("id"));
114+
115+
assertFalse(r.hasNext());
116+
}
117+
118+
public static void assertRagWithVectors(Result r) {
119+
Map<String, Object> row = r.next();
120+
Object value = row.get("value");
121+
assertTrue("The actual value is: " + value, value.toString().contains("Berlin"));
122+
}
123+
124+
public static String ragSetup(GraphDatabaseService db) {
125+
String openAIKey = System.getenv("OPENAI_KEY");;
126+
Assume.assumeNotNull("No OPENAI_KEY environment configured", openAIKey);
127+
db.executeTransactionally("CREATE (:Rag {readID: 'one'}), (:Rag {readID: 'two'})");
128+
return openAIKey;
129+
}
92130
}

0 commit comments

Comments
 (0)