From c4a44b82a7bd719882f9ae144e9bad44d9d32ea1 Mon Sep 17 00:00:00 2001 From: Gregorio Galante Date: Mon, 11 Mar 2024 07:07:04 +0100 Subject: [PATCH] Added remove_texts method on chroma vectorsearch --- lib/langchain/vectorsearch/chroma.rb | 7 +++++++ spec/langchain/vectorsearch/chroma_spec.rb | 10 ++++++++++ 2 files changed, 17 insertions(+) diff --git a/lib/langchain/vectorsearch/chroma.rb b/lib/langchain/vectorsearch/chroma.rb index a9d9b11d5..2d131184b 100644 --- a/lib/langchain/vectorsearch/chroma.rb +++ b/lib/langchain/vectorsearch/chroma.rb @@ -60,6 +60,13 @@ def update_texts(texts:, ids:, metadatas: []) collection.update(embeddings) end + # Remove a list of texts from the index + # @param ids [Array] The list of ids to remove + # @return [Hash] The response from the server + def remove_texts(ids:) + collection.delete(ids) + end + # Create the collection with the default schema # @return [::Chroma::Resources::Collection] Created collection def create_default_schema diff --git a/spec/langchain/vectorsearch/chroma_spec.rb b/spec/langchain/vectorsearch/chroma_spec.rb index ec77b46d2..99ac6e7b7 100644 --- a/spec/langchain/vectorsearch/chroma_spec.rb +++ b/spec/langchain/vectorsearch/chroma_spec.rb @@ -102,6 +102,16 @@ end end + describe "remove_texts" do + before do + allow_any_instance_of(Chroma::Resources::Collection).to receive(:delete).and_return(true) + end + + it "removes texts" do + expect(subject.remove_texts(ids: [1])).to eq(true) + end + end + describe "#collection" do it "returns the collection" do expect(subject.send(:collection)).to be_a(Chroma::Resources::Collection)