From 1a5985134a810c2c964b5c2527a1e374d2dca74f Mon Sep 17 00:00:00 2001
From: Keegan George <kgeorge13@gmail.com>
Date: Tue, 2 Jan 2024 11:24:16 -0800
Subject: [PATCH] FIX: Show illustrate post only if stability API key present
 (#395)

---
 lib/ai_helper/assistant.rb                    | 10 +++-
 spec/lib/modules/ai_helper/assistant_spec.rb  | 46 +++++++++++++++++++
 .../ai_helper/assistant_controller_spec.rb    |  1 +
 3 files changed, 56 insertions(+), 1 deletion(-)

diff --git a/lib/ai_helper/assistant.rb b/lib/ai_helper/assistant.rb
index 7c965fa8..e9cbc96c 100644
--- a/lib/ai_helper/assistant.rb
+++ b/lib/ai_helper/assistant.rb
@@ -5,8 +5,16 @@ module DiscourseAi
     class Assistant
       def available_prompts(name_filter: nil)
         cp = CompletionPrompt
+        prompts = []
 
-        prompts = name_filter ? [cp.enabled_by_name(name_filter)] : cp.where(enabled: true)
+        if name_filter
+          prompts = [cp.enabled_by_name(name_filter)]
+        else
+          prompts = cp.where(enabled: true)
+          # Only show the illustrate_post prompt if the API key is present
+          prompts =
+            prompts.where.not(name: "illustrate_post") if !SiteSetting.ai_stability_api_key.present?
+        end
 
         prompts.map do |prompt|
           translation =
diff --git a/spec/lib/modules/ai_helper/assistant_spec.rb b/spec/lib/modules/ai_helper/assistant_spec.rb
index 8de19228..75085bdd 100644
--- a/spec/lib/modules/ai_helper/assistant_spec.rb
+++ b/spec/lib/modules/ai_helper/assistant_spec.rb
@@ -10,6 +10,52 @@ RSpec.describe DiscourseAi::AiHelper::Assistant do
     defends himself, but instead exclaims: 'You too, my son!' Shakespeare and Quevedo capture the pathetic cry.
   STRING
 
+  describe("#available_prompts") do
+    context "when no name filter is provided" do
+      it "returns all available prompts" do
+        prompts = subject.available_prompts
+
+        expect(prompts.length).to eq(6)
+        expect(prompts.map { |p| p[:name] }).to contain_exactly(
+          "translate",
+          "generate_titles",
+          "proofread",
+          "markdown_table",
+          "custom_prompt",
+          "explain",
+        )
+      end
+    end
+
+    context "when name filter is provided" do
+      it "returns the prompt with the given name" do
+        prompts = subject.available_prompts(name_filter: "translate")
+
+        expect(prompts.length).to eq(1)
+        expect(prompts.first[:name]).to eq("translate")
+      end
+    end
+
+    context "when stability API key is present" do
+      before { SiteSetting.ai_stability_api_key = "foo" }
+
+      it "returns the illustrate_post prompt in the list of all prompts" do
+        prompts = subject.available_prompts
+
+        expect(prompts.length).to eq(7)
+        expect(prompts.map { |p| p[:name] }).to contain_exactly(
+          "translate",
+          "generate_titles",
+          "proofread",
+          "markdown_table",
+          "custom_prompt",
+          "explain",
+          "illustrate_post",
+        )
+      end
+    end
+  end
+
   describe "#generate_and_send_prompt" do
     context "when using a prompt that returns text" do
       let(:mode) { CompletionPrompt::TRANSLATE }
diff --git a/spec/requests/ai_helper/assistant_controller_spec.rb b/spec/requests/ai_helper/assistant_controller_spec.rb
index 67dbc1b2..319844a1 100644
--- a/spec/requests/ai_helper/assistant_controller_spec.rb
+++ b/spec/requests/ai_helper/assistant_controller_spec.rb
@@ -138,6 +138,7 @@ RSpec.describe DiscourseAi::AiHelper::AssistantController do
         sign_in(user)
         user.group_ids = [Group::AUTO_GROUPS[:trust_level_1]]
         SiteSetting.ai_helper_allowed_groups = Group::AUTO_GROUPS[:trust_level_1]
+        SiteSetting.ai_stability_api_key = "foo"
       end
 
       it "returns a list of prompts when no name_filter is provided" do