From f461bd603d554e15139add55c9029e551ec1c828 Mon Sep 17 00:00:00 2001 From: shushengcoder <1937982680@qq.com> Date: Wed, 16 Oct 2024 14:49:58 +0800 Subject: [PATCH] Add support for multiple types in Chroma metadata - Change Chroma collection metadata to Map - Allows integer and other value types to be passed in Fixes #1548 --- .../java/org/springframework/ai/chroma/ChromaApi.java | 4 ++-- .../java/org/springframework/ai/chroma/ChromaApiIT.java | 8 ++++++++ 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java index 2d5c09ef3..ced7e44ff 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/ChromaApi.java @@ -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 metadata) { + public record Collection(String id, String name, Map 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 metadata) { + public record CreateCollectionRequest(String name, Map metadata) { public CreateCollectionRequest(String name) { this(name, new HashMap<>(Map.of("hnsw:space", "cosine"))); } diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java index be9b21288..c1c934df7 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/ChromaApiIT.java @@ -55,6 +55,14 @@ public class ChromaApiIT { chroma.listCollections().stream().forEach(c -> chroma.deleteCollection(c.name())); } + @Test + public void testClientWithMetadata() { + Map 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"));