From 4f2ef9f9fb4c335b713bf9ab649d4cc434d3ae0d Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Mon, 26 May 2025 18:20:11 +0200 Subject: [PATCH 1/3] feat(chroma): improve handling and testing of complex metadata values - Convert non-primitive metadata values to JSON strings in ChromaApi.AddEmbeddingsRequest for compatibility - Add tests to verify metadata conversion and complex metadata handling in Chroma vector store integration - Ensure OpenAiChatModel always returns annotations as a list of maps in metadata - Add test dependency on spring-ai-advisors-vector-store for advisor-related tests Signed-off-by: Christian Tzolov --- .../pom.xml | 7 ++++ .../ChromaVectorStoreAutoConfigurationIT.java | 40 +++++++++++++++++++ .../ai/openai/OpenAiChatModel.java | 2 +- .../ai/chroma/vectorstore/ChromaApi.java | 25 +++++++++++- .../ai/chroma/vectorstore/ChromaApiIT.java | 15 +++++++ 5 files changed, 86 insertions(+), 3 deletions(-) diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/pom.xml b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/pom.xml index f8f55f26dac..eb8333f16bf 100644 --- a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/pom.xml +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/pom.xml @@ -103,5 +103,12 @@ ${project.parent.version} test + + + org.springframework.ai + spring-ai-advisors-vector-store + ${project.parent.version} + test + diff --git a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/src/test/java/org/springframework/ai/vectorstore/chroma/autoconfigure/ChromaVectorStoreAutoConfigurationIT.java b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/src/test/java/org/springframework/ai/vectorstore/chroma/autoconfigure/ChromaVectorStoreAutoConfigurationIT.java index 62c2f87b51e..a96073bfbdf 100644 --- a/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/src/test/java/org/springframework/ai/vectorstore/chroma/autoconfigure/ChromaVectorStoreAutoConfigurationIT.java +++ b/auto-configurations/vector-stores/spring-ai-autoconfigure-vector-store-chroma/src/test/java/org/springframework/ai/vectorstore/chroma/autoconfigure/ChromaVectorStoreAutoConfigurationIT.java @@ -28,6 +28,13 @@ import org.testcontainers.junit.jupiter.Container; import org.testcontainers.junit.jupiter.Testcontainers; +import org.springframework.ai.chat.client.ChatClientRequest; +import org.springframework.ai.chat.client.ChatClientResponse; +import org.springframework.ai.chat.client.advisor.vectorstore.VectorStoreChatMemoryAdvisor; +import org.springframework.ai.chat.messages.AssistantMessage; +import org.springframework.ai.chat.model.ChatResponse; +import org.springframework.ai.chat.model.Generation; +import org.springframework.ai.chat.prompt.Prompt; import org.springframework.ai.chroma.vectorstore.ChromaVectorStore; import org.springframework.ai.document.Document; import org.springframework.ai.embedding.EmbeddingModel; @@ -68,6 +75,39 @@ public class ChromaVectorStoreAutoConfigurationIT { "spring.ai.vectorstore.chroma.client.port=" + chroma.getMappedPort(8000), "spring.ai.vectorstore.chroma.collectionName=TestCollection"); + @Test + public void verifyThatChromaCanHandleComplexMetadataValues() { + this.contextRunner.withPropertyValues("spring.ai.vectorstore.chroma.initializeSchema=true").run(context -> { + + VectorStore vectorStore = context.getBean(VectorStore.class); + + VectorStoreChatMemoryAdvisor advisor = VectorStoreChatMemoryAdvisor.builder(vectorStore) + .defaultTopK(5) + .build(); + + assertThat(advisor.getName()).isEqualTo("VectorStoreChatMemoryAdvisor"); + + var req = ChatClientRequest.builder().prompt(Prompt.builder().content("UserPrompt").build()).build(); + + ChatClientRequest req2 = advisor.before(req, null); + assertThat(req2).isNotNull(); + + var response = ChatClientResponse.builder() + .chatResponse(ChatResponse.builder() + .generations(List + .of(new Generation(new AssistantMessage("AssistantMessage", Map.of("annotations", List.of()))))) + .build()) + .build(); + var res2 = advisor.after(response, null); + assertThat(res2).isNotNull(); + + // Remove all documents from the store + List docs = vectorStore.similaritySearch("UserPrompt, AssistantMessage"); + vectorStore.delete(docs.stream().map(doc -> doc.getId()).toList()); + + }); + } + @Test public void addAndSearchWithFilters() { diff --git a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java index 58a0062eccd..74c776e6dd0 100644 --- a/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java +++ b/models/spring-ai-openai/src/main/java/org/springframework/ai/openai/OpenAiChatModel.java @@ -219,7 +219,7 @@ public ChatResponse internalCall(Prompt prompt, ChatResponse previousChatRespons "index", choice.index(), "finishReason", choice.finishReason() != null ? choice.finishReason().name() : "", "refusal", StringUtils.hasText(choice.message().refusal()) ? choice.message().refusal() : "", - "annotations", choice.message().annotations() != null ? choice.message().annotations() : List.of()); + "annotations", choice.message().annotations() != null ? choice.message().annotations() : List.of(Map.of())); return buildGeneration(choice, metadata, request); }).toList(); // @formatter:on diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java index 9c55e6011c5..bbea537ffdb 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java @@ -30,6 +30,7 @@ import org.springframework.ai.chroma.vectorstore.ChromaApi.QueryRequest.Include; import org.springframework.ai.chroma.vectorstore.common.ChromaApiConstants; +import org.springframework.ai.util.json.JsonParser; import org.springframework.http.HttpHeaders; import org.springframework.http.MediaType; import org.springframework.http.client.support.BasicAuthenticationInterceptor; @@ -253,7 +254,7 @@ public void upsertEmbeddings(String tenantName, String databaseName, String coll this.restClient.post() .uri("/api/v2/tenants/{tenant_name}/databases/{database_name}/collections/{collection_name}/upsert", tenantName, databaseName, collectionId) - .headers(this::httpHeaders) + // .headers(this::httpHeaders) .body(embedding) .retrieve() .toBodilessEntity(); @@ -443,7 +444,27 @@ public record AddEmbeddingsRequest(// @formatter:off @JsonProperty("metadatas") List> metadata, @JsonProperty("documents") List documents) { // @formatter:on - // Convenance for adding a single embedding. + public AddEmbeddingsRequest { + // Process metadata to ensure all values are Integer, Boolean, or String. + // Other types are converted to JSON string using JsonParser.toJson(). + List> processedMetadatas = new ArrayList<>(); + for (Map meta : metadata) { + Map processed = new HashMap<>(); + for (Map.Entry entry : meta.entrySet()) { + Object value = entry.getValue(); + if (value instanceof Number || value instanceof Boolean || value instanceof String) { + processed.put(entry.getKey(), value); + } + else { + processed.put(entry.getKey(), JsonParser.toJson(value)); + } + } + processedMetadatas.add(processed); + } + metadata = processedMetadatas; + } + + // Convenience for adding a single embedding. public AddEmbeddingsRequest(String id, float[] embedding, Map metadata, String document) { this(List.of(id), List.of(embedding), List.of(metadata), List.of(document)); } diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java index 25478fab66d..ae7c4ad123e 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java @@ -275,6 +275,21 @@ void shouldFailWhenCollectionDoesNotExist() { "Collection non-existent doesn't exist and won't be created as the initializeSchema is set to false."); } + @Test + public void testAddEmbeddingsRequestMetadataConversion() { + Map metadata = Map.of("intVal", 42, "boolVal", true, "strVal", "hello", "doubleVal", 3.14, + "listVal", List.of(1, 2, 3), "mapVal", Map.of("a", 1, "b", 2)); + AddEmbeddingsRequest req = new AddEmbeddingsRequest("id", new float[] { 1f, 2f, 3f }, metadata, "doc"); + Map processed = req.metadata().get(0); + + assertThat(processed.get("intVal")).isInstanceOf(Integer.class); + assertThat(processed.get("boolVal")).isInstanceOf(Boolean.class); + assertThat(processed.get("strVal")).isInstanceOf(String.class); + assertThat(processed.get("doubleVal")).isInstanceOf(Number.class).isEqualTo(3.14); + assertThat(processed.get("listVal")).isInstanceOf(String.class).isEqualTo("[1,2,3]"); + assertThat(processed.get("mapVal")).isInstanceOf(String.class); + } + @SpringBootConfiguration public static class Config { From 701971e1e69ab1ebec6b91361d55471dd93b87fb Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Tue, 27 May 2025 09:29:58 +0200 Subject: [PATCH 2/3] Add json-unit-assertj for improved JSON assertions and validating JSON content in tests Signed-off-by: Christian Tzolov --- pom.xml | 1 + vector-stores/spring-ai-chroma-store/pom.xml | 8 ++++++++ .../ai/chroma/vectorstore/ChromaApiIT.java | 3 +++ 3 files changed, 12 insertions(+) diff --git a/pom.xml b/pom.xml index 2cf8b793855..ca54365aa3a 100644 --- a/pom.xml +++ b/pom.xml @@ -309,6 +309,7 @@ 4.12.0 + 4.1.0 0.10.0 diff --git a/vector-stores/spring-ai-chroma-store/pom.xml b/vector-stores/spring-ai-chroma-store/pom.xml index d40a517f1d6..3d415a4e71c 100644 --- a/vector-stores/spring-ai-chroma-store/pom.xml +++ b/vector-stores/spring-ai-chroma-store/pom.xml @@ -95,6 +95,14 @@ ${project.parent.version} test + + + net.javacrumbs.json-unit + json-unit-assertj + ${json-unit-assertj.version} + test + + diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java index ae7c4ad123e..dc8ccde275e 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaApiIT.java @@ -43,6 +43,8 @@ import static org.assertj.core.api.Assertions.assertThat; import static org.assertj.core.api.AssertionsForClassTypes.assertThatNoException; import static org.assertj.core.api.AssertionsForClassTypes.assertThatThrownBy; +import static net.javacrumbs.jsonunit.assertj.JsonAssertions.assertThatJson; +import static net.javacrumbs.jsonunit.assertj.JsonAssertions.json; /** * @author Christian Tzolov @@ -288,6 +290,7 @@ public void testAddEmbeddingsRequestMetadataConversion() { assertThat(processed.get("doubleVal")).isInstanceOf(Number.class).isEqualTo(3.14); assertThat(processed.get("listVal")).isInstanceOf(String.class).isEqualTo("[1,2,3]"); assertThat(processed.get("mapVal")).isInstanceOf(String.class); + assertThatJson(processed.get("mapVal")).isEqualTo("{a:1,b:2}"); } @SpringBootConfiguration From 5b0dce6f85eae8e5afffc5c7a2f4bb5c2b61cc2f Mon Sep 17 00:00:00 2001 From: Christian Tzolov Date: Tue, 27 May 2025 14:23:01 +0200 Subject: [PATCH 3/3] fix a small regression Signed-off-by: Christian Tzolov --- .../org/springframework/ai/chroma/vectorstore/ChromaApi.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java index bbea537ffdb..79166c0c2a1 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaApi.java @@ -254,7 +254,7 @@ public void upsertEmbeddings(String tenantName, String databaseName, String coll this.restClient.post() .uri("/api/v2/tenants/{tenant_name}/databases/{database_name}/collections/{collection_name}/upsert", tenantName, databaseName, collectionId) - // .headers(this::httpHeaders) + .headers(this::httpHeaders) .body(embedding) .retrieve() .toBodilessEntity();