Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -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() {

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
1 change: 1 addition & 0 deletions pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Expand Down
8 changes: 8 additions & 0 deletions vector-stores/spring-ai-chroma-store/pom.xml
Original file line number Diff line number Diff line change
Expand Up @@ -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>
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -443,7 +444,27 @@ public record AddEmbeddingsRequest(// @formatter:off
@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));
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -275,6 +277,22 @@ 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<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);
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can we also verify that the value of mapVal is what we expect it to be?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

good point. It is done

assertThatJson(processed.get("mapVal")).isEqualTo("{a:1,b:2}");
}

@SpringBootConfiguration
public static class Config {

Expand Down