From e052e98405821f526bcda5294ce5952d8280906c Mon Sep 17 00:00:00 2001
From: Andrew Kane <andrew@ankane.org>
Date: Mon, 8 Jul 2024 18:15:08 -0700
Subject: [PATCH] Added extension for Sequel - #3

---
 CHANGELOG.md                      | 4 ++++
 lib/sequel/extensions/pgvector.rb | 5 +++++
 test/sequel_test.rb               | 7 +++++++
 3 files changed, 16 insertions(+)
 create mode 100644 lib/sequel/extensions/pgvector.rb

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6586edb..24f58c3 100644
--- a/CHANGELOG.md
+++ b/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
diff --git a/lib/sequel/extensions/pgvector.rb b/lib/sequel/extensions/pgvector.rb
new file mode 100644
index 0000000..4596530
--- /dev/null
+++ b/lib/sequel/extensions/pgvector.rb
@@ -0,0 +1,5 @@
+require_relative "../plugins/pgvector"
+
+module Sequel
+  Dataset.register_extension(:pgvector, Plugins::Pgvector::DatasetMethods)
+end
diff --git a/test/sequel_test.rb b/test/sequel_test.rb
index f504577..91b04ac 100644
--- a/test/sequel_test.rb
+++ b/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]))