23
23
import java .util .regex .Matcher ;
24
24
import java .util .regex .Pattern ;
25
25
26
+ import com .fasterxml .jackson .annotation .JsonInclude ;
26
27
import com .fasterxml .jackson .annotation .JsonProperty ;
27
28
import com .fasterxml .jackson .core .JsonProcessingException ;
28
29
import com .fasterxml .jackson .databind .ObjectMapper ;
@@ -123,17 +124,6 @@ public Collection createCollection(CreateCollectionRequest createCollectionReque
123
124
.getBody ();
124
125
}
125
126
126
- public Map <String , Object > createCollection2 (CreateCollectionRequest createCollectionRequest ) {
127
-
128
- return this .restClient .post ()
129
- .uri ("/api/v1/collections" )
130
- .headers (this ::httpHeaders )
131
- .body (createCollectionRequest )
132
- .retrieve ()
133
- .toEntity (Map .class )
134
- .getBody ();
135
- }
136
-
137
127
/**
138
128
* Delete a collection with the given name.
139
129
* @param collectionName the name of the collection to delete.
@@ -281,7 +271,11 @@ private String getErrorMessage(HttpStatusCodeException e) {
281
271
* @param name The name of the collection.
282
272
* @param metadata Metadata associated with the collection.
283
273
*/
284
- public record Collection (String id , String name , Map <String , Object > metadata ) {
274
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
275
+ public record Collection (// @formatter:off
276
+ @ JsonProperty ("id" ) String id ,
277
+ @ JsonProperty ("name" ) String name ,
278
+ @ JsonProperty ("metadata" ) Map <String , Object > metadata ) { // @formatter:on
285
279
286
280
}
287
281
@@ -291,7 +285,10 @@ public record Collection(String id, String name, Map<String, Object> metadata) {
291
285
* @param name The name of the collection to create.
292
286
* @param metadata Optional metadata to associate with the collection.
293
287
*/
294
- public record CreateCollectionRequest (String name , Map <String , Object > metadata ) {
288
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
289
+ public record CreateCollectionRequest (// @formatter:off
290
+ @ JsonProperty ("name" ) String name ,
291
+ @ JsonProperty ("metadata" ) Map <String , Object > metadata ) {// @formatter:on
295
292
296
293
public CreateCollectionRequest (String name ) {
297
294
this (name , new HashMap <>(Map .of ("hnsw:space" , "cosine" )));
@@ -300,7 +297,7 @@ public CreateCollectionRequest(String name) {
300
297
}
301
298
302
299
//
303
- // Chroma Collection API (https://docs.trychroma.com/js_reference /Collection)
300
+ // Chroma Collection API (https://docs.trychroma.com/reference/js-client /Collection)
304
301
//
305
302
306
303
/**
@@ -312,14 +309,17 @@ public CreateCollectionRequest(String name) {
312
309
* can filter on this metadata.
313
310
* @param documents The documents contents to associate with the embeddings.
314
311
*/
315
- public record AddEmbeddingsRequest (List <String > ids , List <float []> embeddings ,
316
- @ JsonProperty ("metadatas" ) List <Map <String , Object >> metadata , List <String > documents ) {
312
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
313
+ public record AddEmbeddingsRequest (// @formatter:off
314
+ @ JsonProperty ("ids" ) List <String > ids ,
315
+ @ JsonProperty ("embeddings" ) List <float []> embeddings ,
316
+ @ JsonProperty ("metadatas" ) List <Map <String , Object >> metadata ,
317
+ @ JsonProperty ("documents" ) List <String > documents ) {// @formatter:on
317
318
318
319
// Convenance for adding a single embedding.
319
320
public AddEmbeddingsRequest (String id , float [] embedding , Map <String , Object > metadata , String document ) {
320
321
this (List .of (id ), List .of (embedding ), List .of (metadata ), List .of (document ));
321
322
}
322
-
323
323
}
324
324
325
325
/**
@@ -329,12 +329,14 @@ public AddEmbeddingsRequest(String id, float[] embedding, Map<String, Object> me
329
329
* @param where Condition to filter items to delete based on metadata values.
330
330
* (Optional)
331
331
*/
332
- public record DeleteEmbeddingsRequest (List <String > ids , Map <String , Object > where ) {
332
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
333
+ public record DeleteEmbeddingsRequest (// @formatter:off
334
+ @ JsonProperty ("ids" ) List <String > ids ,
335
+ @ JsonProperty ("where" ) Map <String , Object > where ) {// @formatter:on
333
336
334
337
public DeleteEmbeddingsRequest (List <String > ids ) {
335
- this (ids , Map . of () );
338
+ this (ids , null );
336
339
}
337
-
338
340
}
339
341
340
342
/**
@@ -348,19 +350,24 @@ public DeleteEmbeddingsRequest(List<String> ids) {
348
350
* "metadatas", "documents", "distances". Ids are always included. Defaults to
349
351
* [metadatas, documents, distances].
350
352
*/
351
- public record GetEmbeddingsRequest (List <String > ids , Map <String , Object > where , int limit , int offset ,
352
- List <Include > include ) {
353
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
354
+ public record GetEmbeddingsRequest (// @formatter:off
355
+ @ JsonProperty ("ids" ) List <String > ids ,
356
+ @ JsonProperty ("where" ) Map <String , Object > where ,
357
+ @ JsonProperty ("limit" ) Integer limit ,
358
+ @ JsonProperty ("offset" ) Integer offset ,
359
+ @ JsonProperty ("include" ) List <Include > include ) {// @formatter:on
353
360
354
361
public GetEmbeddingsRequest (List <String > ids ) {
355
- this (ids , Map . of () , 10 , 0 , Include .all );
362
+ this (ids , null , 10 , 0 , Include .all );
356
363
}
357
364
358
365
public GetEmbeddingsRequest (List <String > ids , Map <String , Object > where ) {
359
- this (ids , where , 10 , 0 , Include .all );
366
+ this (ids , CollectionUtils . isEmpty ( where ) ? null : where , 10 , 0 , Include .all );
360
367
}
361
368
362
- public GetEmbeddingsRequest (List <String > ids , Map <String , Object > where , int limit , int offset ) {
363
- this (ids , where , limit , offset , Include .all );
369
+ public GetEmbeddingsRequest (List <String > ids , Map <String , Object > where , Integer limit , Integer offset ) {
370
+ this (ids , CollectionUtils . isEmpty ( where ) ? null : where , limit , offset , Include .all );
364
371
}
365
372
366
373
}
@@ -373,9 +380,12 @@ public GetEmbeddingsRequest(List<String> ids, Map<String, Object> where, int lim
373
380
* @param documents List of document contents. One for each returned document.
374
381
* @param metadata List of document metadata. One for each returned document.
375
382
*/
376
- public record GetEmbeddingResponse (List <String > ids , List <float []> embeddings , List <String > documents ,
377
- @ JsonProperty ("metadatas" ) List <Map <String , String >> metadata ) {
378
-
383
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
384
+ public record GetEmbeddingResponse (// @formatter:off
385
+ @ JsonProperty ("ids" ) List <String > ids ,
386
+ @ JsonProperty ("embeddings" ) List <float []> embeddings ,
387
+ @ JsonProperty ("documents" ) List <String > documents ,
388
+ @ JsonProperty ("metadatas" ) List <Map <String , String >> metadata ) {// @formatter:on
379
389
}
380
390
381
391
/**
@@ -390,18 +400,22 @@ public record GetEmbeddingResponse(List<String> ids, List<float[]> embeddings, L
390
400
* "metadatas", "documents", "distances". Ids are always included. Defaults to
391
401
* [metadatas, documents, distances].
392
402
*/
393
- public record QueryRequest (@ JsonProperty ("query_embeddings" ) List <float []> queryEmbeddings ,
394
- @ JsonProperty ("n_results" ) int nResults , Map <String , Object > where , List <Include > include ) {
403
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
404
+ public record QueryRequest ( // @formatter:off
405
+ @ JsonProperty ("query_embeddings" ) List <float []> queryEmbeddings ,
406
+ @ JsonProperty ("n_results" ) Integer nResults ,
407
+ @ JsonProperty ("where" ) Map <String , Object > where ,
408
+ @ JsonProperty ("include" ) List <Include > include ) { // @formatter:on
395
409
396
410
/**
397
411
* Convenience to query for a single embedding instead of a batch of embeddings.
398
412
*/
399
- public QueryRequest (float [] queryEmbedding , int nResults ) {
400
- this (List .of (queryEmbedding ), nResults , Map . of () , Include .all );
413
+ public QueryRequest (float [] queryEmbedding , Integer nResults ) {
414
+ this (List .of (queryEmbedding ), nResults , null , Include .all );
401
415
}
402
416
403
- public QueryRequest (float [] queryEmbedding , int nResults , Map <String , Object > where ) {
404
- this (List .of (queryEmbedding ), nResults , where , Include .all );
417
+ public QueryRequest (float [] queryEmbedding , Integer nResults , Map <String , Object > where ) {
418
+ this (List .of (queryEmbedding ), nResults , CollectionUtils . isEmpty ( where ) ? null : where , Include .all );
405
419
}
406
420
407
421
public enum Include {
@@ -434,9 +448,13 @@ public enum Include {
434
448
* @param metadata List of list of document metadata. One for each returned document.
435
449
* @param distances List of list of search distances. One for each returned document.
436
450
*/
437
- public record QueryResponse (List <List <String >> ids , List <List <float []>> embeddings , List <List <String >> documents ,
438
- @ JsonProperty ("metadatas" ) List <List <Map <String , Object >>> metadata , List <List <Double >> distances ) {
439
-
451
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
452
+ public record QueryResponse (// @formatter:off
453
+ @ JsonProperty ("ids" ) List <List <String >> ids ,
454
+ @ JsonProperty ("embeddings" ) List <List <float []>> embeddings ,
455
+ @ JsonProperty ("documents" ) List <List <String >> documents ,
456
+ @ JsonProperty ("metadatas" ) List <List <Map <String , Object >>> metadata ,
457
+ @ JsonProperty ("distances" ) List <List <Double >> distances ) {// @formatter:on
440
458
}
441
459
442
460
/**
@@ -448,8 +466,13 @@ public record QueryResponse(List<List<String>> ids, List<List<float[]>> embeddin
448
466
* @param metadata The metadata of the document.
449
467
* @param distances The distance of the document to the query embedding.
450
468
*/
451
- public record Embedding (String id , float [] embedding , String document , Map <String , Object > metadata ,
452
- Double distances ) {
469
+ @ JsonInclude (JsonInclude .Include .NON_NULL )
470
+ public record Embedding (// @formatter:off
471
+ @ JsonProperty ("id" ) String id ,
472
+ @ JsonProperty ("embedding" ) float [] embedding ,
473
+ @ JsonProperty ("document" ) String document ,
474
+ @ JsonProperty ("metadata" ) Map <String , Object > metadata ,
475
+ @ JsonProperty ("distances" ) Double distances ) {// @formatter:on
453
476
454
477
}
455
478
0 commit comments