feat: setup captain limits (#10713)

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.

spec/enterprise/controllers/api/v1/accounts/captain/documents_controller_spec.rb

@@ -1,12 +1,17 @@
 require 'rails_helper'
 
 RSpec.describe 'Api::V1::Accounts::Captain::Documents', type: :request do
-  let(:account) { create(:account) }
+  let(:account) { create(:account, custom_attributes: { plan_name: 'startups' }) }
   let(:admin) { create(:user, account: account, role: :administrator) }
   let(:agent) { create(:user, account: account, role: :agent) }
   let(:assistant) { create(:captain_assistant, account: account) }
   let(:assistant2) { create(:captain_assistant, account: account) }
   let(:document) { create(:captain_document, assistant: assistant, account: account) }
+  let(:captain_limits) do
+    {
+      :startups => { :documents => 1, :responses => 100 }
+    }.with_indifferent_access
+  end
 
   def json_response
     JSON.parse(response.body, symbolize_names: true)
@@ -212,6 +217,21 @@ RSpec.describe 'Api::V1::Accounts::Captain::Documents', type: :request do
           expect(response).to have_http_status(:unprocessable_entity)
         end
       end
+
+      context 'with limits exceeded' do
+        before do
+          create_list(:captain_document, 5, assistant: assistant, account: account)
+
+          create(:installation_config, name: 'CAPTAIN_CLOUD_PLAN_LIMITS', value: captain_limits.to_json)
+          post "/api/v1/accounts/#{account.id}/captain/documents",
+               params: valid_attributes,
+               headers: admin.create_new_auth_token
+        end
+
+        it 'returns an error' do
+          expect(response).to have_http_status(:unprocessable_entity)
+        end
+      end
     end
   end
 

spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb

@@ -0,0 +1,34 @@
+require 'rails_helper'
+
+RSpec.describe Captain::Conversation::ResponseBuilderJob, type: :job 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) }
+
+  describe '#perform' do
+    let(:conversation) { create(:conversation, inbox: inbox, account: account) }
+    let(:mock_llm_chat_service) { instance_double(Captain::Llm::AssistantChatService) }
+
+    before do
+      create(:message, conversation: conversation, content: 'Hello', message_type: :incoming)
+
+      allow(inbox).to receive(:captain_active?).and_return(true)
+      allow(Captain::Llm::AssistantChatService).to receive(:new).and_return(mock_llm_chat_service)
+      allow(mock_llm_chat_service).to receive(:generate_response).and_return({ 'response' => 'Hey, welcome to Captain Specs' })
+    end
+
+    it 'generates and processes response' do
+      described_class.perform_now(conversation, assistant)
+      expect(conversation.messages.count).to eq(2)
+      expect(conversation.messages.outgoing.count).to eq(1)
+      expect(conversation.messages.last.content).to eq('Hey, welcome to Captain Specs')
+    end
+
+    it 'increments usage response' do
+      described_class.perform_now(conversation, assistant)
+      account.reload
+      expect(account.usage_limits[:captain][:responses][:consumed]).to eq(1)
+    end
+  end
+end

spec/enterprise/models/account_spec.rb

@@ -27,12 +27,120 @@ RSpec.describe Account, type: :model do
     end
   end
 
-  describe 'usage_limits' do
+  context 'with usage_limits' do
+    let(:captain_limits) do
+      {
+        :startups => { :documents => 100, :responses => 100 },
+        :business => { :documents => 200, :responses => 300 },
+        :enterprise => { :documents => 300, :responses => 500 }
+      }.with_indifferent_access
+    end
+    let(:account) { create(:account, { custom_attributes: { plan_name: 'startups' } }) }
+    let(:assistant) { create(:captain_assistant, account: account) }
+
     before do
       create(:installation_config, name: 'ACCOUNT_AGENTS_LIMIT', value: 20)
     end
 
-    let!(:account) { create(:account) }
+    describe 'when captain limits are configured' do
+      before do
+        create_list(:captain_document, 3, account: account, assistant: assistant, status: :available)
+        create(:installation_config, name: 'CAPTAIN_CLOUD_PLAN_LIMITS', value: captain_limits.to_json)
+      end
+
+      ## Document
+      it 'updates document count accurately' do
+        account.update_document_usage
+        expect(account.custom_attributes['captain_documents_usage']).to eq(3)
+      end
+
+      it 'handles zero documents' do
+        account.captain_documents.destroy_all
+        account.update_document_usage
+        expect(account.custom_attributes['captain_documents_usage']).to eq(0)
+      end
+
+      it 'reflects document limits' do
+        document_limits = account.usage_limits[:captain][:documents]
+
+        expect(document_limits[:consumed]).to eq 3
+        expect(document_limits[:current_available]).to eq captain_limits[:startups][:documents] - 3
+      end
+
+      ## Responses
+      it 'incrementing responses updates usage_limits' do
+        account.increment_response_usage
+
+        responses_limits = account.usage_limits[:captain][:responses]
+
+        expect(account.custom_attributes['captain_responses_usage']).to eq 1
+        expect(responses_limits[:consumed]).to eq 1
+        expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - 1
+      end
+
+      it 'reseting responses limits updates usage_limits' do
+        account.custom_attributes['captain_responses_usage'] = 30
+        account.save!
+
+        responses_limits = account.usage_limits[:captain][:responses]
+
+        expect(responses_limits[:consumed]).to eq 30
+        expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses] - 30
+
+        account.reset_response_usage
+        responses_limits = account.usage_limits[:captain][:responses]
+
+        expect(account.custom_attributes['captain_responses_usage']).to eq 0
+        expect(responses_limits[:consumed]).to eq 0
+        expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses]
+      end
+
+      it 'returns monthly limit accurately' do
+        %w[startups business enterprise].each do |plan|
+          account.custom_attributes = { 'plan_name': plan }
+          account.save!
+          expect(account.captain_monthly_limit).to eq captain_limits[plan]
+        end
+      end
+
+      it 'current_available is never out of bounds' do
+        account.custom_attributes['captain_responses_usage'] = 3000
+        account.save!
+
+        responses_limits = account.usage_limits[:captain][:responses]
+        expect(responses_limits[:consumed]).to eq 3000
+        expect(responses_limits[:current_available]).to eq 0
+
+        account.custom_attributes['captain_responses_usage'] = -100
+        account.save!
+
+        responses_limits = account.usage_limits[:captain][:responses]
+        expect(responses_limits[:consumed]).to eq 0
+        expect(responses_limits[:current_available]).to eq captain_limits[:startups][:responses]
+      end
+    end
+
+    describe 'when captain limits are not configured' do
+      it 'returns default values' do
+        account.custom_attributes = { 'plan_name': 'unknown' }
+        expect(account.captain_monthly_limit).to eq(
+          { documents: ChatwootApp.max_limit, responses: ChatwootApp.max_limit }.with_indifferent_access
+        )
+      end
+    end
+
+    describe 'when limits are configured for an account' do
+      before do
+        create(:installation_config, name: 'CAPTAIN_CLOUD_PLAN_LIMITS', value: captain_limits.to_json)
+        account.update(limits: { captain_documents: 5555, captain_responses: 9999 })
+      end
+
+      it 'returns limits based on custom attributes' do
+        usage_limits = account.usage_limits
+        expect(usage_limits[:captain][:documents][:total_count]).to eq(5555)
+        expect(usage_limits[:captain][:responses][:total_count]).to eq(9999)
+      end
+    end
 
     describe 'audit logs' do
       it 'returns audit logs' do
@@ -47,54 +155,29 @@ RSpec.describe Account, type: :model do
     end
 
     it 'returns max limits from global config when enterprise version' do
-      expect(account.usage_limits).to eq(
-        {
-          agents: 20,
-          inboxes: ChatwootApp.max_limit
-        }
-      )
+      expect(account.usage_limits[:agents]).to eq(20)
     end
 
     it 'returns max limits from account when enterprise version' do
       account.update(limits: { agents: 10 })
-      expect(account.usage_limits).to eq(
-        {
-          agents: 10,
-          inboxes: ChatwootApp.max_limit
-        }
-      )
+      expect(account.usage_limits[:agents]).to eq(10)
     end
 
     it 'returns limits based on subscription' do
       account.update(limits: { agents: 10 }, custom_attributes: { subscribed_quantity: 5 })
-      expect(account.usage_limits).to eq(
-        {
-          agents: 5,
-          inboxes: ChatwootApp.max_limit
-        }
-      )
+      expect(account.usage_limits[:agents]).to eq(5)
     end
 
     it 'returns max limits from global config if account limit is absent' do
       account.update(limits: { agents: '' })
-      expect(account.usage_limits).to eq(
-        {
-          agents: 20,
-          inboxes: ChatwootApp.max_limit
-        }
-      )
+      expect(account.usage_limits[:agents]).to eq(20)
     end
 
     it 'returns max limits from app limit if account limit and installation config is absent' do
       account.update(limits: { agents: '' })
       InstallationConfig.where(name: 'ACCOUNT_AGENTS_LIMIT').update(value: '')
 
-      expect(account.usage_limits).to eq(
-        {
-          agents: ChatwootApp.max_limit,
-          inboxes: ChatwootApp.max_limit
-        }
-      )
+      expect(account.usage_limits[:agents]).to eq(ChatwootApp.max_limit)
     end
   end
 

spec/enterprise/services/captain/copilot/chat_service_spec.rb

@@ -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

spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb

@@ -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',

spec/models/account_spec.rb

@@ -43,7 +43,8 @@ RSpec.describe Account do
     let(:account) { create(:account) }
 
     it 'returns ChatwootApp.max limits' do
-      expect(account.usage_limits).to eq({ agents: ChatwootApp.max_limit, inboxes: ChatwootApp.max_limit })
+      expect(account.usage_limits[:agents]).to eq(ChatwootApp.max_limit)
+      expect(account.usage_limits[:inboxes]).to eq(ChatwootApp.max_limit)
     end
   end