Add token cost rollups columns to conversation (#499)

app/models/application_record.rb

@@ -1,5 +1,6 @@
 class ApplicationRecord < ActiveRecord::Base
   primary_abstract_class
 
+  include RollupCache
   include Scopes::AutoGenerated
 end

app/models/concerns/rollup_cache.rb

@@ -0,0 +1,189 @@
+module RollupCache
+  extend ActiveSupport::Concern
+
+  # Rails supports counter_cache on associations, e.g.
+  #
+  # Book.rb
+  #   belongs_to :author, counter_cache: true
+  #
+  # And this keeps author.books_count updated each time a new author is created, updated, or destroyed.
+  # However, this option on belongs_to does not support conditional counter caches. For example,
+  # if we only want to cache the published books. rollup_cache was created for this purpose and also
+  # to support other rollups such as sums.
+  #
+  # Example count: book.rb:
+  #
+  # belongs_to :author, inverse_of: :books
+  # rollup_cache :published_books_count, belongs_to: :author, if: :published?
+  # rollup_cache :total_sales, sum: :sales, belongs_to: :author
+  #
+  # if the cache rollup name is _sum then it will be a sum instead of a count.
+  #
+  # Note: if an inverse_of is defined on the belongs_to then the referenced model will get a reset, e.g.:
+  # book.reset_published_books_count!
+
+  class_methods do
+    def rollup_cache(name, opts = {})
+      @_rollup_cache ||= []
+      raise "#{self} defines a rollup_cache #{name} but it's missing a 'belongs_to'" if opts[:belongs_to].nil?
+
+      multiple_belongs_to_models = Array(opts[:belongs_to])
+      multiple_belongs_to_models.each do |model|
+        @_rollup_cache << [ name.to_sym, opts.except(:belongs_to).merge(belongs_to: model) ]
+
+        cache_onto_association = self.reflect_on_all_associations.find { |a| a.name == model }
+        raise "#{self} defines a rollup_cache #{name} which references #{model} but that association cannot be found" if cache_onto_association.nil?
+
+        cache_onto_obj = cache_onto_association.klass
+
+        if opts[:sum] && self.columns.find { |c| c.name == opts[:sum].to_s }.nil?
+          raise "#{self} defines a rollup_cache #{name} which references #{opts[:sum]} but that column does not exist on #{self}"
+        end
+
+        # If we don't have an inverse_of, should we raise or simply not create a reset method? For now, let's just not create a reset method.
+        next  if cache_onto_association.inverse_of.nil?
+        # raise "#{self.to_s} has a rollup_cache :#{name} to :#{model} but the inverse_of :#{model} could not be determined, add inverse_of to :#{model}" if cache_onto_association.inverse_of.nil?
+
+        cache_onto_obj.module_eval <<-STR
+          def reset_#{name}!
+            if #{opts[:sum].present?}
+              if #{opts[:if].present?}
+                new_value = self.send(:#{cache_onto_association.inverse_of.name}).select { |record| record.send(:#{opts[:if] || "object_id"}) }.sum
+              else
+                new_value = self.send(:#{cache_onto_association.inverse_of.name}).sum(:#{opts[:sum]})
+              end
+            else
+              if #{opts[:if].present?}
+                new_value = self.send(:#{cache_onto_association.inverse_of.name}).select { |record| record.send(:#{opts[:if] || "object_id"}) }.length
+              else
+                new_value = self.send(:#{cache_onto_association.inverse_of.name}).count
+              end
+            end
+
+            self.update_column(:#{name}, new_value)
+          end
+        STR
+      end
+    end
+
+    def _rollup_cache
+      @_rollup_cache
+    end
+  end
+
+  included do
+    before_update :_check_rollup_cache_conditions_before_update
+    after_update :_update_rollup_cache_for_update
+    after_create :_update_rollup_cache_for_create
+    after_destroy :_update_rollup_cache_for_destroy
+
+    def _check_rollup_cache_conditions_before_update
+      return unless self.class._rollup_cache
+
+      saved_changes = _revert_changes # so conditions can be checked
+
+      self.class._rollup_cache.each do |name, opts|
+        obj, obj_id, passes_if_now = _rollup_cache_config_for(name, opts)
+        next if obj.nil? || obj_id.nil?
+
+        @_passes_if_before_update ||= {}
+        @_passes_if_before_update[name] = passes_if_now
+      end
+
+      _reapply_changes(saved_changes)
+    end
+
+    def _revert_changes
+      saved_changes = changes
+      changes.each do |key, pair|
+        before, after = pair
+
+        self.send("#{key}=", before)
+      end
+      saved_changes
+    end
+
+    def _reapply_changes(chgs)
+      chgs.each do |key, pair|
+        before, after = pair
+
+        self.send("#{key}=", after)
+      end
+    end
+
+    def _update_rollup_cache_for_update
+      _update_rollup_cache(:update)
+    end
+
+    def _update_rollup_cache_for_create
+      _update_rollup_cache(:create)
+    end
+
+    def _update_rollup_cache_for_destroy
+      _update_rollup_cache(:destroy)
+    end
+
+    def _update_rollup_cache(mode)
+      return unless self.class._rollup_cache
+
+      is_creating = mode == :create
+      is_updating = mode == :update
+      is_destroying = mode == :destroy
+
+      self.class._rollup_cache.each do |name, opts|
+        obj, obj_id, passes_if_now, passes_if_before_update = _rollup_cache_config_for(name, opts)
+        next if obj.nil? || obj_id.nil?
+
+        newly_created_record_that_passes_if = is_creating && passes_if_now
+        destroying_a_record_that_passes_if  = is_destroying && passes_if_now
+        updated_record_to_now_pass_if       = is_updating && !passes_if_before_update && passes_if_now
+        updated_record_to_no_longer_pass_if = is_updating && passes_if_before_update && !passes_if_now
+
+        if opts[:sum].blank? && (newly_created_record_that_passes_if || updated_record_to_now_pass_if)
+          if opts[:callbacks]
+            record = obj.find(obj_id)
+            record.update!(name => record.send(name) + 1)
+          else
+            obj.increment_counter(name, obj_id, touch: true)
+          end
+        end
+
+        if (value_was, value_is = saved_changes[opts[:sum]])
+          record = obj.find(obj_id)
+          record.update!(name => record.send(name) + (value_is - value_was))
+        end
+
+        if destroying_a_record_that_passes_if || updated_record_to_no_longer_pass_if
+          if opts[:sum]
+            record = obj.find(obj_id)
+            record.update!(name => record.send(name) - self.send(opts[:sum]))
+          else
+            if opts[:callbacks]
+              record = obj.find(obj_id)
+              record.update!(name => record.send(name) - 1)
+            else
+              obj.decrement_counter(name, obj_id, touch: true)
+            end
+          end
+        end
+      end
+    end
+
+    def _rollup_cache_config_for(name, opts)
+      @_passes_if_before_update ||= {}
+
+      association = self.class.reflect_on_all_associations.find { |a| a.name == opts[:belongs_to] }
+      raise "#{self} defines a rollup_cache #{name} specified belongs_to of '#{opts[:belongs_to]}' but this association was not found." if association.nil?
+      obj             = association.klass
+      obj_instance    = self.send(opts[:belongs_to])
+      obj_id = self.send(opts[:belongs_to])&.id
+
+      rollup_cache_column = obj.columns.find { |c| c.name == name.to_s }
+      raise "The column #{name} does not exist on #{obj}" if rollup_cache_column.nil?
+
+      passes_if_now = opts[:if].nil? ? true : self.send(opts[:if])
+
+      [obj, obj_id, passes_if_now, @_passes_if_before_update[name]]
+    end
+  end
+end

app/models/conversation.rb

@@ -1,5 +1,5 @@
 class Conversation < ApplicationRecord
-  include Version
+  include Version, Billable
 
   belongs_to :user
   belongs_to :assistant

app/models/conversation/billable.rb

@@ -0,0 +1,10 @@
+module Conversation::Billable
+  extend ActiveSupport::Concern
+  included do
+    attribute :input_token_total_count, :integer, default: 0
+    attribute :output_token_total_count, :integer, default: 0
+
+    attribute :input_token_total_cost, :float, precision: 10, scale: 2, default: 0.0
+    attribute :output_token_total_cost, :float, precision: 10, scale: 2, default: 0.0
+  end
+end

app/models/message.rb

@@ -1,12 +1,14 @@
 class Message < ApplicationRecord
-  include DocumentImage, Version, Cancellable, Toolable, TokenCount
+  include DocumentImage, Version, Cancellable, Toolable
 
   belongs_to :assistant
-  belongs_to :conversation
+  belongs_to :conversation, inverse_of: :messages
   belongs_to :content_document, class_name: "Document", inverse_of: :message, optional: true
   belongs_to :run, optional: true
   has_one :latest_assistant_message_for, class_name: "Conversation", inverse_of: :last_assistant_message, dependent: :nullify
 
+  include Billable
+
   enum role: %w[user assistant tool].index_by(&:to_sym)
 
   delegate :user, to: :conversation

app/models/message/billable.rb

@@ -0,0 +1,16 @@
+module Message::Billable
+  extend ActiveSupport::Concern
+  included do
+    attribute :input_token_count, :integer, default: 0
+    attribute :output_token_count, :integer, default: 0
+
+    rollup_cache :input_token_total_count, sum: :input_token_count, belongs_to: :conversation
+    rollup_cache :output_token_total_count, sum: :output_token_count, belongs_to: :conversation
+
+    attribute :input_token_cost, :decimal, precision: 10, scale: 2, default: 0.0
+    attribute :output_token_cost, :decimal, precision: 10, scale: 2, default: 0.0
+
+    rollup_cache :input_token_total_cost, sum: :input_token_cost, belongs_to: :conversation
+    rollup_cache :output_token_total_cost, sum: :output_token_cost, belongs_to: :conversation
+  end
+end

db/migrate/20240823210939_add_billable_columns.rb

@@ -0,0 +1,12 @@
+class AddBillableColumns < ActiveRecord::Migration[7.1]
+  def change
+    add_column :messages, :input_token_cost, :decimal, precision: 10, scale: 2, default: 0.0, null: false
+    add_column :messages, :output_token_cost, :decimal, precision: 10, scale: 2, default: 0.0, null: false
+
+    add_column :conversations, :input_token_total_cost, :decimal, precision: 10, scale: 2, default: 0.0, null: false
+    add_column :conversations, :output_token_total_cost, :decimal, precision: 10, scale: 2, default: 0.0, null: false
+
+    add_column :conversations, :input_token_total_count, :integer, default: 0, null: false
+    add_column :conversations, :output_token_total_count, :integer, default: 0, null: false
+  end
+end

db/schema.rb

@@ -10,7 +10,7 @@
 #
 # It's strongly recommended that you check this file into your version control system.
 
-ActiveRecord::Schema[7.1].define(version: 2024_08_01_203803) do
+ActiveRecord::Schema[7.1].define(version: 2024_08_23_210939) do
   # These are extensions that must be enabled in order to support this database
   enable_extension "plpgsql"
 
@@ -116,6 +116,10 @@
     t.datetime "updated_at", null: false
     t.bigint "last_assistant_message_id"
     t.text "external_id", comment: "The Backend AI system (e.g OpenAI) Thread Id"
+    t.decimal "input_token_total_cost", precision: 10, scale: 2, default: "0.0", null: false
+    t.decimal "output_token_total_cost", precision: 10, scale: 2, default: "0.0", null: false
+    t.integer "input_token_total_count", default: 0, null: false
+    t.integer "output_token_total_count", default: 0, null: false
     t.index ["assistant_id"], name: "index_conversations_on_assistant_id"
     t.index ["external_id"], name: "index_conversations_on_external_id", unique: true
     t.index ["last_assistant_message_id"], name: "index_conversations_on_last_assistant_message_id"
@@ -200,6 +204,8 @@
     t.string "tool_call_id"
     t.integer "input_token_count", default: 0, null: false
     t.integer "output_token_count", default: 0, null: false
+    t.decimal "input_token_cost", precision: 10, scale: 2, default: "0.0", null: false
+    t.decimal "output_token_cost", precision: 10, scale: 2, default: "0.0", null: false
     t.index ["assistant_id"], name: "index_messages_on_assistant_id"
     t.index ["content_document_id"], name: "index_messages_on_content_document_id"
     t.index ["conversation_id", "index", "version"], name: "index_messages_on_conversation_id_and_index_and_version", unique: true

test/fixtures/conversations.yml

@@ -39,6 +39,10 @@ debugging:
   assistant: rob_gpt4
   title: Verifying Ruby syntax
   last_assistant_message: filter_map_example
+  input_token_total_count: 12
+  output_token_total_count: 36
+  input_token_total_cost: 0.12
+  output_token_total_cost: 0.72
 
 #  0.1
 #  1.1

test/fixtures/messages.yml

@@ -373,6 +373,10 @@ filter_map:
   created_at: 2023-12-30 1:00:00
   index: 0
   version: 1
+  input_token_count: 5
+  output_token_count: 15
+  input_token_cost: 0.05
+  output_token_cost: 0.30
 
 filter_map_example:
   assistant: rob_gpt4
@@ -384,6 +388,10 @@ filter_map_example:
   created_at: 2023-12-30 1:01:00
   index: 1
   version: 1
+  input_token_count: 7
+  output_token_count: 21
+  input_token_cost: 0.07
+  output_token_cost: 0.42
 
 
 # Next conversation