pgvector 0.2.0 → 0.2.2

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: a0729ce3ec22ea08be365fc6c7a5c5e0b038860d5c9a57ccd2554d9d8b57a3f2
-  data.tar.gz: c36eabc18ca6db3dd47204df3c48b58f7ac635e230bb9039269506d7d8afa834
+  metadata.gz: 7485ea4be0d5be0177a972db911c696daf3438a661ddac61b08f4e8b2da3ac51
+  data.tar.gz: 2532ef79f5db88aecb681d9455e38f3e5fc1d30bde015d0a1e9daaa9fe82635e
 SHA512:
-  metadata.gz: ff4dc7d6a7a26b1ff3502d2305b62fc25fff38f3bf314261dcb6aa0f40cc520963cfd9562af03b11678cad7819aef9951b480a61e138ae2ae02459ed668a5298
-  data.tar.gz: 6aa39019d5566f83331f3c95bb6404b065b7f7a29be62e84be2f2b6ea58bc27d20002543d93dfec7a17256c970e32b9676e2d319d07758cf49fc31b89d3d4d7b
+  metadata.gz: be40e4c3e16dd904a200115794a8ffaa850b40c5055330bd873ec7a707164a53b29d22040defbb4dd8a9cff597d4e1ad5c659d37a580ccc201ce30a9eb17fef9
+  data.tar.gz: f5d36289b043d987920911ab08d85d7d1066039f84dcc2a24436701c06b246adcb8fd3c32e3f76f3e1604001403574c14818f91a64357ce4ebd675458e57184c

data/CHANGELOG.md

@@ -1,3 +1,11 @@
+## 0.2.2 (2023-10-03)
+
+- Added `nearest_neighbors` method to datasets with Sequel
+
+## 0.2.1 (2023-06-04)
+
+- Added support for Sequel
+
 ## 0.2.0 (2023-05-11)
 
 - Dropped support for Ruby < 3

data/README.md

@@ -2,7 +2,7 @@
 
 [pgvector](https://github.com/pgvector/pgvector) support for Ruby
 
-Supports the [pg](https://github.com/ged/ruby-pg) gem
+Supports [pg](https://github.com/ged/ruby-pg) and [Sequel](https://github.com/jeremyevans/sequel)
 
 For Rails, check out [Neighbor](https://github.com/ankane/neighbor)
 
@@ -19,9 +19,22 @@ gem "pgvector"
 And follow the instructions for your database library:
 
 - [pg](#pg)
+- [Sequel](#sequel)
+
+Or check out some examples:
+
+- [Embeddings](examples/openai_embeddings.rb) with OpenAI
+- [User-based recommendations](examples/disco_user_recs.rb) with Disco
+- [Item-based recommendations](examples/disco_item_recs.rb) with Disco
 
 ## pg
 
+Enable the extension
+
+```ruby
+conn.exec("CREATE EXTENSION IF NOT EXISTS vector")
+```
+
 Register the vector type with your connection
 
 ```ruby
@@ -43,6 +56,51 @@ Get the nearest neighbors to a vector
 conn.exec_params("SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5", [embedding]).to_a
 ```
 
+## Sequel
+
+Enable the extension
+
+```ruby
+DB.run("CREATE EXTENSION IF NOT EXISTS vector")
+```
+
+Create a table
+
+```ruby
+DB.create_table :items do
+  primary_key :id
+  column :embedding, "vector(3)"
+end
+```
+
+Add the plugin to your model
+
+```ruby
+class Item < Sequel::Model
+  plugin :pgvector, :embedding
+end
+```
+
+Insert a vector
+
+```ruby
+Item.create(embedding: [1, 1, 1])
+```
+
+Get the nearest neighbors to a record
+
+```ruby
+item.nearest_neighbors(:embedding, distance: "euclidean").limit(5)
+```
+
+Also supports `inner_product` and `cosine` distance
+
+Get the nearest neighbors to a vector
+
+```ruby
+Item.nearest_neighbors(:embedding, [1, 1, 1], distance: "euclidean").limit(5)
+```
+
 ## History
 
 View the [changelog](https://github.com/pgvector/pgvector-ruby/blob/master/CHANGELOG.md)

data/lib/pgvector/pg.rb

@@ -20,7 +20,7 @@ module Pgvector
     module TextDecoder
       class Vector < ::PG::SimpleDecoder
         def decode(string, tuple = nil, field = nil)
-          string[1..-2].split(",").map(&:to_f)
+          Pgvector.decode(string)
         end
       end
     end

data/lib/pgvector/version.rb

@@ -1,3 +1,3 @@
 module Pgvector
-  VERSION = "0.2.0"
+  VERSION = "0.2.2"
 end

data/lib/pgvector.rb

@@ -1,5 +1,14 @@
-require "pgvector/version"
+# modules
+require_relative "pgvector/version"
 
 module Pgvector
   autoload :PG, "pgvector/pg"
+
+  def self.encode(data)
+    "[#{data.to_a.map(&:to_f).join(",")}]"
+  end
+
+  def self.decode(string)
+    string[1..-2].split(",").map(&:to_f)
+  end
 end

data/lib/sequel/plugins/pgvector.rb

@@ -0,0 +1,81 @@
+module Sequel
+  module Plugins
+    module Pgvector
+      def self.configure(model, *columns)
+        model.vector_columns ||= {}
+        columns.each do |column|
+          model.vector_columns[column.to_sym] = {}
+        end
+      end
+
+      module DatasetMethods
+        def nearest_neighbors(column, value, distance:)
+          value = ::Pgvector.encode(value) unless value.is_a?(String)
+          quoted_column = quote_identifier(column)
+          distance = distance.to_s
+
+          operator =
+            case distance
+            when "inner_product"
+              "<#>"
+            when "cosine"
+              "<=>"
+            when "euclidean"
+              "<->"
+            end
+
+          raise ArgumentError, "Invalid distance: #{distance}" unless operator
+
+          order = "#{quoted_column} #{operator} ?"
+
+          neighbor_distance =
+            if distance == "inner_product"
+              "(#{order}) * -1"
+            else
+              order
+            end
+
+          select_append(Sequel.lit("#{neighbor_distance} AS neighbor_distance", value))
+            .exclude(column => nil)
+            .order(Sequel.lit(order, value))
+        end
+      end
+
+      module ClassMethods
+        attr_accessor :vector_columns
+
+        Sequel::Plugins.def_dataset_methods(self, :nearest_neighbors)
+
+        Plugins.inherited_instance_variables(self, :@vector_columns => :dup)
+      end
+
+      module InstanceMethods
+        def nearest_neighbors(column, **options)
+          column = column.to_sym
+          # important! check if neighbor attribute before calling send
+          raise ArgumentError, "Invalid column" unless self.class.vector_columns[column]
+
+          self.class
+            .nearest_neighbors(column, self[column], **options)
+            .exclude(primary_key => self[primary_key])
+        end
+
+        def []=(k, v)
+          if self.class.vector_columns.key?(k.to_sym) && !v.is_a?(String)
+            super(k, ::Pgvector.encode(v))
+          else
+            super
+          end
+        end
+
+        def [](k)
+          if self.class.vector_columns.key?(k.to_sym)
+            ::Pgvector.decode(super)
+          else
+            super
+          end
+        end
+      end
+    end
+  end
+end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: pgvector
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.2.2
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2023-05-12 00:00:00.000000000 Z
+date: 2023-10-04 00:00:00.000000000 Z
 dependencies: []
 description:
 email: andrew@ankane.org
@@ -22,6 +22,7 @@ files:
 - lib/pgvector.rb
 - lib/pgvector/pg.rb
 - lib/pgvector/version.rb
+- lib/sequel/plugins/pgvector.rb
 homepage: https://github.com/pgvector/pgvector-ruby
 licenses:
 - MIT