Class: Aoororachain::VectorStores::Chroma

Inherits:
Object
  • Object
show all
Defined in:
lib/aoororachain/vector_stores/chroma.rb

Constant Summary collapse

DEFAULT_RESULTS = 
3

Instance Method Summary collapse

Constructor Details

#initialize(embedder:, options: {}) ⇒ Chroma

Returns a new instance of Chroma.



10
11
12
13
14
15
16
 
# File 'lib/aoororachain/vector_stores/chroma.rb', line 10

def initialize(embedder:, options: {})
  ::Chroma.connect_host = options.delete(:host)
  ::Chroma.logger = options.delete(:logger) || Aoororachain.logger
  ::Chroma.log_level = options.delete(:log_level) || Aoororachain.log_level

  @embedder = embedder
end

Instance Method Details

#add(documents, include_metadata: false) ⇒ Object 



37
38
39
40
41
42
43
44
45
46
 
# File 'lib/aoororachain/vector_stores/chroma.rb', line 37

def add(documents, include_metadata: false)
  documents_embeddings = embed_documents(documents, include_metadata:)

  documents_embeddings.each do |batch|
    Aoororachain::Util.log_debug("Storing embeddings")
    @store.add(batch)
  end

  true
end

#as_retriever(search_type: :similarity, results: DEFAULT_RESULTS) ⇒ Object 



48
49
50
 
# File 'lib/aoororachain/vector_stores/chroma.rb', line 48

def as_retriever(search_type: :similarity, results: DEFAULT_RESULTS)
  Aoororachain::VectorStores::Retriever.new(self, search_type:, similarity:, results:)
end

#from_documents(documents, index:, include_metadata: false) ⇒ Object 



18
19
20
21
22
23
24
25
26
27
28
29
 
# File 'lib/aoororachain/vector_stores/chroma.rb', line 18

def from_documents(documents, index:, include_metadata: false)
  @store = ::Chroma::Resources::Collection.get_or_create(index, {embedder: @embedder.to_s})

  documents_embeddings = embed_documents(documents, include_metadata:)

  documents_embeddings.each do |batch|
    Aoororachain::Util.log_debug("Storing embeddings")
    @store.add(batch)
  end

  true
end

#from_index(index) ⇒ Object 



31
32
33
34
35
 
# File 'lib/aoororachain/vector_stores/chroma.rb', line 31

def from_index(index)
  @store = ::Chroma::Resources::Collection.get_or_create(index, {embedder: @embedder.to_s})

  true
end

#similarity_search_by_vector(query, results: DEFAULT_RESULTS, filter: {}) ⇒ Object 



52
53
54
 
# File 'lib/aoororachain/vector_stores/chroma.rb', line 52

def similarity_search_by_vector(query, results: DEFAULT_RESULTS, filter: {})
  query(query_texts: query, results:, where: filter, include: %w[metadatas documents])
end

#similarity_search_with_score(query, results: DEFAULT_RESULTS, filter: {}) ⇒ Object Also known as: similarity_search 



56
57
58
 
# File 'lib/aoororachain/vector_stores/chroma.rb', line 56

def similarity_search_with_score(query, results: DEFAULT_RESULTS, filter: {})
  query(query_texts: query, results:, where: filter)
end