diff --git a/README.md b/README.md index f330c6b..43585c3 100644 --- a/README.md +++ b/README.md @@ -307,14 +307,12 @@ class AddEmbeddingToItems < ActiveRecord::Migration[7.2] # Rails < 8 execute <<~SQL CREATE VIRTUAL TABLE items USING vec0( - id integer PRIMARY KEY AUTOINCREMENT NOT NULL, embedding float[3] distance_metric=L2 ) SQL # Rails 8+ create_virtual_table :items, :vec0, [ - "id integer PRIMARY KEY AUTOINCREMENT NOT NULL", "embedding float[3] distance_metric=L2" ] end @@ -323,6 +321,16 @@ end Use `distance_metric=cosine` for cosine distance +Create a model with `rowid` as the primary key + +```ruby +class Item < ApplicationRecord + self.primary_key = "rowid" + + has_neighbors :embedding, dimensions: 3 +end +``` + Get the nearest neighbors ```ruby diff --git a/test/sqlite_virtual_test.rb b/test/sqlite_virtual_test.rb index a995c2a..6ee5d0b 100644 --- a/test/sqlite_virtual_test.rb +++ b/test/sqlite_virtual_test.rb @@ -22,6 +22,8 @@ def test_euclidean create_items(SqliteVecItem, :embedding) relation = SqliteVecItem.where("embedding MATCH ?", [1, 1, 1].to_s).order(:distance).limit(3) + assert_equal [1, 3, 2], relation.all.map(&:id) + assert_equal [1, 3, 2], relation.pluck(:rowid) assert_elements_in_delta [0, 1, Math.sqrt(3)], relation.pluck(:distance) assert_match "SCAN vec_items VIRTUAL TABLE INDEX", relation.explain.inspect diff --git a/test/support/sqlite.rb b/test/support/sqlite.rb index aad0681..8cf2479 100644 --- a/test/support/sqlite.rb +++ b/test/support/sqlite.rb @@ -15,13 +15,11 @@ class SqliteRecord < ActiveRecord::Base if ActiveRecord::VERSION::MAJOR >= 8 create_virtual_table :vec_items, :vec0, [ - "id integer PRIMARY KEY AUTOINCREMENT NOT NULL", "embedding float[3] distance_metric=L2" ] else execute <<~SQL CREATE VIRTUAL TABLE vec_items USING vec0( - id integer PRIMARY KEY AUTOINCREMENT NOT NULL, embedding float[3] distance_metric=L2 ) SQL @@ -29,13 +27,11 @@ class SqliteRecord < ActiveRecord::Base if ActiveRecord::VERSION::MAJOR >= 8 create_virtual_table :cosine_items, :vec0, [ - "id integer PRIMARY KEY AUTOINCREMENT NOT NULL", "embedding float[3] distance_metric=cosine" ] else execute <<~SQL CREATE VIRTUAL TABLE cosine_items USING vec0( - id integer PRIMARY KEY AUTOINCREMENT NOT NULL, embedding float[3] distance_metric=cosine ) SQL @@ -51,11 +47,13 @@ class SqliteItem < SqliteRecord class SqliteVecItem < SqliteRecord has_neighbors :embedding, dimensions: 3 + self.primary_key = "rowid" self.table_name = "vec_items" end class SqliteCosineItem < SqliteRecord has_neighbors :embedding, dimensions: 3 + self.primary_key = "rowid" self.table_name = "cosine_items" end