Support embedding generation

CONTRIBUTING.md

@@ -68,10 +68,11 @@ bundle exec rake release
 
 ### In Progress
 
-* Support creating inline embeddings
+* Support creating inline text embeddings
 
 ### Future
 
+* Support creating inline OpenAI embeddings (requires token)
 * Support using text to perform vector search with inline embeddings
 * Support creating vector indexes
 * Support creating embedding jobs

README.md

@@ -7,7 +7,7 @@ No Ruby client is required for `pg` or `Sequel`. For `ActiveRecord` and `Rails`,
 ## Features
 
 - Perform nearest neighbor queries over vectors
-- Create text embeddings using various models
+- Create text embeddings using open-source models
 
 ## Installation
 
@@ -31,8 +31,42 @@ gem install lantern
 
 ## ActiveRecord
 
+### Vector search
+
+The Lantern gem integrates with ActiveRecord to provide vector similarity search capabilities:
+
+```ruby
+class Document < ApplicationRecord
+  # Enable vector similarity search for one or more columns
+  has_neighbors :embedding
+end
+
+# Class-level nearest-neighbor search: find nearest neighbors for a given vector
+Document.nearest_neighbors(:embedding, vector, distance: 'l2')
+
+# Instance-level nearest-neighbor search: find nearest neighbors for a given instance
+document = Document.first
+document.nearest_neighbors(:embedding, distance: 'l2')
+```
+
+Supported distance metrics:
+
+- `l2` (Euclidean distance)
+- `cosine` (Cosine similarity)
+
+### Embedding generation
+
+```ruby
+# Generate embeddings using a specific model
+embedding = Lantern.generate_embedding('BAAI/bge-base-en', 'Your text here')
+```
+
+A full list of supported models can be found [here](lantern.dev/docs/develop/generate).
+
 ## Rails
 
+For Rails, enable the Lantern extension using the provided generator:
+
 ```bash
 rails generate lantern:install
 rails db:migrate

lib/lantern.rb

@@ -1,9 +1,18 @@
 require_relative 'lantern/version'
 require_relative 'lantern/model'
+require_relative 'lantern/embeddings'
 require 'active_support'
 
 module Lantern
   class Error < StandardError; end
+
+  extend Lantern::Embeddings::ClassMethods
+
+  class << self
+    def connection
+      ActiveRecord::Base.connection
+    end
+  end
 end
 
 ActiveSupport.on_load(:active_record) do

lib/lantern/embeddings.rb

@@ -0,0 +1,23 @@
+require 'active_support'
+
+module Lantern
+  module Embeddings
+    extend ActiveSupport::Concern
+
+    class_methods do
+      # Generates a text embedding using the specified model
+      #
+      # @param model [String] The embedding model to use (e.g., 'BAAI/bge-base-en')
+      # @param text [String] The text input to embed
+      # @return [Array<Float>] The generated embedding vector
+      def generate_embedding(model, text)
+        sanitized_model = connection.quote(model)
+        sanitized_text = connection.quote(text)
+        result = connection.select_one("SELECT text_embedding(#{sanitized_model}, #{sanitized_text}) AS embedding")
+        embedding = result['embedding'].tr('{}', '').split(',').map(&:to_f)
+        embedding
+      end
+    end
+  end
+end
+

test/embeddings_test.rb

@@ -0,0 +1,8 @@
+require 'test_helper'
+
+class LanternEmbeddingsTest < Minitest::Test
+  def test_generate_embedding
+    embedding = Lantern.generate_embedding('BAAI/bge-small-en', 'My text input')
+    assert_equal [-0.42607692,-0.6560961,0.24540034,-0.5687317,-0.48004436,-0.3445696,0.15631531,0.44238535,0.6822042,0.11165306,-0.27675378,-0.54439175,0.3238239,-0.018825464,0.15022391,0.16182774,0.0931536,-0.4021465,-0.599053,-0.08376932,0.4803496,0.30632803,-0.24684542,-0.2111539,-0.37804127,0.38152796,-0.3742156,0.021979779,-0.6139418,-0.9352158,0.08769297,0.10611001,0.91340333,0.10319206,-0.43541518,-0.035220027,-0.5422052,-0.13951029,-0.08785111,-0.18059957,0.29017866,-0.52898204,0.08826429,-0.28320208,0.32065412,-0.74469453,0.019126728,-0.1331417,1.1915426,-0.0271228,-0.1791501,0.5045242,0.47735226,0.26709956,0.27072954,0.3015376,0.2857157,0.81967425,0.29811752,0.24612932,0.2402717,0.7456093,-1.9080384,1.1608,-0.24770442,0.30312783,-0.23141631,0.022359706,0.079715215,0.16744456,0.00048166513,0.13478571,-0.32114238,0.9849128,0.21819264,-0.35998306,0.016470224,0.19729656,0.044306144,0.46361348,0.17170845,-0.2316737,-0.310745,0.07062367,0.33972862,-0.10763518,-0.11695275,0.19037059,0.01202409,-0.26588318,-0.2108987,-0.70220965,0.0990309,0.079637624,-0.48364216,-0.021402832,0.3567547,0.08315052,-0.65410924,2.8314016,-0.4592521,0.04489682,0.013682842,-0.49335244,-0.04366027,0.32255328,-0.18192077,-0.156384,-0.6502448,-0.26909065,-0.11918783,-0.42729855,-0.007463634,0.55170995,-0.0718565,0.80208516,0.26283142,0.20096505,0.08236541,-0.28746638,-0.12979397,0.17897911,-0.4580958,0.35778597,0.5124185,-0.5521189,0.7105305,0.95287323,-0.09035744,0.12799554,0.62802976,0.23259547,-0.6584735,-0.24235135,-0.42022616,0.09264918,-0.06097544,-0.3303358,-0.05964969,-0.31162965,-0.62518376,-0.86724585,-0.47793478,0.034938917,-0.15466754,0.6372119,-0.3076179,-0.5150481,0.1450295,-0.45128334,-0.11074954,0.52214897,0.17979234,0.15051602,-0.04626303,0.47930863,0.47703665,-0.2688008,0.27251148,-0.082838476,-0.19751151,0.06756437,-0.15275787,0.08278718,0.30210978,-0.4610209,-0.47056708,-0.11175291,-0.15225048,-0.63996863,0.40118033,0.2311216,-0.38426283,0.27423102,0.91351,0.011803292,-0.11026738,0.11364358,0.30451474,0.28938252,0.36859804,-0.054251574,-0.0122629255,0.29449856,0.4099234,-0.009324573,-0.30096447,-0.9003356,0.0028337762,0.20014527,0.20618732,0.1750438,-0.073773384,0.051853806,-0.83459204,-0.4512414,0.4837396,-0.69534075,-0.038540907,-0.61638665,1.1021618,0.11100495,-0.060204174,-0.18590793,0.26962733,-0.097282946,0.015973601,-0.16622287,0.6560011,-0.37308952,-0.26591653,-0.37483108,1.3029053,-0.47876132,-0.13348062,-0.57279104,-0.091271445,0.30151683,0.10556768,0.1626319,-0.20271748,-0.01064305,-0.39189273,-2.5214515,-0.041893646,0.37145826,-0.34317538,-0.35858908,0.039526075,0.33622658,0.046334162,0.46628746,0.65926677,0.6439412,-0.15697318,-0.069618955,-0.3143116,-0.0854539,-0.021223515,-0.040328622,-0.29221457,0.18152742,0.1023116,-0.48913223,-0.24974301,0.4580552,-0.46850812,0.34635943,-0.10458655,1.4959769,0.732024,0.35794735,-0.13347305,0.680652,0.35506997,-0.21085936,-0.90287113,0.35918146,0.021746077,-0.37929574,0.41729915,-0.39861178,-0.5498347,-0.4037995,0.005103618,-0.032325573,-0.34511617,0.04418596,-0.54290605,-0.35228255,-0.8101264,-0.81166774,0.31071082,0.31730983,0.02845101,0.20515004,0.9182312,0.41709322,-0.16966859,-0.22877869,0.3171808,0.0007266179,-0.2706459,0.107099876,0.24273989,-0.2550255,-0.8434111,0.20187429,0.17889066,-0.15518267,-0.27918166,0.30582142,0.35796964,-0.13066071,-0.32303825,0.41789207,0.49520302,0.36858648,0.16270527,-0.15524365,0.50438327,0.1784228,-0.43594018,-0.11294274,0.24804524,0.51854146,0.3148494,0.028747043,0.19159088,0.53050834,-0.17503415,0.9501747,-0.49234536,-0.06733905,0.22630046,-0.07570651,-0.44593415,0.04914434,-0.51262087,-2.0349698,0.5663595,0.11317264,0.379006,-0.1781563,0.3739446,-0.08112174,-0.3299852,-0.7314233,0.16986695,-0.64576036,0.07652925,-0.53505784,-0.17598443,-0.16785783,-0.17894292,0.064481236,-0.45796177,-0.021510556,-0.54930174,0.4595049,0.27867806,1.6718408,0.025401779,0.24337678,0.32296145,-0.26478213,0.28458935,0.9804398,0.17695138,0.17754546,-0.0919306,0.98911256,0.04851158,0.205428,-0.057857975,-0.2255534,0.153895,-0.18901107,0.093918934,-0.32115656,0.6039499,-0.9937815,-0.37324917,0.4359285,0.0066044778,0.12033479,-0.14657457,-0.5507795,0.025807166,-0.19052581,-0.024470005,-0.2198321,0.44893804,0.29699105,0.70613104,-0.16971812,-0.15046224,-0.007557243,-0.11110652,-0.37352803,-0.34645554,1.0155857,-0.044960663,-0.064731136], embedding
+  end
+end

test/test_helper.rb

@@ -10,5 +10,4 @@
 require "dotenv/load"
 
 ActiveRecord::Base.establish_connection(ENV["DATABASE_URL"])
-
 ActiveRecord::Base.connection.enable_extension("lantern")