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.

enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb

@@ -23,6 +23,8 @@ class Api::V1::Accounts::Captain::DocumentsController < Api::V1::Accounts::BaseC
 
     @document = @assistant.documents.build(document_params)
     @document.save!
+  rescue Captain::Document::LimitExceededError => e
+    render_could_not_create_error(e.message)
   end
 
   def destroy

enterprise/app/controllers/enterprise/api/v1/accounts/conversations_controller.rb

@@ -7,6 +7,7 @@ module Enterprise::Api::V1::Accounts::ConversationsController
   def copilot
     assistant = @conversation.inbox.captain_assistant
     return render json: { message: I18n.t('captain.copilot_error') } unless assistant
+    return render json: { message: I18n.t('captain.copilot_limit') } unless @conversation.inbox.captain_active?
 
     response = Captain::Copilot::ChatService.new(
       assistant,

enterprise/app/controllers/enterprise/super_admin/app_configs_controller.rb

@@ -32,6 +32,7 @@ module Enterprise::SuperAdmin::AppConfigsController
   end
 
   def internal_config_options
-    %w[CHATWOOT_INBOX_TOKEN CHATWOOT_INBOX_HMAC_KEY ANALYTICS_TOKEN CLEARBIT_API_KEY DASHBOARD_SCRIPTS BLOCKED_EMAIL_DOMAINS]
+    %w[CHATWOOT_INBOX_TOKEN CHATWOOT_INBOX_HMAC_KEY ANALYTICS_TOKEN CLEARBIT_API_KEY DASHBOARD_SCRIPTS BLOCKED_EMAIL_DOMAINS
+       CAPTAIN_CLOUD_PLAN_LIMITS]
   end
 end

enterprise/app/helpers/captain/chat_helper.rb

@@ -35,7 +35,6 @@ module Captain::ChatHelper
 
   def handle_response(response)
     message = response.dig('choices', 0, 'message')
-
     if message['tool_calls']
       process_tool_calls(message['tool_calls'])
     else

enterprise/app/jobs/captain/conversation/response_builder_job.rb

@@ -3,6 +3,7 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob
 
   def perform(conversation, assistant)
     @conversation = conversation
+    @inbox = conversation.inbox
     @assistant = assistant
 
     ActiveRecord::Base.transaction do
@@ -25,6 +26,8 @@ class Captain::Conversation::ResponseBuilderJob < ApplicationJob
     return process_action('handoff') if handoff_requested?
 
     create_messages
+    Rails.logger.info("[CAPTAIN][ResponseBuilderJob] Incrementing response usage for #{account.id}")
+    account.increment_response_usage
   end
 
   def collect_previous_messages

enterprise/app/jobs/captain/tools/simple_page_crawl_parser_job.rb

@@ -3,6 +3,13 @@ class Captain::Tools::SimplePageCrawlParserJob < ApplicationJob
 
   def perform(assistant_id:, page_link:)
     assistant = Captain::Assistant.find(assistant_id)
+    account = assistant.account
+
+    if limit_exceeded?(account)
+      Rails.logger.info("Document limit exceeded for #{assistant_id}")
+      return
+    end
+
     crawler = Captain::Tools::SimplePageCrawlService.new(page_link)
 
     page_title = crawler.page_title || ''
@@ -18,4 +25,11 @@ class Captain::Tools::SimplePageCrawlParserJob < ApplicationJob
   rescue StandardError => e
     raise "Failed to parse data: #{page_link} #{e.message}"
   end
+
+  private
+
+  def limit_exceeded?(account)
+    limits = account.usage_limits[:captain][:documents]
+    limits[:current_available].negative? || limits[:current_available].zero?
+  end
 end

enterprise/app/models/captain/document.rb

@@ -20,6 +20,7 @@
 #  index_captain_documents_on_status                          (status)
 #
 class Captain::Document < ApplicationRecord
+  class LimitExceededError < StandardError; end
   self.table_name = 'captain_documents'
 
   belongs_to :assistant, class_name: 'Captain::Assistant'
@@ -35,7 +36,10 @@ class Captain::Document < ApplicationRecord
     available: 1
   }
 
+  before_create :ensure_within_plan_limit
   after_create_commit :enqueue_crawl_job
+  after_create_commit :update_document_usage
+  after_destroy :update_document_usage
   after_commit :enqueue_response_builder_job
   scope :ordered, -> { order(created_at: :desc) }
 
@@ -56,7 +60,16 @@ class Captain::Document < ApplicationRecord
     Captain::Documents::ResponseBuilderJob.perform_later(self)
   end
 
+  def update_document_usage
+    account.update_document_usage
+  end
+
   def ensure_account_id
     self.account_id = assistant&.account_id
   end
+
+  def ensure_within_plan_limit
+    limits = account.usage_limits[:captain][:documents]
+    raise LimitExceededError, 'Document limit exceeded' unless limits[:current_available].positive?
+  end
 end

enterprise/app/models/enterprise/account.rb

@@ -1,11 +1,37 @@
 module Enterprise::Account
+  CAPTAIN_RESPONSES = 'captain_responses'.freeze
+  CAPTAIN_DOCUMENTS = 'captain_documents'.freeze
+  CAPTAIN_RESPONSES_USAGE = 'captain_responses_usage'.freeze
+  CAPTAIN_DOCUMENTS_USAGE = 'captain_documents_usage'.freeze
+
   def usage_limits
     {
       agents: agent_limits.to_i,
-      inboxes: get_limits(:inboxes).to_i
+      inboxes: get_limits(:inboxes).to_i,
+      captain: {
+        documents: get_captain_limits(:documents),
+        responses: get_captain_limits(:responses)
+      }
     }
   end
 
+  def increment_response_usage
+    current_usage = custom_attributes[CAPTAIN_RESPONSES_USAGE].to_i || 0
+    custom_attributes[CAPTAIN_RESPONSES_USAGE] = current_usage + 1
+    save
+  end
+
+  def reset_response_usage
+    custom_attributes[CAPTAIN_RESPONSES_USAGE] = 0
+    save
+  end
+
+  def update_document_usage
+    # this will ensure that the document count is always accurate
+    custom_attributes[CAPTAIN_DOCUMENTS_USAGE] = captain_documents.count
+    save
+  end
+
   def subscribed_features
     plan_features = InstallationConfig.find_by(name: 'CHATWOOT_CLOUD_PLAN_FEATURES')&.value
     return [] if plan_features.blank?
@@ -13,8 +39,58 @@ module Enterprise::Account
     plan_features[plan_name]
   end
 
+  def captain_monthly_limit
+    default_limits = default_captain_limits
+
+    {
+      documents: self[:limits][CAPTAIN_DOCUMENTS] || default_limits['documents'],
+      responses: self[:limits][CAPTAIN_RESPONSES] || default_limits['responses']
+    }.with_indifferent_access
+  end
+
   private
 
+  def get_captain_limits(type)
+    total_count = captain_monthly_limit[type.to_s].to_i
+
+    consumed = if type == :documents
+                 custom_attributes[CAPTAIN_DOCUMENTS_USAGE].to_i || 0
+               else
+                 custom_attributes[CAPTAIN_RESPONSES_USAGE].to_i || 0
+               end
+
+    consumed = 0 if consumed.negative?
+
+    {
+      total_count: total_count,
+      current_available: (total_count - consumed).clamp(0, total_count),
+      consumed: consumed
+    }
+  end
+
+  def default_captain_limits
+    max_limits = { documents: ChatwootApp.max_limit, responses: ChatwootApp.max_limit }.with_indifferent_access
+    zero_limits = { documents: 0, responses: 0 }.with_indifferent_access
+    plan_quota = InstallationConfig.find_by(name: 'CAPTAIN_CLOUD_PLAN_LIMITS')&.value
+
+    # If there are no limits configured, we allow max usage
+    return max_limits if plan_quota.blank?
+
+    # if there is plan_quota configred, but plan_name is not present, we return zero limits
+    return zero_limits if plan_name.blank?
+
+    begin
+      # Now we parse the plan_quota and return the limits for the plan name
+      # but if there's no plan_name present in the plan_quota, we return zero limits
+      plan_quota = JSON.parse(plan_quota) if plan_quota.present?
+      plan_quota[plan_name.downcase] || zero_limits
+    rescue StandardError
+      # if there's any error in parsing the plan_quota, we return max limits
+      # this is to ensure that we don't block the user from using the product
+      max_limits
+    end
+  end
+
   def plan_name
     custom_attributes['plan_name']
   end
@@ -41,7 +117,9 @@ module Enterprise::Account
       'type' => 'object',
       'properties' => {
         'inboxes' => { 'type': 'number' },
-        'agents' => { 'type': 'number' }
+        'agents' => { 'type': 'number' },
+        'captain_responses' => { 'type': 'number' },
+        'captain_documents' => { 'type': 'number' }
       },
       'required' => [],
       'additionalProperties' => false

enterprise/app/models/enterprise/inbox.rb

@@ -6,11 +6,19 @@ module Enterprise::Inbox
   end
 
   def active_bot?
-    super || captain_assistant.present?
+    super || captain_active?
+  end
+
+  def captain_active?
+    captain_assistant.present? && more_responses?
   end
 
   private
 
+  def more_responses?
+    account.usage_limits[:captain][:responses][:current_available].positive?
+  end
+
   def get_agent_ids_over_assignment_limit(limit)
     conversations.open.select(:assignee_id).group(:assignee_id).having("count(*) >= #{limit.to_i}").filter_map(&:assignee_id)
   end

enterprise/app/services/captain/copilot/chat_service.rb

@@ -15,7 +15,11 @@ class Captain::Copilot::ChatService < Captain::Llm::BaseOpenAiService
 
   def generate_response(input)
     @messages << { role: 'user', content: input } if input.present?
-    request_chat_completion
+    response = request_chat_completion
+    Rails.logger.info("[CAPTAIN][CopilotChatService] Incrementing response usage for #{@assistant.account.id}")
+    @assistant.account.increment_response_usage
+
+    response
   end
 
   private

enterprise/app/services/enterprise/billing/handle_stripe_event_service.rb

@@ -22,6 +22,7 @@ class Enterprise::Billing::HandleStripeEventService
     update_account_attributes(subscription, plan)
 
     change_plan_features
+    reset_captain_usage
   end
 
   def update_account_attributes(subscription, plan)
@@ -56,6 +57,10 @@ class Enterprise::Billing::HandleStripeEventService
     account.save!
   end
 
+  def reset_captain_usage
+    account.reset_response_usage
+  end
+
   def ensure_event_context(event)
     @event = event
   end

enterprise/app/services/enterprise/message_templates/hook_execution_service.rb

@@ -2,6 +2,7 @@ module Enterprise::MessageTemplates::HookExecutionService
   def trigger_templates
     super
     return unless should_process_captain_response?
+    return perform_handoff unless inbox.captain_active?
 
     Captain::Conversation::ResponseBuilderJob.perform_later(
       conversation,
@@ -12,4 +13,17 @@ module Enterprise::MessageTemplates::HookExecutionService
   def should_process_captain_response?
     conversation.pending? && message.incoming? && inbox.captain_assistant.present?
   end
+
+  def perform_handoff
+    return unless conversation.pending?
+
+    Rails.logger.info("Captain limit exceeded, performing handoff mid-conversation for conversation: #{conversation.id}")
+    conversation.messages.create!(
+      message_type: :outgoing,
+      account_id: conversation.account.id,
+      inbox_id: conversation.inbox.id,
+      content: 'Transferring to another agent for further assistance.'
+    )
+    conversation.bot_handoff!
+  end
 end

enterprise/listeners/captain_listener.rb

@@ -2,8 +2,7 @@ class CaptainListener < BaseListener
   def conversation_resolved(event)
     conversation = extract_conversation_and_account(event)[0]
     assistant = conversation.inbox.captain_assistant
-
-    return if assistant.blank?
+    return unless conversation.inbox.captain_active?
 
     Captain::Llm::ContactNotesService.new(assistant, conversation).generate_and_update_notes if assistant.config['feature_memory'].present?
     Captain::Llm::ConversationFaqService.new(assistant, conversation).generate_and_deduplicate if assistant.config['feature_faq'].present?