From 9196546f6f6741b855f834a5da4772eda4aa115e Mon Sep 17 00:00:00 2001 From: Sam Date: Thu, 1 May 2025 19:25:38 +1000 Subject: [PATCH] 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 * lint --------- Co-authored-by: Krzysztof Kotlarek --- lib/inference/open_ai_image_generator.rb | 27 ++++++++++++++------ lib/personas/tools/create_image.rb | 20 ++++++++++----- lib/personas/tools/edit_image.rb | 21 +++++++++------ spec/lib/personas/tools/create_image_spec.rb | 20 +++++++++++++++ spec/lib/personas/tools/edit_image_spec.rb | 21 +++++++++++++++ 5 files changed, 86 insertions(+), 23 deletions(-) diff --git a/lib/inference/open_ai_image_generator.rb b/lib/inference/open_ai_image_generator.rb index 63b94fba..f4ff6754 100644 --- a/lib/inference/open_ai_image_generator.rb +++ b/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 - 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 @@ -203,14 +212,16 @@ module ::DiscourseAi ) 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 diff --git a/lib/personas/tools/create_image.rb b/lib/personas/tools/create_image.rb index 7620d426..f57c83ff 100644 --- a/lib/personas/tools/create_image.rb +++ b/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 = - 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 diff --git a/lib/personas/tools/edit_image.rb b/lib/personas/tools/edit_image.rb index 8cda8ac1..b28afd0c 100644 --- a/lib/personas/tools/edit_image.rb +++ b/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 = - 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 diff --git a/spec/lib/personas/tools/create_image_spec.rb b/spec/lib/personas/tools/create_image_spec.rb index 5111d7e2..0aa18fea 100644 --- a/spec/lib/personas/tools/create_image_spec.rb +++ b/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" }] diff --git a/spec/lib/personas/tools/edit_image_spec.rb b/spec/lib/personas/tools/edit_image_spec.rb index 881e3701..4242aec4 100644 --- a/spec/lib/personas/tools/edit_image_spec.rb +++ b/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" }]