Added typecasting for Sequel

pgvector · Jun 4, 2023 · 25fec94 · 25fec94
1 parent 300eb26
commit 25fec94
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 4 deletions.
diff --git a/README.md b/README.md
@@ -72,7 +72,7 @@ end
 Insert a vector
 
 ```ruby
-Item.create(embedding: Pgvector.encode([1, 1, 1]))
+Item.create(embedding: [1, 1, 1])
 ```
 
 Get the nearest neighbors to a record

diff --git a/lib/sequel/plugins/pgvector.rb b/lib/sequel/plugins/pgvector.rb
@@ -52,6 +52,14 @@ def nearest_neighbors(column, **options)
             .nearest_neighbors(column, self[column], **options)
             .exclude(primary_key => self[primary_key])
         end
+
+        def []=(k, v)
+          if self.class.vector_columns.key?(k) && !v.is_a?(String)
+            super(k, ::Pgvector.encode(v))
+          else
+            super
+          end
+        end
       end
     end
   end

diff --git a/test/sequel_test.rb b/test/sequel_test.rb
@@ -27,9 +27,9 @@ def test_dataset
   end
 
   def test_model
-    Item.create(embedding: Pgvector.encode([1, 1, 1]))
-    Item.create(embedding: Pgvector.encode([2, 2, 2]))
-    Item.create(embedding: Pgvector.encode([1, 1, 2]))
+    Item.create(embedding: [1, 1, 1])
+    Item.create(embedding: [2, 2, 2])
+    Item.create(embedding: [1, 1, 2])
 
     results = Item.nearest_neighbors(:embedding, [1, 1, 1], distance: "euclidean").limit(5)
     assert_equal ["[1,1,1]", "[1,1,2]", "[2,2,2]"], results.map(&:embedding)