FIX: also handle error for edit_image

Author: SamSaffron

lib/inference/open_ai_image_generator.rb

@@ -213,14 +213,16 @@ def self.edit_images(
           )
         rescue => e
           attempts += 1
-          sleep 2
-          retry if attempts < 3
-          if Rails.env.development? || Rails.env.test?
+          if !Rails.env.test?
+            sleep 2
+            retry if attempts < 3
+          end
+          if Rails.env.development?
             puts "Error editing image(s) with prompt: #{prompt} #{e}"
             p e
           end
           Discourse.warn_exception(e, message: "Failed to edit image(s) with prompt #{prompt}")
-          nil
+          raise e
         end
       end
 

lib/personas/tools/edit_image.rb

@@ -37,7 +37,7 @@ def prompt
         end
 
         def chain_next_response?
-          false
+          !!@error
         end
 
         def image_urls
@@ -50,19 +50,24 @@ def invoke
           return { prompt: prompt, error: "No valid images provided" } if image_urls.blank?
 
           sha1s = image_urls.map { |url| Upload.sha1_from_short_url(url) }.compact
-
           uploads = Upload.where(sha1: sha1s).order(created_at: :asc).limit(10).to_a
 
           return { prompt: prompt, error: "No valid images provided" } if uploads.blank?
 
-          result =
-            DiscourseAi::Inference::OpenAiImageGenerator.create_edited_upload!(
-              uploads,
-              prompt,
-              user_id: bot_user.id,
-            )
+          begin
+            result =
+              DiscourseAi::Inference::OpenAiImageGenerator.create_edited_upload!(
+                uploads,
+                prompt,
+                user_id: bot_user.id,
+              )
+          rescue => e
+            @error = e
+            return { prompt: prompt, error: e.message }
+          end
 
           if result.blank?
+            @error = true
             return { prompt: prompt, error: "Something went wrong, could not generate image" }
           end
 

spec/lib/personas/tools/edit_image_spec.rb

@@ -35,6 +35,27 @@
   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/edits").to_return(
+        status: 400,
+        body: error_message.to_json,
+      )
+
+      info = edit_image.invoke(&progress_blk).to_json
+      expect(info).to include("Your request was rejected as a result of our safety system.")
+      expect(edit_image.chain_next_response?).to eq(true)
+    end
+
     it "can edit an image with the GPT image model" do
       data = [{ b64_json: base64_image, revised_prompt: "image with rainbow added in background" }]