discourse-ai/lib/ai_bot/playground.rb

# frozen_string_literal: true

module DiscourseAi
  module AiBot
    class Playground
      # An abstraction to manage the bot and topic interactions.
      # The bot will take care of completions while this class updates the topic title
      # and stream replies.

      REQUIRE_TITLE_UPDATE = "discourse-ai-title-update"

      def initialize(bot)
        @bot = bot
      end

      def update_playground_with(post)
        if can_attach?(post) && bot.bot_user
          schedule_playground_titling(post, bot.bot_user)
          schedule_bot_reply(post, bot.bot_user)
        end
      end

      def conversation_context(post)
        # Pay attention to the `post_number <= ?` here.
        # We want to inject the last post as context because they are translated differently.
        context =
          post
            .topic
            .posts
            .includes(:user)
            .joins("LEFT JOIN post_custom_prompts ON post_custom_prompts.post_id = posts.id")
            .where("post_number <= ?", post.post_number)
            .order("post_number desc")
            .where("post_type = ?", Post.types[:regular])
            .limit(50)
            .pluck(:raw, :username, "post_custom_prompts.custom_prompt")

        result = []
        first = true

        context.reverse_each do |raw, username, custom_prompt|
          custom_prompt_translation =
            Proc.new do |message|
              # We can't keep backwards-compatibility for stored functions.
              # Tool syntax requires a tool_call_id which we don't have.
              if message[2] != "function"
                custom_context = {
                  content: message[0],
                  type: message[2].present? ? message[2].to_sym : :model,
                }

                custom_context[:id] = message[1] if custom_context[:type] != :model

                result << custom_context
              end
            end

          if custom_prompt.present?
            if first
              custom_prompt.each(&custom_prompt_translation)
              first = false
            else
              custom_prompt.first(2).each(&custom_prompt_translation)
            end
          else
            context = {
              content: raw,
              type: (available_bot_usernames.include?(username) ? :model : :user),
            }

            context[:id] = username if context[:type] == :user

            result << context
          end
        end

        result
      end

      def title_playground(post)
        context = conversation_context(post)

        bot
          .get_updated_title(context, post.user)
          .tap do |new_title|
            PostRevisor.new(post.topic.first_post, post.topic).revise!(
              bot.bot_user,
              title: new_title.sub(/\A"/, "").sub(/"\Z/, ""),
            )
            post.topic.custom_fields.delete(DiscourseAi::AiBot::EntryPoint::REQUIRE_TITLE_UPDATE)
            post.topic.save_custom_fields
          end
      end

      def reply_to(post)
        reply = +""
        start = Time.now

        context = {
          site_url: Discourse.base_url,
          site_title: SiteSetting.title,
          site_description: SiteSetting.site_description,
          time: Time.zone.now,
          participants: post.topic.allowed_users.map(&:username).join(", "),
          conversation_context: conversation_context(post),
          user: post.user,
        }

        reply_post =
          PostCreator.create!(
            bot.bot_user,
            topic_id: post.topic_id,
            raw: "",
            skip_validations: true,
            skip_jobs: true,
          )

        publish_update(reply_post, { raw: reply_post.cooked })

        redis_stream_key = "gpt_cancel:#{reply_post.id}"
        Discourse.redis.setex(redis_stream_key, 60, 1)

        new_custom_prompts =
          bot.reply(context) do |partial, cancel, placeholder|
            reply << partial
            raw = reply.dup
            raw << "\n\n" << placeholder if placeholder.present?

            if !Discourse.redis.get(redis_stream_key)
              cancel&.call

              reply_post.update!(raw: reply, cooked: PrettyText.cook(reply))
            end

            # Minor hack to skip the delay during tests.
            if placeholder.blank?
              next if (Time.now - start < 0.5) && !Rails.env.test?
              start = Time.now
            end

            Discourse.redis.expire(redis_stream_key, 60)

            publish_update(reply_post, { raw: raw })
          end

        return if reply.blank?

        # land the final message prior to saving so we don't clash
        reply_post.cooked = PrettyText.cook(reply)
        publish_final_update(reply_post)

        reply_post.revise(bot.bot_user, { raw: reply }, skip_validations: true, skip_revision: true)

        # not need to add a custom prompt for a single reply
        if new_custom_prompts.length > 1
          reply_post.post_custom_prompt ||= reply_post.build_post_custom_prompt(custom_prompt: [])
          prompt = reply_post.post_custom_prompt.custom_prompt || []
          prompt.concat(new_custom_prompts)
          reply_post.post_custom_prompt.update!(custom_prompt: prompt)
        end

        reply_post
      ensure
        publish_final_update(reply_post)
      end

      private

      def publish_final_update(reply_post)
        return if @published_final_update
        if reply_post
          publish_update(reply_post, { cooked: reply_post.cooked, done: true })
          # we subscribe at position -2 so we will always get this message
          # moving all cooked on every page load is wasteful ... this means
          # we have a benign message at the end, 2 is set to ensure last message
          # is delivered
          publish_update(reply_post, { noop: true })
          @published_final_update = true
        end
      end

      attr_reader :bot

      def can_attach?(post)
        return false if bot.bot_user.nil?
        return false if post.post_type != Post.types[:regular]
        return false unless post.topic.private_message?
        return false if (SiteSetting.ai_bot_allowed_groups_map & post.user.group_ids).blank?

        true
      end

      def schedule_playground_titling(post, bot_user)
        if post.post_number == 1
          post.topic.custom_fields[REQUIRE_TITLE_UPDATE] = true
          post.topic.save_custom_fields
        end

        ::Jobs.enqueue_in(
          5.minutes,
          :update_ai_bot_pm_title,
          post_id: post.id,
          bot_user_id: bot_user.id,
        )
      end

      def schedule_bot_reply(post, bot_user)
        ::Jobs.enqueue(:create_ai_reply, post_id: post.id, bot_user_id: bot_user.id)
      end

      def context(topic)
        {
          site_url: Discourse.base_url,
          site_title: SiteSetting.title,
          site_description: SiteSetting.site_description,
          time: Time.zone.now,
          participants: topic.allowed_users.map(&:username).join(", "),
        }
      end

      def publish_update(bot_reply_post, payload)
        payload = { post_id: bot_reply_post.id, post_number: bot_reply_post.post_number }.merge(
          payload,
        )
        MessageBus.publish(
          "discourse-ai/ai-bot/topic/#{bot_reply_post.topic_id}",
          payload,
          user_ids: bot_reply_post.topic.allowed_user_ids,
          max_backlog_size: 2,
          max_backlog_age: 60,
        )
      end

      def available_bot_usernames
        @bot_usernames ||= DiscourseAi::AiBot::EntryPoint::BOTS.map(&:second)
      end
    end
  end
end