From abad369c6a9bd2a0d5198ceea942ca18195f69fe Mon Sep 17 00:00:00 2001 From: "Stanko K.R." Date: Wed, 20 Aug 2025 08:25:16 +0200 Subject: [PATCH] Introduce a Quota and Money object --- .../conversations/messages_controller.rb | 2 +- app/models/ai/quota.rb | 35 ++++ app/models/ai/quota/money.rb | 61 +++++++ app/models/ai/quota/money_accessors.rb | 20 +++ app/models/ai/quota/resettable.rb | 26 +++ app/models/conversation.rb | 11 +- app/models/conversation/cost.rb | 30 ---- app/models/conversation/cost_limited.rb | 27 --- app/models/conversation/message.rb | 5 +- .../message/response_generator/response.rb | 14 +- app/models/user.rb | 2 +- app/models/user/ai_quota.rb | 19 +++ db/migrate/20250819101032_create_ai_quotas.rb | 14 ++ db/schema.rb | 12 ++ db/schema_cache.yml | 154 +++++++++++++----- test/fixtures/ai/quotas.yml | 11 ++ test/models/ai/quota/money_test.rb | 40 +++++ test/models/ai/quota_test.rb | 115 +++++++++++++ .../response_generator/response_test.rb | 3 - test/models/conversation_test.rb | 14 +- test/models/user_test.rb | 9 + 21 files changed, 495 insertions(+), 129 deletions(-) create mode 100644 app/models/ai/quota.rb create mode 100644 app/models/ai/quota/money.rb create mode 100644 app/models/ai/quota/money_accessors.rb create mode 100644 app/models/ai/quota/resettable.rb delete mode 100644 app/models/conversation/cost.rb delete mode 100644 app/models/conversation/cost_limited.rb create mode 100644 app/models/user/ai_quota.rb create mode 100644 db/migrate/20250819101032_create_ai_quotas.rb create mode 100644 test/fixtures/ai/quotas.yml create mode 100644 test/models/ai/quota/money_test.rb create mode 100644 test/models/ai/quota_test.rb diff --git a/app/controllers/conversations/messages_controller.rb b/app/controllers/conversations/messages_controller.rb index fc4b46838..007364176 100644 --- a/app/controllers/conversations/messages_controller.rb +++ b/app/controllers/conversations/messages_controller.rb @@ -7,7 +7,7 @@ class Conversations::MessagesController < ApplicationController def create @conversation.ask(question, **message_params) - rescue Conversation::CostExceededError + rescue Ai::Quota::UsageExceedsQuotaError render json: { error: "You've asked too many questions this week" }, status: :too_many_requests rescue Conversation::InvalidStateError render json: { error: "Fizzy is still working on an answer to your last question" }, status: :conflict diff --git a/app/models/ai/quota.rb b/app/models/ai/quota.rb new file mode 100644 index 000000000..6ef4d7dc3 --- /dev/null +++ b/app/models/ai/quota.rb @@ -0,0 +1,35 @@ +class Ai::Quota < ApplicationRecord + class UsageExceedsQuotaError < StandardError; end + + include MoneyAccessors, Resettable + + self.table_name = "ai_quotas" + + money_accessor :used, :limit + + belongs_to :user + + 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) + cost = Money.wrap(cost) + + transaction do + reset_if_due + increment!(:used, cost.in_microcents) + end + end + + def ensure_under_limit + reset_if_due + + if over_limit? + raise UsageExceedsQuotaError + end + end + + def over_limit? + used >= limit + end +end diff --git a/app/models/ai/quota/money.rb b/app/models/ai/quota/money.rb new file mode 100644 index 000000000..2a22ec3b7 --- /dev/null +++ b/app/models/ai/quota/money.rb @@ -0,0 +1,61 @@ +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 + NUMBER_REGEX = /\d+(\.\d+)?/ + + class << self + def wrap(value) + microcents = case value + when nil + raise ArgumentError, "#{self} can't wrap nil" + when self + value.value + when Integer + value + when String + convert_dollars_to_microcents(BigDecimal(value[NUMBER_REGEX])) + else + convert_dollars_to_microcents(value) + end + + new(microcents) + end + + def convert_dollars_to_microcents(dollars) + (dollars.to_d * MICROCENTS_PER_DOLLAR).round.to_i + end + end + + def to_i + in_microcents + end + + def in_microcents + value + end + + 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 new file mode 100644 index 000000000..28497a4d5 --- /dev/null +++ b/app/models/ai/quota/money_accessors.rb @@ -0,0 +1,20 @@ +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 new file mode 100644 index 000000000..de0ee493e --- /dev/null +++ b/app/models/ai/quota/resettable.rb @@ -0,0 +1,26 @@ +module Ai::Quota::Resettable + extend ActiveSupport::Concern + + included do + before_create -> { reset } + scope :due_for_reset, -> { where(reset_at: ...Time.current) } + 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 afdb4b47c..1d7cab876 100644 --- a/app/models/conversation.rb +++ b/app/models/conversation.rb @@ -1,8 +1,7 @@ class Conversation < ApplicationRecord class InvalidStateError < StandardError; end - class CostExceededError < StandardError; end - include Broadcastable, CostLimited + include Broadcastable belongs_to :user, class_name: "User" has_many :messages, dependent: :destroy @@ -10,7 +9,7 @@ class Conversation < ApplicationRecord enum :state, %w[ ready thinking ].index_by(&:itself), default: :ready def ask(question, **attributes) - limit_cost to: "$100", within: 7.days + user.ensure_under_ai_usage_limit create_message_with_state_change(**attributes, role: :user, content: question) do raise(InvalidStateError, "Can't ask questions while thinking") if thinking? @@ -19,10 +18,14 @@ class Conversation < ApplicationRecord end def respond(answer, **attributes) - create_message_with_state_change(**attributes, role: :assistant, content: answer) do + message = create_message_with_state_change(**attributes, role: :assistant, content: answer) do raise(InvalidStateError, "Can't respond when not thinking") unless thinking? ready! end + + user.increment_ai_usage(message.cost) if message.cost + + message end private diff --git a/app/models/conversation/cost.rb b/app/models/conversation/cost.rb deleted file mode 100644 index d42568150..000000000 --- a/app/models/conversation/cost.rb +++ /dev/null @@ -1,30 +0,0 @@ -module Conversation::Cost - CENTS_PER_DOLLAR = 100 - MICROCENTS_PER_CENT = 1_000_000 - MICROCENTS_PER_DOLLAR = CENTS_PER_DOLLAR * MICROCENTS_PER_CENT - NUMBER_REGEX = /\d+(\.\d+)?/ - - extend self - - def convert_to_microcents(value) - case value - when String - decimal_to_microcents(value[NUMBER_REGEX].to_d) - when Float, BigDecimal - decimal_to_microcents(value) - when Numeric - value - else - raise ArgumentError, "Invalid cost value: #{value}" - end - end - - def convert_to_decimal(microcents) - microcents.to_d / MICROCENTS_PER_DOLLAR - end - - private - def decimal_to_microcents(decimal) - (decimal.to_d * MICROCENTS_PER_DOLLAR).round.to_i - end -end diff --git a/app/models/conversation/cost_limited.rb b/app/models/conversation/cost_limited.rb deleted file mode 100644 index 537782b33..000000000 --- a/app/models/conversation/cost_limited.rb +++ /dev/null @@ -1,27 +0,0 @@ -module Conversation::CostLimited - extend ActiveSupport::Concern - - def limit_cost(to:, within: nil) - raise Conversation::CostExceededError, "Cost limit exceeded" if cost_exceeds?(to, within: within) - end - - def cost_exceeds?(amount, within: nil) - cost_microcents(within: within) >= Conversation::Cost.convert_to_microcents(amount) - end - - def cost(...) - Conversation::Cost.convert_to_decimal cost_microcents(...) - end - - def cost_microcents(within: nil) - within = (within.ago...) if within && !within.is_a?(Range) - - scope = if within - messages.where(created_at: within) - else - messages - end - - scope.with_cost.sum(:cost_microcents) - end -end diff --git a/app/models/conversation/message.rb b/app/models/conversation/message.rb index 282ec294d..72ea37b26 100644 --- a/app/models/conversation/message.rb +++ b/app/models/conversation/message.rb @@ -12,7 +12,10 @@ class Conversation::Message < ApplicationRecord validates :client_message_id, presence: true scope :ordered, -> { order(created_at: :asc, id: :asc) } - scope :with_cost, -> { where.not(cost_microcents: nil) } + + def cost + cost_microcents && Ai::Quota::Money.new(cost_microcents) + end def all_emoji? content.to_plain_text.all_emoji? diff --git a/app/models/conversation/message/response_generator/response.rb b/app/models/conversation/message/response_generator/response.rb index e3e0fcd41..4d782bffd 100644 --- a/app/models/conversation/message/response_generator/response.rb +++ b/app/models/conversation/message/response_generator/response.rb @@ -1,6 +1,4 @@ class Conversation::Message::ResponseGenerator::Response - MICROCENTS_PER_DOLLAR = 100_000 - attr_reader :answer, :input_tokens, :output_tokens, :model_id, :tool_calls, :tool_call_id def initialize(answer:, input_tokens:, output_tokens:, model_id:) @@ -10,10 +8,6 @@ class Conversation::Message::ResponseGenerator::Response @model_id = model_id end - def cost - Conversation::Cost.convert_to_decimal(cost_microcents) - end - def cost_microcents input_cost_microcents + output_cost_microcents end @@ -27,7 +21,7 @@ class Conversation::Message::ResponseGenerator::Response def input_token_price_microcents return unless model_info - price_per_million_tokens_to_microcents(model_info.input_price_per_million) + price_per_million_tokens_in_microcents(model_info.input_price_per_million) end def output_cost_microcents @@ -39,7 +33,7 @@ class Conversation::Message::ResponseGenerator::Response def output_token_price_microcents return unless model_info - price_per_million_tokens_to_microcents(model_info.output_price_per_million) + price_per_million_tokens_in_microcents(model_info.output_price_per_million) end def model_info @@ -47,8 +41,8 @@ class Conversation::Message::ResponseGenerator::Response end private - def price_per_million_tokens_to_microcents(price) + def price_per_million_tokens_in_microcents(price) single_token_price = price.to_d / 1_000_000 - Conversation::Cost.convert_to_microcents(single_token_price) + Ai::Quota::Money.wrap(single_token_price).in_microcents end end diff --git a/app/models/user.rb b/app/models/user.rb index bd8cf990d..af3ff9315 100644 --- a/app/models/user.rb +++ b/app/models/user.rb @@ -1,6 +1,6 @@ class User < ApplicationRecord include Accessor, Attachable, Assignee, Mentionable, Named, Role, Searcher, - SignalUser, Staff, Transferable, Conversational + SignalUser, Staff, Transferable, Conversational, AiQuota include Timelined # Depends on Accessor has_one_attached :avatar diff --git a/app/models/user/ai_quota.rb b/app/models/user/ai_quota.rb new file mode 100644 index 000000000..133759c71 --- /dev/null +++ b/app/models/user/ai_quota.rb @@ -0,0 +1,19 @@ +module User::AiQuota + extend ActiveSupport::Concern + + included do + 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")) + end + + def increment_ai_usage(cost) + fetch_or_create_ai_quota.increment_usage(cost) + end + + def ensure_under_ai_usage_limit + fetch_or_create_ai_quota.ensure_under_limit + end +end diff --git a/db/migrate/20250819101032_create_ai_quotas.rb b/db/migrate/20250819101032_create_ai_quotas.rb new file mode 100644 index 000000000..500f00db9 --- /dev/null +++ b/db/migrate/20250819101032_create_ai_quotas.rb @@ -0,0 +1,14 @@ +class CreateAiQuotas < ActiveRecord::Migration[8.1] + def change + create_table :ai_quotas do |t| + t.belongs_to :user, null: false, foreign_key: true + t.integer :limit, null: false + t.integer :used, null: false, default: 0 + t.datetime :reset_at, null: false, index: true + + t.timestamps + end + + add_index :ai_quotas, :user_id, unique: true + end +end diff --git a/db/schema.rb b/db/schema.rb index b31a7dc51..93533668e 100644 --- a/db/schema.rb +++ b/db/schema.rb @@ -73,6 +73,17 @@ ActiveRecord::Schema[8.1].define(version: 2025_08_19_105245) do t.index ["blob_id", "variation_digest"], name: "index_active_storage_variant_records_uniqueness", unique: true end + create_table "ai_quotas", force: :cascade do |t| + t.datetime "created_at", null: false + t.integer "limit", null: false + t.datetime "reset_at", null: false + t.datetime "updated_at", null: false + t.integer "used", default: 0, null: false + t.integer "user_id", null: false + t.index ["reset_at"], name: "index_ai_quotas_on_reset_at" + t.index ["user_id"], name: "index_ai_quotas_on_user_id" + end + create_table "assignees_filters", id: false, force: :cascade do |t| t.integer "assignee_id", null: false t.integer "filter_id", null: false @@ -438,6 +449,7 @@ ActiveRecord::Schema[8.1].define(version: 2025_08_19_105245) do add_foreign_key "active_storage_attachments", "active_storage_blobs", column: "blob_id" add_foreign_key "active_storage_variant_records", "active_storage_blobs", column: "blob_id" + add_foreign_key "ai_quotas", "users" add_foreign_key "card_activity_spikes", "cards" add_foreign_key "card_goldnesses", "cards" add_foreign_key "cards", "workflow_stages", column: "stage_id" diff --git a/db/schema_cache.yml b/db/schema_cache.yml index 58f3518f3..3f0d88ec2 100644 --- a/db/schema_cache.yml +++ b/db/schema_cache.yml @@ -20,7 +20,7 @@ columns: default_function: collation: comment: - - &23 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + - &24 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: collection_id cast_type: &3 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3Adapter::SQLite3Integer @@ -90,7 +90,7 @@ columns: default_function: collation: comment: - - &25 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + - &18 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: user_id cast_type: *3 @@ -258,7 +258,7 @@ columns: collation: comment: - *6 - - &18 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + - &19 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: key cast_type: *7 @@ -301,9 +301,44 @@ columns: default_function: collation: comment: + ai_quotas: + - *5 + - *6 + - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + auto_increment: + name: limit + cast_type: *3 + sql_type_metadata: *4 + 'null': false + default: + default_function: + collation: + comment: + - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + auto_increment: + name: reset_at + cast_type: *1 + sql_type_metadata: *2 + 'null': false + default: + default_function: + collation: + comment: + - *9 + - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + auto_increment: + name: used + cast_type: *3 + sql_type_metadata: *4 + 'null': false + default: 0 + default_function: + collation: + comment: + - *18 ar_internal_metadata: - *5 - - *18 + - *19 - *9 - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: @@ -316,7 +351,7 @@ columns: collation: comment: assignees_filters: - - &20 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + - &21 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: assignee_id cast_type: *3 @@ -326,7 +361,7 @@ columns: default_function: collation: comment: - - &19 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + - &20 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: filter_id cast_type: *3 @@ -337,7 +372,7 @@ columns: collation: comment: assigners_filters: - - &21 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + - &22 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: assigner_id cast_type: *3 @@ -347,11 +382,11 @@ columns: default_function: collation: comment: - - *19 - assignments: - *20 + assignments: - *21 - - &22 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + - *22 + - &23 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: card_id cast_type: *3 @@ -365,7 +400,7 @@ columns: - *6 - *9 card_activity_spikes: - - *22 + - *23 - *5 - *6 - *9 @@ -394,14 +429,14 @@ columns: comment: - *9 card_goldnesses: - - *22 + - *23 - *5 - *6 - *9 cards: - - *23 + - *24 - *5 - - &24 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column + - &25 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: creator_id cast_type: *3 @@ -483,7 +518,7 @@ columns: default_function: collation: comment: - - *19 + - *20 closure_reasons: - *5 - *6 @@ -499,7 +534,7 @@ columns: comment: - *9 closures: - - *22 + - *23 - *5 - *6 - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column @@ -524,7 +559,7 @@ columns: collation: comment: collection_publications: - - *23 + - *24 - *5 - *6 - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column @@ -558,7 +593,7 @@ columns: collation: comment: - *5 - - *24 + - *25 - *6 - *10 - *9 @@ -573,12 +608,12 @@ columns: collation: comment: collections_filters: - - *23 - - *19 - comments: - - *22 - - *5 - *24 + - *20 + comments: + - *23 + - *5 + - *25 - *6 - *9 conversation_messages: @@ -699,10 +734,10 @@ columns: collation: comment: - *9 - - *25 + - *18 creators_filters: - - *24 - - *19 + - *25 + - *20 entropy_configurations: - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: @@ -760,7 +795,7 @@ columns: comment: - *5 - *6 - - *18 + - *19 - *9 events: - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column @@ -773,9 +808,9 @@ columns: default_function: collation: comment: - - *23 - - *5 - *24 + - *5 + - *25 - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: eventable_id @@ -818,7 +853,7 @@ columns: - *9 filters: - *5 - - *24 + - *25 - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: fields @@ -842,7 +877,7 @@ columns: comment: - *9 filters_stages: - - *19 + - *20 - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: stage_id @@ -854,7 +889,7 @@ columns: collation: comment: filters_tags: - - *19 + - *20 - &34 !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: tag_id @@ -935,13 +970,13 @@ columns: - *28 - *29 - *9 - - *25 + - *18 pins: - - *22 + - *23 - *5 - *6 - *9 - - *25 + - *18 push_subscriptions: - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: @@ -986,7 +1021,7 @@ columns: default_function: collation: comment: - - *25 + - *18 reactions: - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: @@ -1103,7 +1138,7 @@ columns: collation: comment: - *9 - - *25 + - *18 search_results: - *5 - *6 @@ -1123,9 +1158,9 @@ columns: comment: - *9 - *30 - - *25 + - *18 steps: - - *22 + - *23 - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: completed @@ -1141,7 +1176,7 @@ columns: - *6 - *9 taggings: - - *22 + - *23 - *5 - *6 - *34 @@ -1207,11 +1242,11 @@ columns: comment: - *9 watches: - - *22 + - *23 - *5 - *6 - *9 - - *25 + - *18 - !ruby/object:ActiveRecord::ConnectionAdapters::SQLite3::Column auto_increment: name: watching @@ -1259,6 +1294,7 @@ primary_keys: active_storage_attachments: id active_storage_blobs: id active_storage_variant_records: id + ai_quotas: id ar_internal_metadata: key assignees_filters: assigners_filters: @@ -1307,6 +1343,7 @@ data_sources: active_storage_attachments: true active_storage_blobs: true active_storage_variant_records: true + ai_quotas: true ar_internal_metadata: true assignees_filters: true assigners_filters: true @@ -1538,6 +1575,39 @@ indexes: nulls_not_distinct: comment: valid: true + ai_quotas: + - !ruby/object:ActiveRecord::ConnectionAdapters::IndexDefinition + table: ai_quotas + name: index_ai_quotas_on_reset_at + unique: false + columns: + - reset_at + lengths: {} + orders: {} + opclasses: {} + where: + type: + using: + include: + nulls_not_distinct: + comment: + valid: true + - !ruby/object:ActiveRecord::ConnectionAdapters::IndexDefinition + table: ai_quotas + name: index_ai_quotas_on_user_id + unique: false + columns: + - user_id + lengths: {} + orders: {} + opclasses: {} + where: + type: + using: + include: + nulls_not_distinct: + comment: + valid: true ar_internal_metadata: [] assignees_filters: - !ruby/object:ActiveRecord::ConnectionAdapters::IndexDefinition diff --git a/test/fixtures/ai/quotas.yml b/test/fixtures/ai/quotas.yml new file mode 100644 index 000000000..66657f117 --- /dev/null +++ b/test/fixtures/ai/quotas.yml @@ -0,0 +1,11 @@ +kevin: + user: kevin + limit: 100_00_000_000 + used: 1053 + reset_at: <%= 7.days.from_now %> + +david: + user: david + limit: 20_00_000_000 + used: 0 + reset_at: <%= 3.days.from_now %> diff --git a/test/models/ai/quota/money_test.rb b/test/models/ai/quota/money_test.rb new file mode 100644 index 000000000..01498a30c --- /dev/null +++ b/test/models/ai/quota/money_test.rb @@ -0,0 +1,40 @@ +require "test_helper" + +class Ai::Quota::MoneyTest < ActiveSupport::TestCase + test "wrapping" do + money = Ai::Quota::Money.wrap("$5.42") + assert_equal 5_42_000_000, money.in_microcents, "Strings with numbers are treated as dollars" + + assert_raises TypeError do + Ai::Quota::Money.wrap("foobar") + end + + money = Ai::Quota::Money.wrap(5.42) + assert_equal 5_42_000_000, money.in_microcents, "Decimals are treated as dollars" + + money = Ai::Quota::Money.wrap(5) + assert_equal 5, money.in_microcents, "Integers are treated as microcents" + + money1 = Ai::Quota::Money.wrap("$5") + money2 = Ai::Quota::Money.wrap(money1) + assert_equal money1, money2, "Money can wrap itself" + + assert_raises ArgumentError do + Ai::Quota::Money.wrap(nil) + end + end + + test "conversions" do + money = Ai::Quota::Money.wrap("$0") + assert_equal 0.0, money.in_dollars + assert_equal 0, money.in_microcents + + money = Ai::Quota::Money.wrap("$1") + assert_equal 1, money.in_dollars + assert_equal 1_00_000_000, money.in_microcents + + money = Ai::Quota::Money.wrap("$5.42") + assert_equal 5.42, money.in_dollars + assert_equal 5_42_000_000, money.in_microcents + end +end diff --git a/test/models/ai/quota_test.rb b/test/models/ai/quota_test.rb new file mode 100644 index 000000000..9d175b80d --- /dev/null +++ b/test/models/ai/quota_test.rb @@ -0,0 +1,115 @@ +require "test_helper" + +class Ai::QuotaTest < ActiveSupport::TestCase + setup do + @quota = Ai::Quota.new(user: users(:jz), limit: "$100") + end + + test "create" do + assert @quota.save + + 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 + 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.reset + assert_equal Ai::Quota::Money.wrap("$0"), @quota.used + assert_in_delta 7.days.from_now, @quota.reset_at, 1.minute + + @quota.increment_usage("$10") + assert_equal Ai::Quota::Money.wrap("$10"), @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 + 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 + 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 = "$300" + assert @quota.over_limit? + + @quota.used = "$0" + @quota.ensure_under_limit + + @quota.used = "$300" + assert_raises Ai::Quota::UsageExceedsQuotaError do + @quota.ensure_under_limit + 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 + end +end diff --git a/test/models/conversation/response_generator/response_test.rb b/test/models/conversation/response_generator/response_test.rb index 606d535cd..40b4c287b 100644 --- a/test/models/conversation/response_generator/response_test.rb +++ b/test/models/conversation/response_generator/response_test.rb @@ -28,8 +28,5 @@ class Conversation::Message::ResponseGenerator::ResponseTest < ActiveSupport::Te # So the total is 594000 + 12000 micro-cents assert_equal 606000, response.cost_microcents - - # If we convert that to a decimal value we get 0.00606 cents - assert_equal "0.00606".to_d, response.cost end end diff --git a/test/models/conversation_test.rb b/test/models/conversation_test.rb index 82de42d49..2f1e108d7 100644 --- a/test/models/conversation_test.rb +++ b/test/models/conversation_test.rb @@ -61,31 +61,25 @@ class ConversationTest < ActiveSupport::TestCase assert conversation.ready?, "The conversation should switch back to ready after a response is made" end - test "cost calculation" do - conversation = conversations(:kevin) - - assert_equal "0.00001053".to_d, conversation.cost - end - test "cost limits" do conversation = conversations(:kevin) conversation.ask("Where does the planning office keep demolition notices?") conversation.respond( "In a locked filing cabinet in a disused lavatory", - cost_microcents: Conversation::Cost.convert_to_microcents("$3") + cost_microcents: Ai::Quota::Money.wrap("$3").in_microcents ) conversation.ask("What's the meaning of life?") - conversation.respond("42", cost_microcents: Conversation::Cost.convert_to_microcents("$120")) + conversation.respond("42", cost_microcents: Ai::Quota::Money.wrap("$120").in_microcents) - assert_raises Conversation::CostExceededError do + assert_raises Ai::Quota::UsageExceedsQuotaError do conversation.ask("Should you leave a house without a towel?") end travel 1.month conversation.ask("Should you leave a house without a towel?") - conversation.respond("Never", cost_microcents: Conversation::Cost.convert_to_microcents("$0.01")) + conversation.respond("Never", cost_microcents: Ai::Quota::Money.wrap("$0.01").in_microcents) end end diff --git a/test/models/user_test.rb b/test/models/user_test.rb index d10fc73bc..e9670d0d9 100644 --- a/test/models/user_test.rb +++ b/test/models/user_test.rb @@ -60,4 +60,13 @@ class UserTest < ActiveSupport::TestCase assert user.conversation assert_equal user.conversation, conversation end + + test "AI quota" do + user = users(:jz) + + assert_nil user.ai_quota + quota = user.fetch_or_create_ai_quota + assert user.ai_quota + assert_equal user.ai_quota, quota + end end