FIX: better LLM feedback for image generation failures (#1306)

* FIX: handle error conditions when generating images gracefully

* FIX: also handle error for edit_image

* Update lib/inference/open_ai_image_generator.rb

Co-authored-by: Krzysztof Kotlarek <kotlarek.krzysztof@gmail.com>

* lint

---------

Co-authored-by: Krzysztof Kotlarek <kotlarek.krzysztof@gmail.com>

lib/inference/open_ai_image_generator.rb

@@ -40,6 +40,10 @@ module ::DiscourseAi
             output_format: output_format,
           )
 
+        raise api_responses[0] if api_responses.all? { |resp| resp.is_a?(StandardError) }
+
+        api_responses = api_responses.filter { |response| !response.is_a?(StandardError) }
+
         create_uploads_from_responses(api_responses, user_id, for_private_message, title)
       end
 
@@ -151,11 +155,16 @@ module ::DiscourseAi
               )
             rescue => e
               attempts += 1
-              sleep 2
+              # to keep tests speedy
-              retry if attempts < 3
+              if !Rails.env.test?
-              Discourse.warn_exception(e, message: "Failed to generate image for prompt #{prompt}")
+                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
@@ -203,14 +212,16 @@ module ::DiscourseAi
           )
         rescue => e
           attempts += 1
-          sleep 2
+          if !Rails.env.test?
-          retry if attempts < 3
+            sleep 2
-          if Rails.env.development? || Rails.env.test?
+            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/create_image.rb

@@ -30,7 +30,7 @@ module DiscourseAi
         end
 
         def chain_next_response?
-          false
+          !!@error
         end
 
         def invoke
@@ -42,14 +42,20 @@ module DiscourseAi
 
           results = nil
 
-          results =
+          begin
-            DiscourseAi::Inference::OpenAiImageGenerator.create_uploads!(
+            results =
-              max_prompts,
+              DiscourseAi::Inference::OpenAiImageGenerator.create_uploads!(
-              model: "gpt-image-1",
+                max_prompts,
-              user_id: bot_user.id,
+                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
 

lib/personas/tools/edit_image.rb

@@ -37,7 +37,7 @@ module DiscourseAi
         end
 
         def chain_next_response?
-          false
+          !!@error
         end
 
         def image_urls
@@ -50,19 +50,24 @@ module DiscourseAi
           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 =
+          begin
-            DiscourseAi::Inference::OpenAiImageGenerator.create_edited_upload!(
+            result =
-              uploads,
+              DiscourseAi::Inference::OpenAiImageGenerator.create_edited_upload!(
-              prompt,
+                uploads,
-              user_id: bot_user.id,
+                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/create_image_spec.rb

@@ -22,6 +22,26 @@ RSpec.describe DiscourseAi::Personas::Tools::CreateImage do
   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" }]
 

spec/lib/personas/tools/edit_image_spec.rb

@@ -35,6 +35,27 @@ RSpec.describe DiscourseAi::Personas::Tools::EditImage do
   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" }]