pgvector 0.1.1 → 0.2.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 4e1249b9ced41c051b06654cea7539a58e2816d794461788bec605eab464cca8
-  data.tar.gz: fc65c3c883a39cd1b63c804eaae65897c2a332836b5d91ba5cf5248dc660fb8f
+  metadata.gz: 215f975cafa2f782f3777ec65dd2fea587b420286f7f5ccf63b6b68376756dfd
+  data.tar.gz: e1979a1aa4fd4157cb04ae73c0c62b378046ca97154a25d4355e827be5abc53b
 SHA512:
-  metadata.gz: b37957502d147a221b3678c373b648090de7f9101fa1199e1721ffc1ac0e04d99318935866f12295a160c2f1d1d7648e3103355865eb1625b11619b9e52981d6
-  data.tar.gz: abdb670ea89993a92982c3913b192c6fed4ac6329beb2c0f1abb5aef39ef47993198ecdcdd5431ea057b84aa3931a532a96f5e85d0ab8a6c4faec48a98fc74de
+  metadata.gz: 1e54d7c41f8750b99262021e402f329252b428678ee99510dd52a1a113aeffdf35628a23b6ca9408ce55350a237cb2d224e66dfe7aafefd645612a29b74529ff
+  data.tar.gz: c8a014bb3708690cf3e463954aa7a846ce79075da039e23b3f94644e50e3c28d05376539699fd8ffe3c3b93d54eecb393fed472bddb60619ba897a5204bc511a

data/CHANGELOG.md

@@ -1,3 +1,11 @@
+## 0.2.1 (2023-06-04)
+
+- Added support for Sequel
+
+## 0.2.0 (2023-05-11)
+
+- Dropped support for Ruby < 3
+
 ## 0.1.1 (2022-02-08)
 
 - Added autoloading

data/LICENSE.txt

@@ -1,6 +1,6 @@
 The MIT License (MIT)
 
-Copyright (c) 2022 Andrew Kane
+Copyright (c) 2022-2023 Andrew Kane
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

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,6 +19,13 @@ 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
 
@@ -33,14 +40,53 @@ conn.type_map_for_results = PG::BasicTypeMapForResults.new(conn, registry: regis
 Insert a vector
 
 ```ruby
-factors = [1, 2, 3]
-conn.exec_params("INSERT INTO items (factors) VALUES ($1)", [factors])
+embedding = [1, 2, 3]
+conn.exec_params("INSERT INTO items (embedding) VALUES ($1)", [embedding])
+```
+
+Get the nearest neighbors to a vector
+
+```ruby
+conn.exec_params("SELECT * FROM items ORDER BY embedding <-> $1 LIMIT 5", [embedding]).to_a
+```
+
+## Sequel
+
+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
-conn.exec_params("SELECT * FROM items ORDER BY factors <-> $1 LIMIT 5", [factors]).to_a
+Item.nearest_neighbors(:embedding, [1, 1, 1], distance: "euclidean").limit(5)
 ```
 
 ## History

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.1.1"
+  VERSION = "0.2.1"
 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,77 @@
+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 ClassMethods
+        attr_accessor :vector_columns
+
+        def nearest_neighbors(column, value, distance:)
+          value = ::Pgvector.encode(value) unless value.is_a?(String)
+          quoted_column = dataset.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
+
+        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.1.1
+  version: 0.2.1
 platform: ruby
 authors:
 - Andrew Kane
 autorequire:
 bindir: bin
 cert_chain: []
-date: 2022-02-08 00:00:00.000000000 Z
+date: 2023-06-05 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
@@ -34,14 +35,14 @@ required_ruby_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
-      version: '2.6'
+      version: '3'
 required_rubygems_version: !ruby/object:Gem::Requirement
   requirements:
   - - ">="
     - !ruby/object:Gem::Version
       version: '0'
 requirements: []
-rubygems_version: 3.3.3
+rubygems_version: 3.4.10
 signing_key:
 specification_version: 4
 summary: pgvector support for Ruby