implement cancel manager for slow ops

Author: SamSaffron

lib/inference/open_ai_image_generator.rb

@@ -21,7 +21,8 @@ def self.create_uploads!(
         moderation: "low",
         output_compression: nil,
         output_format: nil,
-        title: nil
+        title: nil,
+        cancel_manager: nil
       )
         # Get the API responses in parallel threads
         api_responses =
@@ -38,6 +39,7 @@ def self.create_uploads!(
             moderation: moderation,
             output_compression: output_compression,
             output_format: output_format,
+            cancel_manager: cancel_manager,
           )
 
         raise api_responses[0] if api_responses.all? { |resp| resp.is_a?(StandardError) }
@@ -58,7 +60,8 @@ def self.create_edited_upload!(
         user_id:,
         for_private_message: false,
         n: 1,
-        quality: nil
+        quality: nil,
+        cancel_manager: nil
       )
         api_response =
           edit_images(
@@ -70,6 +73,7 @@ def self.create_edited_upload!(
             api_url: api_url,
             n: n,
             quality: quality,
+            cancel_manager: cancel_manager,
           )
 
         create_uploads_from_responses([api_response], user_id, for_private_message).first
@@ -124,7 +128,8 @@ def self.generate_images_in_threads(
         background:,
         moderation:,
         output_compression:,
-        output_format:
+        output_format:,
+        cancel_manager:
       )
         prompts = [prompts] unless prompts.is_a?(Array)
         prompts = prompts.take(4) # Limit to 4 prompts max
@@ -152,18 +157,21 @@ def self.generate_images_in_threads(
                 moderation: moderation,
                 output_compression: output_compression,
                 output_format: output_format,
+                cancel_manager: cancel_manager,
               )
             rescue => e
               attempts += 1
               # to keep tests speedy
-              if !Rails.env.test?
+              if !Rails.env.test? && !cancel_manager&.canceled?
                 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?
+              if !cancel_manager&.canceled?
+                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?
+              end
               e
             end
           end
@@ -181,7 +189,8 @@ def self.edit_images(
         api_key: nil,
         api_url: nil,
         n: 1,
-        quality: nil
+        quality: nil,
+        cancel_manager: nil
       )
         images = [images] if !images.is_a?(Array)
 
@@ -209,8 +218,10 @@ def self.edit_images(
             api_url: api_url,
             n: n,
             quality: quality,
+            cancel_manager: cancel_manager,
           )
         rescue => e
+          raise e if cancel_manager&.canceled?
           attempts += 1
           if !Rails.env.test?
             sleep 2
@@ -238,7 +249,8 @@ def self.perform_generation_api_call!(
         background: nil,
         moderation: nil,
         output_compression: nil,
-        output_format: nil
+        output_format: nil,
+        cancel_manager: nil
       )
         api_key ||= SiteSetting.ai_openai_api_key
         api_url ||= SiteSetting.ai_openai_image_generation_url
@@ -276,6 +288,7 @@ def self.perform_generation_api_call!(
 
         # Store original prompt for upload metadata
         original_prompt = prompt
+        cancel_manager_callback = nil
 
         FinalDestination::HTTP.start(
           uri.host,
@@ -288,6 +301,11 @@ def self.perform_generation_api_call!(
           request = Net::HTTP::Post.new(uri, headers)
           request.body = payload.to_json
 
+          if cancel_manager
+            cancel_manager_callback = lambda { http.finish }
+            cancel_manager.add_callback(cancel_manager_callback)
+          end
+
           json = nil
           http.request(request) do |response|
             if response.code.to_i != 200
@@ -300,6 +318,10 @@ def self.perform_generation_api_call!(
           end
           json
         end
+      ensure
+        if cancel_manager && cancel_manager_callback
+          cancel_manager.remove_callback(cancel_manager_callback)
+        end
       end
 
       def self.perform_edit_api_call!(
@@ -310,7 +332,8 @@ def self.perform_edit_api_call!(
         api_key:,
         api_url:,
         n: 1,
-        quality: nil
+        quality: nil,
+        cancel_manager: nil
       )
         uri = URI(api_url)
 
@@ -403,6 +426,7 @@ def self.perform_edit_api_call!(
 
         # Store original prompt for upload metadata
         original_prompt = prompt
+        cancel_manager_callback = nil
 
         FinalDestination::HTTP.start(
           uri.host,
@@ -415,6 +439,11 @@ def self.perform_edit_api_call!(
           request = Net::HTTP::Post.new(uri.path, headers)
           request.body = body.join
 
+          if cancel_manager
+            cancel_manager_callback = lambda { http.finish }
+            cancel_manager.add_callback(cancel_manager_callback)
+          end
+
           json = nil
           http.request(request) do |response|
             if response.code.to_i != 200
@@ -428,6 +457,9 @@ def self.perform_edit_api_call!(
           json
         end
       ensure
+        if cancel_manager && cancel_manager_callback
+          cancel_manager.remove_callback(cancel_manager_callback)
+        end
         if files_to_delete.present?
           files_to_delete.each { |file| File.delete(file) if File.exist?(file) }
         end

lib/personas/artifact_update_strategies/base.rb

@@ -5,15 +5,24 @@ module ArtifactUpdateStrategies
       class InvalidFormatError < StandardError
       end
       class Base
-        attr_reader :post, :user, :artifact, :artifact_version, :instructions, :llm
+        attr_reader :post, :user, :artifact, :artifact_version, :instructions, :llm, :cancel_manager
 
-        def initialize(llm:, post:, user:, artifact:, artifact_version:, instructions:)
+        def initialize(
+          llm:,
+          post:,
+          user:,
+          artifact:,
+          artifact_version:,
+          instructions:,
+          cancel_manager:
+        )
           @llm = llm
           @post = post
           @user = user
           @artifact = artifact
           @artifact_version = artifact_version
           @instructions = instructions
+          @cancel_manager = cancel_manager
         end
 
         def apply(&progress)
@@ -26,7 +35,7 @@ def apply(&progress)
 
         def generate_changes(&progress)
           response = +""
-          llm.generate(build_prompt, user: user) do |partial|
+          llm.generate(build_prompt, user: user, cancel_manager: cancel_manager) do |partial|
             progress.call(partial) if progress
             response << partial
           end

lib/personas/tools/create_artifact.rb

@@ -151,7 +151,12 @@ def generate_artifact_code(post:, user:)
                 LlmModel.find_by(id: options[:creator_llm].to_i)&.to_llm
             ) || self.llm
 
-          llm.generate(prompt, user: user, feature_name: "create_artifact") do |partial_response|
+          llm.generate(
+            prompt,
+            user: user,
+            feature_name: "create_artifact",
+            cancel_manager: context.cancel_manager,
+          ) do |partial_response|
             response << partial_response
             yield partial_response
           end

lib/personas/tools/create_image.rb

@@ -48,6 +48,7 @@ def invoke
                 max_prompts,
                 model: "gpt-image-1",
                 user_id: bot_user.id,
+                cancel_manager: context.cancel_manager,
               )
           rescue => e
             @error = e

lib/personas/tools/edit_image.rb

@@ -60,6 +60,7 @@ def invoke
                 uploads,
                 prompt,
                 user_id: bot_user.id,
+                cancel_manager: context.cancel_manager,
               )
           rescue => e
             @error = e

lib/personas/tools/update_artifact.rb

@@ -159,6 +159,7 @@ def invoke
                   artifact: artifact,
                   artifact_version: artifact_version,
                   instructions: instructions,
+                  cancel_manager: context.cancel_manager,
                 )
                 .apply do |progress|
                   partial_response << progress