feat: add model based usage helper

This commit is contained in:
Shivam Mishra
2026-01-14 11:31:12 +05:30
parent ce822e1e58
commit c179393a0b
2 changed files with 29 additions and 0 deletions
@@ -21,6 +21,12 @@ module Enterprise::Account::PlanUsageAndLimits
save
end
def increment_response_usage_for_model(model)
credits = using_openai_hook_key? ? 1 : Llm::Models.credit_multiplier_for(model)
Rails.logger.info("[CAPTAIN] Incrementing response usage for account #{id} by #{credits} credits (model: #{model})")
increment_response_usage(credits: credits)
end
def reset_response_usage
custom_attributes[CAPTAIN_RESPONSES_USAGE] = 0
save
@@ -28,4 +28,27 @@ RSpec.describe Enterprise::Account::PlanUsageAndLimits do
expect(account.custom_attributes['captain_responses_usage']).to eq(2)
end
end
describe '#increment_response_usage_for_model' do
it 'increments by model credit multiplier when using system key' do
# gpt-5.1 has credit_multiplier of 2
account.increment_response_usage_for_model('gpt-5.1')
account.reload
expect(account.custom_attributes['captain_responses_usage']).to eq(2)
end
it 'increments by 1 when using hook key' do
create(:integrations_hook, account: account, app_id: 'openai', status: 'enabled', settings: { 'api_key' => 'user-key' })
# gpt-5.1 has credit_multiplier of 2, but hook key overrides to 1
account.increment_response_usage_for_model('gpt-5.1')
account.reload
expect(account.custom_attributes['captain_responses_usage']).to eq(1)
end
it 'defaults to 1 credit for unknown models' do
account.increment_response_usage_for_model('unknown-model')
account.reload
expect(account.custom_attributes['captain_responses_usage']).to eq(1)
end
end
end