From 110a1629aa30052f636f231602b20470814cd8b1 Mon Sep 17 00:00:00 2001
From: Keegan George <kgeorge13@gmail.com>
Date: Thu, 3 Oct 2024 02:36:35 +0900
Subject: [PATCH] DEV: Update rate limits for image captioning (#816)

This PR updates the rate limits for AI helper so that image caption follows a specific rate limit of 20 requests per minute. This should help when uploading multiple files that need to be captioned. This PR also updates the UI so that it shows toast message with the extracted error message instead of having a blocking `popupAjaxError` error dialog.
---------

Co-authored-by: Rafael dos Santos Silva <xfalcox@gmail.com>
Co-authored-by: Penar Musaraj <pmusaraj@gmail.com>
---
 .../ai_helper/assistant_controller.rb         | 21 +++++++++--
 .../initializers/ai-image-caption.js          | 11 ++++--
 .../ai_helper/assistant_controller_spec.rb    | 36 +++++++++++++++++++
 3 files changed, 64 insertions(+), 4 deletions(-)

diff --git a/app/controllers/discourse_ai/ai_helper/assistant_controller.rb b/app/controllers/discourse_ai/ai_helper/assistant_controller.rb
index 3594affe..55e1b0fa 100644
--- a/app/controllers/discourse_ai/ai_helper/assistant_controller.rb
+++ b/app/controllers/discourse_ai/ai_helper/assistant_controller.rb
@@ -6,10 +6,21 @@ module DiscourseAi
       requires_plugin ::DiscourseAi::PLUGIN_NAME
       requires_login
       before_action :ensure_can_request_suggestions
-      before_action :rate_limiter_performed!, except: %i[prompts]
+      before_action :rate_limiter_performed!
 
       include SecureUploadEndpointHelpers
 
+      RATE_LIMITS = {
+        "default" => {
+          amount: 6,
+          interval: 3.minutes,
+        },
+        "caption_image" => {
+          amount: 20,
+          interval: 1.minute,
+        },
+      }.freeze
+
       def suggest
         input = get_text_param!
         force_default_locale = params[:force_default_locale] || false
@@ -161,7 +172,13 @@ module DiscourseAi
       end
 
       def rate_limiter_performed!
-        RateLimiter.new(current_user, "ai_assistant", 6, 3.minutes).performed!
+        action_rate_limit = RATE_LIMITS[action_name] || RATE_LIMITS["default"]
+        RateLimiter.new(
+          current_user,
+          "ai_assistant",
+          action_rate_limit[:amount],
+          action_rate_limit[:interval],
+        ).performed!
       end
 
       def ensure_can_request_suggestions
diff --git a/assets/javascripts/initializers/ai-image-caption.js b/assets/javascripts/initializers/ai-image-caption.js
index b70bec65..245da8a7 100644
--- a/assets/javascripts/initializers/ai-image-caption.js
+++ b/assets/javascripts/initializers/ai-image-caption.js
@@ -1,5 +1,5 @@
 import { ajax } from "discourse/lib/ajax";
-import { popupAjaxError } from "discourse/lib/ajax-error";
+import { extractError, popupAjaxError } from "discourse/lib/ajax-error";
 import { apiInitializer } from "discourse/lib/api";
 import { getUploadMarkdown, isImage } from "discourse/lib/uploads";
 import I18n from "discourse-i18n";
@@ -111,7 +111,13 @@ export default apiInitializer("1.25.0", (api) => {
       });
       return response.caption;
     } catch (error) {
-      popupAjaxError(error);
+      toasts.error({
+        class: "ai-image-caption-error-toast",
+        duration: 3000,
+        data: {
+          message: extractError(error),
+        },
+      });
     }
   }
 
@@ -129,6 +135,7 @@ export default apiInitializer("1.25.0", (api) => {
     return;
   }
 
+  const toasts = api.container.lookup("service:toasts");
   // Automatically caption uploaded images
   api.addComposerUploadMarkdownResolver(async (upload) => {
     const autoCaptionEnabled = currentUser.get(
diff --git a/spec/requests/ai_helper/assistant_controller_spec.rb b/spec/requests/ai_helper/assistant_controller_spec.rb
index 9df109fb..21890411 100644
--- a/spec/requests/ai_helper/assistant_controller_spec.rb
+++ b/spec/requests/ai_helper/assistant_controller_spec.rb
@@ -100,6 +100,23 @@ RSpec.describe DiscourseAi::AiHelper::AssistantController do
           expect(response.parsed_body["diff"]).to eq(expected_diff)
         end
       end
+
+      context "when performing numerous requests" do
+        it "rate limits" do
+          RateLimiter.enable
+          rate_limit = described_class::RATE_LIMITS["default"]
+          amount = rate_limit[:amount]
+
+          amount.times do
+            post "/discourse-ai/ai-helper/suggest", params: { mode: mode, text: text_to_proofread }
+            expect(response.status).to eq(200)
+          end
+          DiscourseAi::Completions::Llm.with_prepared_responses([proofread_text]) do
+            post "/discourse-ai/ai-helper/suggest", params: { mode: mode, text: text_to_proofread }
+            expect(response.status).to eq(429)
+          end
+        end
+      end
     end
   end
 
@@ -258,6 +275,25 @@ RSpec.describe DiscourseAi::AiHelper::AssistantController do
           end
         end
       end
+
+      context "when performing numerous requests" do
+        it "rate limits" do
+          RateLimiter.enable
+
+          rate_limit = described_class::RATE_LIMITS["caption_image"]
+          amount = rate_limit[:amount]
+
+          amount.times do
+            request_caption({ image_url: image_url, image_url_type: "long_url" }) do |r|
+              expect(r.status).to eq(200)
+            end
+          end
+
+          request_caption({ image_url: image_url, image_url_type: "long_url" }) do |r|
+            expect(r.status).to eq(429)
+          end
+        end
+      end
     end
   end
 end