chromable 0.1.5 → 0.3.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9df743def1991fb00b56fdfb0878642a2e5568136d4078dcb77084cf47c6ce14
-  data.tar.gz: 6972ef16ef5c9f6f2f39de7e18c7028399d73b917001243c272e3b94d41839dc
+  metadata.gz: d33a705c79d42fb0495bb63e4ad7d04e8e3eea823400ea213f362af3e4ce09c8
+  data.tar.gz: d18bc615f1a6b22530d76b06b471e07688be422a9af55a92d81419ae48eb4d84
 SHA512:
-  metadata.gz: 07adc75fffe609a6e493f30091745093c3bb9164cd8e64c1f81519f8c447c4a6142075de97d461062903c28ac4d1a15d79bc501c57731e0d167a8562a696650b
-  data.tar.gz: 7818ee8d5a931f5e7936a0e27bb3bc716f349415e01f3cda33ff4cd72cf3ee237a5df157709a562238469659a828121e7d55b772b531bd06c1ddc2c7bc499204
+  metadata.gz: cfc65982c4f7952170630630eba0c6bf2bdc7708fae70727cb013d1cc1fe2b88b55e2bf08d9e0231864a0790373db2dfb1db95584c8ed328f0131c9e6a4c8d10
+  data.tar.gz: 6a0b4a7af862ada818d888ed6ae672a23d85057726df9da40f3c715cd1c80fa71ebfac10e00833dd0cf9908b678d3746d3ead9d9cc6f73fc4ae23436a1e3a13e

data/.rubocop.yml

@@ -2,7 +2,7 @@ require:
   - rubocop-rspec
 
 AllCops:
-  TargetRubyVersion: 2.6
+  TargetRubyVersion: 3.0
   NewCops: enable
 
 Layout/LineLength:

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,33 +32,54 @@ 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
 
-  def embed
-    # Call OpenAI now please :)
+    if options[:is_query]
+      # Call OpenAI API to embed `text` as a search query.
+    else
+      # Call OpenAI API to embed `text` as a post content.
+    end
   end
 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 returns the embedding representation for the current instance.
+- `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.
 
-All `chromable` method arguments are optional.
+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.
+
+To interact with the ChromaDB collection, `chromable` provides `Model.query` method to query the collection and `Model.collection` method to access the collection directly.
+
+```ruby
+puts Post.collection.count # Gets the number of documents inside the collection. Should always match Post.count.
+
+Post.query(
+  query: params[:query],
+  results: 20,
+  where: chroma_search_filters,
+  is_query: true # `is_query` here will be passed to `Post.embed` as an option.
+)
+```
 
-At this point, `chromable` will create, update, and destroy the ChromaDB embeddings for your objects based on Rails callbacks.
+`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.
 
-To interact with the ChromaDB collection, `chromable` provides `Model.collection` method to retrieve the collection instance.
 Also, `chromable` provides the following methods for each model instance:
 
 - `embedding`: Retrieves the instance's ChromaDB embedding object.
 - `upsert_embedding`: Creates or updates the instance's ChromaDB embedding object.
 - `destroy_embedding`: Destroys the instance's ChromaDB embedding object.
 
-All these methods (including `Model.collection`) are available with `chroma_` prefix, if you have similar methods defined in your model.
+All these methods (including `Model.query` and `Model.collection`) are available with `chroma_` prefix, if you have similar methods defined in your model.
 
 ## Development
 

data/lib/chromable/version.rb

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

data/lib/chromable.rb

@@ -4,52 +4,116 @@ require_relative 'chromable/version'
 
 # Ruby on Rails integration for ChromaDB.
 module Chromable
-  def self.included(base) 
+  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(collection_name: nil, document: nil, metadata: nil, embedder: 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
 
-    alias_method :collection, :chroma_collection unless method_defined? :collection
-  end
+    def chroma_delete_collection
+      Chroma::Resources::Collection.delete(chromable_settings.collection_name)
+    end
 
-  def chroma_embedding
-    self.class.chroma_collection.get(ids: [id])[0]
+    def chroma_query( # rubocop:disable Metrics/ParameterLists
+      text:,
+      results: 10,
+      where: {},
+      where_document: {},
+      include: %w[metadatas documents distances],
+      **embedder_options
+    )
+      find(chroma_collection.query(
+        query_embeddings: [send(chromable_settings.embedder, text, **embedder_options)],
+        results: results,
+        where: where,
+        where_document: where_document,
+        include: include
+      ).map(&:id))
+    end
   end
 
-  def chroma_upsert_embedding
-    self.class.chroma_collection.upsert(
+  # 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_embedding
+      self.class.chroma_collection.get(ids: [id])[0]
+    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
+
+    private
+
+    def build_embedding
       Chroma::Resources::Embedding.new(
         id: id,
-        document: self.class.document ? send(self.class.document) : nil,
-        embedding: self.class.embedder ? send(self.class.embedder) : nil,
-        metadata: self.class.metadata ? self.class.metadata.index_with { |attribute| send(attribute) } : nil
+        document: document_to_embed,
+        embedding: document_embedding,
+        metadata: embedding_metadata
       )
-    )
-  end
+    end
 
-  def chroma_destroy_embedding
-    self.class.chroma_collection.delete(ids: [id])
-  end
+    def document_to_embed
+      chromable_settings.keep_document ? send(chromable_settings.document) : nil
+    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
+    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.1.5
+  version: 0.3.0
 platform: ruby
 authors:
 - Ali Hamdi Ali Fadel
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2023-12-09 00:00:00.000000000 Z
+date: 2023-12-12 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: chroma-db