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