Added extension for Sequel - #3

CHANGELOG.md

@@ -1,3 +1,7 @@
+## 0.3.1 (unreleased)
+
+- Added extension for Sequel
+
 ## 0.3.0 (2024-06-25)
 
 - Added support for `halfvec` and `sparsevec` types

lib/sequel/extensions/pgvector.rb

@@ -0,0 +1,5 @@
+require_relative "../plugins/pgvector"
+
+module Sequel
+  Dataset.register_extension(:pgvector, Plugins::Pgvector::DatasetMethods)
+end

test/sequel_test.rb

@@ -36,6 +36,13 @@ def test_dataset
     assert_equal [[1, 1, 1], [1, 1, 2], [2, 2, 2]], results.map { |r| Pgvector.decode(r[:embedding]) }
   end
 
+  def test_extension
+    items.insert(embedding: Pgvector.encode([1, 1, 1]))
+    items.multi_insert([{embedding: "[2,2,2]"}, {embedding: "[1,1,2]"}])
+    results = items.extension(:pgvector).nearest_neighbors(:embedding, [1, 1, 1], distance: "euclidean").limit(5)
+    assert_equal [[1, 1, 1], [1, 1, 2], [2, 2, 2]], results.map { |r| Pgvector.decode(r[:embedding]) }
+  end
+
   def test_model
     Item.create(id: 1, embedding: [1, 1, 1], half_embedding: [1, 1, 1], binary_embedding: "000", sparse_embedding: Pgvector::SparseVector.new([1, 1, 1]))
     Item.create(id: 2, embedding: [2, 2, 2], half_embedding: [2, 2, 2], binary_embedding: "101", sparse_embedding: Pgvector::SparseVector.new([2, 2, 2]))