discourse/plugins/discourse-ai/lib/agents/tools/tool.rb

# frozen_string_literal: true

module DiscourseAi
  module Agents
    module Tools
      class Tool
        # Why 30 mega bytes?
        # This general limit is mainly a security feature to avoid tools
        # forcing infinite downloads or causing memory exhaustion.
        # The limit is somewhat arbitrary and can be increased in future if needed.
        MAX_RESPONSE_BODY_LENGTH = 30.megabytes

        class << self
          def signature
            raise NotImplemented
          end

          def name
            raise NotImplemented
          end

          def custom?
            false
          end

          def accepted_options
            []
          end

          def option(name, type:, values: nil, default: nil)
            Option.new(tool: self, name: name, type: type, values: values, default: default)
          end

          def help
            I18n.t("discourse_ai.ai_bot.tool_help.#{signature[:name]}")
          end

          def custom_system_message
            nil
          end

          def allow_partial_tool_calls?
            false
          end

          def requires_approval?
            false
          end

          def inject_prompt(prompt:, context:, agent:)
          end

          def available_custom_image_tools
            image_tool_ids = AiTool.where(enabled: true, is_image_generation_tool: true).pluck(:id)
            image_tool_ids.map do |tool_id|
              DiscourseAi::Agents::Tools::Custom.class_instance(tool_id)
            end
          end
        end

        # llm being public makes it a bit easier to test
        attr_accessor :custom_raw, :parameters, :llm, :provider_data
        attr_reader :tool_call_id, :agent_options, :bot_user, :context, :agent

        def initialize(
          parameters,
          tool_call_id: "",
          agent_options: {},
          bot_user:,
          llm:,
          context: nil,
          provider_data: {},
          agent: nil
        )
          @parameters = parameters
          @tool_call_id = tool_call_id
          @agent_options =
            agent_options.is_a?(Hash) ? agent_options.with_indifferent_access : agent_options
          @bot_user = bot_user
          @llm = llm
          @context = context.nil? ? DiscourseAi::Agents::BotContext.new(messages: []) : context
          @provider_data = provider_data.is_a?(Hash) ? provider_data.deep_symbolize_keys : {}
          @agent = agent
          if !@context.is_a?(DiscourseAi::Agents::BotContext)
            raise ArgumentError, "context must be a DiscourseAi::Agents::Context"
          end
        end

        def name
          self.class.name
        end

        def summary
          I18n.t("discourse_ai.ai_bot.tool_summary.#{name}")
        end

        def details
          I18n.t("discourse_ai.ai_bot.tool_description.#{name}", description_args)
        end

        def help
          I18n.t("discourse_ai.ai_bot.tool_help.#{name}")
        end

        def options
          result = ActiveSupport::HashWithIndifferentAccess.new
          self.class.accepted_options.each do |option|
            val = @agent_options[option.name]
            if val
              case option.type
              when :boolean
                val = (val.to_s == "true")
              when :integer
                val = val.to_i
              when :enum
                val = val.to_s
                val = option.default if option.values && !option.values.include?(val)
              end
              result[option.name] = val
            elsif val.nil?
              result[option.name] = option.default
            end
          end
          result
        end

        def chain_next_response?
          true
        end

        protected

        def fetch_default_branch(repo)
          github_client.get("https://api.github.com/repos/#{repo}")["default_branch"]
        rescue Discourse::GithubApi::Error
          "main"
        end

        def send_http_request(
          url,
          headers: {},
          follow_redirects: false,
          method: :get,
          body: nil,
          &blk
        )
          self.class.send_http_request(
            url,
            headers: headers,
            follow_redirects: follow_redirects,
            method: method,
            body: body,
            &blk
          )
        end

        def github_client
          ::Discourse::GithubApi.for(token: SiteSetting.ai_bot_github_access_token)
        end

        def self.send_http_request(
          url,
          headers: {},
          follow_redirects: false,
          method: :get,
          body: nil
        )
          raise "Expecting caller to use a block" if !block_given?

          uri = nil
          url = UrlHelper.normalized_encode(url)
          uri =
            begin
              URI.parse(url)
            rescue StandardError
              nil
            end

          return if !uri

          if follow_redirects
            fd =
              FinalDestination.new(
                url,
                validate_uri: true,
                max_redirects: 5,
                follow_canonical: true,
              )

            uri = fd.resolve
          end

          return if uri.blank?

          request = nil
          if method == :get
            request = FinalDestination::HTTP::Get.new(uri)
          elsif method == :post
            request = FinalDestination::HTTP::Post.new(uri)
          elsif method == :put
            request = FinalDestination::HTTP::Put.new(uri)
          elsif method == :patch
            request = FinalDestination::HTTP::Patch.new(uri)
          elsif method == :delete
            request = FinalDestination::HTTP::Delete.new(uri)
          end

          raise ArgumentError, "Invalid method: #{method}" if !request

          request.body = body if body

          request["User-Agent"] = DiscourseAi::AiBot::USER_AGENT
          headers.each { |k, v| request[k] = v }

          FinalDestination::HTTP.start(uri.hostname, uri.port, use_ssl: uri.port != 80) do |http|
            http.request(request) { |response| yield response, uri }
          end
        end

        def self.read_response_body(response, max_length: nil)
          max_length ||= MAX_RESPONSE_BODY_LENGTH

          body = +""
          response.read_body do |chunk|
            body << chunk
            break if body.bytesize > max_length
          end

          if body.bytesize > max_length
            body[0...max_length].scrub
          else
            body.scrub
          end
        end

        def read_response_body(response, max_length: nil)
          self.class.read_response_body(response, max_length: max_length)
        end

        def truncate(text, llm:, percent_length: nil, max_length: nil)
          if !percent_length && !max_length
            raise ArgumentError, "You must provide either percent_length or max_length"
          end

          target = llm.max_prompt_tokens
          target = (target * percent_length).to_i if percent_length

          if max_length
            target = max_length if target > max_length
          end

          llm.tokenizer.truncate(text, target, strict: SiteSetting.ai_strict_token_counting)
        end

        def accepted_options
          []
        end

        def option(name, type:)
          Option.new(tool: self, name: name, type: type)
        end

        def acting_user
          bot_user || Discourse.system_user
        end

        def guardian
          Guardian.new(context.user || acting_user)
        end

        def reason
          parameters[:reason].to_s.strip
        end

        def error_response(message)
          { status: "error", error: message }
        end

        def description_args
          {}
        end

        def format_results(rows, column_names = nil, args: nil)
          rows = rows&.map { |row| yield row } if block_given?

          if !column_names
            index = -1
            column_indexes = {}

            rows =
              rows&.map do |data|
                new_row = []
                data.each do |key, value|
                  found_index = column_indexes[key.to_s] ||= (index += 1)
                  new_row[found_index] = value
                end
                new_row
              end
            column_names = column_indexes.keys
          end

          # this is not the most efficient format
          # however this is needed cause GPT 3.5 / 4 was steered using JSON
          result = { column_names: column_names, rows: rows }
          result[:args] = args if args
          result
        end
      end
    end
  end
end