boxcars 0.2.11 → 0.2.12

data/lib/boxcars/vector_store/hnswlib/save_to_hnswlib.rb

@@ -12,67 +12,77 @@ module Boxcars
 
         # @param document_embeddings [Array] An array of hashes containing the document id, document text, and embedding.
         # @param index_file_path [String] The path to the index file.
-        # @param hnswlib_config [Boxcars::VectorStore::Hnswlib::Config] The config object for the hnswlib index.
         # @option json_doc_file_path [String] Optional. The path to the json file containing the document text.
-        def initialize(document_embeddings:, index_file_path:, hnswlib_config:, json_doc_file_path: nil)
-          @document_embeddings = document_embeddings
-          @index_file_path = index_file_path
-          @json_doc_file_path = json_doc_file_path || index_file_path.gsub(/\.bin$/, '.json')
-
-          @hnswlib_config = hnswlib_config
-          @index = ::Hnswlib::HnswIndex.new(
-            n_features: hnswlib_config.dim,
-            max_item: hnswlib_config.max_item,
-            metric: hnswlib_config.metric
-          )
-        end
+        def initialize(hnsw_vectors_array)
+          @metadata = hnsw_vectors_array&.first&.metadata
+          validate_params(hnsw_vectors_array, metadata)
 
-        def call
-          validate_params
-          document_texts = []
+          @vectors = hnsw_vectors_array
+          @index_file_path = metadata[:index_file_path]
+          @json_doc_file_path = metadata[:json_doc_file_path] || @index_file_path.gsub(/\.bin$/, '.json')
 
-          document_embeddings.each do |embedding|
-            index.add_item(embedding[:doc_id], embedding[:embedding])
+          @metric = metadata[:metric] || "l2"
+          @dim = metadata[:dim]
+          @max_item = metadata[:max_item] || 10000
 
-            document_texts << { doc_id: embedding[:doc_id], embedding: embedding[:embedding], document: embedding[:document] }
-          end
+          @index = ::Hnswlib::HierarchicalNSW.new(
+            space: @metric,
+            dim: @dim
+          )
+          @index.init_index(max_elements: @max_item)
+        end
 
+        def call
+          document_texts = add_vectors_to_index
           write_files(index, document_texts)
         end
 
         private
 
-        def write_files(index, document_texts)
-          FileUtils.mkdir_p(File.dirname(json_doc_file_path))
-          File.write(json_doc_file_path, document_texts.to_json)
+        attr_reader :metadata, :index, :vectors, :index_file_path, :json_doc_file_path, :metric, :dim, :max_item
 
-          FileUtils.mkdir_p(File.dirname(index_file_path))
-          File.write("#{File.dirname(index_file_path)}/hnswlib_config.json", hnswlib_config.to_json)
+        def validate_params(hnsw_vectors_array, metadata)
+          raise_argument_error('argument must be an array') unless hnsw_vectors_array.is_a?(Array)
+          raise_argument_error('missing data') if hnsw_vectors_array.empty?
+          raise_error('missing metadata') unless metadata || metadata.empty?
 
-          index.save(index_file_path)
+          raise_argument_error("dim must be an integer") unless metadata[:dim].is_a?(Integer)
+          raise_argument_error('missing dim') unless metadata[:dim]
+          raise_argument_error('missing index_file_path') unless metadata[:index_file_path]
+
+          check_parent_directory(metadata[:index_file_path])
+          check_parent_directory(metadata[:json_doc_file_path])
         end
 
-        attr_reader :index, :document_embeddings, :index_file_path, :json_doc_file_path, :hnswlib_config
+        def add_vectors_to_index
+          document_texts = []
 
-        def validate_params
-          raise_error("document_embeddings must be an array") unless document_embeddings.is_a?(Array)
-          raise_error("dim must be an integer") unless hnswlib_config.dim.is_a?(Integer)
-          raise_error("index_file_path must be a string") unless index_file_path.is_a?(String)
+          vectors.each do |item|
+            index.add_point(item.embedding, item.metadata[:doc_id])
 
-          [index_file_path, json_doc_file_path].each do |path|
-            check_parent_directory(path)
+            document_texts << {
+              doc_id: item.metadata[:doc_id],
+              embedding: item.embedding,
+              document: item.content,
+              metadata: item.metadata
+            }
           end
+          document_texts
+        end
+
+        def write_files(index, document_texts)
+          FileUtils.mkdir_p(File.dirname(json_doc_file_path))
+          File.write(json_doc_file_path, document_texts.to_json)
+
+          FileUtils.mkdir_p(File.dirname(index_file_path))
+          index.save_index(index_file_path)
         end
 
         def check_parent_directory(path)
           return unless path
 
           parent_dir = File.dirname(path)
-          raise_error('parent directory must exist') unless File.directory?(parent_dir)
-        end
-
-        def raise_error(message)
-          raise ::Boxcars::ValueError, message
+          raise_argument_error('parent directory must exist') unless File.directory?(parent_dir)
         end
       end
     end

data/lib/boxcars/vector_store/hnswlib/search.rb

@@ -0,0 +1,70 @@
+# frozen_string_literal: true
+
+require 'hnswlib'
+require 'json'
+
+module Boxcars
+  module VectorStore
+    module Hnswlib
+      class Search
+        include VectorStore
+
+        def initialize(params)
+          validate_params(params[:vector_documents])
+          @vector_documents = params[:vector_documents]
+          @search_index = load_index(params[:vector_documents])
+        end
+
+        def call(query_vector:, count: 1)
+          search(query_vector, count)
+        end
+
+        private
+
+        attr_reader :vector_documents, :vector_store, :json_doc, :search_index, :metadata
+
+        def validate_params(vector_documents)
+          raise_argument_error('vector_documents is nil') unless vector_documents
+          raise_arugment_error('vector_documents must be a hash') unless vector_documents.is_a?(Hash)
+          raise_arugment_error('type must be hnswlib') unless vector_documents[:type] == :hnswlib
+          raise_arugment_error('vector_store is nil') unless vector_documents[:vector_store]
+          raise_arugment_error('vector_store must be an array') unless vector_documents[:vector_store].is_a?(Array)
+
+          unless vector_documents[:vector_store].all? { |doc| doc.is_a?(Document) }
+            raise_arugment_error('vector_store must be an array of Document objects')
+          end
+
+          true
+        end
+
+        def load_index(vector_documents)
+          @metadata = vector_documents[:vector_store].first.metadata
+          @json_doc = @metadata[:json_doc_file_path]
+
+          search_index = ::Hnswlib::HierarchicalNSW.new(
+            space: metadata[:metric],
+            dim: metadata[:dim]
+          )
+          search_index.load_index(metadata[:index_file_path])
+          @search_index = search_index
+          @vector_store = vector_documents[:vector_store]
+
+          search_index
+        end
+
+        def search(query_vector, num_neighbors)
+          raw_results = search_index.search_knn(query_vector, num_neighbors)
+          raw_results.map { |doc_id, distance| lookup_embedding(doc_id, distance) }.compact
+        rescue StandardError => e
+          raise_argument_error("Error searching for #{query_vector}: #{e.message}")
+        end
+
+        def lookup_embedding(doc_id, distance)
+          return unless vector_store[doc_id]
+
+          { document: vector_store[doc_id], distance: distance }
+        end
+      end
+    end
+  end
+end

data/lib/boxcars/vector_store/in_memory/build_from_document_array.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+module Boxcars
+  module VectorStore
+    module InMemory
+      class BuildFromDocumentArray
+        include VectorStore
+
+        def initialize(embedding_tool: :openai, documents: nil)
+          validate_params(embedding_tool, documents)
+          @embedding_tool = embedding_tool
+          @documents = documents
+          @memory_vectors = []
+        end
+
+        def call
+          texts = documents
+          vectors = generate_vectors(texts)
+          add_vectors(vectors, documents)
+          {
+            type: :in_memory,
+            vector_store: memory_vectors
+          }
+        end
+
+        private
+
+        attr_reader :documents, :memory_vectors
+
+        def validate_params(embedding_tool, documents)
+          raise_argument_error('documents is nil') unless documents
+          return if %i[openai tensorflow].include?(embedding_tool)
+
+          raise_argument_error('embedding_tool is invalid')
+        end
+
+        # returns array of documents with vectors
+        def add_vectors(vectors, documents)
+          vectors.zip(documents).each do |vector, doc|
+            memory_vector = Document.new(
+              content: doc[:content],
+              embedding: vector[:embedding],
+              metadata: doc[:metadata].merge(dim: vector[:dim])
+            )
+            @memory_vectors << memory_vector
+          end
+        end
+      end
+    end
+  end
+end

data/lib/boxcars/vector_store/in_memory/build_from_files.rb

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module Boxcars
+  module VectorStore
+    module InMemory
+      class BuildFromFiles
+        include VectorStore
+
+        def initialize(params)
+          @split_chunk_size = params[:split_chunk_size] || 2000
+          @training_data_path = File.absolute_path(params[:training_data_path])
+          @embedding_tool = params[:embedding_tool] || :openai
+
+          validate_params(embedding_tool, training_data_path)
+          @memory_vectors = []
+        end
+
+        def call
+          data = load_data_files(training_data_path)
+          texts = split_text_into_chunks(data)
+          vectors = generate_vectors(texts)
+          add_vectors(vectors, texts)
+
+          {
+            type: :in_memory,
+            vector_store: memory_vectors
+          }
+        end
+
+        private
+
+        attr_reader :split_chunk_size, :training_data_path, :embedding_tool, :memory_vectors
+
+        def validate_params(embedding_tool, training_data_path)
+          training_data_dir = File.dirname(training_data_path.gsub(/\*{1,2}/, ''))
+
+          raise_argument_error('training_data_path parent directory must exist') unless File.directory?(training_data_dir)
+          raise_argument_error('No files found at the training_data_path pattern') if Dir.glob(training_data_path).empty?
+
+          return if %i[openai tensorflow].include?(embedding_tool)
+
+          raise_argument_error('embedding_tool is invalid')
+        end
+
+        def add_vectors(vectors, texts)
+          vectors.map.with_index do |vector, index|
+            memory_vector = Document.new(
+              content: texts[index],
+              embedding: vector[:embedding],
+              metadata: {
+                doc_id: index,
+                training_data_path: training_data_path
+              }
+            )
+            memory_vectors << memory_vector
+          end
+        end
+      end
+    end
+  end
+end

data/lib/boxcars/vector_store/in_memory/search.rb

@@ -1,72 +1,52 @@
 # frozen_string_literal: true
 
-# require 'openai'
-#
-# documents = [
-#     { page_content: "hello", metadata: { a: 1 } },
-#     { page_content: "hi", metadata: { a: 1 } },
-#     { page_content: "bye", metadata: { a: 1 } },
-#     { page_content: "what's this", metadata: { a: 1 } },
-# ]
-#
-# vector_documents = Boxcars::VectorStore::InMemory::AddDocuments.call(embedding_tool: :openai, documents: documents)
-#
-# result = Boxcars::VectorStore::InMemory::Search.call(vecotr_documents: vector_documents, query: "hello")
-#
-# expect(result).to eq(Boxcars::VectorStore::Document.new({ page_content: "hello", metadata: { a: 1 } }))
-
 module Boxcars
   module VectorStore
     module InMemory
       class Search
         include VectorStore
-        def initialize(vector_documents:, query:, embedding_tool: :openai)
-          validate_params(vector_documents, query, embedding_tool)
-          @vector_documents = vector_documents
-          @query = query
-          @embedding_tool = embedding_tool
+
+        def initialize(params)
+          validate_params(params[:vector_documents])
+          @vector_documents = params[:vector_documents]
         end
 
-        def call
-          results = @vector_documents.map do |doc|
-            {
-              document: doc,
-              similarity: cosine_similarity(query_vector, doc[:vector])
-            }
-          end
-          results.min_by { |result| -result[:similarity] }[:document]
+        def call(query_vector:, count: 1)
+          raise ::Boxcars::ArgumentError, 'query_vector is empty' if query_vector.empty?
+
+          search(query_vector, count)
         end
 
         private
 
-        def validate_params(vector_documents, query, embedding_tool)
-          raise ::Boxcars::ArgumentError, 'query is empty' if query.to_s.empty?
-          raise ::Boxcars::ArgumentError, 'embedding_tool is invalid' unless %i[openai tensorflow].include?(embedding_tool)
+        attr_reader :vector_documents
 
-          unless vector_documents.is_a?(Array) && vector_documents.all? do |doc|
-                   doc.is_a?(Hash) && doc.key?(:document) && doc.key?(:vector)
-                 end
-            raise ::Boxcars::ArgumentError, "vector_documents is not valid"
-          end
-        end
+        def validate_params(vector_documents)
+          return if valid_vector_store?(vector_documents)
 
-        def query_vector
-          embeddings_method(@embedding_tool)[:klass].call(
-            texts: [@query], client: embeddings_method(@embedding_tool)[:client]
-          ).first
+          raise ::Boxcars::ArgumentError, "vector_documents is not valid"
         end
 
-        def openai_client(openai_access_token: nil)
-          @openai_client ||= Openai.open_ai_client(openai_access_token: openai_access_token)
+        def valid_vector_store?(vector_documents)
+          vector_documents && vector_documents[:type] == :in_memory &&
+            vector_documents[:vector_store].is_a?(Array) &&
+            vector_documents[:vector_store].all? do |doc|
+              doc.is_a?(Boxcars::VectorStore::Document)
+            end
         end
 
-        def embeddings_method(embedding_tool)
-          case embedding_tool
-          when :openai
-            { klass: Boxcars::VectorStore::EmbedViaOpenAI, client: openai_client }
-          when :tensorflow
-            { klass: Boxcars::VectorStore::EmbedViaTensorflow, client: nil }
+        def search(query_vector, num_neighbors)
+          results = vector_documents[:vector_store].map do |doc|
+            {
+              document: doc,
+              similarity: cosine_similarity(query_vector, doc.embedding)
+            }
           end
+          results.sort_by { |result| -result[:similarity] }
+                 .first(num_neighbors)
+        rescue StandardError => e
+          raise_argument_error("Error searching for #{query_vector}: #{e.message}")
+          raise_error
         end
 
         def cosine_similarity(vector1, vector2)

data/lib/boxcars/vector_store/pgvector/build_from_array.rb

@@ -0,0 +1,95 @@
+# frozen_string_literal: true
+
+module Boxcars
+  module VectorStore
+    # install pgvector: https://github.com/pgvector/pgvector#installation-notes
+    module Pgvector
+      class BuildFromArray
+        include VectorStore
+
+        # params =  {
+        #   embedding_tool: embedding_tool,
+        #   input_array: input_array,
+        #   database_url: db_url,
+        #   table_name: table_name,
+        #   embedding_column_name: embedding_column_name,
+        #   content_column_name: content_column_name,
+        #   metadata_column_name: metadata_column_name
+        # }
+        def initialize(params)
+          @embedding_tool = params[:embedding_tool] || :openai
+
+          validate_params(embedding_tool, params[:input_array])
+
+          @database_url = params[:database_url]
+          @table_name = params[:table_name]
+          @embedding_column_name = params[:embedding_column_name]
+          @content_column_name = params[:content_column_name]
+          @metadata_column_name = params[:metadata_column_name]
+
+          @input_array = params[:input_array]
+          @pg_vectors = []
+        end
+
+        def call
+          texts = input_array
+          vectors = generate_vectors(texts)
+          add_vectors(vectors, texts)
+          documents = save_vector_store
+
+          {
+            type: :pgvector,
+            vector_store: documents
+          }
+        end
+
+        private
+
+        attr_reader :input_array, :embedding_tool, :pg_vectors, :database_url,
+                    :table_name, :embedding_column_name, :content_column_name,
+                    :metadata_column_name
+
+        def validate_params(embedding_tool, input_array)
+          raise_argument_error('input_array is nil') unless input_array
+          return if %i[openai tensorflow].include?(embedding_tool)
+
+          raise_argument_error('embedding_tool is invalid') unless %i[openai tensorflow].include?(embedding_tool)
+
+          input_array.each do |item|
+            next if item.key?(:content) && item.key?(:metadata)
+
+            return raise_argument_error('embedding_tool is invalid')
+          end
+        end
+
+        def add_vectors(vectors, texts)
+          raise_argument_error("vectors are nil") unless vectors
+          raise_argument_error("vectors and texts are not the same size") unless vectors.size == texts.size
+
+          vectors.zip(texts) do |vector, doc|
+            pg_vector = Document.new(
+              content: doc[:content],
+              embedding: vector[:embedding],
+              metadata: doc[:metadata]
+            )
+            @pg_vectors << pg_vector
+          end
+        end
+
+        def save_vector_store
+          result = Boxcars::VectorStore::Pgvector::SaveToDatabase.call(
+            pg_vectors: pg_vectors,
+            database_url: database_url,
+            table_name: table_name,
+            embedding_column_name: embedding_column_name,
+            content_column_name: content_column_name,
+            metadata_column_name: metadata_column_name
+          )
+          raise_argument_error('Error saving vector store to database.') unless result
+
+          result
+        end
+      end
+    end
+  end
+end

data/lib/boxcars/vector_store/pgvector/build_from_files.rb

@@ -0,0 +1,97 @@
+# frozen_string_literal: true
+
+require 'pgvector'
+require 'fileutils'
+require 'json'
+
+module Boxcars
+  module VectorStore
+    module Pgvector
+      class BuildFromFiles
+        include VectorStore
+
+        # params = {
+        #   training_data_path: training_data_path,
+        #   split_chunk_size: 200,
+        #   embedding_tool: embedding_tool,
+        #   database_url: db_url,
+        #   table_name: table_name,
+        #   embedding_column_name: embedding_column_name,
+        #   content_column_name: content_column_name
+        # }
+        def initialize(params)
+          @split_chunk_size = params[:split_chunk_size] || 2000
+          @training_data_path = File.absolute_path(params[:training_data_path])
+          @embedding_tool = params[:embedding_tool] || :openai
+
+          validate_params(embedding_tool, training_data_path)
+
+          @database_url = params[:database_url]
+          @table_name = params[:table_name]
+          @embedding_column_name = params[:embedding_column_name]
+          @content_column_name = params[:content_column_name]
+          @metadata_column_name = params[:metadata_column_name]
+
+          @pg_vectors = []
+        end
+
+        def call
+          data = load_data_files(training_data_path)
+          texts = split_text_into_chunks(data)
+          embeddings = generate_vectors(texts)
+          add_vectors(embeddings, texts)
+          documents = save_vector_store
+
+          {
+            type: :pgvector,
+            vector_store: documents
+          }
+        end
+
+        private
+
+        attr_reader :split_chunk_size, :training_data_path, :embedding_tool, :database_url,
+                    :table_name, :embedding_column_name, :content_column_name,
+                    :metadata_column_name, :pg_vectors
+
+        def validate_params(embedding_tool, training_data_path)
+          training_data_dir = File.dirname(training_data_path.gsub(/\*{1,2}/, ''))
+
+          raise_argument_error('training_data_path parent directory must exist') unless File.directory?(training_data_dir)
+          raise_argument_error('No files found at the training_data_path pattern') if Dir.glob(training_data_path).empty?
+          return if %i[openai tensorflow].include?(embedding_tool)
+
+          raise_argument_error('embedding_tool is invalid')
+        end
+
+        def add_vectors(vectors, texts)
+          vectors.map.with_index do |vector, index|
+            pg_vector = Document.new(
+              content: texts[index],
+              embedding: vector[:embedding],
+              metadata: {
+                doc_id: index,
+                training_data_path: training_data_path
+              }
+            )
+            pg_vectors << pg_vector
+          end
+        end
+
+        def save_vector_store
+          result = Boxcars::VectorStore::Pgvector::SaveToDatabase.call(
+            pg_vectors: pg_vectors,
+            database_url: database_url,
+            table_name: table_name,
+            embedding_column_name: embedding_column_name,
+            content_column_name: content_column_name,
+            metadata_column_name: metadata_column_name
+          )
+          raise_argument_error('Error saving vector store to database.') unless result
+
+          result
+        end
+      end
+    end
+  end
+end