boxcars 0.7.7 → 0.8.0

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

@@ -1,96 +1,122 @@
 # frozen_string_literal: true
 
-require 'pgvector'
-require 'fileutils'
-require 'json'
-
-module Boxcars
-  module VectorStore
-    module Pgvector
-      class BuildFromFiles
-        include VectorStore
-
-        # @param training_data_path [String] path to training data files
-        # @param split_chunk_size [Integer] number of characters to split the text into
-        # @param embedding_tool [Symbol] embedding tool to use
-        # @param database_url [String] database url
-        # @param table_name [String] table name
-        # @param embedding_column_name [String] embedding column name
-        # @param content_column_name [String] content column name
-        # @param metadata_column_name [String] metadata column name
-        # @return [Hash] vector_store: array of hashes with :content, :metadata, and :embedding keys
-        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
+if Gem.loaded_specs.key?('pgvector') && Gem.loaded_specs.key?('pg')
+  require 'pgvector'
+  require 'fileutils'
+  require 'json'
+
+  module Boxcars
+    module VectorStore
+      module Pgvector
+        class BuildFromFiles
+          include VectorStore
+
+          # @param training_data_path [String] path to training data files
+          # @param split_chunk_size [Integer] number of characters to split the text into
+          # @param embedding_tool [Symbol] embedding tool to use
+          # @param database_url [String] database url
+          # @param table_name [String] table name
+          # @param embedding_column_name [String] embedding column name
+          # @param content_column_name [String] content column name
+          # @param metadata_column_name [String] metadata column name
+          # @return [Hash] vector_store: array of hashes with :content, :metadata, and :embedding keys
+          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
 
-        # @return [Hash] vector_store: array of Inventor::VectorStore::Document
-        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
+          # @return [Hash] vector_store: array of Inventor::VectorStore::Document
+          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
+          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
+          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}/, ''))
+          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 Dir.exist?(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('training_data_path parent directory must exist') unless Dir.exist?(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
+            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 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
-              }
+          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
             )
-            pg_vectors << pg_vector
+            raise_argument_error('Error saving vector store to database.') unless result
+
+            result
           end
         end
+      end
+    end
+  end
+else
+  # Define placeholder modules/classes that raise an error if pgvector is not available
+  module Boxcars
+    module VectorStore
+      module Pgvector
+        class PgvectorNotAvailableError < StandardError
+          DEFAULT_MESSAGE = "The 'pgvector' and 'pg' gems are required. Please add them to your Gemfile."
+          def initialize(message = DEFAULT_MESSAGE)
+            super
+          end
+        end
+
+        class BuildFromFiles
+          def initialize(*_args)
+            raise PgvectorNotAvailableError
+          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
+          def call(*_args)
+            raise PgvectorNotAvailableError
+          end
         end
       end
     end

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

@@ -1,148 +1,174 @@
 # frozen_string_literal: true
 
-require 'pg'
-require 'pgvector'
-
-module Boxcars
-  module VectorStore
-    module Pgvector
-      class SaveToDatabase
-        include VectorStore
-
-        # @param pg_vectors [Array] array of Boxcars::VectorStore::Document
-        # @param database_url [String] database url
-        # @param table_name [String] table name
-        # @param embedding_column_name [String] embedding column name
-        # @param content_column_name [String] content column name
-        # @param metadata_column_name [String] metadata column name
-        # @return [Array] array of Boxcars::VectorStore::Document
-        def initialize(params)
-          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
+if Gem.loaded_specs.key?('pgvector') && Gem.loaded_specs.key?('pg')
+  require 'pg'
+  require 'pgvector'
+
+  module Boxcars
+    module VectorStore
+      module Pgvector
+        class SaveToDatabase
+          include VectorStore
+
+          # @param pg_vectors [Array] array of Boxcars::VectorStore::Document
+          # @param database_url [String] database url
+          # @param table_name [String] table name
+          # @param embedding_column_name [String] embedding column name
+          # @param content_column_name [String] content column name
+          # @param metadata_column_name [String] metadata column name
+          # @return [Array] array of Boxcars::VectorStore::Document
+          def initialize(params)
+            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
 
-        # @return [Array] array of Boxcars::VectorStore::Document
-        def call
-          add_vectors_to_database
-        end
+          # @return [Array] array of Boxcars::VectorStore::Document
+          def call
+            add_vectors_to_database
+          end
 
-        private
+          private
 
-        attr_reader :database_url, :pg_vectors, :db_connection, :table_name,
-                    :embedding_column_name, :content_column_name,
-                    :metadata_column_name
+          attr_reader :database_url, :pg_vectors, :db_connection, :table_name,
+                      :embedding_column_name, :content_column_name,
+                      :metadata_column_name
 
-        def validate_param_types(params)
-          pg_vectors = params[:pg_vectors]
+          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
+            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)
+          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
-        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 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
+          def check_db_connection(conn)
+            return if conn.status == PG::CONNECTION_OK
 
-          raise_argument_error("PostgreSQL connection is not ok")
-        end
+            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?
+          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
+            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
 
-        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
 
-          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]
 
-        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
 
-          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
 
-            raise_argument_error("Column '#{column_name}' does not exist in table '#{table_name}'")
+          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
+else
+  # Define placeholder modules/classes that raise an error if pgvector is not available
+  module Boxcars
+    module VectorStore
+      module Pgvector
+        class PgvectorNotAvailableError < StandardError
+          DEFAULT_MESSAGE = "The 'pgvector' and 'pg' gems are required. Please add them to your Gemfile."
+          def initialize(message = DEFAULT_MESSAGE)
+            super
           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
+        class SaveToDatabase
+          def initialize(*_args)
+            raise PgvectorNotAvailableError
+          end
+
+          def call(*_args)
+            raise PgvectorNotAvailableError
           end
-        rescue PG::Error => e
-          raise_argument_error(e.message)
         end
       end
     end