Add support for multiple types in Chroma metadata

- Change Chroma collection metadata to Map<String, Object>
- Allows integer and other value types to be passed in

Fixes #1548
This commit is contained in:
shushengcoder
2024-10-16 14:49:58 +08:00
committed by Mark Pollack
parent d0c53392aa
commit f461bd603d
2 changed files with 10 additions and 2 deletions

View File

@@ -106,7 +106,7 @@ public class ChromaApi {
* @param name The name of the collection.
* @param metadata Metadata associated with the collection.
*/
public record Collection(String id, String name, Map<String, String> metadata) {
public record Collection(String id, String name, Map<String, Object> metadata) {
}
/**
@@ -115,7 +115,7 @@ public class ChromaApi {
* @param name The name of the collection to create.
* @param metadata Optional metadata to associate with the collection.
*/
public record CreateCollectionRequest(String name, Map<String, String> metadata) {
public record CreateCollectionRequest(String name, Map<String, Object> metadata) {
public CreateCollectionRequest(String name) {
this(name, new HashMap<>(Map.of("hnsw:space", "cosine")));
}

View File

@@ -55,6 +55,14 @@ public class ChromaApiIT {
chroma.listCollections().stream().forEach(c -> chroma.deleteCollection(c.name()));
}
@Test
public void testClientWithMetadata() {
Map<String, Object> metadata = Map.of("hnsw:space", "cosine", "hnsw:M", 5);
var newCollection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection", metadata));
assertThat(newCollection).isNotNull();
assertThat(newCollection.name()).isEqualTo("TestCollection");
}
@Test
public void testClient() {
var newCollection = chroma.createCollection(new ChromaApi.CreateCollectionRequest("TestCollection"));