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
- Add json-unit-assertj for improved JSON assertions and validating JSON content in tests

Signed-off-by: Christian Tzolov <christian.tzolov@broadcom.com>
This commit is contained in:
Christian Tzolov
2025-05-26 18:20:11 +02:00
committed by Ilayaperumal Gopinathan
parent f6d73a910e
commit aba837a6cc
7 changed files with 97 additions and 2 deletions

View File

@@ -103,5 +103,12 @@
<version>${project.parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>org.springframework.ai</groupId>
<artifactId>spring-ai-advisors-vector-store</artifactId>
<version>${project.parent.version}</version>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View File

@@ -28,6 +28,13 @@ import org.testcontainers.chromadb.ChromaDBContainer;
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<Document> docs = vectorStore.similaritySearch("UserPrompt, AssistantMessage");
vectorStore.delete(docs.stream().map(doc -> doc.getId()).toList());
});
}
@Test
public void addAndSearchWithFilters() {

View File

@@ -219,7 +219,7 @@ public class OpenAiChatModel implements ChatModel {
"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

View File

@@ -309,6 +309,7 @@
<!-- testing dependencies -->
<okhttp3.version>4.12.0</okhttp3.version>
<json-unit-assertj.version>4.1.0</json-unit-assertj.version>
<!-- MCP-->
<mcp.sdk.version>0.10.0</mcp.sdk.version>

View File

@@ -95,6 +95,14 @@
<version>${project.parent.version}</version>
<scope>test</scope>
</dependency>
<dependency>
<groupId>net.javacrumbs.json-unit</groupId>
<artifactId>json-unit-assertj</artifactId>
<version>${json-unit-assertj.version}</version>
<scope>test</scope>
</dependency>
</dependencies>
</project>

View File

@@ -30,6 +30,7 @@ import com.fasterxml.jackson.databind.ObjectMapper;
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;
@@ -443,7 +444,27 @@ public class ChromaApi {
@JsonProperty("metadatas") List<Map<String, Object>> metadata,
@JsonProperty("documents") List<String> 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<Map<String, Object>> processedMetadatas = new ArrayList<>();
for (Map<String, Object> meta : metadata) {
Map<String, Object> processed = new HashMap<>();
for (Map.Entry<String, Object> 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<String, Object> metadata, String document) {
this(List.of(id), List.of(embedding), List.of(metadata), List.of(document));
}

View File

@@ -43,6 +43,8 @@ import org.springframework.context.annotation.Bean;
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
@@ -275,6 +277,22 @@ public class ChromaApiIT {
"Collection non-existent doesn't exist and won't be created as the initializeSchema is set to false.");
}
@Test
public void testAddEmbeddingsRequestMetadataConversion() {
Map<String, Object> 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<String, Object> 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);
assertThatJson(processed.get("mapVal")).isEqualTo("{a:1,b:2}");
}
@SpringBootConfiguration
public static class Config {