LeadMagnet/leadchat
1
0
Fork 0
Code Issues Pull Requests Actions 3 Packages Projects Releases Wiki Activity

feat: setup captain limits (#10713)

Browse Source 
This pull request introduces several changes to implement and manage
usage limits for the Captain AI service. The key changes include adding
configuration for plan limits, updating error messages, modifying
controllers and models to handle usage limits, and updating tests to
ensure the new functionality works correctly.

## Implementation Checklist

- [x] Ability to configure captain limits per check
- [x] Update response for `usage_limits` to include captain limits
- [x] Methods to increment or reset captain responses limits in the
`limits` column for the `Account` model
- [x] Check documents limit using a count query
- [x] Ensure Captain hand-off if a limit is reached
- [x] Ensure limits are enforced for Copilot Chat
- [x] Ensure limits are reset when stripe webhook comes in 
- [x] Increment usage for FAQ generation and Contact notes
- [x] Ensure documents limit is enforced

These changes ensure that the Captain AI service operates within the defined usage limits for different subscription plans, providing appropriate error messages and handling when limits are exceeded.
This commit is contained in:
Shivam Mishra
2025-01-23 01:23:18 +05:30
committed by GitHub
parent 52362ec1ea
commit 3b366f43e6
22 changed files with 394 additions and 57 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
spec/enterprise/services/captain/copilot/chat_service_spec.rb Normal file
View File

@@ -0,0 +1,34 @@
require 'rails_helper'
RSpec.describe Captain::Copilot::ChatService do
let(:account) { create(:account, custom_attributes: { plan_name: 'startups' }) }
let(:inbox) { create(:inbox, account: account) }
let(:assistant) { create(:captain_assistant, account: account) }
let(:captain_inbox_association) { create(:captain_inbox, captain_assistant: assistant, inbox: inbox) }
let(:mock_captain_agent) { instance_double(Captain::Agent) }
let(:mock_captain_tool) { instance_double(Captain::Tool) }
let(:mock_openai_client) { instance_double(OpenAI::Client) }
describe '#execute' do
before do
create(:installation_config) { create(:installation_config, name: 'CAPTAIN_OPEN_AI_API_KEY', value: 'test-key') }
allow(OpenAI::Client).to receive(:new).and_return(mock_openai_client)
allow(mock_openai_client).to receive(:chat).and_return({ choices: [{ message: { content: '{ "result": "Hey" }' } }] }.with_indifferent_access)
allow(Captain::Agent).to receive(:new).and_return(mock_captain_agent)
allow(mock_captain_agent).to receive(:execute).and_return(true)
allow(mock_captain_agent).to receive(:register_tool).and_return(true)
allow(Captain::Tool).to receive(:new).and_return(mock_captain_tool)
allow(mock_captain_tool).to receive(:register_method).and_return(true)
allow(account).to receive(:increment_response_usage).and_return(true)
end
it 'increments usage' do
described_class.new(assistant, { previous_messages: ['Hello'], conversation_history: 'Hi' }).generate_response('Hey')
expect(account).to have_received(:increment_response_usage).once
end
end
end

43
spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb
View File

@@ -37,19 +37,34 @@ describe Enterprise::Billing::HandleStripeEventService do
end
describe '#perform' do
it 'handle customer.subscription.updated' do
allow(event).to receive(:type).and_return('customer.subscription.updated')
allow(subscription).to receive(:customer).and_return('cus_123')
stripe_event_service.new.perform(event: event)
expect(account.reload.custom_attributes).to eq({
'stripe_customer_id' => 'cus_123',
'stripe_price_id' => 'test',
'stripe_product_id' => 'plan_id',
'plan_name' => 'Hacker',
'subscribed_quantity' => '10',
'subscription_ends_on' => Time.zone.at(1_686_567_520).as_json,
'subscription_status' => 'active'
})
context 'when it gets customer.subscription.updated event' do
it 'updates subscription attributes' do
allow(event).to receive(:type).and_return('customer.subscription.updated')
allow(subscription).to receive(:customer).and_return('cus_123')
stripe_event_service.new.perform(event: event)
expect(account.reload.custom_attributes).to eq({
'captain_responses_usage' => 0,
'stripe_customer_id' => 'cus_123',
'stripe_price_id' => 'test',
'stripe_product_id' => 'plan_id',
'plan_name' => 'Hacker',
'subscribed_quantity' => '10',
'subscription_ends_on' => Time.zone.at(1_686_567_520).as_json,
'subscription_status' => 'active'
})
end
it 'resets captain usage' do
5.times { account.increment_response_usage }
expect(account.custom_attributes['captain_responses_usage']).to eq(5)
allow(event).to receive(:type).and_return('customer.subscription.updated')
allow(subscription).to receive(:customer).and_return('cus_123')
stripe_event_service.new.perform(event: event)
expect(account.reload.custom_attributes['captain_responses_usage']).to eq(0)
end
end
it 'disable features on customer.subscription.updated for default plan' do
@@ -57,6 +72,7 @@ describe Enterprise::Billing::HandleStripeEventService do
allow(subscription).to receive(:customer).and_return('cus_123')
stripe_event_service.new.perform(event: event)
expect(account.reload.custom_attributes).to eq({
'captain_responses_usage' => 0,
'stripe_customer_id' => 'cus_123',
'stripe_price_id' => 'test',
'stripe_product_id' => 'plan_id',
@@ -96,6 +112,7 @@ describe Enterprise::Billing::HandleStripeEventService do
allow(subscription).to receive(:customer).and_return('cus_123')
stripe_event_service.new.perform(event: event)
expect(account.reload.custom_attributes).to eq({
'captain_responses_usage' => 0,
'stripe_customer_id' => 'cus_123',
'stripe_price_id' => 'test',
'stripe_product_id' => 'plan_id_2',