chatwoot/lib/captain/base_task_service.rb

class Captain::BaseTaskService
  include Integrations::LlmInstrumentation
  include Captain::ToolInstrumentation
  include Llm::ExceptionTrackable

  # gpt-4o-mini supports 128,000 tokens
  # 1 token is approx 4 characters
  # sticking with 120000 to be safe
  # 120000 * 4 = 480,000 characters (rounding off downwards to 400,000 to be safe)
  TOKEN_LIMIT = 400_000
  GPT_MODEL = Llm::Config::DEFAULT_MODEL

  # Prepend enterprise module to subclasses when they're defined.
  # This ensures the enterprise perform wrapper is applied even when
  # subclasses define their own perform method, since prepend puts
  # the module before the class in the ancestor chain.
  def self.inherited(subclass)
    super
    subclass.prepend_mod_with('Captain::BaseTaskService')
  end

  pattr_initialize [:account!, { conversation_display_id: nil }]

  private

  def event_name
    raise NotImplementedError, "#{self.class} must implement #event_name"
  end

  def conversation
    @conversation ||= account.conversations.find_by(display_id: conversation_display_id)
  end

  def api_base = Llm::Config.api_base_for(llm_provider)

  def make_api_call(messages:, model: nil, feature: nil, schema: nil, tools: [])
    llm_route = resolved_llm_route(model: model, feature: feature)

    # Community edition prerequisite checks
    # Enterprise module handles these with more specific error messages (cloud vs self-hosted)
    return { error: I18n.t('captain.disabled'), error_code: 403 } unless captain_tasks_enabled?
    return { error: I18n.t('captain.api_key_missing'), error_code: 401 } unless api_key_configured?(llm_route[:provider])

    @llm_provider = llm_route[:provider]
    model = llm_route[:model]
    request_tools = Llm::Config.supports_tools_and_schema?(llm_route[:provider]) ? tools : []
    instrumentation_params = build_instrumentation_params(model, messages, llm_route[:provider])
    instrumentation_method = request_tools.any? ? :instrument_tool_session : :instrument_llm_call

    response = send(instrumentation_method, instrumentation_params) do
      execute_ruby_llm_request(llm_route: llm_route, messages: messages, schema: schema, tools: request_tools)
    end

    return response unless build_follow_up_context? && response[:message].present?

    response.merge(follow_up_context: build_follow_up_context(messages, response))
  end

  def resolved_llm_route(model:, feature:)
    return explicit_model_route(model) if feature.blank?

    route = Llm::FeatureRouter.resolve(feature: feature, account: account)
    resolved_model = model.present? && route[:source] == :default ? model : route[:model]

    route.merge(model: resolved_model, provider: provider_for_model(resolved_model, route[:provider]))
  end

  def explicit_model_route(model)
    resolved_model = model.presence || GPT_MODEL
    { model: resolved_model, provider: provider_for_model(resolved_model), source: :explicit }
  end

  def provider_for_model(model, fallback_provider = Llm::Config::DEFAULT_PROVIDER) = Llm::Models.provider_for(model) || fallback_provider

  def execute_ruby_llm_request(llm_route:, messages:, schema: nil, tools: [])
    provider = llm_route[:provider]
    credential = llm_credential(provider)

    Llm::Config.with_provider(provider: provider, config_values: credential[:config_values]) do |context|
      chat = build_chat(context, llm_route: llm_route, messages: messages, schema: schema, tools: tools)

      conversation_messages = messages.reject { |m| m[:role] == 'system' }
      return { error: 'No conversation messages provided', error_code: 400, request_messages: messages } if conversation_messages.empty?

      add_messages_if_needed(chat, conversation_messages)
      build_ruby_llm_response(chat.ask(conversation_messages.last[:content]), messages)
    end
  rescue StandardError => e
    capture_llm_exception(e, credential: credential)
    { error: e.message, request_messages: messages }
  end

  def build_chat(context, llm_route:, messages:, schema: nil, tools: [])
    model = llm_route[:model]
    provider = llm_route[:provider]
    chat = Llm::ProviderChat.new(context.chat(model: model, provider: provider, assume_model_exists: true), provider: provider)
    system_msg = messages.find { |m| m[:role] == 'system' }
    chat.with_instructions(system_msg[:content]) if system_msg
    chat.with_schema(schema) if schema

    if tools.any?
      tools.each { |tool| chat = chat.with_tool(tool) }
      chat.on_end_message { |message| record_generation(chat, message, model, provider) }
    end

    chat
  end

  def add_messages_if_needed(chat, conversation_messages)
    return if conversation_messages.length == 1

    conversation_messages[0...-1].each do |msg|
      chat.add_message(role: msg[:role].to_sym, content: msg[:content])
    end
  end

  def build_ruby_llm_response(response, messages)
    {
      message: response.content,
      usage: {
        'prompt_tokens' => response.input_tokens,
        'completion_tokens' => response.output_tokens,
        'total_tokens' => (response.input_tokens || 0) + (response.output_tokens || 0)
      },
      request_messages: messages
    }
  end

  def build_instrumentation_params(model, messages, provider)
    {
      span_name: "llm.#{event_name}",
      account_id: account.id,
      conversation_id: conversation&.display_id,
      feature_name: event_name,
      model: model,
      provider: provider,
      messages: messages,
      temperature: nil,
      metadata: instrumentation_metadata
    }
  end

  def instrumentation_metadata
    {
      channel_type: conversation&.inbox&.channel_type
    }.compact
  end

  def conversation_messages(start_from: 0)
    messages = []
    character_count = start_from

    conversation.messages
                .where(message_type: [:incoming, :outgoing])
                .where(private: false)
                .reorder('id desc')
                .each do |message|
      content = message.content_for_llm
      next if content.blank?
      break if character_count + content.length > TOKEN_LIMIT

      messages.prepend({ role: (message.incoming? ? 'user' : 'assistant'), content: content })
      character_count += content.length
    end

    messages
  end

  def captain_tasks_enabled? = account.feature_enabled?('captain_tasks')

  # Extension point consulted by the Enterprise quota wrapper. Subclasses
  # whose calls should not consume captain_responses should override this to
  # return false. When false, the wrapper neither blocks the call on an
  # exhausted captain_responses quota nor decrements it on success — the call
  # participates in the quota system in neither direction.
  def counts_toward_usage?
    llm_credential&.dig(:source) != :hook
  end

  def api_key_configured?(provider = llm_provider) = llm_credential(provider).present?

  def api_key = llm_credential&.dig(:api_key)

  def llm_provider = @llm_provider || Llm::Config::DEFAULT_PROVIDER

  def llm_credential(provider = llm_provider)
    @llm_credentials ||= {}
    @llm_credentials[provider.to_s] ||= Llm::CredentialResolver.new(provider: provider, openai_hook: resolved_openai_hook(provider)).resolve
  end

  def resolved_openai_hook(provider) = use_account_openai_hook? && Llm::Config.openai_provider?(provider) ? openai_hook : nil

  def use_account_openai_hook? = false

  def system_llm_credential(provider = llm_provider) = Llm::CredentialResolver.new(provider: provider).resolve

  def openai_hook
    @openai_hook ||= account.hooks.find_by(app_id: 'openai', status: 'enabled')
  end

  def exception_tracking_account
    account
  end

  def prompt_from_file(file_name)
    Rails.root.join('lib/integrations/openai/openai_prompts', "#{file_name}.liquid").read
  end

  # Follow-up context for client-side refinement
  def build_follow_up_context?
    # FollowUpService should return its own updated context
    !is_a?(Captain::FollowUpService)
  end

  def build_follow_up_context(messages, response)
    {
      event_name: event_name,
      original_context: extract_original_context(messages),
      last_response: response[:message],
      conversation_history: [],
      channel_type: conversation&.inbox&.channel_type
    }
  end

  def extract_original_context(messages)
    # Get the most recent user message for follow-up context
    user_msg = messages.reverse.find { |m| m[:role] == 'user' }
    user_msg ? user_msg[:content] : nil
  end
end
Captain::BaseTaskService.prepend_mod_with('Captain::BaseTaskService')