Class: BxBuilderChain::Vectorsearch::Pgvector

Inherits:
Base
  • Object
show all
Defined in:
lib/bx_builder_chain/vectorsearch/pgvector.rb

Constant Summary collapse

OPERATORS =

The operators supported by the PostgreSQL vector search adapter

{
  "cosine_distance" => "cosine",
  "euclidean_distance" => "euclidean"
}
DEFAULT_OPERATOR = 
"cosine_distance"

Constants inherited from Base

Base::DEFAULT_METRIC

Instance Attribute Summary collapse

Attributes inherited from Base

#client, #llm

Instance Method Summary collapse

Methods inherited from Base

#generate_prompt, #get_default_schema, logger_options

Methods included from DependencyHelper

#depends_on

Constructor Details

#initialize(llm:, namespaces: [BxBuilderChain.configuration.public_namespace] || ['public']) ⇒ Pgvector

Returns a new instance of Pgvector.

Parameters:

  • url (String) —

    The URL of the PostgreSQL database

  • table_name (String) —

    The name of the table to use for the index

  • llm (Object) —

    The LLM client to use

  • namespace (String) —

    The namespace to use for the index when inserting/querying



25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 25

def initialize(llm:, namespaces: [BxBuilderChain.configuration.public_namespace] || ['public'])
  depends_on "sequel"
  require "sequel"
  
  @db = create_sequel_connection
  @table_name = "bx_builder_chain_embeddings"
  @namespace_column = "namespace"
  set_namespaces(namespaces)
  @threshold = BxBuilderChain.configuration.threshold

  validate_threshold(@threshold)

  @operator = OPERATORS[DEFAULT_OPERATOR]

  super(llm: llm)
end

Instance Attribute Details

#db ⇒ Object (readonly)

Returns the value of attribute db.



19
20
21
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 19

def db
  @db
end

#documents_table ⇒ Object (readonly)

Returns the value of attribute documents_table.



19
20
21
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 19

def documents_table
  @documents_table
end

#namespace_column ⇒ Object (readonly)

Returns the value of attribute namespace_column.



19
20
21
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 19

def namespace_column
  @namespace_column
end

#namespaces ⇒ Object (readonly)

Returns the value of attribute namespaces.



19
20
21
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 19

def namespaces
  @namespaces
end

#operator ⇒ Object (readonly)

Returns the value of attribute operator.



19
20
21
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 19

def operator
  @operator
end

#table_name ⇒ Object (readonly)

Returns the value of attribute table_name.



19
20
21
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 19

def table_name
  @table_name
end

Instance Method Details

#add_data(paths:) ⇒ Object

Raises:

  • (ArgumentError) 


181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 181

def add_data(paths:)
  raise ArgumentError, "Paths must be provided" if Array(paths).empty?

  all_added_chunk_ids = []

  @db.transaction do  # Start the transaction
    paths.each do |file_n_path|
      path, file = extract_path_and_file(file_n_path)

      texts = BxBuilderChain::Loader.new(path)&.load&.chunks.map { |chunk| chunk[:text] }
      
      texts.flatten!

      added_chunk_ids_for_current_path = add_texts(texts: texts)
      
      all_added_chunk_ids.concat(added_chunk_ids_for_current_path)

      document_record_id = @db[:bx_builder_chain_documents].insert(
                                                              name: file, 
                                                              namespace: namespaces[0],
                                                              created_at: Time.now.utc,
                                                              updated_at: Time.now.utc
                                                            )

      document_chunks_data = added_chunk_ids_for_current_path.map do |chunk_id|
        {document_id: document_record_id, embedding_id: chunk_id}
      end
      @db[:bx_builder_chain_document_chunks].multi_insert(document_chunks_data)
    end
  end  # End the transaction

  all_added_chunk_ids
end

#add_texts(texts:, ids: nil) ⇒ Array<Integer>

Add a list of texts to the index

Parameters:

  • texts (Array<String>) —

    The texts to add to the index

  • ids (Array<String>) (defaults to: nil) —

    The ids to add to the index, in the same order as the texts

Returns:

  • (Array<Integer>) —

    The the ids of the added texts.



70
71
72
73
74
75
76
77
78
79
80
81
82
83
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 70

def add_texts(texts:, ids: nil)
  if ids.nil? || ids.empty?
    Async do
      data = texts.map do |text|
        Async do
          {content: text, vectors: llm.embed(text: text).to_s, namespace: namespaces[0]}
        end
      end
    end
    @db[@table_name.to_sym].multi_insert(data, return: :primary_key)
  else
    upsert_texts(texts: texts, ids: ids)
  end
end

#ask(question:, context_results: 4, prompt_template: nil) ⇒ String

Ask a question and return the answer

Parameters:

  • question (String) —

    The question to ask

Returns:

  • (String) —

    The answer to the question



168
169
170
171
172
173
174
175
176
177
178
179
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 168

def ask(question:, context_results: 4, prompt_template: nil)
  search_results = similarity_search(query: question, k: context_results)

  context = search_results.map do |result|
    result.content.to_s
  end
  context = context.join("\n---\n")

  prompt = generate_prompt(question: question, context: context, prompt_template: nil)

  llm.chat(prompt: prompt)
end

#create_default_schema ⇒ Object 



93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 93

def create_default_schema
  db.run "CREATE EXTENSION IF NOT EXISTS vector"
  
  namespace_column = @namespace_column
  vector_dimension = llm.default_dimension || 1000
  
  # bx_builder_chain_embeddings table
  db.create_table? :bx_builder_chain_embeddings do
    primary_key :id
    text :content
    column :vectors, "vector(#{vector_dimension})"
    text namespace_column.to_sym, default: 'public'
    
    index namespace_column.to_sym
  end

  # bx_builder_chain_documents table
  db.create_table? :bx_builder_chain_documents do
    primary_key :id
    text :name
    text namespace_column.to_sym, default: 'public'
    timestamp :created_at
    timestamp :updated_at
    
    index [:name, namespace_column.to_sym], unique: true
  end
  
  # bx_builder_chain_document_chunks table
  db.create_table? :bx_builder_chain_document_chunks do
    primary_key :id
    foreign_key :document_id, :bx_builder_chain_documents, null: false, on_delete: :cascade
    foreign_key :embedding_id, :bx_builder_chain_embeddings, null: false, on_delete: :cascade

    unique [:document_id, :embedding_id]
  end
end

#destroy_default_schema ⇒ Object

Destroy default schema



132
133
134
135
136
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 132

def destroy_default_schema
  db.drop_table? :bx_builder_chain_document_chunks
  db.drop_table? :bx_builder_chain_documents
  db.drop_table? :bx_builder_chain_embeddings
end

#documents_model ⇒ Object 



42
43
44
45
46
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 42

def documents_model
  Class.new(Sequel::Model(@table_name.to_sym)) do
    plugin :pgvector, :vectors
  end
end

#similarity_search(query:, k: 4) ⇒ Array<Hash>

Search for similar texts in the index

Parameters:

  • query (String) —

    The text to search for

  • k (Integer) (defaults to: 4) —

    The number of top results to return

Returns:

  • (Array<Hash>) —

    The results of the search



142
143
144
145
146
147
148
149
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 142

def similarity_search(query:, k: 4)
  embedding = llm.embed(text: query)

  similarity_search_by_vector(
    embedding: embedding,
    k: k
  )
end

#similarity_search_by_vector(embedding:, k: 4) ⇒ Array<Hash>

Search for similar texts in the index by the passed in vector. You must generate your own vector using the same LLM that generated the embeddings stored in the Vectorsearch DB.

Parameters:

  • embedding (Array<Float>) —

    The vector to search for

  • k (Integer) (defaults to: 4) —

    The number of top results to return

Returns:

  • (Array<Hash>) —

    The results of the search



156
157
158
159
160
161
162
163
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 156

def similarity_search_by_vector(embedding:, k: 4)
  db.transaction do # BEGIN
    documents_model
      .nearest_neighbors(:vectors, embedding, distance: operator, threshold: @threshold)
      .where(@namespace_column.to_sym => namespaces)
      .limit(k)
  end
end

#update_texts(texts:, ids:) ⇒ Array<Integer>

Update a list of ids and corresponding texts to the index

Parameters:

  • texts (Array<String>) —

    The texts to add to the index

  • ids (Array<String>) —

    The ids to add to the index, in the same order as the texts

Returns:

  • (Array<Integer>) —

    The ids of the updated texts.



89
90
91
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 89

def update_texts(texts:, ids:)
  upsert_texts(texts: texts, ids: ids)
end

#upsert_texts(texts:, ids:) ⇒ PG::Result

Upsert a list of texts to the index the added or updated texts.

Parameters:

  • texts (Array<String>) —

    The texts to add to the index

  • ids (Array<Integer>) —

    The ids of the objects to add to the index, in the same order as the texts

Returns:

  • (PG::Result) —

    The response from the database including the ids of



53
54
55
56
57
58
59
60
61
62
63
64
 
# File 'lib/bx_builder_chain/vectorsearch/pgvector.rb', line 53

def upsert_texts(texts:, ids:)
  data = texts.zip(ids).flat_map do |(text, id)|
    {id: id, content: text, vectors: llm.embed(text: text).to_s, namespace: namespaces[0]}
  end
  # @db[table_name.to_sym].multi_insert(data, return: :primary_key)
  @db[@table_name.to_sym]
    .insert_conflict(
      target: :id,
      update: {content: Sequel[:excluded][:content], vectors: Sequel[:excluded][:vectors]}
    )
    .multi_insert(data, return: :primary_key)
end