Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add PineconeDataPlaneClient Interface #76

Merged
merged 2 commits into from
Mar 12, 2024
Merged

Add PineconeDataPlaneClient Interface #76

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4336aeb
add new PineconeDataPlaneClient interface and implement in PineconeBl…
austin-denoble Mar 8, 2024
eed989e
rename PineconeDataPlaneClient to PineconeDataPlaneInterface and move…
austin-denoble Mar 11, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
216 changes: 56 additions & 160 deletions src/main/java/io/pinecone/clients/PineconeBlockingDataPlaneClient.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,15 +1,15 @@
package io.pinecone.clients;

import com.google.protobuf.Struct;
import io.pinecone.commons.PineconeDataPlaneInterface;
import io.pinecone.exceptions.PineconeValidationException;
import io.pinecone.proto.*;
import io.pinecone.unsigned_indices_model.QueryResponseWithUnsignedIndices;

import java.util.List;

import static io.pinecone.utils.SparseIndicesConverter.convertUnsigned32IntToSigned32Int;

public class PineconeBlockingDataPlaneClient {
public class PineconeBlockingDataPlaneClient implements PineconeDataPlaneInterface<UpsertResponse,
QueryResponseWithUnsignedIndices, FetchResponse, UpdateResponse, DeleteResponse, DescribeIndexStatsResponse> {

private final VectorServiceGrpc.VectorServiceBlockingStub blockingStub;

Expand All @@ -21,216 +21,124 @@ public PineconeBlockingDataPlaneClient(VectorServiceGrpc.VectorServiceBlockingSt
this.blockingStub = blockingStub;
}

@Override
public UpsertResponse upsert(String id,
List<Float> values) {
return upsert(id, values, null, null, null, null);
}

@Override
public UpsertResponse upsert(String id,
List<Float> values,
String namespace) {
return upsert(id, values, null, null, null, namespace);
}

@Override
public UpsertResponse upsert(String id,
List<Float> values,
List<Long> sparseIndices,
List<Float> sparseValues,
com.google.protobuf.Struct metadata,
String namespace) {
UpsertRequest.Builder upsertRequest = UpsertRequest.newBuilder();

if (id == null || id.isEmpty() || values == null || values.isEmpty()) {
throw new PineconeValidationException("Invalid upsert request. Please ensure that both id and values are provided.");
}

Vector.Builder vectorBuilder = Vector.newBuilder()
.setId(id)
.addAllValues(values);

if ((sparseIndices != null && sparseValues == null) || (sparseIndices == null && sparseValues != null)) {
throw new PineconeValidationException("Invalid upsert request. Please ensure that both sparse indices and values are present.");
}
if (sparseIndices != null) {
if (sparseIndices.size() != sparseValues.size()) {
throw new PineconeValidationException("Invalid upsert request. Please ensure that both sparse indices and values are of the same length.");
}
vectorBuilder.setSparseValues(SparseValues.newBuilder()
.addAllIndices(convertUnsigned32IntToSigned32Int(sparseIndices))
.addAllValues(sparseValues)
.build());
}
UpsertRequest upsertRequest = validateUpsertRequest(id, values, sparseIndices, sparseValues, metadata,
namespace);

if (metadata != null) {
vectorBuilder.setMetadata(metadata);
}

upsertRequest.addVectors(vectorBuilder.build());

if (namespace != null) {
upsertRequest.setNamespace(namespace);
}

return blockingStub.upsert(upsertRequest.build());
return blockingStub.upsert(upsertRequest);
}

@Override
public QueryResponseWithUnsignedIndices query(int topK,
List<Float> vector,
List<Long> sparseIndices,
List<Float> sparseValues,
String id,
String namespace,
Struct filter,
boolean includeValues,
boolean includeMetadata) {
QueryRequest.Builder queryRequest = QueryRequest.newBuilder();

if (id != null && !id.isEmpty() && vector != null && !vector.isEmpty()) {
throw new PineconeValidationException("Invalid query request. Cannot query with both vector id and vector values.");
}

if (id != null && !id.isEmpty()) {
queryRequest.setId(id);
}

if (namespace != null) {
queryRequest.setNamespace(namespace);
}

queryRequest.setTopK(topK)
.setIncludeValues(includeValues)
.setIncludeMetadata(includeMetadata);

if (filter != null) {
queryRequest.setFilter(filter);
}

if (vector != null && !vector.isEmpty()) {
queryRequest.addAllVector(vector);
}

if ((sparseIndices != null && sparseValues == null) || (sparseIndices == null && sparseValues != null)) {
throw new PineconeValidationException("Invalid upsert request. Please ensure that both sparse indices and values are present.");
}

if (sparseIndices != null) {
if (sparseIndices.size() != sparseValues.size()) {
throw new PineconeValidationException("Invalid upsert request. Please ensure that both sparse indices and values are of the same length.");
}
queryRequest.setSparseVector(SparseValues.newBuilder()
.addAllIndices(convertUnsigned32IntToSigned32Int(sparseIndices))
.addAllValues(sparseValues)
.build());
}
List<Float> vector,
List<Long> sparseIndices,
List<Float> sparseValues,
String id,
String namespace,
Struct filter,
boolean includeValues,
boolean includeMetadata) {
QueryRequest queryRequest = validateQueryRequest(topK, vector, sparseIndices, sparseValues, id, namespace,
filter, includeValues, includeMetadata);

return new QueryResponseWithUnsignedIndices(blockingStub.query(queryRequest.build()));
return new QueryResponseWithUnsignedIndices(blockingStub.query(queryRequest));
}

@Override
public QueryResponseWithUnsignedIndices queryByVectorId(int topK,
String id,
String namespace,
Struct filter,
boolean includeValues,
boolean includeMetadata) {
String id,
String namespace,
Struct filter,
boolean includeValues,
boolean includeMetadata) {
return query(topK, null, null, null, id, namespace, filter, includeValues, includeMetadata);
}

@Override
public QueryResponseWithUnsignedIndices queryByVectorId(int topK,
String id,
String namespace,
Struct filter) {
String id,
String namespace,
Struct filter) {
return query(topK, null, null, null, id, namespace, filter, false, false);
}

@Override
public QueryResponseWithUnsignedIndices queryByVectorId(int topK,
String id,
String namespace) {
String id,
String namespace) {
return query(topK, null, null, null, id, namespace, null, false, false);
}

@Override
public QueryResponseWithUnsignedIndices queryByVectorId(int topK,
String id) {
String id) {
return query(topK, null, null, null, id, null, null, false, false);
}

@Override
public FetchResponse fetch(List<String> ids) {
return fetch(ids, null);
}

@Override
public FetchResponse fetch(List<String> ids,
String namespace) {
FetchRequest.Builder fetchRequest = FetchRequest.newBuilder();
FetchRequest fetchRequest = validateFetchRequest(ids, namespace);

if (ids == null || ids.isEmpty()) {
throw new PineconeValidationException("Invalid fetch request. Vector ids must be present");
}

fetchRequest.addAllIds(ids);

if (namespace != null) {
fetchRequest.setNamespace(namespace);
}

return blockingStub.fetch(fetchRequest.build());
return blockingStub.fetch(fetchRequest);
}

@Override
public UpdateResponse update(String id,
List<Float> values) {
return update(id, values, null, null, null, null);
}

@Override
public UpdateResponse update(String id,
List<Float> values,
String namespace) {
return update(id, values, null, namespace, null, null);
}

@Override
public UpdateResponse update(String id,
List<Float> values,
Struct metadata,
String namespace,
List<Long> sparseIndices,
List<Float> sparseValues) {
UpdateRequest.Builder updateRequest = UpdateRequest.newBuilder();

if (id == null) {
throw new PineconeValidationException("Invalid update request. Vector id must be present");
}
updateRequest.setId(id);

if (values != null) {
updateRequest.addAllValues(values);
}

if (metadata != null) {
updateRequest.setSetMetadata(metadata);
}

if (namespace != null) {
updateRequest.setNamespace(namespace);
}
UpdateRequest updateRequest = validateUpdateRequest(id, values, metadata, namespace, sparseIndices,
sparseValues);

if ((sparseIndices != null && sparseValues == null) || (sparseIndices == null && sparseValues != null)) {
throw new PineconeValidationException("Invalid upsert request. Please ensure that both sparse indices and values are present.");
}

if (sparseIndices != null) {
if (sparseIndices.size() != sparseValues.size()) {
throw new PineconeValidationException("Invalid upsert request. Please ensure that both sparse indices and values are of the same length.");
}
updateRequest.setSparseValues(SparseValues.newBuilder()
.addAllIndices(convertUnsigned32IntToSigned32Int(sparseIndices))
.addAllValues(sparseValues)
.build());
}

return blockingStub.update(updateRequest.build());
return blockingStub.update(updateRequest);
}

@Override
public DeleteResponse deleteByIds(List<String> ids, String namespace) {
return delete(ids, false, namespace, null);
}

@Override
public DeleteResponse deleteByIds(List<String> ids) {
return delete(ids, false, null, null);
}
Expand All @@ -239,42 +147,30 @@ public DeleteResponse deleteByFilter(Struct filter, String namespace) {
return delete(null, false, namespace, filter);
}

@Override
public DeleteResponse deleteByFilter(Struct filter) {
return delete(null, false, null, filter);
}

@Override
public DeleteResponse deleteAll(String namespace) {
return delete(null, true, namespace, null);
}

@Override
public DeleteResponse delete(List<String> ids,
boolean deleteAll,
String namespace,
Struct filter) {
DeleteRequest.Builder deleteRequest = DeleteRequest.newBuilder().setDeleteAll(deleteAll);

if (ids != null && !ids.isEmpty()) {
deleteRequest.addAllIds(ids);
}
DeleteRequest deleteRequest = validateDeleteRequest(ids, deleteAll, namespace, filter);

if (namespace != null) {
deleteRequest.setNamespace(namespace);
}

if (filter != null) {
deleteRequest.setFilter(filter);
}

return blockingStub.delete(deleteRequest.build());
return blockingStub.delete(deleteRequest);
}

@Override
public DescribeIndexStatsResponse describeIndexStats(Struct filter) {
DescribeIndexStatsRequest.Builder describeIndexStatsRequest = DescribeIndexStatsRequest.newBuilder();

if (filter != null) {
describeIndexStatsRequest.setFilter(filter);
}
DescribeIndexStatsRequest describeIndexStatsRequest = validateDescribeIndexStatsRequest(filter);

return blockingStub.describeIndexStats(describeIndexStatsRequest.build());
return blockingStub.describeIndexStats(describeIndexStatsRequest);
}
}
Loading
Loading