From 3b366f43e65efff1bc57a338b5ee830d720c61bf Mon Sep 17 00:00:00 2001 From: Shivam Mishra Date: Thu, 23 Jan 2025 01:23:18 +0530 Subject: [PATCH] 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. --- app/fields/enterprise/account_limits_field.rb | 2 +- config/installation_config.yml | 6 + config/locales/en.yml | 3 +- .../accounts/captain/documents_controller.rb | 2 + .../v1/accounts/conversations_controller.rb | 1 + .../super_admin/app_configs_controller.rb | 3 +- enterprise/app/helpers/captain/chat_helper.rb | 1 - .../conversation/response_builder_job.rb | 3 + .../tools/simple_page_crawl_parser_job.rb | 14 ++ enterprise/app/models/captain/document.rb | 13 ++ enterprise/app/models/enterprise/account.rb | 82 +++++++++- enterprise/app/models/enterprise/inbox.rb | 10 +- .../services/captain/copilot/chat_service.rb | 6 +- .../billing/handle_stripe_event_service.rb | 5 + .../hook_execution_service.rb | 14 ++ enterprise/listeners/captain_listener.rb | 3 +- .../captain/documents_controller_spec.rb | 22 ++- .../conversation/response_builder_job_spec.rb | 34 ++++ spec/enterprise/models/account_spec.rb | 147 ++++++++++++++---- .../captain/copilot/chat_service_spec.rb | 34 ++++ .../handle_stripe_event_service_spec.rb | 43 +++-- spec/models/account_spec.rb | 3 +- 22 files changed, 394 insertions(+), 57 deletions(-) create mode 100644 spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb create mode 100644 spec/enterprise/services/captain/copilot/chat_service_spec.rb diff --git a/app/fields/enterprise/account_limits_field.rb b/app/fields/enterprise/account_limits_field.rb index 5833952ad..b014435c6 100644 --- a/app/fields/enterprise/account_limits_field.rb +++ b/app/fields/enterprise/account_limits_field.rb @@ -2,6 +2,6 @@ require 'administrate/field/base' class Enterprise::AccountLimitsField < Administrate::Field::Base def to_s - data.present? ? data.to_json : { agents: nil, inboxes: nil }.to_json + data.present? ? data.to_json : { agents: nil, inboxes: nil, captain_responses: nil, captain_documents: nil }.to_json end end diff --git a/config/installation_config.yml b/config/installation_config.yml index 2b4241697..a3b15a191 100644 --- a/config/installation_config.yml +++ b/config/installation_config.yml @@ -142,6 +142,12 @@ display_title: 'OpenAI Model' description: 'The OpenAI model configured for use in Captain AI. Default: gpt-4o-mini' locked: false +- name: CAPTAIN_CLOUD_PLAN_LIMITS + display_title: 'Captain Cloud Plan Limits' + description: 'The limits for the Captain AI service for different plans' + value: + type: code + # End of Captain Config # ------- Chatwoot Internal Config for Cloud ----# diff --git a/config/locales/en.yml b/config/locales/en.yml index 310235d5e..c525074cf 100644 --- a/config/locales/en.yml +++ b/config/locales/en.yml @@ -230,7 +230,8 @@ en: name: 'Linear' description: 'Create issues in Linear directly from your conversation window. Alternatively, link existing Linear issues for a more streamlined and efficient issue tracking process.' captain: - copilot_error: 'Please connect an assistant to this inbox to use copilot' + copilot_error: 'Please connect an assistant to this inbox to use Copilot' + copilot_limit: 'You are out of Copilot credits. You can buy more credits from the billing section.' public_portal: search: search_placeholder: Search for article by title or body... diff --git a/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb b/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb index 9743c3892..594aa0642 100644 --- a/enterprise/app/controllers/api/v1/accounts/captain/documents_controller.rb +++ b/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 diff --git a/enterprise/app/controllers/enterprise/api/v1/accounts/conversations_controller.rb b/enterprise/app/controllers/enterprise/api/v1/accounts/conversations_controller.rb index f22812f19..e9113f4d0 100644 --- a/enterprise/app/controllers/enterprise/api/v1/accounts/conversations_controller.rb +++ b/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, diff --git a/enterprise/app/controllers/enterprise/super_admin/app_configs_controller.rb b/enterprise/app/controllers/enterprise/super_admin/app_configs_controller.rb index 8cf8e7258..b11aab487 100644 --- a/enterprise/app/controllers/enterprise/super_admin/app_configs_controller.rb +++ b/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 diff --git a/enterprise/app/helpers/captain/chat_helper.rb b/enterprise/app/helpers/captain/chat_helper.rb index 258328f9d..a10c9e797 100644 --- a/enterprise/app/helpers/captain/chat_helper.rb +++ b/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 diff --git a/enterprise/app/jobs/captain/conversation/response_builder_job.rb b/enterprise/app/jobs/captain/conversation/response_builder_job.rb index 6b6570f0b..c8c511a57 100644 --- a/enterprise/app/jobs/captain/conversation/response_builder_job.rb +++ b/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 diff --git a/enterprise/app/jobs/captain/tools/simple_page_crawl_parser_job.rb b/enterprise/app/jobs/captain/tools/simple_page_crawl_parser_job.rb index fa658dd80..ccf34adf2 100644 --- a/enterprise/app/jobs/captain/tools/simple_page_crawl_parser_job.rb +++ b/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 diff --git a/enterprise/app/models/captain/document.rb b/enterprise/app/models/captain/document.rb index 11338d409..8331d291b 100644 --- a/enterprise/app/models/captain/document.rb +++ b/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 diff --git a/enterprise/app/models/enterprise/account.rb b/enterprise/app/models/enterprise/account.rb index bc07e1171..a1b0d0449 100644 --- a/enterprise/app/models/enterprise/account.rb +++ b/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 diff --git a/enterprise/app/models/enterprise/inbox.rb b/enterprise/app/models/enterprise/inbox.rb index 5f89a92a8..4cd627e41 100644 --- a/enterprise/app/models/enterprise/inbox.rb +++ b/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 diff --git a/enterprise/app/services/captain/copilot/chat_service.rb b/enterprise/app/services/captain/copilot/chat_service.rb index 6bcda66fa..019136540 100644 --- a/enterprise/app/services/captain/copilot/chat_service.rb +++ b/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 diff --git a/enterprise/app/services/enterprise/billing/handle_stripe_event_service.rb b/enterprise/app/services/enterprise/billing/handle_stripe_event_service.rb index 1ddf00d7a..c3c570596 100644 --- a/enterprise/app/services/enterprise/billing/handle_stripe_event_service.rb +++ b/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 diff --git a/enterprise/app/services/enterprise/message_templates/hook_execution_service.rb b/enterprise/app/services/enterprise/message_templates/hook_execution_service.rb index 858b5e903..92ccba553 100644 --- a/enterprise/app/services/enterprise/message_templates/hook_execution_service.rb +++ b/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 diff --git a/enterprise/listeners/captain_listener.rb b/enterprise/listeners/captain_listener.rb index da77bfbaa..adf156562 100644 --- a/enterprise/listeners/captain_listener.rb +++ b/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? diff --git a/spec/enterprise/controllers/api/v1/accounts/captain/documents_controller_spec.rb b/spec/enterprise/controllers/api/v1/accounts/captain/documents_controller_spec.rb index f9385af1c..ef0aa6b2e 100644 --- a/spec/enterprise/controllers/api/v1/accounts/captain/documents_controller_spec.rb +++ b/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 diff --git a/spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb b/spec/enterprise/jobs/captain/conversation/response_builder_job_spec.rb new file mode 100644 index 000000000..1e4a6e824 --- /dev/null +++ b/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 diff --git a/spec/enterprise/models/account_spec.rb b/spec/enterprise/models/account_spec.rb index 8e2681a9b..e36a7217d 100644 --- a/spec/enterprise/models/account_spec.rb +++ b/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 diff --git a/spec/enterprise/services/captain/copilot/chat_service_spec.rb b/spec/enterprise/services/captain/copilot/chat_service_spec.rb new file mode 100644 index 000000000..c4b759c85 --- /dev/null +++ b/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 diff --git a/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb b/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb index 36c163f79..d6d2e187e 100644 --- a/spec/enterprise/services/enterprise/billing/handle_stripe_event_service_spec.rb +++ b/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', diff --git a/spec/models/account_spec.rb b/spec/models/account_spec.rb index aa0556492..4b200fb9a 100644 --- a/spec/models/account_spec.rb +++ b/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