boxcars 0.2.11 → 0.2.12

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

@@ -0,0 +1,152 @@
+# frozen_string_literal: true
+
+require 'pg'
+require 'pgvector'
+
+module Boxcars
+  module VectorStore
+    module Pgvector
+      class SaveToDatabase
+        include VectorStore
+
+        # params = {
+        #   pg_vectors: pg_vectors,
+        #   database_url: db_url,
+        #   table_name: table_name,
+        #   embedding_column_name: embedding_column_name,
+        #   content_column_name: content_column_name
+        # }
+        def initialize(params)
+          @errors = []
+          validate_param_types(params)
+          @db_connection = test_db_params(params)
+
+          @table_name = params[:table_name]
+          @content_column_name = params[:content_column_name]
+          @embedding_column_name = params[:embedding_column_name]
+          @metadata_column_name = params[:metadata_column_name]
+
+          @pg_vectors = params[:pg_vectors]
+        end
+
+        def call
+          return { success: false, error: errors } unless errors.empty?
+
+          add_vectors_to_database
+        end
+
+        private
+
+        attr_reader :database_url, :pg_vectors, :db_connection, :table_name,
+                    :embedding_column_name, :content_column_name,
+                    :metadata_column_name, :errors
+
+        def validate_param_types(params)
+          pg_vectors = params[:pg_vectors]
+
+          raise_argument_error('pg_vectors must be an array') unless pg_vectors.is_a?(Array)
+          raise_argument_error('missing data') if pg_vectors.empty?
+          raise_argument_error('invalid vector_store') unless valid_vector_store?(pg_vectors)
+          @database_url = params[:database_url]
+          raise_argument_error('missing database_url argument') if @database_url.to_s.empty?
+        end
+
+        def valid_vector_store?(pg_vectors)
+          pg_vectors.all? do |doc|
+            doc.is_a?(Boxcars::VectorStore::Document)
+          end
+        rescue TypeError => e
+          raise_argument_error(e.message)
+        end
+
+        def test_db_params(params)
+          conn = ::PG::Connection.new(@database_url)
+
+          check_db_connection(conn)
+          check_vector_extension(conn)
+          check_table_exists(conn, params[:table_name])
+          check_column_exists(conn, params)
+
+          registry = PG::BasicTypeRegistry.new.define_default_types
+          ::Pgvector::PG.register_vector(registry)
+          conn.type_map_for_queries = PG::BasicTypeMapForQueries.new(conn, registry: registry)
+          conn.type_map_for_results = PG::BasicTypeMapForResults.new(conn, registry: registry)
+          conn
+        rescue PG::Error, NameError => e
+          raise_argument_error(e.message)
+        end
+
+        def check_db_connection(conn)
+          return if conn.status == PG::CONNECTION_OK
+
+          raise_argument_error("PostgreSQL connection is not ok")
+        end
+
+        def check_vector_extension(conn)
+          return if conn.exec("SELECT 1 FROM pg_extension WHERE extname = 'vector'").any?
+
+          raise_argument_error("PostgreSQL 'vector' extension is not installed")
+        end
+
+        def check_table_exists(conn, table_name)
+          table_exists = conn.exec_params(
+            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)", [table_name]
+          ).getvalue(0, 0) == "t"
+          return if table_exists
+
+          raise_argument_error("Table '#{table_name}' does not exist")
+        end
+
+        def check_column_exists(conn, params)
+          column_names = %i[embedding_column_name content_column_name metadata_column_name]
+          table_name = params[:table_name]
+
+          column_names.each do |target|
+            column_name = params[target]
+            column_exists = conn.exec_params(
+              "SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = $2)",
+              [table_name, column_name]
+            ).getvalue(0, 0) == "t"
+            next if column_exists
+
+            raise_argument_error("Column '#{column_name}' does not exist in table '#{table_name}'")
+          end
+        end
+
+        def add_vectors_to_database
+          pg_vectors.each do |document|
+            embedding = document.embedding.map(&:to_f)
+            content = document.content
+            metadata = document.metadata.to_json
+
+            if document.metadata[:id]
+              id = document.metadata[:id]
+              # directly inserting table_name, embedding_column_name, and content_column_name
+              # into the SQL command. If these values are coming from an untrusted source,
+              # there is a risk of SQL injection
+              sql = <<-SQL
+                INSERT INTO #{table_name} (id, #{embedding_column_name}, #{content_column_name}, #{metadata_column_name})
+                VALUES ($1, $2, $3, $4)
+                ON CONFLICT (id) DO UPDATE
+                SET #{embedding_column_name} = EXCLUDED.#{embedding_column_name},
+                    #{content_column_name} = EXCLUDED.#{content_column_name},
+                    #{metadata_column_name} = EXCLUDED.#{metadata_column_name}
+              SQL
+              # parameters are given separately from the SQL command,
+              # there's no risk of them being interpreted as part of the command.
+              db_connection.exec_params(sql, [id, embedding, content, metadata])
+            else
+              sql = <<-SQL
+                INSERT INTO #{table_name} (#{embedding_column_name}, #{content_column_name}, #{metadata_column_name})
+                VALUES ($1, $2, $3)
+              SQL
+              db_connection.exec_params(sql, [embedding, content, metadata])
+            end
+          end
+        rescue PG::Error => e
+          raise_argument_error(e.message)
+        end
+      end
+    end
+  end
+end

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

@@ -0,0 +1,144 @@
+# frozen_string_literal: true
+
+require 'pg'
+require 'json'
+
+module Boxcars
+  module VectorStore
+    module Pgvector
+      class Search
+        include VectorStore
+
+        # required params:
+        # {
+        #   type: :pgvector,
+        #   vector_store: {
+        #     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
+        #   }
+        # }
+        def initialize(params)
+          vector_store = validate_params(params)
+          db_url = validate_vector_store(vector_store)
+          @db_connection = test_db(db_url)
+
+          @vector_documents = params[:vector_documents]
+        end
+
+        def call(query_vector:, count: 1)
+          raise ::Boxcars::ArgumentError, 'query_vector is empty' if query_vector.empty?
+
+          search(query_vector, count)
+        end
+
+        private
+
+        attr_reader :vector_documents, :vector_store, :db_connection,
+                    :table_name, :embedding_column_name, :content_column_name
+
+        def validate_params(params)
+          @vector_documents = 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 pgvector') unless vector_documents[:type] == :pgvector
+
+          @vector_store = vector_documents[:vector_store]
+          @vector_store
+        end
+
+        def validate_vector_store(vector_store)
+          raise_arugment_error('vector_store is nil') unless vector_store
+          raise_arugment_error('vector_store must be a hash') unless vector_store.is_a?(Hash)
+          raise_arugment_error('vector_store must have a table_name') unless vector_store[:table_name]
+          raise_arugment_error('vector_store must have a embedding_column_name') unless vector_store[:embedding_column_name]
+          raise_arugment_error('vector_store must have a content_column_name') unless vector_store[:content_column_name]
+          raise_argument_error('missing DATABASE_URL') unless vector_store[:database_url]
+
+          vector_store[:database_url]
+        end
+
+        def test_db(db_url)
+          conn = ::PG::Connection.new(db_url)
+
+          check_db_connection(conn)
+          check_vector_extension(conn)
+          check_table_exists(conn, vector_store[:table_name])
+          check_column_exists(conn)
+
+          @table_name = vector_store[:table_name]
+          @embedding_column_name = vector_store[:embedding_column_name]
+          @content_column_name = vector_store[:content_column_name]
+
+          conn
+        rescue PG::Error, PG::UndefinedTable, NameError => e
+          raise_argument_error(e.message)
+        end
+
+        def check_db_connection(conn)
+          return if conn.status == PG::CONNECTION_OK
+
+          raise_argument_error("PostgreSQL connection is not ok")
+        end
+
+        def check_vector_extension(conn)
+          return if conn.exec("SELECT 1 FROM pg_extension WHERE extname = 'vector'").any?
+
+          raise_argument_error("PostgreSQL 'vector' extension is not installed")
+        end
+
+        def check_table_exists(conn, table_name)
+          table_exists = conn.exec_params(
+            "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = $1)", [table_name]
+          ).getvalue(0, 0) == "t"
+          return if table_exists
+
+          raise_argument_error("Table '#{table_name}' does not exist")
+        end
+
+        def check_column_exists(conn)
+          column_names = %i[embedding_column_name content_column_name]
+          table_name = vector_store[:table_name]
+
+          column_names.each do |target|
+            column_name = vector_store[target]
+            column_exists = conn.exec_params(
+              "SELECT EXISTS (SELECT 1 FROM information_schema.columns WHERE table_name = $1 AND column_name = $2)",
+              [table_name, column_name]
+            ).getvalue(0, 0) == "t"
+            next if column_exists
+
+            raise_argument_error("Column '#{column_name}' does not exist in table '#{table_name}'")
+          end
+        end
+
+        def search(query_vector, num_neighbors)
+          sql = <<-SQL
+            SELECT *, #{embedding_column_name} <-> $1 AS distance FROM #{table_name}
+            ORDER BY #{embedding_column_name} <-> $1
+            LIMIT #{num_neighbors}
+          SQL
+          result = db_connection.exec_params(sql, [query_vector.to_s]).to_a
+          return [] if result.empty?
+
+          result.map { |hash| hash.transform_keys(&:to_sym) }
+                .map do |item|
+                  {
+                    document: Boxcars::VectorStore::Document.new(
+                      content: item[:content],
+                      embedding: JSON.parse(item[:embedding]),
+                      metadata: JSON.parse(item[:metadata], symbolize_names: true)
+                    ),
+                    distance: item[:distance].to_f
+                  }
+                end
+        rescue StandardError => e
+          raise_argument_error("Error searching for #{query_vector}: #{e.message}")
+        end
+      end
+    end
+  end
+end

data/lib/boxcars/vector_store/split_text.rb

@@ -6,14 +6,11 @@ module Boxcars
     class SplitText
       include VectorStore
 
-      attr_reader :separator, :chunk_size, :chunk_overlap, :text
-
       # @param separator [String] The string to use to split the text.
       # @param chunk_size [Integer] The size of each chunk.
       # @param chunk_overlap [Integer] The amount of overlap between chunks.
       # @param text [String] The text to split.
       def initialize(separator: "Search", chunk_size: 7, chunk_overlap: 3, text: "")
-        # require 'debugger'; debugger
         validate_params(separator, chunk_size, chunk_overlap, text)
 
         @separator = separator
@@ -31,6 +28,8 @@ module Boxcars
 
       private
 
+      attr_reader :separator, :chunk_size, :chunk_overlap, :text
+
       def validate_params(separator, chunk_size, chunk_overlap, text)
         raise_error("separator must be a string") unless separator.is_a?(String)
         raise_error("chunk_size must be an integer") unless chunk_size.is_a?(Integer)

data/lib/boxcars/vector_store.rb

@@ -13,10 +13,72 @@ module Boxcars
 
     def self.included(base)
       base.extend(ClassMethods)
+    end
+
+    private
+
+    attr_reader :embedding_tool
+
+    def generate_vectors(texts)
+      @embedding_tool = embedding_tool || :openai
+
+      embeddings_method[:klass]
+        .call(
+          texts: texts, client: embeddings_method[:client]
+        )
+        .map { |item| item.transform_keys(&:to_sym) }
+    end
+
+    def embeddings_method
+      case @embedding_tool
+      when :openai
+        { klass: Boxcars::VectorStore::EmbedViaOpenAI, client: openai_client }
+      when :tensorflow
+        { klass: Boxcars::VectorStore::EmbedViaTensorflow, client: nil }
+      end
+    end
+
+    # Get the OpenAI client
+    # @param openai_access_token [String] the OpenAI access token
+    # @return [OpenAI::Client]
+    def openai_client(openai_access_token: nil)
+      @openai_client ||= Openai.open_ai_client(openai_access_token: openai_access_token)
+    end
+
+    def raise_argument_error(message)
+      raise ::Boxcars::ArgumentError, message
+    end
+
+    def parse_json_file(file_path)
+      return [] if file_path.nil?
+
+      file_content = File.read(file_path)
+      JSON.parse(file_content, symbolize_names: true)
+    rescue JSON::ParserError => e
+      raise_argument_error("Error parsing #{file_path}: #{e.message}")
+    end
+
+    def load_data_files(training_data_path)
+      data = []
+      files = Dir.glob(training_data_path)
+      raise_error "No files found at #{training_data_path}" if files.empty?
+
+      files.each do |file|
+        data << File.read(file)
+      end
+      puts "Added #{files.length} files to data. Splitting text into chunks..."
+      data
+    end
 
-      class << base
-        private :new
+    def split_text_into_chunks(data)
+      docs = []
+      data.each do |chunk|
+        doc_output = Boxcars::VectorStore::SplitText.call(
+          separator: "\n", chunk_size: split_chunk_size, chunk_overlap: 0, text: chunk
+        )
+        docs.concat(doc_output)
       end
+      docs
     end
   end
 end
@@ -25,10 +87,14 @@ require_relative "vector_store/document"
 require_relative "vector_store/embed_via_open_ai"
 require_relative "vector_store/embed_via_tensorflow"
 require_relative "vector_store/split_text"
-require_relative "vector_store/similarity_search"
-require_relative "vector_store/hnswlib/hnswlib_config"
+require_relative "vector_store/hnswlib/load_from_disk"
 require_relative "vector_store/hnswlib/save_to_hnswlib"
-require_relative "vector_store/hnswlib/build_vector_store"
-require_relative "vector_store/hnswlib/hnswlib_search"
-require_relative "vector_store/in_memory/add_documents"
+require_relative "vector_store/hnswlib/build_from_files"
+require_relative "vector_store/hnswlib/search"
+require_relative "vector_store/in_memory/build_from_files"
+require_relative "vector_store/in_memory/build_from_document_array"
 require_relative "vector_store/in_memory/search"
+require_relative "vector_store/pgvector/build_from_files"
+require_relative "vector_store/pgvector/build_from_array"
+require_relative "vector_store/pgvector/save_to_database"
+require_relative "vector_store/pgvector/search"

data/lib/boxcars/version.rb

@@ -2,5 +2,5 @@
 
 module Boxcars
   # The current version of the gem.
-  VERSION = "0.2.11"
+  VERSION = "0.2.12"
 end

data/lib/boxcars.rb

@@ -58,7 +58,7 @@ module Boxcars
               # override with kwargs if present
               kwargs[key]
             elsif (provided_val = instance_variable_get("@#{key}"))
-              # use saved value if present. Set using Boxcars::configuration.the_key = "abcde"
+              # use saved value if present. Set using Boxcars.configuration.the_key = "abcde"
               provided_val
             else
               # otherwise, dig out of the environment

metadata

@@ -1,7 +1,7 @@
 --- !ruby/object:Gem::Specification
 name: boxcars
 version: !ruby/object:Gem::Version
-  version: 0.2.11
+  version: 0.2.12
 platform: ruby
 authors:
 - Francis Sullivan
@@ -9,7 +9,7 @@ authors:
 autorequire:
 bindir: exe
 cert_chain: []
-date: 2023-05-05 00:00:00.000000000 Z
+date: 2023-05-22 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: debug
@@ -87,14 +87,14 @@ dependencies:
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '3.0'
+        version: '4.0'
   type: :runtime
   prerelease: false
   version_requirements: !ruby/object:Gem::Requirement
     requirements:
     - - "~>"
       - !ruby/object:Gem::Version
-        version: '3.0'
+        version: '4.0'
 description: You simply set an OpenAI key, give a number of Boxcars to a Train, and
   magic ensues when you run it.
 email:
@@ -144,13 +144,17 @@ files:
 - lib/boxcars/vector_store/document.rb
 - lib/boxcars/vector_store/embed_via_open_ai.rb
 - lib/boxcars/vector_store/embed_via_tensorflow.rb
-- lib/boxcars/vector_store/hnswlib/build_vector_store.rb
-- lib/boxcars/vector_store/hnswlib/hnswlib_config.rb
-- lib/boxcars/vector_store/hnswlib/hnswlib_search.rb
+- lib/boxcars/vector_store/hnswlib/build_from_files.rb
+- lib/boxcars/vector_store/hnswlib/load_from_disk.rb
 - lib/boxcars/vector_store/hnswlib/save_to_hnswlib.rb
-- lib/boxcars/vector_store/in_memory/add_documents.rb
+- lib/boxcars/vector_store/hnswlib/search.rb
+- lib/boxcars/vector_store/in_memory/build_from_document_array.rb
+- lib/boxcars/vector_store/in_memory/build_from_files.rb
 - lib/boxcars/vector_store/in_memory/search.rb
-- lib/boxcars/vector_store/similarity_search.rb
+- lib/boxcars/vector_store/pgvector/build_from_array.rb
+- lib/boxcars/vector_store/pgvector/build_from_files.rb
+- lib/boxcars/vector_store/pgvector/save_to_database.rb
+- lib/boxcars/vector_store/pgvector/search.rb
 - lib/boxcars/vector_store/split_text.rb
 - lib/boxcars/version.rb
 homepage: https://github.com/BoxcarsAI/boxcars
@@ -169,7 +173,7 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: 2.6.0
+      version: '3.0'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="

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

@@ -1,157 +0,0 @@
-# frozen_string_literal: true
-
-require 'fileutils'
-require 'hnswlib'
-require 'json'
-
-module Boxcars
-  module VectorStore
-    module Hnswlib
-      class BuildVectorStore
-        include VectorStore
-
-        # This class is responsible for building the vector store for the hnswlib similarity search.
-        # It will load the training data, generate the embeddings, and save the vector store.
-        # It will also load the vector store into memory.
-        # For later use, it will save the splitted document with index numbers to a json file.
-        #
-        # @param training_data_path [String] The path to the training data. Can be a glob pattern.
-        # @param index_file_path [String] The path to the index file.
-        # @param split_chunk_size [Integer] The number of documents to split the text into. default 2000
-        # @option json_doc_file_path [String]. The json file containing the document text.
-        #                                      if nil, it will reuse index file name.
-        # @option force_rebuild [Boolean] Optional. If true, will rebuild the index even if it already exists.
-        def initialize(
-          training_data_path:,
-          index_file_path:,
-          split_chunk_size: 2000,
-          json_doc_file_path: nil,
-          force_rebuild: true
-        )
-          @training_data_path = training_data_path
-          @index_file_path = index_file_path
-          @split_chunk_size = split_chunk_size
-          @json_doc_file_path = json_doc_file_path || index_file_path.gsub(/\.bin$/, '.json')
-          @force_rebuild = force_rebuild
-        end
-
-        def call
-          validate_params
-          data = load_files
-          documents = split_text_into_chunks(data)
-          embeddings_with_config = generate_embeddings(documents)
-          save_vector_store(embeddings_with_config)
-          load_hnsw
-        end
-
-        private
-
-        attr_reader :training_data_path, :index_file_path, :split_chunk_size, :json_doc_file_path, :force_rebuild
-
-        def validate_params
-          training_data_dir = File.dirname(training_data_path.gsub(/\*{1,2}/, ''))
-          raise_error('training_data_path parent directory must exist') unless File.directory?(training_data_dir)
-          raise_error('No files found at the training_data_path pattern') if Dir.glob(training_data_path).empty?
-
-          index_dir = File.dirname(index_file_path)
-          raise_error('index_file_path parent directory must exist') unless File.directory?(index_dir)
-
-          raise_error('split_chunk_size must be an integer') unless split_chunk_size.is_a?(Integer)
-        end
-
-        def load_files
-          data = []
-          files = Dir.glob(training_data_path)
-          raise_error "No files found at #{training_data_path}" if files.empty?
-
-          files.each do |file|
-            data << File.read(file)
-          end
-          puts "Added #{files.length} files to data. Splitting text into chunks..."
-          data
-        end
-
-        def split_text_into_chunks(data)
-          return true unless rebuild_required?
-
-          docs = []
-          data.each do |chunk|
-            doc_output = Boxcars::VectorStore::SplitText.call(
-              separator: "\n", chunk_size: split_chunk_size, chunk_overlap: 0, text: chunk
-            )
-            docs.concat(doc_output)
-          end
-          docs
-        end
-
-        def rebuild_required?
-          hnswlib_config_json = "#{File.dirname(index_file_path)}/hnswlib_config.json"
-          return true unless File.exist?(index_file_path)
-          return true if File.exist?(index_file_path) && !File.exist?(hnswlib_config_json)
-          return true if force_rebuild
-
-          false
-        end
-
-        def generate_embeddings(documents)
-          return true unless rebuild_required?
-
-          puts "Initializing Store..."
-          openai_client = Openai.open_ai_client
-          embeddings_with_dim = Boxcars::VectorStore::EmbedViaOpenAI.call(texts: documents, client: openai_client)
-          document_embeddings = embeddings_with_dim.map.with_index do |item, index|
-            { doc_id: index, embedding: item[:embedding], document: documents[index] }
-          end
-
-          { document_embeddings: document_embeddings, dim: embeddings_with_dim.first[:dim] }
-        end
-
-        def save_vector_store(embeddings_with_config)
-          return true unless rebuild_required?
-
-          puts "Saving Vectorstore"
-          Boxcars::VectorStore::Hnswlib::SaveToHnswlib.call(
-            document_embeddings: embeddings_with_config[:document_embeddings],
-            index_file_path: index_file_path,
-            json_doc_file_path: json_doc_file_path,
-            hnswlib_config: hnswlib_config(embeddings_with_config[:dim])
-          )
-          puts "VectorStore saved"
-        end
-
-        def hnswlib_config(dim)
-          # dim: length of datum point vector that will be indexed.
-          Boxcars::VectorStore::Hnswlib::HnswlibConfig.new(
-            metric: "l2", max_item: 10000, dim: dim
-          )
-        end
-
-        def load_hnsw
-          puts "Loading Hnswlib"
-
-          config_file = "#{File.dirname(index_file_path)}/hnswlib_config.json"
-          json_config = parse_json_file(config_file)
-          document_embeddings = parse_json_file(json_doc_file_path)
-
-          search_index = ::Hnswlib::HierarchicalNSW.new(space: json_config[:metric], dim: json_config[:dim])
-          search_index.load_index(index_file_path)
-
-          { vector_store: search_index, document_embeddings: document_embeddings }
-        end
-
-        def parse_json_file(file_path)
-          return [] if file_path.nil?
-
-          file_content = File.read(file_path)
-          JSON.parse(file_content, symbolize_names: true)
-        rescue JSON::ParserError => e
-          raise_error("Error parsing hnswlib_config.json: #{e.message}")
-        end
-
-        def raise_error(message)
-          raise ::Boxcars::Error, message
-        end
-      end
-    end
-  end
-end