diff --git a/app/assets/stylesheets/conversation.css b/app/assets/stylesheets/conversation.css index 95641849c..f5b0bbe8a 100644 --- a/app/assets/stylesheets/conversation.css +++ b/app/assets/stylesheets/conversation.css @@ -20,6 +20,7 @@ .conversation__message { display: flex; + flex-direction: column; .action-text-content { border-radius: 0.5rem; @@ -51,7 +52,7 @@ } .conversation__message--assistant { - justify-content: flex-start; + align-items: flex-start; .action-text-content { margin-inline-end: 15%; @@ -67,8 +68,15 @@ color: var(--color-negative); } + .conversation__message--failed[data-error]::after { + color: var(--color-negative); + content: attr(data-error); + display: block; + padding: var(--inline-space); + } + .conversation__message--user { - justify-content: flex-end; + align-items: flex-end; .action-text-content { background-color: var(--color-selected); diff --git a/app/controllers/conversations/messages_controller.rb b/app/controllers/conversations/messages_controller.rb index 0a6ec8b31..b7585ac72 100644 --- a/app/controllers/conversations/messages_controller.rb +++ b/app/controllers/conversations/messages_controller.rb @@ -7,6 +7,10 @@ class Conversations::MessagesController < ApplicationController def create @conversation.ask(question, **message_params) + rescue Ai::Quota::UsageExceedsQuotaError + render json: { error: "You've depleted your quota" }, status: :too_many_requests + rescue Conversation::InvalidStateError + render json: { error: "Fizzy is still working on an answer to your last question" }, status: :conflict end private diff --git a/app/javascript/controllers/conversation/composer_controller.js b/app/javascript/controllers/conversation/composer_controller.js index 368716846..0f7e83223 100644 --- a/app/javascript/controllers/conversation/composer_controller.js +++ b/app/javascript/controllers/conversation/composer_controller.js @@ -27,7 +27,10 @@ export default class extends Controller { if (event.detail.success) { this.#reset() } else { - this.conversationMessagesOutlet.failPendingMessage(this.clientMessageIdInputTarget.value) + this.#failPendingMessage( + this.clientMessageIdInputTarget.value, + event.detail.fetchResponse + ) } } @@ -60,4 +63,18 @@ export default class extends Controller { this.textInputTarget.value = "" this.clientMessageIdInputTarget.value = "" } + + async #failPendingMessage(clientMessageId, response) { + let errorMessage = null + + if (response?.contentType?.includes('application/json')) { + const error = JSON.parse(await response.responseText) + errorMessage = error.error + } + + this.conversationMessagesOutlet.failPendingMessage( + this.clientMessageIdInputTarget.value, + errorMessage + ) + } } diff --git a/app/javascript/controllers/conversation/messages_controller.js b/app/javascript/controllers/conversation/messages_controller.js index 4d484ee2d..32e3e58e9 100644 --- a/app/javascript/controllers/conversation/messages_controller.js +++ b/app/javascript/controllers/conversation/messages_controller.js @@ -12,11 +12,15 @@ export default class extends Controller { this.scrollToBottom() } - failPendingMessage(clientMessageId) { + failPendingMessage(clientMessageId, errorMessage) { const message = this.element.querySelector(`#message_${clientMessageId}`) if (message) { message.classList.add(this.failedMessageClass) + + if (errorMessage) { + message.dataset.error = errorMessage + } } } 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 new file mode 100644 index 000000000..955a855af --- /dev/null +++ b/app/models/ai/quota.rb @@ -0,0 +1,50 @@ +class Ai::Quota < ApplicationRecord + class UsageExceedsQuotaError < StandardError; end + + 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 spend(cost) + cost = Money.wrap(cost) + + transaction do + reset_if_due + increment!(:used, cost.in_microcents) + end + end + + def ensure_not_depleted + reset_if_due + + if depleted? + raise UsageExceedsQuotaError + end + 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 new file mode 100644 index 000000000..f229c8def --- /dev/null +++ b/app/models/ai/quota/money.rb @@ -0,0 +1,42 @@ +class Ai::Quota::Money < Data.define(:value) + 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 + + private + 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 +end diff --git a/app/models/conversation.rb b/app/models/conversation.rb index 184394fb1..bf3a4a0ab 100644 --- a/app/models/conversation.rb +++ b/app/models/conversation.rb @@ -8,16 +8,9 @@ class Conversation < ApplicationRecord enum :state, %w[ ready thinking ].index_by(&:itself), default: :ready - def cost - messages.where.not(cost_microcents: nil).sum(:cost_microcents).to_d / 100_000 - end - - def clear - messages.delete_all - touch - end - def ask(question, **attributes) + 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? thinking! @@ -25,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.spend_ai_quota(message.cost) if message.cost + + message end private diff --git a/app/models/conversation/message.rb b/app/models/conversation/message.rb index dc817c500..72ea37b26 100644 --- a/app/models/conversation/message.rb +++ b/app/models/conversation/message.rb @@ -13,6 +13,10 @@ class Conversation::Message < ApplicationRecord scope :ordered, -> { order(created_at: :asc, id: :asc) } + def cost + cost_microcents && Ai::Quota::Money.new(cost_microcents) + end + def all_emoji? content.to_plain_text.all_emoji? end diff --git a/app/models/conversation/message/response_generator/response.rb b/app/models/conversation/message/response_generator/response.rb index f404835a3..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 - cost_microcents.to_d / MICROCENTS_PER_DOLLAR - 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 - single_token_price * MICROCENTS_PER_DOLLAR + 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..b34879fb9 --- /dev/null +++ b/app/models/user/ai_quota.rb @@ -0,0 +1,22 @@ +module User::AiQuota + extend ActiveSupport::Concern + + DEFAULT_QUOTA = Ai::Quota::Money.wrap("$100").in_microcents + + included do + has_one :ai_quota, class_name: "Ai::Quota" + end + + def spend_ai_quota(cost) + fetch_or_create_ai_quota.spend(cost) + end + + def ensure_ai_quota_not_depleted + fetch_or_create_ai_quota.ensure_not_depleted + end + + private + def fetch_or_create_ai_quota + ai_quota || create_ai_quota!(limit: DEFAULT_QUOTA) + 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/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..b37e60ee9 --- /dev/null +++ b/test/models/ai/quota_test.rb @@ -0,0 +1,63 @@ +require "test_helper" + +class Ai::QuotaTest < ActiveSupport::TestCase + setup do + @quota = Ai::Quota.new(user: users(:jz), limit: Ai::Quota::Money.wrap("$100").in_microcents) + 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").in_microcents, @quota.limit + end + + test "increment usage" do + @quota.save + + @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.used = 0 + + @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.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.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 + @quota.ensure_not_depleted + + @quota.used = Ai::Quota::Money.wrap("$300").in_microcents + assert_raises Ai::Quota::UsageExceedsQuotaError do + @quota.ensure_not_depleted + end + + travel 10.days + @quota.ensure_not_depleted + end +end diff --git a/test/models/conversation/response_generator/response_test.rb b/test/models/conversation/response_generator/response_test.rb index c33798b63..40b4c287b 100644 --- a/test/models/conversation/response_generator/response_test.rb +++ b/test/models/conversation/response_generator/response_test.rb @@ -13,23 +13,20 @@ class Conversation::Message::ResponseGenerator::ResponseTest < ActiveSupport::Te # and 60 USD per million output tokens # That's 0.00003 cents per input token and 0.00006 cents # per output token - # Which is 3 micro-cents per input token and 6 micro-cents + # Which is 3000 micro-cents per input token and 6000 micro-cents # per output token - assert_equal "3.0".to_d, response.input_token_price_microcents - assert_equal "6.0".to_d, response.output_token_price_microcents + assert_equal "3000.0".to_d, response.input_token_price_microcents + assert_equal "6000.0".to_d, response.output_token_price_microcents # We've got 198 input tokens, so that's - # 193 * 3 = 594 - assert_equal 594, response.input_cost_microcents + # 198 * 3000 = 594000 + assert_equal 594000, response.input_cost_microcents # We've got 2 output tokens, so that's - # 2 * 6 = 12 - assert_equal 12, response.output_cost_microcents + # 2 * 6000 = 12 + assert_equal 12000, response.output_cost_microcents - # So the total is 594 + 12 micro-cents - assert_equal 606, response.cost_microcents - - # If we convert that to a decimal value we get 0.00606 cents - assert_equal "0.00606".to_d, response.cost + # So the total is 594000 + 12000 micro-cents + assert_equal 606000, response.cost_microcents end end diff --git a/test/models/conversation_test.rb b/test/models/conversation_test.rb index 5691e7363..2f1e108d7 100644 --- a/test/models/conversation_test.rb +++ b/test/models/conversation_test.rb @@ -61,21 +61,25 @@ class ConversationTest < ActiveSupport::TestCase assert conversation.ready?, "The conversation should switch back to ready after a response is made" end - test "clearing conversation messages" do + test "cost limits" do conversation = conversations(:kevin) - assert conversation.messages.any?, "The conversation should have messages before clearing" + conversation.ask("Where does the planning office keep demolition notices?") + conversation.respond( + "In a locked filing cabinet in a disused lavatory", + cost_microcents: Ai::Quota::Money.wrap("$3").in_microcents + ) - original_updated_at = conversation.updated_at - conversation.clear - assert conversation.updated_at > original_updated_at, "The conversation's updated_at timestamp should change after clearing messages" + conversation.ask("What's the meaning of life?") + conversation.respond("42", cost_microcents: Ai::Quota::Money.wrap("$120").in_microcents) - assert conversation.messages.empty?, "All messages should be deleted when clearing the conversation" - end + assert_raises Ai::Quota::UsageExceedsQuotaError do + conversation.ask("Should you leave a house without a towel?") + end - test "cost calculation" do - conversation = conversations(:kevin) + travel 1.month - assert_equal "0.01053".to_d, conversation.cost + conversation.ask("Should you leave a house without a towel?") + conversation.respond("Never", cost_microcents: Ai::Quota::Money.wrap("$0.01").in_microcents) end end