FIX: handle error conditions when generating images gracefully

Author: SamSaffron

lib/inference/open_ai_image_generator.rb

@@ -40,6 +40,11 @@ def self.create_uploads!(
             output_format: output_format,
           )
 
+        raise api_responses[0] if api_responses.all? { |resp| resp.is_a?(StandardError) }
+
+        api_responses =
+          api_responses.filter_map { |response| response.is_a?(StandardError) ? nil : response }
+
         create_uploads_from_responses(api_responses, user_id, for_private_message, title)
       end
 
@@ -151,11 +156,16 @@ def self.generate_images_in_threads(
               )
             rescue => e
               attempts += 1
-              sleep 2
-              retry if attempts < 3
-              Discourse.warn_exception(e, message: "Failed to generate image for prompt #{prompt}")
+              # to keep tests speedy
+              if !Rails.env.test?
+                retry if attempts < 3
+              end
+              Discourse.warn_exception(
+                e,
+                message: "Failed to generate image for prompt #{prompt}\n",
+              )
               puts "Error generating image for prompt: #{prompt} #{e}" if Rails.env.development?
-              nil
+              e
             end
           end
         end

lib/personas/tools/create_image.rb

@@ -30,7 +30,7 @@ def prompts
         end
 
         def chain_next_response?
-          false
+          !!@error
         end
 
         def invoke
@@ -42,14 +42,20 @@ def invoke
 
           results = nil
 
-          results =
-            DiscourseAi::Inference::OpenAiImageGenerator.create_uploads!(
-              max_prompts,
-              model: "gpt-image-1",
-              user_id: bot_user.id,
-            )
+          begin
+            results =
+              DiscourseAi::Inference::OpenAiImageGenerator.create_uploads!(
+                max_prompts,
+                model: "gpt-image-1",
+                user_id: bot_user.id,
+              )
+          rescue => e
+            @error = e
+            return { prompts: max_prompts, error: e.message }
+          end
 
           if results.blank?
+            @error = true
             return { prompts: max_prompts, error: "Something went wrong, could not generate image" }
           end
 

spec/lib/personas/tools/create_image_spec.rb

@@ -22,6 +22,26 @@
   end
 
   describe "#process" do
+    it "can reject generation of images and return a proper error to llm" do
+      error_message = {
+        error: {
+          message:
+            "Your request was rejected as a result of our safety system. Your request may contain content that is not allowed by our safety system.",
+          type: "user_error",
+          param: nil,
+          code: "moderation_blocked",
+        },
+      }
+
+      WebMock.stub_request(:post, "https://api.openai.com/v1/images/generations").to_return(
+        status: 400,
+        body: error_message.to_json,
+      )
+
+      info = create_image.invoke(&progress_blk).to_json
+      expect(info).to include("Your request was rejected as a result of our safety system.")
+      expect(create_image.chain_next_response?).to eq(true)
+    end
     it "can generate images with gpt-image-1 model" do
       data = [{ b64_json: base64_image, revised_prompt: "a watercolor painting of flowers" }]