Merge pull request #158 from merefield/paint_aspect_fixes

FIX: image aspect ratio rendering and edits

Author: merefield

config/locales/server.en.yml

@@ -183,6 +183,7 @@ en:
             creates a picture based on the provided direction and the last image in the conversation
           parameters:
             description: a description of the edit to the picture the user wants to make
+            aspect_ratio: optional image shape override. if omitted, keep the original image aspect ratio
           no_image_error: there was no image in the conversation to edit
           error: there was an issue creating the new picture.
         forum_search:

lib/discourse_chatbot/functions/paint_edit_function.rb

@@ -1,13 +1,13 @@
 # frozen_string_literal: true
 
-require_relative '../function'
+require_relative "../function"
 
 module DiscourseChatbot
   class PaintEditFunction < Function
-    TOKEN_COST = 1000000 # 1M tokens per request based on cost of dall-e-3 model vs gpt-4o-mini
+    TOKEN_COST = 1_000_000 # 1M tokens per request based on cost of dall-e-3 model vs gpt-4o-mini
 
     def name
-      'paint_edit_picture'
+      "paint_edit_picture"
     end
 
     def description
@@ -16,12 +16,22 @@ def description
 
     def parameters
       [
-        { name: "description", type: String, description: I18n.t("chatbot.prompt.function.paint_edit.parameters.description") } ,
+        {
+          name: "description",
+          type: String,
+          description: I18n.t("chatbot.prompt.function.paint_edit.parameters.description"),
+        },
+        {
+          name: "aspect_ratio",
+          type: String,
+          enum: PaintFunction::ASPECT_RATIO_OPTIONS,
+          description: I18n.t("chatbot.prompt.function.paint_edit.parameters.aspect_ratio"),
+        },
       ]
     end
 
     def required
-      ['description']
+      ["description"]
     end
 
     def process(args, opts)
@@ -31,39 +41,49 @@ def process(args, opts)
 
         description = args[parameters[0][:name]]
 
-        client = OpenAI::Client.new do |f|
-          f.response :logger, Logger.new($stdout), bodies: true if SiteSetting.chatbot_enable_verbose_console_logging
-          if SiteSetting.chatbot_enable_verbose_rails_logging != "off"
-            case SiteSetting.chatbot_verbose_rails_logging_destination_level
+        client =
+          OpenAI::Client.new do |f|
+            if SiteSetting.chatbot_enable_verbose_console_logging
+              f.response :logger, Logger.new($stdout), bodies: true
+            end
+            if SiteSetting.chatbot_enable_verbose_rails_logging != "off"
+              case SiteSetting.chatbot_verbose_rails_logging_destination_level
               when "warn"
                 f.response :logger, Rails.logger, bodies: true, log_level: :warn
               else
                 f.response :logger, Rails.logger, bodies: true, log_level: :info
+              end
             end
           end
+
+        type = opts[:type]
+
+        last_image_upload =
+          (
+            if type == ::DiscourseChatbot::POST
+              last_post_image_upload(opts[:reply_to_message_or_post_id])
+            else
+              last_message_image_upload(opts[:reply_to_message_or_post_id])
+            end
+          )
+
+        if last_image_upload.nil?
+          return(
+            { answer: I18n.t("chatbot.prompt.function.paint_edit.no_image_error"), token_usage: 0 }
+          )
         end
 
-        size = "1536x1024"
+        aspect_ratio = resolved_aspect_ratio(args[parameters[1][:name]], last_image_upload)
+        size =
+          PaintFunction.size_for(SiteSetting.chatbot_support_picture_creation_model, aspect_ratio)
         quality = "auto"
 
         options = {
           model: SiteSetting.chatbot_support_picture_creation_model,
           prompt: description,
           size: size,
           quality: quality,
-       }
-
-        type = opts[:type]
-
-        last_image_upload = type == ::DiscourseChatbot::POST ? last_post_image_upload(opts[:reply_to_message_or_post_id]) : last_message_image_upload(opts[:reply_to_message_or_post_id])
-
-        return {
-          answer: I18n.t("chatbot.prompt.function.paint_edit.no_image_error"),
-          token_usage: 0
-        } if last_image_upload.nil?
-
-        file_path = path = Discourse.store.path_for(last_image_upload)
-        base64_encoded_data = Base64.strict_encode64(File.read(file_path))
+        }
 
         file_path = Discourse.store.path_for(last_image_upload)
         extension = last_image_upload.extension
@@ -73,16 +93,17 @@ def process(args, opts)
         f.binmode
         f.write(File.binread(file_path))
         f.rewind
-        
+
         # Specify the file with MIME type
         upload_io = Faraday::Multipart::FilePart.new(f, mime_type)
-        
+
         options[:image] = upload_io
-        
-        response = client.images.edit(parameters: options)
-        
-        f.close
-        f.unlink
+        begin
+          response = client.images.edit(parameters: options)
+        ensure
+          f.close
+          f.unlink
+        end
 
         if response.dig("error")
           error_text = "ERROR when trying to call paint API: #{response.dig("error", "message")}"
@@ -91,32 +112,30 @@ def process(args, opts)
 
         tokens_used = response.dig("usage", "total_tokens")
 
-        artifacts = response.dig("data")
-          .to_a
-          .map { |art| art["b64_json"] }
+        artifacts = response.dig("data").to_a.map { |art| art["b64_json"] }
 
         bot_username = SiteSetting.chatbot_bot_user
         bot_user = ::User.find_by(username: bot_username)
 
         thumbnails = base64_to_image(artifacts, description, bot_user.id)
-        short_url = thumbnails.first.short_url
-        markdown = "![#{description}|690x460](#{short_url})"
-
-        {
-          answer: markdown,
-          token_usage: tokens_used
-        }
+        markdown =
+          PaintFunction.markdown_for(
+            upload: thumbnails.first,
+            description: description,
+            fallback_size: size,
+          )
+
+        { answer: markdown, token_usage: tokens_used }
       rescue => e
         Rails.logger.error("Chatbot: Error in paint edit function: #{e}")
         if e.respond_to?(:response)
           status = e.response[:status]
           message = e.response[:body]["error"]["message"]
-          Rails.logger.error("Chatbot: There was a problem with Image call: status: #{status}, message: #{message}")
+          Rails.logger.error(
+            "Chatbot: There was a problem with Image call: status: #{status}, message: #{message}",
+          )
         end
-        {
-          answer: I18n.t("chatbot.prompt.function.paint_edit.error"),
-          token_usage: TOKEN_COST
-        }
+        { answer: I18n.t("chatbot.prompt.function.paint_edit.error"), token_usage: TOKEN_COST }
       end
     end
 
@@ -130,34 +149,39 @@ def base64_to_image(artifacts, description, user_id)
         f.binmode
         f.write(Base64.decode64(art))
         f.rewind
-        upload = UploadCreator.new(f, attribution).create_for(user_id)
-        f.unlink
+        begin
+          upload = UploadCreator.new(f, attribution).create_for(user_id)
+        ensure
+          f.close
+          f.unlink
+        end
 
-        UploadSerializer.new(upload, root: false)
+        upload
       end
     end
 
+    def resolved_aspect_ratio(aspect_ratio, upload)
+      return PaintFunction.normalized_aspect_ratio(aspect_ratio) if aspect_ratio.present?
+
+      PaintFunction.aspect_ratio_for_upload(upload)
+    end
+
     def last_post_image_upload(post_id)
       post_collection = ::DiscourseChatbot::PostPromptUtils.collect_past_interactions(post_id)
 
       return nil if post_collection.empty?
-    
+
       upload_id = post_collection.map(&:image_upload_id).compact.max
-      return Upload.find_by(id: upload_id)
-    
-      nil
+      Upload.find_by(id: upload_id)
     end
 
     def last_message_image_upload(message_id)
-      message_collection = ::DiscourseChatbot::MessagePromptUtils.collect_past_interactions(message_id)
+      message_collection =
+        ::DiscourseChatbot::MessagePromptUtils.collect_past_interactions(message_id)
       uploads = []
 
       message_collection.each do |cm|
-        cm.uploads.each do |ul|
-          if %w[png webp jpg jpeg].include?(ul.extension)
-            uploads << ul
-          end
-        end
+        cm.uploads.each { |ul| uploads << ul if %w[png webp jpg jpeg].include?(ul.extension) }
       end
 
       return nil if uploads.empty?