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:
committed by
Ilayaperumal Gopinathan
parent
f6d73a910e
commit
aba837a6cc
@@ -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>
|
||||
|
||||
@@ -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() {
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
1
pom.xml
1
pom.xml
@@ -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>
|
||||
|
||||
@@ -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>
|
||||
|
||||
@@ -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));
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
|
||||
Reference in New Issue
Block a user