lanterndb 0.0.0

checksums.yaml

@@ -0,0 +1,7 @@
+---
+SHA256:
+  metadata.gz: 33b108e60126444cecd60d7b8bb093edd3a260993842974d138e70628622dd89
+  data.tar.gz: 554d9f355237b15b6d32ba694753369891cec7e721cd99d330330dd1d1e3e176
+SHA512:
+  metadata.gz: f58f15546955bfd16c01a665c9da82106bd417accca3117e8eba2a2e318a713297d07c28a5a1a935a71c2b1c6c3d9508cb49ef1a64fba89f1dc302b71abe23f0
+  data.tar.gz: 9224859157287a5180bed90f7bc9bca8216b021d2129f7fe93c3b8465b48554222d73dfa060449296a4a6b62a64b79db755197d37c62c39c0ab764754172218a

data/LICENSE.txt

@@ -0,0 +1,21 @@
+The MIT License (MIT)
+
+Copyright (c) 2024 Lantern Systems, Inc.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.

data/README.md

@@ -0,0 +1,150 @@
+# Lantern Ruby Client
+
+[![codecov](https://codecov.io/gh/lanterndata/lantern-ruby/branch/main/graph/badge.svg)](https://codecov.io/gh/lanterndata/lantern-ruby)
+
+No Ruby client is required for `pg` or `Sequel`. For `ActiveRecord` and `Rails`, you can use the `lantern` gem.
+
+## Features
+
+- Perform nearest neighbor queries over vectors using vectors or text
+- Create text embeddings using OpenAI, Cophere, and open-source models
+
+## Installation
+
+Add this line to your application's Gemfile:
+
+```ruby
+gem 'lanterndb'
+```
+
+And then execute:
+
+```bash
+bundle install
+```
+
+Or install it yourself as:
+
+```bash
+gem install lanterndb
+```
+
+## ActiveRecord
+
+### Connect to the database
+
+```ruby
+require 'active_record'
+require 'lantern'
+ActiveRecord::Base.establish_connection("DATABASE_URL")
+ActiveRecord::Base.connection.enable_extension("lantern")
+conn = ActiveRecord::Base.connection
+```
+
+### Create a model
+
+```ruby
+ActiveRecord::Migration.create_table :movies do |t|
+  t.column :movie_embedding, :real, array: true
+end
+conn.execute("INSERT INTO movies (movie_embedding) VALUES ('{0,1,0}'), ('{3,2,4}')")
+```
+
+### Embedding generation
+
+```ruby
+embedding1 = Lantern.text_embedding('BAAI/bge-base-en', 'Your text here')
+
+Lantern.set_api_token(openai_token: 'your_openai_token')
+embedding2 = Lantern.openai_embedding('text-embedding-3-small', 'Hello')
+
+Lantern.set_api_token(cohere_token: 'your_cohere_token')
+embedding3 = Lantern.cohere_embedding('embed-english-v3.0', 'Hello')
+```
+
+A full list of supported models can be found [here](lantern.dev/docs/develop/generate).
+
+### Vector search
+
+This gem provides several ways to perform vector search. We support the following distance metrics:
+
+- `l2` (Euclidean distance)
+- `cosine` (Cosine similarity)
+
+Using pre-computed vectors:
+
+```ruby
+class Document < ApplicationRecord
+  has_neighbors :embedding
+end
+
+# Find 5 nearest neighbors using L2 distance 
+Document.nearest_neighbors(:embedding, [0.1, 0.2, 0.3], distance: 'l2').limit(5)
+
+# Given a document, find 5 nearest neighbors using cosine distance
+document = Document.first
+document.nearest_neighbors(:embedding, distance: 'cosine').limit(5)
+```
+
+Using text:
+
+```ruby
+class Book < ApplicationRecord
+  has_neighbors :embedding
+end
+
+# Find 5 nearest neighbors using open-source model
+Book.nearest_neighbors(:embedding, 'The quick brown fox', model: 'BAAI/bge-small-en', distance: 'l2').limit(5)
+
+# Find 5 nearest neighbors using OpenAI
+Lantern.set_api_token(openai_token: 'your_openai_token')
+Book.nearest_neighbors(:embedding, 'The quick brown fox', model: 'openai/text-embedding-3-small', distance: 'cosine').limit(5)
+```
+
+### Vector index
+
+To speed up vector search queries, you can add an HNSW index to your model:
+
+```ruby
+class CreateVectorIndex < ActiveRecord::Migration[7.0]
+  def up
+    add_index :books, :embedding, using: :lantern_hnsw, opclass: :dist_l2sq_ops, name: 'book_embedding_index'
+  end
+  def down
+    remove_index :books, name: 'book_embedding_index'
+  end
+end
+```
+
+Note: This does not support `WITH` parameters (e.g., `ef_construction`, `ef`, `m`, `dim`).
+
+To specify `WITH` parameters, you can pass them as options with raw SQL:
+
+```ruby
+class CreateHnswIndex < ActiveRecord::Migration[7.0]
+  def up
+    execute <<-SQL
+      CREATE INDEX movie_embedding_hnsw_idx 
+      ON movies 
+      USING lantern_hnsw (movie_embedding dist_l2sq_ops) 
+      WITH (
+        ef = 15,
+        m = 16,
+        ef_construction = 64
+      )
+    SQL
+  end
+  def down
+    drop index movie_embedding_hnsw_idx
+  end
+end
+```
+
+## Rails
+
+For Rails, enable the Lantern extension using the provided generator:
+
+```bash
+rails generate lantern:install
+rails db:migrate
+```

data/lib/generators/lantern/lantern_generator.rb

@@ -0,0 +1,21 @@
+require 'rails/generators'
+require 'rails/generators/active_record'
+
+module Lantern
+  module Generators
+    class LanternGenerator < Rails::Generators::Base
+      include ActiveRecord::Generators::Migration
+      source_root File.join(__dir__, 'templates')
+
+      def copy_migration
+        migration_template 'lantern.rb.tt', 'db/migrate/install_lantern.rb', migration_version: migration_version
+      end
+
+      private
+
+      def migration_version
+        "[#{ActiveRecord::VERSION::MAJOR}.#{ActiveRecord::VERSION::MINOR}]"
+      end
+    end
+  end
+end

data/lib/generators/lantern/templates/lantern.rb.tt

@@ -0,0 +1,5 @@
+class <%= migration_class_name %> < ActiveRecord::Migration<%= migration_version %>
+  def change
+    enable_extension "lantern"
+  end
+end

data/lib/lantern/embeddings.rb

@@ -0,0 +1,42 @@
+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 text_embedding(model, text)
+        generate_embedding(model, text, 'text_embedding')
+      end
+      def openai_embedding(model, text, dim = nil)
+        Lantern.ensure_token!(:openai)
+        generate_embedding(model, text, 'openai_embedding', dim)
+      end
+      def cohere_embedding(model, text, input_type = nil)
+        Lantern.ensure_token!(:cohere)
+        generate_embedding(model, text, 'cohere_embedding', input_type)
+      end
+
+      private
+
+      def generate_embedding(model, text, embedding_function, other = nil)
+        sanitized_model = connection.quote(model)
+        sanitized_text = connection.quote(text)
+        if other
+          query = "SELECT #{embedding_function}(#{sanitized_model}, #{sanitized_text}, #{other}) AS embedding"
+        else
+          query = "SELECT #{embedding_function}(#{sanitized_model}, #{sanitized_text}) AS embedding"
+        end
+        result = connection.select_one(query)
+        embedding = result['embedding'].tr('{}', '').split(',').map(&:to_f)
+        embedding
+      end
+    end
+  end
+end
+  

data/lib/lantern/model.rb

@@ -0,0 +1,60 @@
+module Lantern
+  module Model
+    def has_neighbors(*attribute_names)
+      attribute_names.map!(&:to_sym)
+
+      scope :nearest_neighbors, ->(attribute_name, vector_or_text, model: nil, openai_dim: nil, cohere_input_type: nil, distance: 'l2') {
+        attribute_name = attribute_name.to_sym
+
+        # Distance operator
+        distance_operators = {
+          'l2' => '<->',
+          'cosine' => '<=>'
+        }
+        operator = distance_operators[distance]
+        unless operator
+          raise ArgumentError, "Invalid distance metric. Allowed metrics are #{distance_operators.keys.join(', ')}"
+        end
+
+        # Vector order by expression
+        order_expression = "#{quoted_table_name}.#{connection.quote_column_name(attribute_name)} #{operator}" + if model
+          # Generate vector from text
+          text = vector_or_text
+          embedding_function = case model
+                               when /^openai/ then 'openai_embedding'
+                               when /^cohere/ then 'cohere_embedding'
+                               else 'text_embedding'
+                               end
+          other = model.start_with?('openai') ? openai_dim : model.start_with?('cohere') ? cohere_input_type : nil
+          sanitized_model = connection.quote(model)
+          sanitized_text = connection.quote(text)
+          if other
+            "#{embedding_function}(#{sanitized_model}, #{sanitized_text}, #{other})"
+          else
+            "#{embedding_function}(#{sanitized_model}, #{sanitized_text})"
+          end
+        else
+          # Vector-based search
+          vector = vector_or_text
+          unless vector.all? { |v| v.is_a?(Integer) }
+            vector = vector.map(&:to_f)
+          end
+          vector_literal = vector.join(',')
+          "ARRAY[#{vector_literal}]"
+        end
+
+        select_columns = select_values.any? ? [] : column_names
+        select(select_columns, "#{order_expression} AS distance")
+          .where.not(attribute_name => nil)
+          .order(Arel.sql(order_expression))
+      }
+
+      define_method :nearest_neighbors do |attribute_name, distance: 'l2'|
+        attribute_name = attribute_name.to_sym
+        self.class
+          .where.not(id: id)
+          .nearest_neighbors(attribute_name, self[attribute_name], distance: distance)
+      end
+    end
+  end
+end

data/lib/lantern/railtie.rb

@@ -0,0 +1,11 @@
+require 'rails/railtie'
+
+module Lantern
+  class Railtie < Rails::Railtie
+    railtie_name :lantern
+
+    generators do
+      require 'generators/lantern/lantern_generator'
+    end
+  end
+end

data/lib/lantern/version.rb

@@ -0,0 +1,3 @@
+module Lantern
+  VERSION = '0.0.0'
+end

data/lib/lantern.rb

@@ -0,0 +1,54 @@
+require_relative 'lantern/version'
+require_relative 'lantern/model'
+require_relative 'lantern/embeddings'
+require 'active_support'
+
+module Lantern
+  class Error < StandardError; end
+  class MissingTokenError < StandardError; end
+
+  extend Lantern::Embeddings::ClassMethods
+
+  class << self
+    def connection
+      ActiveRecord::Base.connection
+    end
+
+    def current_user
+      connection.select_value("SELECT CURRENT_USER")
+    end
+
+    def set_api_token(openai_token: nil, cohere_token: nil)
+      if openai_token
+        connection.execute("ALTER ROLE #{current_user} SET lantern_extras.openai_token = #{connection.quote(openai_token)};")
+      end
+      if cohere_token
+        connection.execute("ALTER ROLE #{current_user} SET lantern_extras.cohere_token = #{connection.quote(cohere_token)};")
+      end
+      connection.execute("SELECT pg_reload_conf();")
+    end
+
+    def openai_token
+      connection.select_value("SHOW lantern_extras.openai_token")
+    end
+
+    def cohere_token
+      connection.select_value("SHOW lantern_extras.cohere_token")
+    end
+
+    def ensure_token!(provider)
+      case provider
+      when :openai
+        raise MissingTokenError, "OpenAI token is required for OpenAI embedding generation" if openai_token.blank?
+      when :cohere
+        raise MissingTokenError, "Cohere token is required for Cohere embedding generation" if cohere_token.blank?
+      end
+    end
+  end
+end
+
+ActiveSupport.on_load(:active_record) do
+  extend Lantern::Model
+end
+
+require_relative 'lantern/railtie' if defined?(Rails::Railtie)

metadata

@@ -0,0 +1,196 @@
+--- !ruby/object:Gem::Specification
+name: lanterndb
+version: !ruby/object:Gem::Version
+  version: 0.0.0
+platform: ruby
+authors:
+- Di Qi
+autorequire:
+bindir: exe
+cert_chain: []
+date: 2024-11-10 00:00:00.000000000 Z
+dependencies:
+- !ruby/object:Gem::Dependency
+  name: activerecord
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '7.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '7.0'
+- !ruby/object:Gem::Dependency
+  name: activesupport
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '7.0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '7.0'
+- !ruby/object:Gem::Dependency
+  name: pg
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.2'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.2'
+- !ruby/object:Gem::Dependency
+  name: dotenv
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '2.7'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '2.7'
+- !ruby/object:Gem::Dependency
+  name: minitest
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '5.14'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '5.14'
+- !ruby/object:Gem::Dependency
+  name: rake
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '13.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '13.0'
+- !ruby/object:Gem::Dependency
+  name: rubocop
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '1.0'
+- !ruby/object:Gem::Dependency
+  name: railties
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '7.0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '7.0'
+- !ruby/object:Gem::Dependency
+  name: simplecov
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+- !ruby/object:Gem::Dependency
+  name: simplecov-cobertura
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :development
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+description:
+email: support@lantern.dev
+executables: []
+extensions: []
+extra_rdoc_files: []
+files:
+- LICENSE.txt
+- README.md
+- lib/generators/lantern/lantern_generator.rb
+- lib/generators/lantern/templates/lantern.rb.tt
+- lib/lantern.rb
+- lib/lantern/embeddings.rb
+- lib/lantern/model.rb
+- lib/lantern/railtie.rb
+- lib/lantern/version.rb
+homepage: https://lantern.dev
+licenses:
+- MIT
+metadata:
+  homepage_uri: https://lantern.dev
+  source_code_uri: https://github.com/lanterndata/lantern-ruby
+  bug_tracker_uri: https://github.com/lanterndata/lantern-ruby/issues
+  documentation_uri: https://lantern.dev/docs
+  changelog_uri: https://github.com/lanterndata/lantern-ruby/blob/main/CHANGELOG.md
+post_install_message:
+rdoc_options: []
+require_paths:
+- lib
+required_ruby_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '3.1'
+required_rubygems_version: !ruby/object:Gem::Requirement
+  requirements:
+  - - ">="
+    - !ruby/object:Gem::Version
+      version: '0'
+requirements: []
+rubygems_version: 3.5.23
+signing_key:
+specification_version: 4
+summary: Lantern Rails Client
+test_files: []