discourse-ai/spec/lib/modules/ai_bot/tools/dall_e_spec.rb

#frozen_string_literal: true

RSpec.describe DiscourseAi::AiBot::Tools::DallE do
  let(:prompts) { ["a pink cow", "a red cow"] }

  fab!(:gpt_35_turbo) { Fabricate(:llm_model, name: "gpt-3.5-turbo") }

  before do
    SiteSetting.ai_bot_enabled = true
    toggle_enabled_bots(bots: [gpt_35_turbo])
    SiteSetting.ai_openai_api_key = "abc"
  end

  let(:bot_user) { DiscourseAi::AiBot::EntryPoint.find_user_from_model("gpt-3.5-turbo") }
  let(:llm) { DiscourseAi::Completions::Llm.proxy("open_ai:gpt-3.5-turbo") }
  let(:progress_blk) { Proc.new {} }

  let(:dall_e) do
    described_class.new({ prompts: prompts }, llm: llm, bot_user: bot_user, context: {})
  end

  let(:base64_image) do
    "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mP8z8BQDwAEhQGAhKmMIQAAAABJRU5ErkJggg=="
  end

  describe "#process" do
    it "can generate tall images" do
      generator =
        described_class.new(
          { prompts: ["a cat"], aspect_ratio: "tall" },
          llm: llm,
          bot_user: bot_user,
          context: {
          },
        )

      data = [{ b64_json: base64_image, revised_prompt: "a tall cat" }]

      WebMock
        .stub_request(:post, "https://api.openai.com/v1/images/generations")
        .with do |request|
          json = JSON.parse(request.body, symbolize_names: true)

          expect(json[:prompt]).to eq("a cat")
          expect(json[:size]).to eq("1024x1792")
          true
        end
        .to_return(status: 200, body: { data: data }.to_json)

      info = generator.invoke(&progress_blk).to_json
      expect(JSON.parse(info)).to eq("prompts" => ["a tall cat"])
    end

    it "can generate correct info with azure" do
      _post = Fabricate(:post)

      SiteSetting.ai_openai_dall_e_3_url = "https://test.azure.com/some_url"

      data = [{ b64_json: base64_image, revised_prompt: "a pink cow 1" }]

      WebMock
        .stub_request(:post, SiteSetting.ai_openai_dall_e_3_url)
        .with do |request|
          json = JSON.parse(request.body, symbolize_names: true)

          expect(prompts).to include(json[:prompt])
          expect(request.headers["Api-Key"]).to eq("abc")
          true
        end
        .to_return(status: 200, body: { data: data }.to_json)

      info = dall_e.invoke(&progress_blk).to_json

      expect(JSON.parse(info)).to eq("prompts" => ["a pink cow 1", "a pink cow 1"])
      expect(dall_e.custom_raw).to include("upload://")
      expect(dall_e.custom_raw).to include("[grid]")
      expect(dall_e.custom_raw).to include("a pink cow 1")
    end

    it "can generate correct info" do
      data = [{ b64_json: base64_image, revised_prompt: "a pink cow 1" }]

      WebMock
        .stub_request(:post, "https://api.openai.com/v1/images/generations")
        .with do |request|
          json = JSON.parse(request.body, symbolize_names: true)
          expect(prompts).to include(json[:prompt])
          true
        end
        .to_return(status: 200, body: { data: data }.to_json)

      info = dall_e.invoke(&progress_blk).to_json

      expect(JSON.parse(info)).to eq("prompts" => ["a pink cow 1", "a pink cow 1"])
      expect(dall_e.custom_raw).to include("upload://")
      expect(dall_e.custom_raw).to include("[grid]")
      expect(dall_e.custom_raw).to include("a pink cow 1")
    end
  end
end