boxcars 0.2.10 → 0.2.12

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

@@ -0,0 +1,90 @@
+# frozen_string_literal: true
+
+require 'hnswlib'
+require 'json'
+require 'fileutils'
+
+module Boxcars
+  module VectorStore
+    module Hnswlib
+      class SaveToHnswlib
+        include VectorStore
+
+        # @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.
+        # @option json_doc_file_path [String] Optional. The path to the json file containing the document text.
+        def initialize(hnsw_vectors_array)
+          @metadata = hnsw_vectors_array&.first&.metadata
+          validate_params(hnsw_vectors_array, metadata)
+
+          @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')
+
+          @metric = metadata[:metric] || "l2"
+          @dim = metadata[:dim]
+          @max_item = metadata[:max_item] || 10000
+
+          @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
+
+        attr_reader :metadata, :index, :vectors, :index_file_path, :json_doc_file_path, :metric, :dim, :max_item
+
+        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?
+
+          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
+
+        def add_vectors_to_index
+          document_texts = []
+
+          vectors.each do |item|
+            index.add_point(item.embedding, item.metadata[:doc_id])
+
+            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_argument_error('parent directory must exist') unless File.directory?(parent_dir)
+        end
+      end
+    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

@@ -0,0 +1,61 @@
+# frozen_string_literal: true
+
+module Boxcars
+  module VectorStore
+    module InMemory
+      class Search
+        include VectorStore
+
+        def initialize(params)
+          validate_params(params[:vector_documents])
+          @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
+
+        def validate_params(vector_documents)
+          return if valid_vector_store?(vector_documents)
+
+          raise ::Boxcars::ArgumentError, "vector_documents is not valid"
+        end
+
+        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 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)
+          dot_product = vector1.zip(vector2).reduce(0) { |sum, (a, b)| sum + (a * b) }
+          magnitude1 = Math.sqrt(vector1.reduce(0) { |sum, a| sum + (a**2) })
+          magnitude2 = Math.sqrt(vector2.reduce(0) { |sum, b| sum + (b**2) })
+          dot_product / (magnitude1 * magnitude2)
+        end
+      end
+    end
+  end
+end

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

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