From af0179f49ab6f1cd6d82aace5601ef0a0422ed32 Mon Sep 17 00:00:00 2001 From: Ilayaperumal Gopinathan Date: Wed, 13 Nov 2024 07:20:11 +0000 Subject: [PATCH] GH-1529 Fix ChromaVectorStore delete Ids operation - As a result of the the change:https://github.com/chroma-core/chroma/pull/2880 introduced in chroma:0.5.13,the delete operation doesn't return any values. Hence, we need to check the Http Client Response with the status code instead of return value. - Check the status code for successful delete operation - Update test Resolves #1529 --- .../org/springframework/ai/ollama/BaseOllamaIT.java | 1 - .../org/springframework/ai/chroma/ChromaApi.java | 12 +++++------- .../ai/vectorstore/ChromaVectorStore.java | 5 ++--- .../java/org/springframework/ai/ChromaImage.java | 2 +- .../ai/vectorstore/ChromaVectorStoreIT.java | 4 +++- 5 files changed, 11 insertions(+), 13 deletions(-) diff --git a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/BaseOllamaIT.java b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/BaseOllamaIT.java index 419811f3069..bb5a7f6d048 100644 --- a/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/BaseOllamaIT.java +++ b/models/spring-ai-ollama/src/test/java/org/springframework/ai/ollama/BaseOllamaIT.java @@ -19,7 +19,6 @@ import java.time.Duration; import org.junit.jupiter.api.AfterAll; -import org.junit.jupiter.api.condition.EnabledIfEnvironmentVariable; import org.testcontainers.junit.jupiter.Testcontainers; import org.testcontainers.ollama.OllamaContainer; 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 8d4e5b0369c..c1974e71d7e 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 @@ -28,7 +28,6 @@ import com.fasterxml.jackson.databind.ObjectMapper; import org.springframework.ai.chroma.ChromaApi.QueryRequest.Include; -import org.springframework.core.ParameterizedTypeReference; import org.springframework.http.HttpHeaders; import org.springframework.http.MediaType; import org.springframework.http.client.SimpleClientHttpRequestFactory; @@ -188,17 +187,16 @@ public void upsertEmbeddings(String collectionId, AddEmbeddingsRequest embedding .toBodilessEntity(); } - public List deleteEmbeddings(String collectionId, DeleteEmbeddingsRequest deleteRequest) { + public String deleteEmbeddings(String collectionId, DeleteEmbeddingsRequest deleteRequest) { - return this.restClient.post() + return String.valueOf(this.restClient.post() .uri("/api/v1/collections/{collection_id}/delete", collectionId) .headers(this::httpHeaders) .body(deleteRequest) .retrieve() - .toEntity(new ParameterizedTypeReference>() { - - }) - .getBody(); + .toEntity(String.class) + .getStatusCode() + .value()); } public Long countEmbeddings(String collectionId) { diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java index ba1166927c6..9f434afe612 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/vectorstore/ChromaVectorStore.java @@ -162,9 +162,8 @@ public void doAdd(List documents) { @Override public Optional doDelete(List idList) { Assert.notNull(idList, "Document id list must not be null"); - List deletedIds = this.chromaApi.deleteEmbeddings(this.collectionId, - new DeleteEmbeddingsRequest(idList)); - return Optional.of(deletedIds.size() == idList.size()); + String status = this.chromaApi.deleteEmbeddings(this.collectionId, new DeleteEmbeddingsRequest(idList)); + return Optional.of(status.equals("200")); } @Override diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/ChromaImage.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/ChromaImage.java index ca51285ddec..a2b56266bfc 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/ChromaImage.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/ChromaImage.java @@ -23,7 +23,7 @@ */ public final class ChromaImage { - public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ghcr.io/chroma-core/chroma:0.5.11"); + public static final DockerImageName DEFAULT_IMAGE = DockerImageName.parse("ghcr.io/chroma-core/chroma:0.5.16"); private ChromaImage() { diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java index 966eb06baf2..f24d8c62856 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/vectorstore/ChromaVectorStoreIT.java @@ -19,6 +19,7 @@ import java.util.Collections; import java.util.List; import java.util.Map; +import java.util.Optional; import java.util.UUID; import org.junit.jupiter.api.Test; @@ -83,7 +84,8 @@ public void addAndSearch() { assertThat(resultDoc.getMetadata()).containsKeys("meta2", "distance"); // Remove all documents from the store - vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList()); + assertThat(vectorStore.delete(this.documents.stream().map(doc -> doc.getId()).toList())) + .isEqualTo(Optional.of(Boolean.TRUE)); List results2 = vectorStore.similaritySearch(SearchRequest.query("Great").withTopK(1)); assertThat(results2).hasSize(0);