diff --git a/app/models/ai.rb b/app/models/ai.rb new file mode 100644 index 000000000..4b4c17743 --- /dev/null +++ b/app/models/ai.rb @@ -0,0 +1,5 @@ +module Ai + def self.table_name_prefix + "ai_" + end +end diff --git a/app/models/ai/quota.rb b/app/models/ai/quota.rb index 6ef4d7dc3..955a855af 100644 --- a/app/models/ai/quota.rb +++ b/app/models/ai/quota.rb @@ -1,18 +1,14 @@ class Ai::Quota < ApplicationRecord class UsageExceedsQuotaError < StandardError; end - include MoneyAccessors, Resettable - - self.table_name = "ai_quotas" - - money_accessor :used, :limit - belongs_to :user + before_create -> { reset } + validates :limit, presence: true, numericality: { only_integer: true, greater_than_or_equal_to: 0 } validates :used, presence: true, numericality: { only_integer: true, greater_than_or_equal_to: 0 } - def increment_usage(cost) + def spend(cost) cost = Money.wrap(cost) transaction do @@ -21,15 +17,34 @@ class Ai::Quota < ApplicationRecord end end - def ensure_under_limit + def ensure_not_depleted reset_if_due - if over_limit? + if depleted? raise UsageExceedsQuotaError end end - def over_limit? - used >= limit - end + private + def reset_if_due + reset if due_for_reset? + end + + def reset + attributes = { used: 0, reset_at: 7.days.from_now } + + if persisted? + update(**attributes) + else + assign_attributes(**attributes) + end + end + + def due_for_reset? + reset_at.before?(Time.current) + end + + def depleted? + used >= limit + end end diff --git a/app/models/ai/quota/money.rb b/app/models/ai/quota/money.rb index 2a22ec3b7..f229c8def 100644 --- a/app/models/ai/quota/money.rb +++ b/app/models/ai/quota/money.rb @@ -1,6 +1,4 @@ class Ai::Quota::Money < Data.define(:value) - include Comparable - CENTS_PER_DOLLAR = 100 MICROCENTS_PER_CENT = 1_000_000 MICROCENTS_PER_DOLLAR = CENTS_PER_DOLLAR * MICROCENTS_PER_CENT @@ -24,9 +22,10 @@ class Ai::Quota::Money < Data.define(:value) new(microcents) end - def convert_dollars_to_microcents(dollars) - (dollars.to_d * MICROCENTS_PER_DOLLAR).round.to_i - end + private + def convert_dollars_to_microcents(dollars) + (dollars.to_d * MICROCENTS_PER_DOLLAR).round.to_i + end end def to_i @@ -40,22 +39,4 @@ class Ai::Quota::Money < Data.define(:value) def in_dollars in_microcents.to_d / MICROCENTS_PER_DOLLAR end - - def <=>(other) - in_microcents <=> self.class.wrap(other).in_microcents - end - - def +(other) - other = self.class.wrap(other) - self.class.new(in_microcents + other.in_microcents) - end - - def -(other) - other = self.class.wrap(other) - self.class.new(in_microcents - other.in_microcents) - end - - def abs - self.class.new(in_microcents.abs) - end end diff --git a/app/models/ai/quota/money_accessors.rb b/app/models/ai/quota/money_accessors.rb deleted file mode 100644 index 28497a4d5..000000000 --- a/app/models/ai/quota/money_accessors.rb +++ /dev/null @@ -1,20 +0,0 @@ -module Ai::Quota::MoneyAccessors - extend ActiveSupport::Concern - - class_methods do - def money_accessor(*attribute_names) - attribute_names.each do |name| - define_method(name) do - if super() - Ai::Quota::Money.new(super()) - end - end - - define_method("#{name}=") do |value| - value = Ai::Quota::Money.wrap(value).in_microcents if value - super(value) - end - end - end - end -end diff --git a/app/models/ai/quota/resettable.rb b/app/models/ai/quota/resettable.rb deleted file mode 100644 index 497ea8b93..000000000 --- a/app/models/ai/quota/resettable.rb +++ /dev/null @@ -1,25 +0,0 @@ -module Ai::Quota::Resettable - extend ActiveSupport::Concern - - included do - before_create -> { reset } - end - - def reset_if_due - reset if due_for_reset? - end - - def reset - attributes = { used: 0, reset_at: 7.days.from_now } - - if persisted? - update(**attributes) - else - assign_attributes(**attributes) - end - end - - def due_for_reset? - reset_at.before?(Time.current) - end -end diff --git a/app/models/conversation.rb b/app/models/conversation.rb index 1d7cab876..bf3a4a0ab 100644 --- a/app/models/conversation.rb +++ b/app/models/conversation.rb @@ -9,7 +9,7 @@ class Conversation < ApplicationRecord enum :state, %w[ ready thinking ].index_by(&:itself), default: :ready def ask(question, **attributes) - user.ensure_under_ai_usage_limit + user.ensure_ai_quota_not_depleted create_message_with_state_change(**attributes, role: :user, content: question) do raise(InvalidStateError, "Can't ask questions while thinking") if thinking? @@ -23,7 +23,7 @@ class Conversation < ApplicationRecord ready! end - user.increment_ai_usage(message.cost) if message.cost + user.spend_ai_quota(message.cost) if message.cost message end diff --git a/app/models/user/ai_quota.rb b/app/models/user/ai_quota.rb index 133759c71..01be7ae1c 100644 --- a/app/models/user/ai_quota.rb +++ b/app/models/user/ai_quota.rb @@ -5,15 +5,16 @@ module User::AiQuota has_one :ai_quota, class_name: "Ai::Quota" end - def fetch_or_create_ai_quota - ai_quota || create_ai_quota!(limit: Ai::Quota::Money.wrap("$100")) + def spend_ai_quota(cost) + fetch_or_create_ai_quota.spend(cost) end - def increment_ai_usage(cost) - fetch_or_create_ai_quota.increment_usage(cost) + def ensure_ai_quota_not_depleted + fetch_or_create_ai_quota.ensure_not_depleted end - def ensure_under_ai_usage_limit - fetch_or_create_ai_quota.ensure_under_limit - end + private + def fetch_or_create_ai_quota + ai_quota || create_ai_quota!(limit: Ai::Quota::Money.wrap("$100")) + end end diff --git a/config/initializers/inflections.rb b/config/initializers/inflections.rb index 3d2359e52..62ffe294f 100644 --- a/config/initializers/inflections.rb +++ b/config/initializers/inflections.rb @@ -13,4 +13,7 @@ # These inflection rules are supported but not enabled by default: ActiveSupport::Inflector.inflections(:en) do |inflect| inflect.acronym "SQLite" + + inflect.singular "quotas", "quota" + inflect.plural "quota", "quotas" end diff --git a/test/models/ai/quota_test.rb b/test/models/ai/quota_test.rb index 9d175b80d..b37e60ee9 100644 --- a/test/models/ai/quota_test.rb +++ b/test/models/ai/quota_test.rb @@ -2,7 +2,7 @@ require "test_helper" class Ai::QuotaTest < ActiveSupport::TestCase setup do - @quota = Ai::Quota.new(user: users(:jz), limit: "$100") + @quota = Ai::Quota.new(user: users(:jz), limit: Ai::Quota::Money.wrap("$100").in_microcents) end test "create" do @@ -10,106 +10,54 @@ class Ai::QuotaTest < ActiveSupport::TestCase assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute assert_equal 0, @quota.used - assert_equal Ai::Quota::Money.wrap("$100"), @quota.limit - end - - test "accessors" do - @quota.limit = nil - assert_nil @quota.limit - - @quota.limit = "$100" - assert_kind_of Ai::Quota::Money, @quota.limit - assert_equal Ai::Quota::Money.wrap("$100"), @quota.limit - - @quota.used = nil - assert_nil @quota.used - - @quota.used = "$100" - assert_kind_of Ai::Quota::Money, @quota.used - assert_equal Ai::Quota::Money.wrap("$100"), @quota.used + assert_equal Ai::Quota::Money.wrap("$100").in_microcents, @quota.limit end test "increment usage" do @quota.save - @quota.increment_usage("$100") - assert_equal Ai::Quota::Money.wrap("$100"), @quota.used - @quota.increment_usage("$500") - assert_equal Ai::Quota::Money.wrap("$600"), @quota.used - @quota.increment_usage("$1000") - assert_equal Ai::Quota::Money.wrap("$1600"), @quota.used - @quota.increment_usage("$5000") - assert_equal Ai::Quota::Money.wrap("$6600"), @quota.used - @quota.increment_usage("$10000") - assert_equal Ai::Quota::Money.wrap("$16600"), @quota.used + @quota.spend("$100") + assert_equal Ai::Quota::Money.wrap("$100").in_microcents, @quota.used + @quota.spend("$500") + assert_equal Ai::Quota::Money.wrap("$600").in_microcents, @quota.used + @quota.spend("$1000") + assert_equal Ai::Quota::Money.wrap("$1600").in_microcents, @quota.used + @quota.spend("$5000") + assert_equal Ai::Quota::Money.wrap("$6600").in_microcents, @quota.used + @quota.spend("$10000") + assert_equal Ai::Quota::Money.wrap("$16600").in_microcents, @quota.used - @quota.reset - assert_equal Ai::Quota::Money.wrap("$0"), @quota.used - assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute + @quota.used = 0 - @quota.increment_usage("$10") - assert_equal Ai::Quota::Money.wrap("$10"), @quota.used + @quota.spend("$10") + assert_equal Ai::Quota::Money.wrap("$10").in_microcents, @quota.used assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute travel 2.days - @quota.increment_usage("$5") - assert_equal Ai::Quota::Money.wrap("$15"), @quota.used + @quota.spend("$5") + assert_equal Ai::Quota::Money.wrap("$15").in_microcents, @quota.used assert_in_delta 5.days.from_now, @quota.reset_at, 1.minute travel 8.days - @quota.increment_usage("$5") - assert_equal Ai::Quota::Money.wrap("$5"), @quota.used + @quota.spend("$5") + assert_equal Ai::Quota::Money.wrap("$5").in_microcents, @quota.used assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute end test "limit checks" do @quota.save - @quota.used = "$0" - assert_not @quota.over_limit? + @quota.used = 0 + @quota.ensure_not_depleted - @quota.used = "$300" - assert @quota.over_limit? - - @quota.used = "$0" - @quota.ensure_under_limit - - @quota.used = "$300" + @quota.used = Ai::Quota::Money.wrap("$300").in_microcents assert_raises Ai::Quota::UsageExceedsQuotaError do - @quota.ensure_under_limit + @quota.ensure_not_depleted end travel 10.days - @quota.ensure_under_limit - end - - test "reset" do - @quota.used = "$15" - @quota.reset_at = 3.days.from_now - - @quota.reset - - assert_equal Ai::Quota::Money.wrap("$0"), @quota.used - assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute - - @quota.reset_at = 10.minutes.from_now - assert_not @quota.due_for_reset? - - @quota.reset_at = 10.minutes.ago - assert @quota.due_for_reset? - - @quota.reset_at = 10.minutes.from_now - @quota.used = "$15" - @quota.reset_if_due - assert_equal Ai::Quota::Money.wrap("$15"), @quota.used - assert_in_delta 10.minutes.from_now, @quota.reset_at, 1.minute - - @quota.reset_at = 10.minutes.ago - @quota.used = "$15" - @quota.reset_if_due - assert_equal Ai::Quota::Money.wrap("$0"), @quota.used - assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute + @quota.ensure_not_depleted end end