Working version

Author: dqii

CONTRIBUTING.md

@@ -57,3 +57,14 @@ Run a specific test:
 ```bash
 bundle exec rake test TEST=test/path/to/your_test_file.rb
 ```
+
+## Roadmap
+
+### In Progress
+
+* Support nearest neighbor queries
+
+### Future
+
+* Support creating vector indexes
+* Support creating embedding jobs

README.md

@@ -32,4 +32,3 @@ rails db:migrate
 After checking out the repo, run `bin/setup` to install dependencies. Then, run `rake spec` to run the tests. You can also run `bin/console` for an interactive prompt that will allow you to experiment.
 
 To install this gem onto your local machine, run `bundle exec rake install`. To release a new version, update the version number in `version.rb`, and then run `bundle exec rake release`, which will create a git tag for the version, push git commits and tags, and push the `.gem` file to [rubygems.org](https://rubygems.org).
-

lib/generators/lantern/lantern_generator.rb

@@ -3,14 +3,16 @@
 
 module Lantern
   module Generators
-    class InstallGenerator < Rails::Generators::Base
+    class LanternGenerator < Rails::Generators::Base
       include ActiveRecord::Generators::Migration
       source_root File.join(__dir__, 'templates')
 
       def copy_migration
         migration_template 'lantern.rb', 'db/migrate/install_lantern.rb', migration_version: migration_version
       end
 
+      private
+
       def migration_version
         "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
       end

lib/lantern.rb

@@ -1,7 +1,13 @@
 require_relative 'lantern/version'
+require_relative 'lantern/model'
+require 'active_support'
 
 module Lantern
   class Error < StandardError; end
 end
 
+ActiveSupport.on_load(:active_record) do
+  extend Lantern::Model
+end
+
 require_relative 'lantern/railtie' if defined?(Rails::Railtie)

lib/lantern/model.rb

@@ -0,0 +1,29 @@
+module Lantern
+  module Model
+    def has_neighbors(*attribute_names)
+      attribute_names.map!(&:to_sym)
+
+      scope :nearest_neighbors, ->(attribute_name, vector) {
+        attribute_name = attribute_name.to_sym
+
+        unless vector.all? { |v| v.is_a?(Integer) }
+            vector = vector.map(&:to_f)
+        end
+
+        vector_literal = vector.join(',')
+        order_expression = "#{quoted_table_name}.#{connection.quote_column_name(attribute_name)} <-> ARRAY[#{vector_literal}]"
+
+        select("#{quoted_table_name}.*, #{order_expression} AS distance")
+          .where.not(attribute_name => nil)
+          .order(Arel.sql(order_expression))
+      }
+
+      define_method :nearest_neighbors do |attribute_name|
+        attribute_name = attribute_name.to_sym
+        self.class
+          .where.not(id: id)
+          .nearest_neighbors(attribute_name, self[attribute_name])
+      end
+    end
+  end
+end

test/lantern_generator_test.rb

@@ -2,7 +2,7 @@
 require "generators/lantern/lantern_generator"
 
 class LanternGeneratorTest < Rails::Generators::TestCase
-  tests Lantern::Generators::InstallGenerator
+  tests Lantern::Generators::LanternGenerator
   destination File.expand_path('../tmp', __dir__)
   setup :prepare_destination
 

test/model_test.rb

@@ -0,0 +1,62 @@
+require "test_helper"
+
+class Item < ActiveRecord::Base
+  has_neighbors :embedding
+end
+
+class LanternTest < Minitest::Test
+  def setup
+    @conn = ActiveRecord::Base.connection
+
+    @conn.execute "DROP TABLE IF EXISTS items"
+    @conn.execute <<~SQL
+      CREATE TABLE items (
+        id serial PRIMARY KEY,
+        embedding REAL[]
+      )
+    SQL
+
+    @vectors = [
+      [1.0, 0.0, 0.0],
+      [0.0, 1.0, 0.0],
+      [0.0, 0.0, 1.0]
+    ]
+
+    @vectors.each do |vector|
+      @conn.execute "INSERT INTO items (embedding) VALUES (ARRAY#{vector})"
+    end
+  end
+
+  def test_instance_nearest_neighbors
+    item = Item.first
+    neighbors = item.nearest_neighbors(:embedding).limit(2)
+  
+    assert_equal 2, neighbors.length
+    assert_equal [2, 3], neighbors.pluck(:id)
+
+    distances = neighbors.map { |neighbor| neighbor.distance }
+    expected_distances = [2.0, 2.0]
+    assert_equal expected_distances, distances
+  end
+  
+  def test_class_nearest_neighbors
+    neighbors = Item.nearest_neighbors(:embedding, [0.5, 0.0, 0.5]).limit(2)
+
+    assert 2, neighbors.length
+    assert_equal [1, 3], neighbors.pluck(:id)
+
+    distances = neighbors.map { |neighbor| neighbor.distance }
+    expected_distances = [0.5, 0.5]
+    assert_equal expected_distances, distances
+  end
+
+  def test_invalid_dimensions
+    assert_raises(ActiveRecord::StatementInvalid) do
+      Item.nearest_neighbors(:embedding, [1.0, 0.0]).limit(2).to_a
+    end
+  end
+
+  def teardown
+    @conn.execute "DROP TABLE IF EXISTS items"
+  end
+end

test/test_helper.rb

@@ -2,3 +2,8 @@
 Bundler.require(:default)
 require "minitest/autorun"
 require "active_record"
+require "dotenv/load"
+
+ActiveRecord::Base.establish_connection(ENV["DATABASE_URL"])
+
+ActiveRecord::Base.connection.enable_extension("lantern")