Skip to content

Commit be96c18

Browse files
vga91vga91
authored
[NOID] Fixes #4091: Update RAG docs with vector db examples (#4116) (#4272)
* [NOID] Fixes #4091: Update RAG docs with vector db examples (#4116) * [NOID] test fixes
1 parent dc72f76 commit be96c18

File tree

7 files changed

+180
-7
lines changed

7 files changed

+180
-7
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
@@ -156,6 +156,22 @@ CALL apoc.vectordb.chroma.query($host, '<collection_id>',
156156
----
157157

158158

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

160176
which returns a string that answers the `$question` by leveraging the embeddings of the db vector.
161177

‎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,
@@ -218,7 +217,15 @@ For example, by executing a `CALL apoc.vectordb.qdrant.query(...) YIELD metadata
218217
so that we do not return the other values that we do not need.
219218
====
220219

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

222+
[source,cypher]
223+
----
224+
CALL apoc.vectordb.qdrant.getAndUpdate($host, $collection, [<id1>, <id2>], $conf) YIELD node, metadata, id, vector
225+
WITH collect(node) as paths
226+
CALL apoc.ml.rag(paths, $attributes, $question, $confPrompt) YIELD value
227+
RETURN value
228+
----
222229

223230
which returns a string that answers the `$question` by leveraging the embeddings of the db vector.
224231

‎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,
@@ -235,7 +234,15 @@ For example, by executing a `CALL apoc.vectordb.weaviate.query(...) YIELD metada
235234
so that we do not return the other values that we do not need.
236235
====
237236

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

239+
[source,cypher]
240+
----
241+
CALL apoc.vectordb.weaviate.getAndUpdate($host, $collection, [<id1>, <id2>], $conf) YIELD score, node, metadata, id, vector
242+
WITH collect(node) as paths
243+
CALL apoc.ml.rag(paths, $attributes, $question, $confPrompt) YIELD value
244+
RETURN value
245+
----
239246

240247
which returns a string that answers the `$question` by leveraging the embeddings of the db vector.
241248

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

Lines changed: 28 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package apoc.full.it.vectordb;
22

3+
import static apoc.ml.Prompt.API_KEY_CONF;
34
import static apoc.util.MapUtil.map;
45
import static apoc.util.TestUtil.testCall;
56
import static apoc.util.TestUtil.testResult;
@@ -10,6 +11,7 @@
1011
import static apoc.vectordb.VectorDbTestUtil.assertNodesCreated;
1112
import static apoc.vectordb.VectorDbTestUtil.assertRelsCreated;
1213
import static apoc.vectordb.VectorDbTestUtil.dropAndDeleteAll;
14+
import static apoc.vectordb.VectorDbTestUtil.ragSetup;
1315
import static apoc.vectordb.VectorDbUtil.ERROR_READONLY_MAPPING;
1416
import static apoc.vectordb.VectorEmbeddingConfig.ALL_RESULTS_KEY;
1517
import static apoc.vectordb.VectorEmbeddingConfig.MAPPING_KEY;
@@ -21,9 +23,11 @@
2123
import static org.neo4j.configuration.GraphDatabaseSettings.DEFAULT_DATABASE_NAME;
2224
import static org.neo4j.configuration.GraphDatabaseSettings.SYSTEM_DATABASE_NAME;
2325

26+
import apoc.ml.Prompt;
2427
import apoc.util.TestUtil;
2528
import apoc.vectordb.ChromaDb;
2629
import apoc.vectordb.VectorDb;
30+
import apoc.vectordb.VectorDbTestUtil;
2731
import java.util.List;
2832
import java.util.Map;
2933
import java.util.concurrent.atomic.AtomicReference;
@@ -60,9 +64,9 @@ public static void setUp() throws Exception {
6064
sysDb = databaseManagementService.database(SYSTEM_DATABASE_NAME);
6165

6266
CHROMA_CONTAINER.start();
63-
HOST = CHROMA_CONTAINER.getEndpoint();
6467

65-
TestUtil.registerProcedure(db, ChromaDb.class, VectorDb.class);
68+
HOST = CHROMA_CONTAINER.getEndpoint();
69+
TestUtil.registerProcedure(db, ChromaDb.class, VectorDb.class, Prompt.class);
6670

6771
testCall(
6872
db,
@@ -452,4 +456,26 @@ public void queryVectorsWithSystemDbStorage() {
452456

453457
assertNodesCreated(db);
454458
}
459+
460+
@Test
461+
public void queryVectorsWithRag() {
462+
String openAIKey = ragSetup(db);
463+
464+
Map<String, Object> conf = map(
465+
ALL_RESULTS_KEY, true, MAPPING_KEY, map(NODE_LABEL, "Rag", ENTITY_KEY, "readID", METADATA_KEY, "foo"));
466+
467+
testResult(
468+
db,
469+
"CALL apoc.vectordb.chroma.getAndUpdate($host, $collection, ['1', '2'], $conf) YIELD node, metadata, id, vector\n"
470+
+ "WITH collect(node) as paths\n"
471+
+ "CALL apoc.ml.rag(paths, $attributes, \"Which city has foo equals to one?\", $confPrompt) YIELD value\n"
472+
+ "RETURN value",
473+
map(
474+
"host", HOST,
475+
"conf", conf,
476+
"collection", COLL_ID.get(),
477+
"confPrompt", map(API_KEY_CONF, openAIKey),
478+
"attributes", List.of("city", "foo")),
479+
VectorDbTestUtil::assertRagWithVectors);
480+
}
455481
}

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

Lines changed: 45 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package apoc.full.it.vectordb;
22

3+
import static apoc.ml.Prompt.API_KEY_CONF;
34
import static apoc.ml.RestAPIConfig.HEADERS_KEY;
45
import static apoc.util.MapUtil.map;
56
import static apoc.util.TestUtil.testCall;
@@ -22,17 +23,21 @@
2223
import static org.junit.Assert.assertEquals;
2324
import static org.junit.Assert.assertNotNull;
2425
import static org.junit.Assert.assertNull;
26+
import static org.junit.Assert.assertTrue;
2527
import static org.junit.Assert.fail;
2628
import static org.neo4j.configuration.GraphDatabaseSettings.*;
2729

30+
import apoc.ml.Prompt;
2831
import apoc.util.TestUtil;
2932
import apoc.util.Util;
3033
import apoc.vectordb.Qdrant;
3134
import apoc.vectordb.VectorDb;
3235
import apoc.vectordb.VectorDbTestUtil;
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;
@@ -72,9 +77,9 @@ public static void setUp() throws Exception {
7277
sysDb = databaseManagementService.database(SYSTEM_DATABASE_NAME);
7378

7479
QDRANT_CONTAINER.start();
75-
HOST = QDRANT_CONTAINER.getHost() + ":" + QDRANT_CONTAINER.getMappedPort(6333);
7680

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

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

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

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

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

3+
import static apoc.ml.Prompt.API_KEY_CONF;
34
import static apoc.ml.RestAPIConfig.HEADERS_KEY;
45
import static apoc.util.TestUtil.testCall;
56
import static apoc.util.TestUtil.testCallEmpty;
@@ -8,6 +9,16 @@
89
import static apoc.vectordb.VectorDbHandler.Type.WEAVIATE;
910
import static apoc.vectordb.VectorDbTestUtil.*;
1011
import static apoc.vectordb.VectorDbTestUtil.EntityType.*;
12+
import static apoc.vectordb.VectorDbTestUtil.EntityType.FALSE;
13+
import static apoc.vectordb.VectorDbTestUtil.EntityType.NODE;
14+
import static apoc.vectordb.VectorDbTestUtil.EntityType.REL;
15+
import static apoc.vectordb.VectorDbTestUtil.assertBerlinResult;
16+
import static apoc.vectordb.VectorDbTestUtil.assertLondonResult;
17+
import static apoc.vectordb.VectorDbTestUtil.assertNodesCreated;
18+
import static apoc.vectordb.VectorDbTestUtil.assertRelsCreated;
19+
import static apoc.vectordb.VectorDbTestUtil.dropAndDeleteAll;
20+
import static apoc.vectordb.VectorDbTestUtil.getAuthHeader;
21+
import static apoc.vectordb.VectorDbTestUtil.ragSetup;
1122
import static apoc.vectordb.VectorDbUtil.ERROR_READONLY_MAPPING;
1223
import static apoc.vectordb.VectorEmbeddingConfig.ALL_RESULTS_KEY;
1324
import static apoc.vectordb.VectorEmbeddingConfig.FIELDS_KEY;
@@ -23,6 +34,7 @@
2334
import static org.neo4j.configuration.GraphDatabaseSettings.DEFAULT_DATABASE_NAME;
2435
import static org.neo4j.configuration.GraphDatabaseSettings.SYSTEM_DATABASE_NAME;
2536

37+
import apoc.ml.Prompt;
2638
import apoc.util.MapUtil;
2739
import apoc.util.TestUtil;
2840
import apoc.vectordb.VectorDb;
@@ -84,7 +96,7 @@ public static void setUp() throws Exception {
8496
WEAVIATE_CONTAINER.start();
8597
HOST = WEAVIATE_CONTAINER.getHttpHostAddress();
8698

87-
TestUtil.registerProcedure(db, Weaviate.class, VectorDb.class);
99+
TestUtil.registerProcedure(db, Weaviate.class, VectorDb.class, Prompt.class);
88100

89101
testCall(
90102
db,
@@ -593,4 +605,33 @@ private static void assertQueryVectorsWithSystemDbStorage(String keyConfig, Stri
593605
});
594606
assertNodesCreated(db);
595607
}
608+
609+
@Test
610+
public void queryVectorsWithRag() {
611+
String openAIKey = ragSetup(db);
612+
613+
Map<String, Object> conf = MapUtil.map(
614+
FIELDS_KEY,
615+
FIELDS,
616+
ALL_RESULTS_KEY,
617+
true,
618+
HEADERS_KEY,
619+
READONLY_AUTHORIZATION,
620+
MAPPING_KEY,
621+
MapUtil.map(EMBEDDING_KEY, "vect", NODE_LABEL, "Rag", ENTITY_KEY, "readID", METADATA_KEY, "foo"));
622+
623+
testResult(
624+
db,
625+
"CALL apoc.vectordb.weaviate.getAndUpdate($host, 'TestCollection', [$id1], $conf) YIELD score, node, metadata, id, vector\n"
626+
+ "WITH collect(node) as paths\n"
627+
+ "CALL apoc.ml.rag(paths, $attributes, \"Which city has foo equals to one?\", $confPrompt) YIELD value\n"
628+
+ "RETURN value",
629+
MapUtil.map(
630+
"host", HOST,
631+
"id1", ID_1,
632+
"conf", conf,
633+
"confPrompt", MapUtil.map(API_KEY_CONF, openAIKey),
634+
"attributes", List.of("city", "foo")),
635+
VectorDbTestUtil::assertRagWithVectors);
636+
}
596637
}

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

Lines changed: 33 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,12 @@
44
import static apoc.util.Util.map;
55
import static org.junit.Assert.assertEquals;
66
import static org.junit.Assert.assertFalse;
7+
import static org.junit.Assert.assertNotNull;
78
import static org.junit.Assert.assertTrue;
89

10+
import apoc.util.MapUtil;
911
import java.util.Map;
12+
import org.junit.Assume;
1013
import org.neo4j.graphdb.Entity;
1114
import org.neo4j.graphdb.GraphDatabaseService;
1215
import org.neo4j.graphdb.ResourceIterator;
@@ -89,4 +92,34 @@ private static void assertBerlinProperties(Map props) {
8992
public static Map<String, String> getAuthHeader(String key) {
9093
return map("Authorization", "Bearer " + key);
9194
}
95+
96+
public static void assertReadOnlyProcWithMappingResults(Result r, String node) {
97+
Map<String, Object> row = r.next();
98+
Map<String, Object> props = ((Entity) row.get(node)).getAllProperties();
99+
assertEquals(MapUtil.map("readID", "one"), props);
100+
assertNotNull(row.get("vector"));
101+
assertNotNull(row.get("id"));
102+
103+
row = r.next();
104+
props = ((Entity) row.get(node)).getAllProperties();
105+
assertEquals(MapUtil.map("readID", "two"), props);
106+
assertNotNull(row.get("vector"));
107+
assertNotNull(row.get("id"));
108+
109+
assertFalse(r.hasNext());
110+
}
111+
112+
public static void assertRagWithVectors(Result r) {
113+
Map<String, Object> row = r.next();
114+
Object value = row.get("value");
115+
assertTrue("The actual value is: " + value, value.toString().contains("Berlin"));
116+
}
117+
118+
public static String ragSetup(GraphDatabaseService db) {
119+
String openAIKey = System.getenv("OPENAI_KEY");
120+
;
121+
Assume.assumeNotNull("No OPENAI_KEY environment configured", openAIKey);
122+
db.executeTransactionally("CREATE (:Rag {readID: 'one'}), (:Rag {readID: 'two'})");
123+
return openAIKey;
124+
}
92125
}

0 commit comments

Comments
 (0)