diff --git a/app/models/concerns/captain_featurable.rb b/app/models/concerns/captain_featurable.rb index af73fded3..7b51daa13 100644 --- a/app/models/concerns/captain_featurable.rb +++ b/app/models/concerns/captain_featurable.rb @@ -27,6 +27,26 @@ module CaptainFeaturable }.with_indifferent_access end + # TODO: Once we support multiple LLM providers, ensure provider hook is correctly looked up + # Currently works because the only model/provider is OpenAI + def openai_hook + @openai_hook ||= hooks.find_by(app_id: 'openai', status: 'enabled') + end + + # TODO: Once we support multiple LLM providers, ensure provider key is correctly looked up + # Currently works because the only model/provider is OpenAI + def using_openai_hook_key? + openai_hook&.settings&.dig('api_key').present? + end + + def captain_api_key + openai_hook&.settings&.dig('api_key') || captain_system_api_key + end + + def captain_system_api_key + InstallationConfig.find_by(name: 'CAPTAIN_OPEN_AI_API_KEY')&.value + end + private def captain_models_with_defaults diff --git a/enterprise/lib/enterprise/captain/base_task_service.rb b/enterprise/lib/enterprise/captain/base_task_service.rb index bda4708af..c3514758e 100644 --- a/enterprise/lib/enterprise/captain/base_task_service.rb +++ b/enterprise/lib/enterprise/captain/base_task_service.rb @@ -28,7 +28,7 @@ module Enterprise::Captain::BaseTaskService end def increment_usage - credits = using_hook_key? ? 1 : Llm::Models.credit_multiplier_for(configured_model) + credits = account.using_openai_hook_key? ? 1 : Llm::Models.credit_multiplier_for(configured_model) Rails.logger.info("[CAPTAIN][#{self.class.name}] Incrementing response usage for account #{account.id} by #{credits} credits") account.increment_response_usage(credits: credits) end diff --git a/lib/captain/base_task_service.rb b/lib/captain/base_task_service.rb index e14048451..a3583c8f3 100644 --- a/lib/captain/base_task_service.rb +++ b/lib/captain/base_task_service.rb @@ -59,7 +59,7 @@ class Captain::BaseTaskService end def execute_ruby_llm_request(model:, messages:) - Llm::Config.with_api_key(api_key, api_base: api_base) do |context| + Llm::Config.with_api_key(account.captain_api_key, api_base: api_base) do |context| chat = context.chat(model: model) system_msg = messages.find { |m| m[:role] == 'system' } chat.with_instructions(system_msg[:content]) if system_msg @@ -132,26 +132,6 @@ class Captain::BaseTaskService messages end - def api_key - # TODO: Once we support multiple LLM providers, ensure provider key is correctly looked up - # Currently works because the only model/provider is OpenAI - @api_key ||= openai_hook&.settings&.dig('api_key') || system_api_key - end - - def openai_hook - @openai_hook ||= account.hooks.find_by(app_id: 'openai', status: 'enabled') - end - - def using_hook_key? - # TODO: Once we support multiple LLM providers, ensure provider key is correctly looked up - # Currently works because the only model/provider is OpenAI - openai_hook&.settings&.dig('api_key').present? - end - - def system_api_key - @system_api_key ||= InstallationConfig.find_by(name: 'CAPTAIN_OPEN_AI_API_KEY')&.value - end - def prompt_from_file(file_name) Rails.root.join('lib/integrations/openai/openai_prompts', "#{file_name}.liquid").read end