chromable 0.2.0 → 0.3.1

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: d44fa49957eefd26d00d3fb2315b51f928bd6ecdac7cdc4f849537cac23f036b
-  data.tar.gz: ab51d49d9b558e6d6ba936640026af66258567c0750f2ce11c7b6e01a0b8fea1
+  metadata.gz: ad740d4c1d70642aa0e03518c74cb5e2b77a68df923a9e46e0797eb42908c312
+  data.tar.gz: 4cdb45e1599c5f579082cac30ff9c3887adaee9921d2bf73f20b6f96401769cd
 SHA512:
-  metadata.gz: 2a6be5bd9ec59535eb33a59fc543baa23954ac43c4bc69b8a6e6d19eaab231d313180ba714781f85342f838bfdd4fd70fcfa32c785f663b880128d6a9425a69e
-  data.tar.gz: f63ffea813494e3143c16ed900e190d52866fa11348bac39db894c09ad885e66115b0cb8d0e492f2fb18846b29d8fcb76c4dbdbf3c271e3fd4f6afb417350ee9
+  metadata.gz: 19ee5d2e13a7b61962442766faa52f52930acc221b346ac91609eab01c1c9cb6aad01bfd96c444451925434c39ede78d7b949923a0824ef7ae89276b11c49869
+  data.tar.gz: 8de31a2f2b16402acef0d391e39a1049622dcbc3a69968ce0e8485492e1acbb7dae77ab88ae782bc1ef07ae42a43102929e5a47791a8c5443dd67f758e054c4b

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,38 +47,43 @@ 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.
 
-To interact with the ChromaDB collection, `chromable` provides `Model.query` method to query the collection and `Model.collection` method to access the collection directly.
+To interact with the ChromaDB collection, `chromable` provides:
+- `Model.query` to query the collection using similar API used in `chroma-db` gem, except for accepting `text:` instead of `query_embeddings:`. Extra arguments will be passed to the `embedder:` as `options`. Behind the scenes, `Model.query` will embed the given `text:`, then query the collection, and return the closest `results:` records.
+- `Model.collection` to access the collection directly.
+- `Model.delete_collection` to delete the entire collection.
 
 ```ruby
 puts Post.collection.count # Gets the number of documents inside the collection. Should always match Post.count.
 
 Post.query(
-  query: params[:query],
+  text: 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.
+Post.first.neighbors(results: 20) # => [#<Post:0x0000ffff9e0b5f10 id: "0beb0f98, ...>, ...]
+```
 
 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.
+- `neighbors`: Gets the closest `results:` records to the current record.
 
-All these methods (including `Model.query` and `Model.collection`) are available with `chroma_` prefix, if you have similar methods defined in your model.
+All these methods (including `Model.query`, `Model.collection`, and `Model.delete_collection`) are available with `chroma_` prefix, if you have similar methods defined in your model.
 
 ## Development
 
@@ -86,7 +93,7 @@ To install this gem onto your local machine, run `bundle install`. To release a
 
 ## Contributing
 
-Bug reports and pull requests are welcome on GitHub at https://github.com/AliOsm/chromable. This project is intended to be a safe, welcoming space for collaboration, and contributors are expected to adhere to the [code of conduct](https://github.com/[USERNAME]/chromable/blob/main/CODE_OF_CONDUCT.md).
+Bug reports and pull requests are welcome on GitHub at https://github.com/AliOsm/chromable. This project is intended to be a safe, welcoming space for collaboration, and contributors are expected to adhere to the [code of conduct](https://github.com/AliOsm/chromable/blob/main/CODE_OF_CONDUCT.md).
 
 ## License
 
@@ -94,4 +101,4 @@ The gem is available as open source under the terms of the [MIT License](https:/
 
 ## Code of Conduct
 
-Everyone interacting in the Chromable project's codebases, issue trackers, chat rooms and mailing lists is expected to follow the [code of conduct](https://github.com/[USERNAME]/chromable/blob/main/CODE_OF_CONDUCT.md).
+Everyone interacting in the Chromable project's codebases, issue trackers, chat rooms and mailing lists is expected to follow the [code of conduct](https://github.com/AliOsm/chromable/blob/main/CODE_OF_CONDUCT.md).

data/lib/chromable/version.rb

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

data/lib/chromable.rb

@@ -6,75 +6,122 @@ 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_query( # rubocop:disable Metrics/ParameterLists
-      query:,
-      results: 10,
-      where: {},
-      where_document: {},
-      include: %w[metadatas documents distances],
-      **embedder_options
-    )
+    def chroma_delete_collection
+      Chroma::Resources::Collection.delete(chromable_settings.collection_name)
+    end
+
+    def chroma_query(text:, results: 10, where: {}, where_document: {}, **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
+        alias_method :neighbors, :chroma_neighbors unless method_defined? :neighbors
+        # 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
 
-  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 chroma_destroy_embedding
+      self.class.chroma_collection.delete(ids: [id])
+    end
+
+    def chroma_neighbors(results: 10, where: {}, where_document: {})
+      collection = self.class.chroma_collection
+
+      embedding = collection.get(ids: [id], include: [:embeddings])[0].embedding
+
+      self.class.find(collection.query(
+        query_embeddings: [embedding],
+        results: results,
+        where: where,
+        where_document: where_document,
+        include: include
+      ).map(&:id))
+    end
+
+    private
+
+    def build_embedding
+      Chroma::Resources::Embedding.new(
+        id: id,
+        document: document_to_embed,
+        embedding: document_embedding,
+        metadata: embedding_metadata
+      )
+    end
 
-  private
+    def document_to_embed
+      chromable_settings.keep_document ? send(chromable_settings.document) : nil
+    end
 
-  def build_embedding
-    document = send(self.class.document)
+    def document_embedding
+      chromable_settings.embedder && self.class.send(chromable_settings.embedder, send(chromable_settings.document))
+    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 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.1
 platform: ruby
 authors:
 - Ali Hamdi Ali Fadel
 autorequire: 
 bindir: exe
 cert_chain: []
-date: 2023-12-11 00:00:00.000000000 Z
+date: 2023-12-13 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: chroma-db