chromable 0.1.5 → 0.2.0

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 9df743def1991fb00b56fdfb0878642a2e5568136d4078dcb77084cf47c6ce14
-  data.tar.gz: 6972ef16ef5c9f6f2f39de7e18c7028399d73b917001243c272e3b94d41839dc
+  metadata.gz: d44fa49957eefd26d00d3fb2315b51f928bd6ecdac7cdc4f849537cac23f036b
+  data.tar.gz: ab51d49d9b558e6d6ba936640026af66258567c0750f2ce11c7b6e01a0b8fea1
 SHA512:
-  metadata.gz: 07adc75fffe609a6e493f30091745093c3bb9164cd8e64c1f81519f8c447c4a6142075de97d461062903c28ac4d1a15d79bc501c57731e0d167a8562a696650b
-  data.tar.gz: 7818ee8d5a931f5e7936a0e27bb3bc716f349415e01f3cda33ff4cd72cf3ee237a5df157709a562238469659a828121e7d55b772b531bd06c1ddc2c7bc499204
+  metadata.gz: 2a6be5bd9ec59535eb33a59fc543baa23954ac43c4bc69b8a6e6d19eaab231d313180ba714781f85342f838bfdd4fd70fcfa32c785f663b880128d6a9425a69e
+  data.tar.gz: f63ffea813494e3143c16ed900e190d52866fa11348bac39db894c09ad885e66115b0cb8d0e492f2fb18846b29d8fcb76c4dbdbf3c271e3fd4f6afb417350ee9

data/.rubocop.yml

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

data/README.md

@@ -32,8 +32,14 @@ class Post < ApplicationRecord
 
   chromable document: :content, metadata: %i[author category], embedder: :embed
 
-  def embed
-    # Call OpenAI now please :)
+  def self.embed(text, **options)
+    options[:is_query] ||= false
+
+    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
 ```
@@ -41,22 +47,36 @@ 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.
+- `embedder:` is a callable defined in the model that returns the embedding representation for the given `text` and `options`.
 
 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` method is `document:`.
 
-At this point, `chromable` will create, update, and destroy the ChromaDB embeddings for your objects based on Rails callbacks.
+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,
+  type: 'query' # `type` 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.
 
-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.2.0'
 end

data/lib/chromable.rb

@@ -4,19 +4,20 @@ 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.after_save :chroma_upsert_embedding
     base.after_destroy :chroma_destroy_embedding
   end
 
+  # Methods to be added to the model class.
   module ClassMethods
-    def chromable(collection_name: nil, document: nil, metadata: nil, embedder: nil)
+    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
@@ -27,7 +28,25 @@ module Chromable
       Chroma::Resources::Collection.get_or_create(collection_name)
     end
 
-    alias_method :collection, :chroma_collection unless method_defined? :collection
+    def chroma_query( # rubocop:disable Metrics/ParameterLists
+      query:,
+      results: 10,
+      where: {},
+      where_document: {},
+      include: %w[metadatas documents distances],
+      **embedder_options
+    )
+      find(chroma_collection.query(
+        query_embeddings: [send(embedder, query, **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
@@ -35,14 +54,7 @@ module Chromable
   end
 
   def chroma_upsert_embedding
-    self.class.chroma_collection.upsert(
-      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
-      )
-    )
+    self.class.chroma_collection.upsert(build_embedding)
   end
 
   def chroma_destroy_embedding
@@ -52,4 +64,17 @@ module Chromable
   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
+
+  def build_embedding
+    document = send(self.class.document)
+
+    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) }
+    )
+  end
 end

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: chromable
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  version: 0.2.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-11 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: chroma-db