chromable 0.2.0 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d44fa49957eefd26d00d3fb2315b51f928bd6ecdac7cdc4f849537cac23f036b
-  data.tar.gz: ab51d49d9b558e6d6ba936640026af66258567c0750f2ce11c7b6e01a0b8fea1
+  metadata.gz: d33a705c79d42fb0495bb63e4ad7d04e8e3eea823400ea213f362af3e4ce09c8
+  data.tar.gz: d18bc615f1a6b22530d76b06b471e07688be422a9af55a92d81419ae48eb4d84
 SHA512:
-  metadata.gz: 2a6be5bd9ec59535eb33a59fc543baa23954ac43c4bc69b8a6e6d19eaab231d313180ba714781f85342f838bfdd4fd70fcfa32c785f663b880128d6a9425a69e
-  data.tar.gz: f63ffea813494e3143c16ed900e190d52866fa11348bac39db894c09ad885e66115b0cb8d0e492f2fb18846b29d8fcb76c4dbdbf3c271e3fd4f6afb417350ee9
+  metadata.gz: cfc65982c4f7952170630630eba0c6bf2bdc7708fae70727cb013d1cc1fe2b88b55e2bf08d9e0231864a0790373db2dfb1db95584c8ed328f0131c9e6a4c8d10
+  data.tar.gz: 6a0b4a7af862ada818d888ed6ae672a23d85057726df9da40f3c715cd1c80fa71ebfac10e00833dd0cf9908b678d3746d3ead9d9cc6f73fc4ae23436a1e3a13e

data/README.md

@@ -2,6 +2,8 @@
 
 Ruby on Rails integration for ChromaDB based on `chroma-db` gem.
 
+`chromable` were tested with Ruby 3.2.2 and Rails 7.1.2.
+
 ## Installation
 
 Install `chromable` and add it to the application's Gemfile by executing:
@@ -30,7 +32,7 @@ Then, include `Chromable` module in your model and initialize it:
 class Post < ApplicationRecord
   include Chromable
 
-  chromable document: :content, metadata: %i[author category], embedder: :embed
+  chromable document: :content, metadata: %i[author category], embedder: :embed, keep_document: false
 
   def self.embed(text, **options)
     options[:is_query] ||= false
@@ -45,13 +47,14 @@ end
 ```
 
 Where:
-- `document:` is a callable represents the text content you want to embed and store in ChromaDB.
-- `metadata:` is the list of attributes to be passed to ChromaDB as metadata to be used to filter.
-- `embedder:` is a callable defined in the model that returns the embedding representation for the given `text` and `options`.
+- `document:` is a callable represents the text content you want to embed and store in ChromaDB (e.g. Could be a model attribute).
+- `metadata:` is a list of callables to be evaluated and passed to ChromaDB as metadata to be used to filter (e.g. Could be an instance method).
+- `embedder:` is a callable defined at the model level that returns the embedding representation for the given `text` based on some `options`.
+- `keep_document:` tells chromable to pass the `document:` to ChromaDB and save it or not. It is useful if you just want to have the embeddings in ChromaDB and the rest of the data in your Rails application database to reduce memory footprint. `keep_document:` is `true` by default.
 
 Optionaly you can pass `collection_name:`. If not passed, the plural form of the model name will be used.
 
-The only required option for `chromable` method is `document:`.
+The only required option for `chromable` is `document:`.
 
 At this point, `chromable` will create, update, and destroy the ChromaDB embeddings for your objects based on Rails `after_save` and `after_destroy` callbacks.
 
@@ -64,11 +67,11 @@ Post.query(
   query: params[:query],
   results: 20,
   where: chroma_search_filters,
-  type: 'query' # `type` here will be passed to `Post.embed` as an option.
+  is_query: true # `is_query` here will be passed to `Post.embed` as an option.
 )
 ```
 
-`Model.query` accepts the same arguments accepted by `chroma-db` gem `query` method. Extra arguments will be passed to the `embedder:`. Behind the scene, `Model.query` will embed the given `query:` text, then query the collection, and return the closest `results:` records.
+`Model.query` accepts the same arguments accepted by `chroma-db` gem `query` method. Extra arguments will be passed to the `embedder:` as `options`. Behind the scenes, `Model.query` will embed the given `query:` text, then query the collection, and return the closest `results:` records.
 
 Also, `chromable` provides the following methods for each model instance:
 

data/lib/chromable/version.rb

@@ -1,5 +1,5 @@
 # frozen_string_literal: true
 
 module Chromable
-  VERSION = '0.2.0'
+  VERSION = '0.3.0'
 end

data/lib/chromable.rb

@@ -6,30 +6,53 @@ require_relative 'chromable/version'
 module Chromable
   def self.included(base)
     base.extend ClassMethods
-    base.class_attribute :collection_name
-    base.class_attribute :document
-    base.class_attribute :metadata
-    base.class_attribute :embedder
+    base.include InstanceMethods
 
     base.after_save :chroma_upsert_embedding
     base.after_destroy :chroma_destroy_embedding
   end
 
+  # Chromable settings class to hide them from Rails models.
+  class Settings
+    attr_accessor :document, :collection_name, :metadata, :embedder, :keep_document
+
+    def initialize(document:, metadata: nil, embedder: nil, collection_name: nil, keep_document: true)
+      @collection_name = collection_name
+      @document = document
+      @metadata = metadata
+      @embedder = embedder
+      @keep_document = keep_document
+    end
+  end
+
   # Methods to be added to the model class.
   module ClassMethods
-    def chromable(document:, metadata: nil, embedder: nil, collection_name: nil)
-      self.collection_name = (collection_name.presence || name.underscore.pluralize)
-      self.document = document
-      self.metadata = metadata
-      self.embedder = embedder
+    def self.extended(base)
+      class << base
+        alias_method :collection, :chroma_collection unless method_defined? :collection
+        alias_method :delete_collection, :chroma_delete_collection unless method_defined? :delete_collection
+        alias_method :query, :chroma_query unless method_defined? :query
+      end
+
+      base.cattr_accessor :chromable_settings
+    end
+
+    def chromable(**options)
+      options[:collection_name] ||= name.underscore.pluralize
+
+      self.chromable_settings = Settings.new(**options)
     end
 
     def chroma_collection
-      Chroma::Resources::Collection.get_or_create(collection_name)
+      Chroma::Resources::Collection.get_or_create(chromable_settings.collection_name)
+    end
+
+    def chroma_delete_collection
+      Chroma::Resources::Collection.delete(chromable_settings.collection_name)
     end
 
     def chroma_query( # rubocop:disable Metrics/ParameterLists
-      query:,
+      text:,
       results: 10,
       where: {},
       where_document: {},
@@ -37,44 +60,60 @@ module Chromable
       **embedder_options
     )
       find(chroma_collection.query(
-        query_embeddings: [send(embedder, query, **embedder_options)],
+        query_embeddings: [send(chromable_settings.embedder, text, **embedder_options)],
         results: results,
         where: where,
         where_document: where_document,
         include: include
       ).map(&:id))
     end
-
-    alias collection chroma_collection unless method_defined? :collection
-    alias query chroma_query unless method_defined? :query
   end
 
-  def chroma_embedding
-    self.class.chroma_collection.get(ids: [id])[0]
-  end
+  # Methods to be added to the model instances.
+  module InstanceMethods
+    def self.included(base)
+      base.instance_eval do
+        # rubocop:disable Style/Alias
+        alias_method :embedding, :chroma_embedding unless method_defined? :embedding
+        alias_method :upsert_embedding, :chroma_upsert_embedding unless method_defined? :upsert_embedding
+        alias_method :destroy_embedding, :chroma_destroy_embedding unless method_defined? :destroy_embedding
+        # rubocop:enable Style/Alias
+      end
+    end
 
-  def chroma_upsert_embedding
-    self.class.chroma_collection.upsert(build_embedding)
-  end
+    def chroma_embedding
+      self.class.chroma_collection.get(ids: [id])[0]
+    end
 
-  def chroma_destroy_embedding
-    self.class.chroma_collection.delete(ids: [id])
-  end
+    def chroma_upsert_embedding
+      self.class.chroma_collection.upsert(build_embedding)
+    end
+
+    def chroma_destroy_embedding
+      self.class.chroma_collection.delete(ids: [id])
+    end
 
-  alias embedding chroma_embedding unless method_defined? :embedding
-  alias upsert_embedding chroma_upsert_embedding unless method_defined? :upsert_embedding
-  alias destroy_embedding chroma_destroy_embedding unless method_defined? :destroy_embedding
+    private
 
-  private
+    def build_embedding
+      Chroma::Resources::Embedding.new(
+        id: id,
+        document: document_to_embed,
+        embedding: document_embedding,
+        metadata: embedding_metadata
+      )
+    end
 
-  def build_embedding
-    document = send(self.class.document)
+    def document_to_embed
+      chromable_settings.keep_document ? send(chromable_settings.document) : nil
+    end
 
-    Chroma::Resources::Embedding.new(
-      id: id,
-      document: document,
-      embedding: self.class.embedder && self.class.send(self.class.embedder, document),
-      metadata: self.class.metadata&.index_with { |attribute| send(attribute) }
-    )
+    def document_embedding
+      chromable_settings.embedder && self.class.send(chromable_settings.embedder, send(chromable_settings.document))
+    end
+
+    def embedding_metadata
+      chromable_settings.metadata&.index_with { |attribute| send(attribute) }
+    end
   end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: chromable
 version: !ruby/object:Gem::Version
-  version: 0.2.0
+  version: 0.3.0
 platform: ruby
 authors:
 - Ali Hamdi Ali Fadel
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2023-12-11 00:00:00.000000000 Z
+date: 2023-12-12 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: chroma-db