Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
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
48 changes: 17 additions & 31 deletions src/main/java/io/endee/client/Collection.java
Original file line number Diff line number Diff line change
Expand Up @@ -67,8 +67,7 @@ private Map<String, Map<String, Object>> fieldMap() {
Map<String, Map<String, Object>> idx = new LinkedHashMap<>();
for (Map<String, Object> f : fields) {
@SuppressWarnings("unchecked")
Map<String, Object> params =
(Map<String, Object>) f.getOrDefault("params", Map.of());
Map<String, Object> params = (Map<String, Object>) f.getOrDefault("params", Map.of());
Map<String, Object> entry = new HashMap<>();
entry.put("type", f.getOrDefault("type", "vector"));
entry.put("space_type", params.getOrDefault("space_type", "cosine"));
Expand All @@ -95,8 +94,7 @@ public Map<String, Object> upsert(List<ObjectItem> objects) {
"Cannot upsert more than " + MAX_BATCH_SIZE + " objects at a time");
}

List<String> ids =
objects.stream().map(ObjectItem::getId).collect(Collectors.toList());
List<String> ids = objects.stream().map(ObjectItem::getId).collect(Collectors.toList());
ValidationUtils.validateObjectIds(ids);

Map<String, Map<String, Object>> fMap = fieldMap();
Expand Down Expand Up @@ -125,7 +123,10 @@ public Map<String, Object> upsert(List<ObjectItem> objects) {

String ftype = (String) cfg.get("type");
String space = (String) cfg.getOrDefault("space_type", "cosine");
int dim = cfg.get("dimension") instanceof Number ? ((Number) cfg.get("dimension")).intValue() : 0;
int dim =
cfg.get("dimension") instanceof Number
? ((Number) cfg.get("dimension")).intValue()
: 0;

if ("vector".equals(ftype)) {
double[] vec = (double[]) fdata;
Expand All @@ -151,10 +152,7 @@ public Map<String, Object> upsert(List<ObjectItem> objects) {
validateVectorValues(vecs[i], item.getId());
if (dim > 0 && vecs[i].length != dim) {
throw new IllegalArgumentException(
"Field '"
+ fname
+ "': every multi_vector must have dimension "
+ dim);
"Field '" + fname + "': every multi_vector must have dimension " + dim);
}
normalizedVecs[i] = normalizeDense(vecs[i], space);
if ("cosine".equals(space)) {
Expand Down Expand Up @@ -187,8 +185,7 @@ public Map<String, Object> upsert(List<ObjectItem> objects) {
byte[] payload = MessagePackUtils.packObjects(wireObjects);

try {
HttpRequest request =
buildPostMsgpackRequest("/collection/" + name + "/objects", payload);
HttpRequest request = buildPostMsgpackRequest("/collection/" + name + "/objects", payload);
HttpResponse<String> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofString());

Expand Down Expand Up @@ -379,14 +376,12 @@ public Map<String, List<SearchHit>> search(

/** Convenience overload with defaults: efSearch=128, no filter tuning. */
public Map<String, List<SearchHit>> search(
Map<String, Map<String, Object>> queryFields,
List<Map<String, Object>> filter) {
Map<String, Map<String, Object>> queryFields, List<Map<String, Object>> filter) {
return search(queryFields, filter, 128, null, null);
}

/** Convenience overload: no filter. */
public Map<String, List<SearchHit>> search(
Map<String, Map<String, Object>> queryFields) {
public Map<String, List<SearchHit>> search(Map<String, Map<String, Object>> queryFields) {
return search(queryFields, null, 128, null, null);
}

Expand Down Expand Up @@ -460,8 +455,7 @@ public List<ObjectInfo> getObjects(List<String> ids) {
if (sparsesRaw != null) {
for (Map.Entry<String, Object[]> se : sparsesRaw.entrySet()) {
sparses.put(
se.getKey(),
new SparseData((int[]) se.getValue()[0], (double[]) se.getValue()[1]));
se.getKey(), new SparseData((int[]) se.getValue()[0], (double[]) se.getValue()[1]));
}
}
info.setSparses(sparses);
Expand Down Expand Up @@ -502,8 +496,7 @@ public List<ObjectInfo> getObjects(List<String> ids) {
/** Deletes a single object by ID. */
public Map<String, Object> deleteObject(String id) {
try {
HttpRequest request =
buildDeleteRequest("/collection/" + name + "/objects/" + id);
HttpRequest request = buildDeleteRequest("/collection/" + name + "/objects/" + id);
HttpResponse<String> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofString());

Expand Down Expand Up @@ -531,8 +524,7 @@ public Map<String, Object> deleteByFilter(List<Map<String, Object>> filter) {

try {
String jsonBody = JsonUtils.toJson(Map.of("filter", filter));
HttpRequest request =
buildDeleteJsonRequest("/collection/" + name + "/objects", jsonBody);
HttpRequest request = buildDeleteJsonRequest("/collection/" + name + "/objects", jsonBody);
HttpResponse<String> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofString());

Expand Down Expand Up @@ -569,8 +561,7 @@ public Map<String, Object> updateFilters(List<UpdateFilterParams> updates) {

try {
String jsonBody = JsonUtils.toJson(Map.of("updates", payload));
HttpRequest request =
buildPostJsonRequest("/collection/" + name + "/filters", jsonBody);
HttpRequest request = buildPostJsonRequest("/collection/" + name + "/filters", jsonBody);
HttpResponse<String> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofString());

Expand Down Expand Up @@ -685,8 +676,7 @@ public Map<String, Object> rebuildStatus() {
/** Defragments the collection's storage in place. */
public Map<String, Object> shrink() {
try {
HttpRequest request =
buildPostJsonRequest("/collection/" + name + "/shrink", "{}");
HttpRequest request = buildPostJsonRequest("/collection/" + name + "/shrink", "{}");
HttpResponse<String> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofString());

Expand All @@ -713,8 +703,7 @@ public Map<String, Object> createBackup(String backupName) {

try {
String jsonBody = JsonUtils.toJson(Map.of("name", backupName));
HttpRequest request =
buildPostJsonRequest("/collection/" + name + "/backup", jsonBody);
HttpRequest request = buildPostJsonRequest("/collection/" + name + "/backup", jsonBody);
HttpResponse<String> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofString());

Expand Down Expand Up @@ -815,10 +804,7 @@ private HttpRequest buildPostMsgpackRequest(String path, byte[] body) {

private HttpRequest buildDeleteRequest(String path) {
HttpRequest.Builder builder =
HttpRequest.newBuilder()
.uri(URI.create(baseUrl + path))
.timeout(DEFAULT_TIMEOUT)
.DELETE();
HttpRequest.newBuilder().uri(URI.create(baseUrl + path)).timeout(DEFAULT_TIMEOUT).DELETE();

if (token != null && !token.isBlank()) {
builder.header("Authorization", token);
Expand Down
53 changes: 29 additions & 24 deletions src/main/java/io/endee/client/Endee.java
Original file line number Diff line number Diff line change
Expand Up @@ -40,8 +40,7 @@
public class Endee {
private static final Logger logger = LoggerFactory.getLogger(Endee.class);
private static final Duration DEFAULT_TIMEOUT = Duration.ofSeconds(30);
private static final Set<String> VALID_DB_TYPES =
Set.of("starter", "pro", "scale", "enterprise");
private static final Set<String> VALID_DB_TYPES = Set.of("starter", "pro", "scale", "enterprise");
private static final Set<String> VALID_TOKEN_TYPES = Set.of("rw", "r");

private String token;
Expand Down Expand Up @@ -247,8 +246,7 @@ public String createToken(String dbName, String name) {
@SuppressWarnings("unchecked")
public List<Map<String, Object>> listTokens(String dbName) {
requireNonEmpty(dbName, "db_name");
Map<String, Object> result =
call("GET", "/admin/dbs/" + dbName + "/tokens", null, Set.of(200));
Map<String, Object> result = call("GET", "/admin/dbs/" + dbName + "/tokens", null, Set.of(200));
Object t = result.get("tokens");
return t instanceof List ? (List<Map<String, Object>>) t : List.of();
}
Expand Down Expand Up @@ -350,21 +348,23 @@ public String downloadBackup(String backupName, String destPath, String dbName)
requireNonEmpty(backupName, "backup_name");
requireNonEmpty(destPath, "dest_path");

StringBuilder url = new StringBuilder(baseUrl)
.append("/backup/")
.append(backupName)
.append("/download?token=")
.append(URLEncoder.encode(token != null ? token : "", StandardCharsets.UTF_8));
StringBuilder url =
new StringBuilder(baseUrl)
.append("/backup/")
.append(backupName)
.append("/download?token=")
.append(URLEncoder.encode(token != null ? token : "", StandardCharsets.UTF_8));
if (dbName != null && !dbName.isEmpty()) {
url.append("&db=").append(URLEncoder.encode(dbName, StandardCharsets.UTF_8));
}

try {
HttpRequest request = HttpRequest.newBuilder()
.uri(URI.create(url.toString()))
.timeout(DEFAULT_TIMEOUT)
.GET()
.build();
HttpRequest request =
HttpRequest.newBuilder()
.uri(URI.create(url.toString()))
.timeout(DEFAULT_TIMEOUT)
.GET()
.build();
HttpResponse<byte[]> response =
httpClient.send(request, HttpResponse.BodyHandlers.ofByteArray());
if (response.statusCode() != 200) {
Expand Down Expand Up @@ -406,11 +406,12 @@ public Map<String, Object> uploadBackup(String filePath) {

byte[] multipartBody = buildMultipartBody(boundary, "backup", fileName, fileBytes);

HttpRequest.Builder builder = HttpRequest.newBuilder()
.uri(URI.create(baseUrl + "/backup/upload"))
.timeout(DEFAULT_TIMEOUT)
.header("Content-Type", "multipart/form-data; boundary=" + boundary)
.POST(HttpRequest.BodyPublishers.ofByteArray(multipartBody));
HttpRequest.Builder builder =
HttpRequest.newBuilder()
.uri(URI.create(baseUrl + "/backup/upload"))
.timeout(DEFAULT_TIMEOUT)
.header("Content-Type", "multipart/form-data; boundary=" + boundary)
.POST(HttpRequest.BodyPublishers.ofByteArray(multipartBody));

if (token != null && !token.isEmpty()) {
builder.header("Authorization", token);
Expand Down Expand Up @@ -443,8 +444,14 @@ private static byte[] buildMultipartBody(
String CRLF = "\r\n";
var baos = new java.io.ByteArrayOutputStream();
baos.write(("--" + boundary + CRLF).getBytes(StandardCharsets.UTF_8));
baos.write(("Content-Disposition: form-data; name=\"" + fieldName + "\"; filename=\""
+ fileName + "\"" + CRLF).getBytes(StandardCharsets.UTF_8));
baos.write(
("Content-Disposition: form-data; name=\""
+ fieldName
+ "\"; filename=\""
+ fileName
+ "\""
+ CRLF)
.getBytes(StandardCharsets.UTF_8));
baos.write(("Content-Type: application/x-tar" + CRLF).getBytes(StandardCharsets.UTF_8));
baos.write(CRLF.getBytes(StandardCharsets.UTF_8));
baos.write(fileBytes);
Expand Down Expand Up @@ -484,9 +491,7 @@ private Map<String, Object> call(

private HttpRequest buildRequest(String method, String path, Map<String, Object> json) {
HttpRequest.Builder builder =
HttpRequest.newBuilder()
.uri(URI.create(baseUrl + path))
.timeout(DEFAULT_TIMEOUT);
HttpRequest.newBuilder().uri(URI.create(baseUrl + path)).timeout(DEFAULT_TIMEOUT);

if (token != null && !token.isEmpty()) {
builder.header("Authorization", token);
Expand Down
3 changes: 1 addition & 2 deletions src/main/java/io/endee/client/Reranker.java
Original file line number Diff line number Diff line change
Expand Up @@ -107,8 +107,7 @@ public static List<SearchHit> rerank(
}

/** Convenience: rerank with uniform weights. */
public static List<SearchHit> rerank(
Map<String, List<SearchHit>> searchResults, int limit) {
public static List<SearchHit> rerank(Map<String, List<SearchHit>> searchResults, int limit) {
return rerank(searchResults, limit, null, DEFAULT_RRF_K);
}
}
3 changes: 2 additions & 1 deletion src/main/java/io/endee/client/types/SearchHit.java
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,8 @@ public class SearchHit {

public SearchHit() {}

public SearchHit(String id, double similarity, Map<String, Object> meta, Map<String, Object> filter) {
public SearchHit(
String id, double similarity, Map<String, Object> meta, Map<String, Object> filter) {
this.id = id;
this.similarity = similarity;
this.meta = meta;
Expand Down
1 change: 0 additions & 1 deletion src/main/java/io/endee/client/util/MessagePackUtils.java
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,6 @@
import org.msgpack.core.MessagePack;
import org.msgpack.core.MessageUnpacker;
import org.msgpack.value.Value;
import org.msgpack.value.ValueType;

/**
* MessagePack serialization utilities for the v2 wire format.
Expand Down
Loading