Skip to content

Commit

Permalink
small code refactoring
Browse files Browse the repository at this point in the history
  • Loading branch information
@vga91
vga91 committed May 27, 2024
1 parent 226a739 commit 5db7fcf
Show file tree
Hide file tree
Showing 6 changed files with 38 additions and 24 deletions.
22 changes: 17 additions & 5 deletions docs/asciidoc/modules/ROOT/pages/database-integration/vectordb/milvus.adoc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,16 +8,28 @@ Here is a list of all available Milvus procedures:
| name | description
| apoc.vectordb.milvus.createCollection(hostOrKey, collection, similarity, size, $config) |
Creates a collection, with the name specified in the 2nd parameter, and with the specified `similarity` and `size`.
The default endpoint is `<hostOrKey param>/v2/vectordb/collections/create`.
| apoc.vectordb.milvus.deleteCollection(hostOrKey, collection, $config) |
Deletes a collection with the name specified in the 2nd parameter
Deletes a collection with the name specified in the 2nd parameter.
The default endpoint is `<hostOrKey param>/v2/vectordb/collections/drop`.
| apoc.vectordb.milvus.upsert(hostOrKey, collection, vectors, $config) |
Upserts, in the collection with the name specified in the 2nd parameter, the vectors [{id: 'id', vector: '<vectorDb>', medatada: '<metadata>'}]
Upserts, in the collection with the name specified in the 2nd parameter, the vectors [{id: 'id', vector: '<vectorDb>', medatada: '<metadata>'}].
The default endpoint is `<hostOrKey param>/v2/vectordb/entities/upsert`.
| apoc.vectordb.milvus.delete(hostOrKey, collection, ids, $config) |
Delete the vectors with the specified `ids`.
The default endpoint is `<hostOrKey param>/v2/vectordb/entities/delete`.
| apoc.vectordb.milvus.get(hostOrKey, collection, ids, $config) |
Get the vectors with the specified `ids`.
The default endpoint is `<hostOrKey param>/v2/vectordb/entities/get`.
| apoc.vectordb.milvus.query(hostOrKey, collection, vector, filter, limit, $config) |
Retrieve closest vectors the the defined `vector`, `limit` of results, in the collection with the name specified in the 2nd parameter.
The default endpoint is `<hostOrKey param>/v2/vectordb/entities/search`.
| apoc.vectordb.milvus.getAndUpdate(hostOrKey, collection, ids, $config) |
Get the vectors with the specified `ids`.
The default endpoint is `<hostOrKey param>/v2/vectordb/entities/get`, and optionally creates/updates neo4j entities.
| apoc.vectordb.milvus.queryAndUpdate(hostOrKey, collection, vector, filter, limit, $config) |
Retrieve closest vectors the the defined `vector`, `limit` of results, in the collection with the name specified in the 2nd parameter, and optionally creates/updates neo4j entities.
The default endpoint is `<hostOrKey param>/v2/vectordb/entities/search`.
|===

where the 1st parameter can be a key defined by the apoc config `apoc.milvus.<key>.host=myHost`.
Expand Down Expand Up @@ -118,7 +130,7 @@ we can populate some existing nodes (i.e. `(:Test {myId: 'one'})` and `(:Test {m

[source,cypher]
----
CALL apoc.vectordb.milvus.query('http://localhost:19531', 'test_collection',
CALL apoc.vectordb.milvus.queryAndUpdate('http://localhost:19531', 'test_collection',
[0.2, 0.1, 0.9, 0.7],
{},
5,
Expand All @@ -139,7 +151,7 @@ Or else, we can create a node if not exists, via `create: true`:

[source,cypher]
----
CALL apoc.vectordb.milvus.query('http://localhost:19531', 'test_collection',
CALL apoc.vectordb.milvus.queryAndUpdate('http://localhost:19531', 'test_collection',
[0.2, 0.1, 0.9, 0.7],
{},
5,
Expand All @@ -160,7 +172,7 @@ Or, we can populate an existing relationship (i.e. `(:Start)-[:TEST {myId: 'one'

[source,cypher]
----
CALL apoc.vectordb.milvus.query('http://localhost:19531', 'test_collection',
CALL apoc.vectordb.milvus.queryAndUpdate('http://localhost:19531', 'test_collection',
[0.2, 0.1, 0.9, 0.7],
{},
5,
Expand Down
4 changes: 2 additions & 2 deletions docs/asciidoc/modules/ROOT/pages/database-integration/vectordb/pinecone.adoc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,10 +25,10 @@ Here is a list of all available Pinecone procedures:
Get the vectors with the specified `ids`.
The default endpoint is `<hostOrKey param>/vectors/fetch`.
| apoc.vectordb.pinecone.query(hostOrKey, index, vector, filter, limit, $config) |
Retrieve closest vectors the the defined `vector`, `limit` of results, in the index with the name specified in the 2nd parameter.
Retrieve closest vectors the the defined `vector`, `limit` of results, in the index with the name specified in the 2nd parameter, and optionally creates/updates neo4j entities.
The default endpoint is `<hostOrKey param>/query`.
| apoc.vectordb.pinecone.queryAndUpdate(hostOrKey, index, vector, filter, limit, $config) |
Retrieve closest vectors the the defined `vector`, `limit` of results, in the index with the name specified in the 2nd parameter.
Retrieve closest vectors the the defined `vector`, `limit` of results, in the index with the name specified in the 2nd parameter, and optionally creates/updates neo4j entities.
The default endpoint is `<hostOrKey param>/query`.
|===

Expand Down
2 changes: 1 addition & 1 deletion extended-it/src/test/java/apoc/vectordb/MilvusTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -382,7 +382,7 @@ public void queryVectorsWithSystemDbStorage() {
"keyConfig", keyConfig,
"databaseName", DEFAULT_DATABASE_NAME,
"conf", map(
"host", HOST,
"host", HOST + "/v2/vectordb",
"mapping", mapping
)
)
Expand Down
22 changes: 11 additions & 11 deletions extended/src/main/java/apoc/vectordb/Milvus.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,7 +26,8 @@

@Extended
public class Milvus {

public static final VectorDbHandler DB_HANDLER = MILVUS.get();

@Context
public ProcedureCallContext procedureCallContext;

Expand All @@ -46,7 +47,7 @@ public Stream<MapResult> createCollection(@Name("hostOrKey") String hostOrKey,
@Name("similarity") String similarity,
@Name("size") Long size,
@Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception {
String url = "%s/v2/vectordb/collections/create";
String url = "%s/collections/create";
Map<String, Object> config = getVectorDbInfo(hostOrKey, collection, configuration, url);
config.putIfAbsent(METHOD_KEY, "POST");

Expand All @@ -67,7 +68,7 @@ public Stream<MapResult> deleteCollection(
@Name("collection") String collection,
@Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception {

String url = "%s/v2/vectordb/collections/drop";
String url = "%s/collections/drop";
Map<String, Object> config = getVectorDbInfo(hostOrKey, collection, configuration, url);
config.putIfAbsent(METHOD_KEY, "POST");
Map<String, Object> additionalBodies = Map.of("collectionName", collection);
Expand All @@ -86,7 +87,7 @@ public Stream<MapResult> upsert(
@Name("vectors") List<Map<String, Object>> vectors,
@Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception {

String url = "%s/v2/vectordb/entities/upsert";
String url = "%s/entities/upsert";

Map<String, Object> config = getVectorDbInfo(hostOrKey, collection, configuration, url);
config.putIfAbsent(METHOD_KEY, "POST");
Expand Down Expand Up @@ -114,12 +115,11 @@ public Stream<MapResult> delete(
@Name("vectors") List<Object> ids,
@Name(value = "configuration", defaultValue = "{}") Map<String, Object> configuration) throws Exception {

String url = "%s/v2/vectordb/entities/delete";
String url = "%s/entities/delete";
Map<String, Object> config = getVectorDbInfo(hostOrKey, collection, configuration, url);
config.putIfAbsent(METHOD_KEY, "POST");

String filter = "id in " + ids;
System.out.println("filter = " + filter);
Map<String, Object> additionalBodies = Map.of("collectionName", collection, "filter", filter);
RestAPIConfig apiConfig = new RestAPIConfig(config, Map.of(), additionalBodies);
return executeRequest(apiConfig, urlAccessChecker)
Expand All @@ -146,14 +146,14 @@ public Stream<EmbeddingResult> getAndUpdate(@Name("hostOrKey") String hostOrKey,
}

private Stream<EmbeddingResult> getCommon(String hostOrKey, String collection, List<Object> ids, Map<String, Object> configuration, boolean readOnly) throws Exception {
String url = "%s/v2/vectordb/entities/get";
String url = "%s/entities/get";
Map<String, Object> config = getVectorDbInfo(hostOrKey, collection, configuration, url);

if (readOnly) {
checkMappingConf(configuration, "apoc.vectordb.milvus.getAndUpdate");
}

VectorEmbeddingConfig apiConfig = MILVUS.get().getEmbedding().fromGet(config, procedureCallContext, ids, collection);
VectorEmbeddingConfig apiConfig = DB_HANDLER.getEmbedding().fromGet(config, procedureCallContext, ids, collection);
return getEmbeddingResultStream(apiConfig, procedureCallContext, urlAccessChecker, tx,
v -> getMapStream((Map) v));
}
Expand Down Expand Up @@ -199,20 +199,20 @@ private Stream<Map> getMapStream(Map v) {
}

private Stream<EmbeddingResult> queryCommon(String hostOrKey, String collection, List<Double> vector, Object filter, long limit, Map<String, Object> configuration, boolean readOnly) throws Exception {
String url = "%s/v2/vectordb/entities/search";
String url = "%s/entities/search";
Map<String, Object> config = getVectorDbInfo(hostOrKey, collection, configuration, url);

if (readOnly) {
checkMappingConf(configuration, "apoc.vectordb.milvus.queryAndUpdate");
}

VectorEmbeddingConfig apiConfig = MILVUS.get().getEmbedding().fromQuery(config, procedureCallContext, vector, filter, limit, collection);
VectorEmbeddingConfig apiConfig = DB_HANDLER.getEmbedding().fromQuery(config, procedureCallContext, vector, filter, limit, collection);
return getEmbeddingResultStream(apiConfig, procedureCallContext, urlAccessChecker, tx,
v -> getMapStream((Map) v));
}

private Map<String, Object> getVectorDbInfo(
String hostOrKey, String collection, Map<String, Object> configuration, String templateUrl) {
return getCommonVectorDbInfo(hostOrKey, collection, configuration, templateUrl, MILVUS.get());
return getCommonVectorDbInfo(hostOrKey, collection, configuration, templateUrl, DB_HANDLER);
}
}
3 changes: 2 additions & 1 deletion extended/src/main/java/apoc/vectordb/MilvusHandler.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,8 @@ public class MilvusHandler implements VectorDbHandler {

@Override
public String getUrl(String hostOrKey) {
return new UrlResolver("http", "localhost", 19530).getUrl("milvus", hostOrKey);
String url = new UrlResolver("http", "localhost", 19530).getUrl("milvus", hostOrKey);
return url + "/v2/vectordb";
}

@Override
Expand Down
9 changes: 5 additions & 4 deletions extended/src/main/java/apoc/vectordb/Pinecone.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,7 +27,8 @@

@Extended
public class Pinecone {

public static final VectorDbHandler DB_HANDLER = PINECONE.get();

@Context
public ProcedureCallContext procedureCallContext;

Expand Down Expand Up @@ -152,7 +153,7 @@ private Stream<VectorDbUtil.EmbeddingResult> getCommon(String hostOrKey, String
checkMappingConf(configuration, "apoc.vectordb.pinecone.getAndUpdate");
}

VectorEmbeddingConfig apiConfig = PINECONE.get().getEmbedding().fromGet(config, procedureCallContext, ids, collection);
VectorEmbeddingConfig apiConfig = DB_HANDLER.getEmbedding().fromGet(config, procedureCallContext, ids, collection);
return getEmbeddingResultStream(apiConfig, procedureCallContext, urlAccessChecker, tx,
v -> {
Object vectors = ((Map) v).get("vectors");
Expand Down Expand Up @@ -191,7 +192,7 @@ private Stream<VectorDbUtil.EmbeddingResult> queryCommon(String hostOrKey, Strin
checkMappingConf(configuration, "apoc.vectordb.pinecone.queryAndUpdate");
}

VectorEmbeddingConfig apiConfig = PINECONE.get().getEmbedding().fromQuery(config, procedureCallContext, vector, filter, limit, collection);
VectorEmbeddingConfig apiConfig = DB_HANDLER.getEmbedding().fromQuery(config, procedureCallContext, vector, filter, limit, collection);
return getEmbeddingResultStream(apiConfig, procedureCallContext, urlAccessChecker, tx,
v -> {
Map map = (Map) v;
Expand All @@ -202,6 +203,6 @@ private Stream<VectorDbUtil.EmbeddingResult> queryCommon(String hostOrKey, Strin

private Map<String, Object> getVectorDbInfo(
String hostOrKey, String collection, Map<String, Object> configuration, String templateUrl) {
return getCommonVectorDbInfo(hostOrKey, collection, configuration, templateUrl, PINECONE.get());
return getCommonVectorDbInfo(hostOrKey, collection, configuration, templateUrl, DB_HANDLER);
}
}

0 comments on commit 5db7fcf

Please # to comment.