discourse/plugins/discourse-ai/spec/requests/ai_bot/bot_controller_spec.rb

# frozen_string_literal: true

RSpec.describe DiscourseAi::AiBot::BotController do
  fab!(:user)
  fab!(:pm_topic, :private_message_topic)
  fab!(:pm_post) { Fabricate(:post, topic: pm_topic) }
  fab!(:pm_post2) { Fabricate(:post, topic: pm_topic) }
  fab!(:pm_post3) { Fabricate(:post, topic: pm_topic) }

  before do
    enable_current_plugin
    sign_in(user)
  end

  describe "#show_debug_info" do
    fab!(:debug_bot_model) { Fabricate(:llm_model, name: "debug-bot-model") }
    fab!(:debug_bot_user) do
      enable_current_plugin
      toggle_enabled_bots(bots: [debug_bot_model])
      debug_bot_model.reload.user
    end

    fab!(:pm_topic) { Fabricate(:private_message_topic, user: user, recipient: debug_bot_user) }
    fab!(:pm_post) { Fabricate(:post, topic: pm_topic, user: debug_bot_user) }
    fab!(:pm_post2) { Fabricate(:post, topic: pm_topic, user: debug_bot_user) }
    fab!(:pm_post3) { Fabricate(:post, topic: pm_topic, user: user) }

    before { SiteSetting.ai_bot_enabled = true }

    it "returns a 403 when the user cannot debug the AI bot conversation" do
      get "/discourse-ai/ai-bot/post/#{pm_post.id}/show-debug-info"
      expect(response.status).to eq(403)
    end

    it "does not disclose hidden whisper audit logs" do
      debug_group = Fabricate(:group)
      debug_group.add(user)
      SiteSetting.ai_bot_debugging_allowed_groups = debug_group.id.to_s
      SiteSetting.whispers_allowed_groups = Group::AUTO_GROUPS[:staff].to_s

      visible_bot_post = Fabricate(:post, user: debug_bot_user)
      hidden_whisper_post =
        Fabricate(
          :post,
          topic: visible_bot_post.topic,
          user: Fabricate(:admin),
          post_type: Post.types[:whisper],
        )
      later_regular_post = Fabricate(:post, topic: visible_bot_post.topic, user: user)

      visible_bot_log =
        AiApiAuditLog.create!(
          post_id: visible_bot_post.id,
          provider_id: 1,
          topic_id: visible_bot_post.topic_id,
          feature_name: "ai_bot",
          raw_request_payload: "bot request",
          raw_response_payload: "bot response",
          request_tokens: 1,
          response_tokens: 2,
          created_at: 2.minutes.ago,
        )

      hidden_bot_log =
        AiApiAuditLog.create!(
          post_id: hidden_whisper_post.id,
          provider_id: 1,
          topic_id: hidden_whisper_post.topic_id,
          feature_name: "ai_bot",
          raw_request_payload: "hidden bot request",
          raw_response_payload: "hidden bot response",
          request_tokens: 3,
          response_tokens: 4,
          created_at: 1.minute.ago,
        )

      hidden_translation_log =
        AiApiAuditLog.create!(
          post_id: hidden_whisper_post.id,
          provider_id: 1,
          topic_id: hidden_whisper_post.topic_id,
          feature_name: "translation",
          raw_request_payload: "hidden translation request",
          raw_response_payload: "hidden translation response",
          request_tokens: 5,
          response_tokens: 6,
          created_at: Time.zone.now,
        )

      get "/discourse-ai/ai-bot/post/#{later_regular_post.id}/show-debug-info"
      show_debug_info_status = response.status
      show_debug_info_body = response.parsed_body.deep_dup

      get "/discourse-ai/ai-bot/show-debug-info/#{hidden_translation_log.id}"
      hidden_translation_status = response.status
      hidden_translation_body = response.parsed_body.deep_dup

      get "/discourse-ai/ai-bot/show-debug-info/#{hidden_bot_log.id}"
      hidden_bot_status = response.status
      hidden_bot_body = response.parsed_body.deep_dup

      aggregate_failures do
        expect(show_debug_info_status).to eq(200)
        expect(show_debug_info_body["id"]).to eq(visible_bot_log.id)
        expect(show_debug_info_body["raw_request_payload"]).to eq("bot request")
        expect(show_debug_info_body["raw_response_payload"]).to eq("bot response")

        expect(hidden_translation_status).to eq(403)
        expect(hidden_translation_body["errors"]).to include(I18n.t("invalid_access"))

        expect(hidden_bot_status).to eq(403)
        expect(hidden_bot_body["errors"]).to include(I18n.t("invalid_access"))
      end
    end

    it "returns debug info if the user can debug the AI bot conversation" do
      user = pm_topic.topic_allowed_users.first.user
      sign_in(user)

      log1 =
        AiApiAuditLog.create!(
          provider_id: 1,
          topic_id: pm_topic.id,
          feature_name: "ai_bot",
          raw_request_payload: "request",
          raw_response_payload: "response",
          request_tokens: 1,
          response_tokens: 2,
        )

      log2 =
        AiApiAuditLog.create!(
          post_id: pm_post.id,
          provider_id: 1,
          topic_id: pm_topic.id,
          feature_name: "ai_bot",
          raw_request_payload: "request",
          raw_response_payload: "response",
          request_tokens: 1,
          response_tokens: 2,
        )

      log3 =
        AiApiAuditLog.create!(
          post_id: pm_post2.id,
          provider_id: 1,
          topic_id: pm_topic.id,
          feature_name: "ai_bot",
          raw_request_payload: "request",
          raw_response_payload: "response",
          request_tokens: 1,
          response_tokens: 2,
        )

      Group.refresh_automatic_groups!
      SiteSetting.ai_bot_debugging_allowed_groups = user.groups.first.id.to_s

      get "/discourse-ai/ai-bot/post/#{pm_post.id}/show-debug-info"
      expect(response.status).to eq(200)

      expect(response.parsed_body["id"]).to eq(log2.id)
      expect(response.parsed_body["next_log_id"]).to eq(log3.id)
      expect(response.parsed_body["prev_log_id"]).to eq(log1.id)
      expect(response.parsed_body["topic_id"]).to eq(pm_topic.id)

      expect(response.parsed_body["request_tokens"]).to eq(1)
      expect(response.parsed_body["response_tokens"]).to eq(2)
      expect(response.parsed_body["raw_request_payload"]).to eq("request")
      expect(response.parsed_body["raw_response_payload"]).to eq("response")

      get "/discourse-ai/ai-bot/post/#{pm_post3.id}/show-debug-info"
      expect(response.status).to eq(200)
      expect(response.parsed_body["request_tokens"]).to eq(1)
      expect(response.parsed_body["response_tokens"]).to eq(2)

      get "/discourse-ai/ai-bot/show-debug-info/#{log1.id}"
      expect(response.status).to eq(200)
      expect(response.parsed_body["id"]).to eq(log1.id)
    end

    context "with conversation totals and spending" do
      fab!(:llm_model) do
        Fabricate(
          :llm_model,
          input_cost: 3.0,
          output_cost: 6.0,
          cached_input_cost: 1.0,
          cache_write_cost: 2.0,
        )
      end

      fab!(:other_llm_model) do
        Fabricate(
          :llm_model,
          name: "other-model",
          input_cost: 10.0,
          output_cost: 20.0,
          cached_input_cost: 5.0,
          cache_write_cost: 8.0,
        )
      end

      let(:allowed_user) { pm_topic.topic_allowed_users.first.user }

      before do
        sign_in(allowed_user)
        Group.refresh_automatic_groups!
        SiteSetting.ai_bot_debugging_allowed_groups = allowed_user.groups.first.id.to_s
      end

      it "computes per-turn spending and conversation totals across multiple models" do
        log_a =
          AiApiAuditLog.create!(
            post_id: pm_post.id,
            provider_id: 1,
            topic_id: pm_topic.id,
            llm_id: llm_model.id,
            feature_name: "ai_bot",
            raw_request_payload: "req",
            raw_response_payload: "res",
            request_tokens: 1_000_000,
            response_tokens: 500_000,
            cache_read_tokens: 200_000,
            cache_write_tokens: 100_000,
          )

        AiApiAuditLog.create!(
          post_id: pm_post2.id,
          provider_id: 1,
          topic_id: pm_topic.id,
          llm_id: other_llm_model.id,
          feature_name: "ai_bot",
          raw_request_payload: "req",
          raw_response_payload: "res",
          request_tokens: 2_000_000,
          response_tokens: 1_000_000,
        )

        get "/discourse-ai/ai-bot/show-debug-info/#{log_a.id}"
        expect(response.status).to eq(200)

        body = response.parsed_body
        expected_turn_spending =
          (1_000_000 * 3.0 + 500_000 * 6.0 + 200_000 * 1.0 + 100_000 * 2.0) / 1_000_000.0
        expect(body["spending"]).to be_within(0.000001).of(expected_turn_spending)

        expect(body["conversation_request_tokens"]).to eq(3_000_000)
        expect(body["conversation_response_tokens"]).to eq(1_500_000)
        expect(body["conversation_cache_read_tokens"]).to eq(200_000)
        expect(body["conversation_cache_write_tokens"]).to eq(100_000)

        expected_conversation_spending =
          expected_turn_spending + (2_000_000 * 10.0 + 1_000_000 * 20.0) / 1_000_000.0
        expect(body["conversation_spending"]).to be_within(0.000001).of(
          expected_conversation_spending,
        )
      end

      it "returns nil spending when the log has no llm_model" do
        log_a =
          AiApiAuditLog.create!(
            post_id: pm_post.id,
            provider_id: 1,
            topic_id: pm_topic.id,
            feature_name: "ai_bot",
            raw_request_payload: "req",
            raw_response_payload: "res",
            request_tokens: 100,
            response_tokens: 200,
          )

        get "/discourse-ai/ai-bot/show-debug-info/#{log_a.id}"
        expect(response.status).to eq(200)

        body = response.parsed_body
        expect(body["spending"]).to be_nil
        expect(body["conversation_request_tokens"]).to eq(100)
        expect(body["conversation_response_tokens"]).to eq(200)
        expect(body["conversation_spending"]).to be_nil
      end

      it "returns nil conversation fields when topic_id is absent" do
        log_a =
          AiApiAuditLog.create!(
            provider_id: 1,
            llm_id: llm_model.id,
            feature_name: "ai_bot",
            raw_request_payload: "req",
            raw_response_payload: "res",
            request_tokens: 100,
            response_tokens: 200,
          )
        SiteSetting.ai_bot_debugging_allowed_groups = Group::AUTO_GROUPS[:admins].to_s
        sign_in(Fabricate(:admin))

        get "/discourse-ai/ai-bot/show-debug-info/#{log_a.id}"
        # no topic means the endpoint 404s, so assert via serializer directly
        serialized = AiApiAuditLogSerializer.new(log_a, root: false).as_json
        expect(serialized[:conversation_request_tokens]).to be_nil
        expect(serialized[:conversation_response_tokens]).to be_nil
        expect(serialized[:conversation_cache_read_tokens]).to be_nil
        expect(serialized[:conversation_cache_write_tokens]).to be_nil
        expect(serialized[:conversation_spending]).to be_nil
      end
    end

    it "returns topic-level debug info for any feature for a post in the conversation" do
      user = pm_topic.topic_allowed_users.first.user
      sign_in(user)

      log =
        AiApiAuditLog.create!(
          provider_id: 1,
          topic_id: pm_topic.id,
          feature_name: "translation",
          raw_request_payload: "request",
          raw_response_payload: "response",
          request_tokens: 1,
          response_tokens: 2,
        )

      Group.refresh_automatic_groups!
      SiteSetting.ai_bot_debugging_allowed_groups = user.groups.first.id.to_s

      get "/discourse-ai/ai-bot/post/#{pm_post.id}/show-debug-info"

      expect(response.status).to eq(200)
      expect(response.parsed_body["id"]).to eq(log.id)

      get "/discourse-ai/ai-bot/show-debug-info/#{log.id}"

      expect(response.status).to eq(200)
      expect(response.parsed_body["id"]).to eq(log.id)
    end
  end

  describe "#stop_streaming_response" do
    let(:redis_stream_key) { "gpt_cancel:#{pm_post.id}" }

    before do
      SiteSetting.ai_bot_enabled = true
      Discourse.redis.setex(redis_stream_key, 60, 1)
    end

    it "returns a 403 when the user cannot see the PM" do
      post "/discourse-ai/ai-bot/post/#{pm_post.id}/stop-streaming"

      expect(response.status).to eq(403)
    end

    it "returns a 403 when the user can see the PM but is not in ai_bot_allowed_groups" do
      allowed_user = pm_topic.topic_allowed_users.first.user
      sign_in(allowed_user)

      SiteSetting.ai_bot_allowed_groups = Group::AUTO_GROUPS[:staff].to_s

      post "/discourse-ai/ai-bot/post/#{pm_post.id}/stop-streaming"

      expect(response.status).to eq(403)
      expect(Discourse.redis.get(redis_stream_key)).to eq("1")
    end

    it "returns a 403 when ai_bot_enabled is false" do
      allowed_user = pm_topic.topic_allowed_users.first.user
      sign_in(allowed_user)

      SiteSetting.ai_bot_enabled = false

      post "/discourse-ai/ai-bot/post/#{pm_post.id}/stop-streaming"

      expect(response.status).to eq(403)
      expect(Discourse.redis.get(redis_stream_key)).to eq("1")
    end

    it "deletes the key using to track the streaming" do
      allowed_user = pm_topic.topic_allowed_users.first.user
      sign_in(allowed_user)

      Group.refresh_automatic_groups!
      SiteSetting.ai_bot_allowed_groups = allowed_user.groups.first.id.to_s

      post "/discourse-ai/ai-bot/post/#{pm_post.id}/stop-streaming"

      expect(response.status).to eq(200)
      expect(Discourse.redis.get(redis_stream_key)).to be_nil
    end
  end

  describe "#retry_response" do
    fab!(:bot_user, :user)
    let!(:llm_model) { Fabricate(:llm_model, user: bot_user) }
    let!(:ai_agent) do
      Fabricate(
        :ai_agent,
        user: bot_user,
        default_llm_id: llm_model.id,
        allowed_group_ids: [Group::AUTO_GROUPS[:trust_level_0]],
      )
    end
    let(:agent) { ai_agent.class_instance.new }
    let(:bot) { DiscourseAi::Agents::Bot.as(bot_user, agent: agent) }

    let!(:prompt_post) do
      Fabricate(:post, topic: pm_topic, user: user, raw: "Hello @#{bot_user.username}")
    end

    let!(:reply_post) do
      DiscourseAi::Completions::Llm.with_prepared_responses(["first try"]) do
        DiscourseAi::AiBot::Playground.new(bot).reply_to(prompt_post)
      end

      pm_topic.reload.posts.last
    end

    before do
      Group.refresh_automatic_groups!
      SiteSetting.ai_bot_enabled = true
      SiteSetting.ai_bot_allowed_groups = Group::AUTO_GROUPS[:trust_level_0].to_s
      AiAgent.agent_cache.flush!

      tl0_group =
        Group.find_by(name: "trust_level_0") || Group.find(Group::AUTO_GROUPS[:trust_level_0])
      GroupUser.find_or_create_by!(user: user, group: tl0_group)
      user.reload

      pm_topic.topic_allowed_users.find_or_create_by!(user: user)

      unless pm_topic.topic_allowed_users.exists?(user: bot_user)
        pm_topic.topic_allowed_users.create!(user: bot_user)
      end
    end

    it "returns 403 when ai_bot_enabled is false" do
      SiteSetting.ai_bot_enabled = false

      post "/discourse-ai/ai-bot/post/#{reply_post.id}/retry"

      expect(response.status).to eq(403)
    end

    it "returns 403 when user is not in ai_bot_allowed_groups" do
      SiteSetting.ai_bot_allowed_groups = Group::AUTO_GROUPS[:admins].to_s

      post "/discourse-ai/ai-bot/post/#{reply_post.id}/retry"

      expect(response.status).to eq(403)
    end

    it "streams a replacement into the existing bot reply" do
      retry_text = "second attempt"
      messages = nil

      DiscourseAi::Completions::Llm.with_prepared_responses([retry_text]) do
        post "/discourse-ai/ai-bot/post/#{reply_post.id}/retry"

        messages =
          MessageBus.track_publish("discourse-ai/ai-bot/topic/#{reply_post.topic_id}") do
            Jobs::CreateAiReply.new.execute(
              post_id: prompt_post.id,
              bot_user_id: bot_user.id,
              agent_id: reply_post.custom_fields[DiscourseAi::AiBot::POST_AI_AGENT_ID_FIELD].to_i,
              reply_post_id: reply_post.id,
            )
          end
      end

      expect(response.status).to eq(200)
      expect(reply_post.reload.raw).to eq(retry_text)
      expect(messages.first.data).to include(post_id: reply_post.id, raw: "")
      expect(reply_post.custom_fields[DiscourseAi::AiBot::POST_AI_AGENT_ID_FIELD].to_i).to eq(
        agent.id,
      )
    end

    it "returns a 404 when there is no previous non-bot prompt" do
      lone_topic = Fabricate(:private_message_topic)
      lone_topic.topic_allowed_users.create!(user: user)
      lone_topic.topic_allowed_users.create!(user: bot_user)
      bot_only_post = Fabricate(:post, topic: lone_topic, user: bot_user)

      post "/discourse-ai/ai-bot/post/#{bot_only_post.id}/retry"

      expect(response.status).to eq(404)
    end

    it "allows retrying if LLM model has a negative id (seeded)" do
      seeded_llm_model = Fabricate(:llm_model, id: -9999, user: bot_user, name: "second-model")

      bot = DiscourseAi::Agents::Bot.as(bot_user, agent: agent, model: seeded_llm_model)
      DiscourseAi::Completions::Llm.with_prepared_responses(["first try"], llm: seeded_llm_model) do
        DiscourseAi::AiBot::Playground.new(bot).reply_to(prompt_post)
      end

      reply = pm_topic.reload.posts.last
      original_llm_id = reply.custom_fields[DiscourseAi::AiBot::POST_AI_LLM_MODEL_ID_FIELD]

      expect(original_llm_id.to_i).to eq(-9999)

      retry_text = "retry with original model"

      DiscourseAi::Completions::Llm.with_prepared_responses([retry_text], llm: seeded_llm_model) do
        post "/discourse-ai/ai-bot/post/#{reply.id}/retry"

        job_args = Jobs::CreateAiReply.jobs.last["args"].first
        expect(job_args["llm_model_id"]).to eq(-9999)

        Jobs::CreateAiReply.new.execute(job_args.symbolize_keys)
      end

      expect(response.status).to eq(200)
      expect(reply.reload.raw).to eq(retry_text)
    end

    it "uses the original LLM model when retrying even if agent default changed" do
      second_bot_user = Fabricate(:user)
      second_llm_model = Fabricate(:llm_model, user: second_bot_user, name: "second-model")

      pm_topic.topic_allowed_users.find_or_create_by!(user: second_bot_user)

      original_llm_name = reply_post.custom_fields[DiscourseAi::AiBot::POST_AI_LLM_NAME_FIELD]
      original_llm_id = reply_post.custom_fields[DiscourseAi::AiBot::POST_AI_LLM_MODEL_ID_FIELD]

      expect(original_llm_name).to be_present
      expect(original_llm_id.to_i).to eq(llm_model.id)

      ai_agent.update!(default_llm_id: second_llm_model.id)
      AiAgent.agent_cache.flush!

      retry_text = "retry with original model"

      DiscourseAi::Completions::Llm.with_prepared_responses([retry_text], llm: llm_model) do
        post "/discourse-ai/ai-bot/post/#{reply_post.id}/retry"

        job_args = Jobs::CreateAiReply.jobs.last["args"].first
        expect(job_args["llm_model_id"]).to eq(llm_model.id)

        Jobs::CreateAiReply.new.execute(job_args.symbolize_keys)
      end

      expect(response.status).to eq(200)
      expect(reply_post.reload.raw).to eq(retry_text)
      expect(reply_post.custom_fields[DiscourseAi::AiBot::POST_AI_LLM_NAME_FIELD]).to eq(
        original_llm_name,
      )
      expect(reply_post.custom_fields[DiscourseAi::AiBot::POST_AI_LLM_MODEL_ID_FIELD].to_i).to eq(
        original_llm_id.to_i,
      )
    end
  end

  describe "#show_bot_username" do
    it "returns the username_lower of the selected bot" do
      gpt_35_bot = Fabricate(:llm_model, name: "gpt-3.5-turbo")

      SiteSetting.ai_bot_enabled = true
      toggle_enabled_bots(bots: [gpt_35_bot])

      expected_username =
        DiscourseAi::AiBot::EntryPoint.find_user_from_model("gpt-3.5-turbo").username_lower

      get "/discourse-ai/ai-bot/bot-username", params: { username: gpt_35_bot.name }

      expect(response.status).to eq(200)
      expect(response.parsed_body["bot_username"]).to eq(expected_username)
    end
  end
end