Skip to content

Commit

Permalink
wip - finising chroma procs
Browse files Browse the repository at this point in the history
  • Loading branch information
@vga91
vga91 committed Apr 16, 2024
1 parent bba15c8 commit 04e549a
Show file tree
Hide file tree
Showing 7 changed files with 280 additions and 183 deletions.
39 changes: 20 additions & 19 deletions extended-it/src/test/java/apoc/vectordb/ChromaDbTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -37,19 +37,20 @@ public class ChromaDbTest {
@ClassRule
public static DbmsRule db = new ImpermanentDbmsRule();

private static ChromaDBContainer qdrant = new ChromaDBContainer("chromadb/chroma:0.4.25.dev137");
private static ChromaDBContainer chroma = new ChromaDBContainer("chromadb/chroma:0.4.25.dev137");
public static String HOST;
private static AtomicReference<String> collId = new AtomicReference<>();

@BeforeClass
public static void setUp() throws Exception {
qdrant.start();
chroma.start();

HOST = "localhost:" + qdrant.getMappedPort(8000);
TestUtil.registerProcedure(db, VectorDb.class, Qdrant.class);
HOST = "localhost:" + chroma.getMappedPort(8000);
TestUtil.registerProcedure(db, VectorDb.class, ChromaDb.class);

apocConfig().setProperty(APOC_IMPORT_FILE_ENABLED, true);
apocConfig().setProperty(APOC_EXPORT_FILE_ENABLED, true);
AtomicReference<String> id = new AtomicReference<>();

testCall(db, """
CALL apoc.vectordb.custom({
endpoint: $endpoint,
Expand All @@ -63,7 +64,7 @@ public static void setUp() throws Exception {
Map.of("endpoint", "http://" + HOST + "/api/v1/collections"),
r -> {
Map value = (Map) r.get("value");
id.set((String) value.get("id"));
collId.set((String) value.get("id"));
});

testCall(db, """
Expand All @@ -85,7 +86,7 @@ public static void setUp() throws Exception {
payload: {city: "London", foo: "two"}
}
]*/
}, method: 'POST'})""", Map.of("endpoint", "http://" + HOST + "/api/v1/collections/%s/add".formatted(id.get())),
}, method: 'POST'})""", Map.of("endpoint", "http://" + HOST + "/api/v1/collections/%s/add".formatted(collId.get())),
r -> {
assertEquals(true, r.get("value"));
});
Expand All @@ -103,8 +104,8 @@ public void after() {

@Test
public void getEmbeddings() {
testResult(db, "CALL apoc.vectordb.qdrant.get($host, 'test_collection', [1]) ",
Map.of("host", HOST),
testResult(db, "CALL apoc.vectordb.chroma.get($host, $collection, [1]) ",
Map.of("host", HOST, "collection", collId.get()),
r -> {
System.out.println("r = " + r.next());
});
Expand All @@ -115,8 +116,8 @@ public void getEmbedding() {
// String filter = System.getenv("PINECONE_FILTER");
// Assume.assumeNotNull("No PINECONE_FILTER environment configured", host);
// todo -> nResults: 10, ovvero limit, come parametro opzionale
testResult(db, "CALL apoc.vectordb.qdrant.query($host, 'test_collection', [0.2, 0.1, 0.9, 0.7], {}, 5)",
Map.of("host", HOST, /*"filter", filter, */"conf", emptyMap()),
testResult(db, "CALL apoc.vectordb.chroma.query($host, $collection, [0.2, 0.1, 0.9, 0.7], {}, 5)",
Map.of("host", HOST, "collection", collId.get(), /*"filter", filter, */"conf", emptyMap()),
r -> {
System.out.println("r = " + r.next());
System.out.println("r = " + r.next());
Expand All @@ -128,8 +129,8 @@ public void getEmbeddingWithYield() {
// String filter = System.getenv("PINECONE_FILTER");
// Assume.assumeNotNull("No PINECONE_FILTER environment configured", host);
// todo -> nResults: 10, ovvero limit, come parametro opzionale
testResult(db, "CALL apoc.vectordb.qdrant.query($host, 'test_collection', [0.2, 0.1, 0.9, 0.7], {}, 5) YIELD metadata, id",
Map.of("host", HOST, /*"filter", filter, */"conf", emptyMap()),
testResult(db, "CALL apoc.vectordb.chroma.query($host, $collection, [0.2, 0.1, 0.9, 0.7], {}, 5) YIELD metadata, id",
Map.of("host", HOST, "collection", collId.get(), /*"filter", filter, */"conf", emptyMap()),
r -> {
System.out.println("r = " + r.next());
System.out.println("r = " + r.next());
Expand All @@ -144,8 +145,8 @@ public void getEmbeddingWithCreateIndex() {
"prop", "myId",
"id", "foo",
"create", true));
testResult(db, "CALL apoc.vectordb.qdrant.query($host, 'test_collection', [0.2, 0.1, 0.9, 0.7], {}, 5, $conf)",
Map.of("host", HOST, "conf", conf),
testResult(db, "CALL apoc.vectordb.chroma.query($host, $collection, [0.2, 0.1, 0.9, 0.7], {}, 5, $conf)",
Map.of("host", HOST, "collection", collId.get(), "conf", conf),
r -> {
System.out.println("r = " + r.next());
System.out.println("r = " + r.next());
Expand Down Expand Up @@ -178,8 +179,8 @@ public void getEmbeddingWithCreateExistingNode() {
"label", "Test",
"prop", "myId",
"id", "foo"));
testResult(db, "CALL apoc.vectordb.qdrant.query($host, 'test_collection', [0.2, 0.1, 0.9, 0.7], {}, 5, $conf)",
Map.of("host", HOST, "conf", conf),
testResult(db, "CALL apoc.vectordb.chroma.query($host, $collection, [0.2, 0.1, 0.9, 0.7], {}, 5, $conf)",
Map.of("host", HOST, "collection", collId.get(), "conf", conf),
r -> {
System.out.println("r = " + r.next());
System.out.println("r = " + r.next());
Expand Down Expand Up @@ -213,8 +214,8 @@ public void getEmbeddingWithCreateRelIndex() {
"prop", "myId",
"id", "foo",
"create", true));
testResult(db, "CALL apoc.vectordb.qdrant.query($host, 'test_collection', [0.2, 0.1, 0.9, 0.7], {}, 5, $conf)",
Map.of("host", HOST, "conf", conf),
testResult(db, "CALL apoc.vectordb.chroma.query($host, $collection, [0.2, 0.1, 0.9, 0.7], {}, 5, $conf)",
Map.of("host", HOST, "collection", collId.get(), "conf", conf),
r -> {
System.out.println("r = " + r.next());
System.out.println("r = " + r.next());
Expand Down
148 changes: 0 additions & 148 deletions extended/src/main/java/apoc/vectordb/Chroma.java
View file

This file was deleted.

Loading

0 comments on commit 04e549a

Please # to comment.