diff --git a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java index 37e7ed2f91..c233fd42b7 100644 --- a/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java +++ b/vector-stores/spring-ai-chroma-store/src/main/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStore.java @@ -63,6 +63,8 @@ */ public class ChromaVectorStore extends AbstractObservationVectorStore implements InitializingBean { + private static final int DEFAULT_PURGE_BATCH_SIZE = 1000; + private final ChromaApi chromaApi; private final String tenantName; @@ -175,6 +177,44 @@ public void doDelete(List idList) { new DeleteEmbeddingsRequest(idList)); } + /** + * Deletes all embeddings from the current collection while keeping the collection + * itself intact. Embeddings are deleted in batches by repeatedly fetching the first + * page of ids until the collection is empty. + *

+ * Note: This method does not support concurrency control. Concurrent writes or + * deletions may result in incomplete deletions or duplicate deletion attempts. + * @return the number of embeddings requested for deletion + * @since 2.0.0-SNAPSHOT + */ + public int purgeEmbeddings() { + int deleteCount = 0; + String collectionId = this.requireCollectionId(); + try { + while (true) { + var response = this.chromaApi.getEmbeddings(this.tenantName, this.databaseName, collectionId, + new ChromaApi.GetEmbeddingsRequest(List.of(), Map.of(), DEFAULT_PURGE_BATCH_SIZE, 0, + List.of())); + + if (response == null || CollectionUtils.isEmpty(response.ids())) { + return deleteCount; + } + List ids = response.ids(); + this.chromaApi.deleteEmbeddings(this.tenantName, this.databaseName, collectionId, + new DeleteEmbeddingsRequest(ids)); + deleteCount += ids.size(); + + if (ids.size() < DEFAULT_PURGE_BATCH_SIZE) { + return deleteCount; + } + } + } + catch (Exception e) { + logger.error("Purge failed: {}", e.getMessage()); + throw new IllegalStateException("Failed to purge chroma collection", e); + } + } + @Override protected void doDelete(Filter.Expression expression) { Assert.notNull(expression, "Filter expression must not be null"); diff --git a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java index 64d6512ae0..1a0fd731ed 100644 --- a/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java +++ b/vector-stores/spring-ai-chroma-store/src/test/java/org/springframework/ai/chroma/vectorstore/ChromaVectorStoreIT.java @@ -272,6 +272,17 @@ public void searchThresholdTest() { }); } + @Test + public void purgeEmbeddingsCollection() { + this.contextRunner.run(context -> { + VectorStore vectorStore = context.getBean(VectorStore.class); + vectorStore.add(this.documents); + + int deletedCount = ((ChromaVectorStore) vectorStore).purgeEmbeddings(); + assertThat(deletedCount).isEqualTo(this.documents.size()); + }); + } + @SpringBootConfiguration public static class TestApplication {