quota update logic

Author: SamSaffron

app/controllers/discourse_ai/admin/ai_llms_controller.rb

@@ -56,6 +56,25 @@ def create
       def update
         llm_model = LlmModel.find(params[:id])
 
+        if params[:ai_llm].key?(:llm_quotas)
+          if quota_params
+            existing_quota_group_ids = llm_model.llm_quotas.pluck(:group_id)
+            new_quota_group_ids = quota_params.map { |q| q[:group_id] }
+
+            llm_model
+              .llm_quotas
+              .where(group_id: existing_quota_group_ids - new_quota_group_ids)
+              .destroy_all
+
+            quota_params.each do |quota_param|
+              quota = llm_model.llm_quotas.find_or_initialize_by(group_id: quota_param[:group_id])
+              quota.update!(quota_param)
+            end
+          else
+            llm_model.llm_quotas.destroy_all
+          end
+        end
+
         if llm_model.seeded?
           return render_json_error(I18n.t("discourse_ai.llm.cannot_edit_builtin"), status: 403)
         end

spec/requests/admin/ai_llms_controller_spec.rb

@@ -256,6 +256,71 @@
     context "with valid update params" do
       let(:update_attrs) { { provider: "anthropic" } }
 
+      context "with quotas" do
+        it "updates quotas correctly" do
+          group1 = Fabricate(:group)
+          group2 = Fabricate(:group)
+          group3 = Fabricate(:group)
+
+          _quota1 =
+            Fabricate(
+              :llm_quota,
+              llm_model: llm_model,
+              group: group1,
+              max_tokens: 1000,
+              max_usages: 10,
+              duration_seconds: 86_400,
+            )
+          _quota2 =
+            Fabricate(
+              :llm_quota,
+              llm_model: llm_model,
+              group: group2,
+              max_tokens: 2000,
+              max_usages: 20,
+              duration_seconds: 86_400,
+            )
+
+          put "/admin/plugins/discourse-ai/ai-llms/#{llm_model.id}.json",
+              params: {
+                ai_llm: {
+                  llm_quotas: [
+                    {
+                      group_id: group1.id,
+                      max_tokens: 1500,
+                      max_usages: 15,
+                      duration_seconds: 43_200,
+                    },
+                    {
+                      group_id: group3.id,
+                      max_tokens: 3000,
+                      max_usages: 30,
+                      duration_seconds: 86_400,
+                    },
+                  ],
+                },
+              }
+
+          expect(response.status).to eq(200)
+
+          llm_model.reload
+          expect(llm_model.llm_quotas.count).to eq(2)
+
+          updated_quota1 = llm_model.llm_quotas.find_by(group: group1)
+          expect(updated_quota1.max_tokens).to eq(1500)
+          expect(updated_quota1.max_usages).to eq(15)
+          expect(updated_quota1.duration_seconds).to eq(43_200)
+
+          expect(llm_model.llm_quotas.find_by(group: group2)).to be_nil
+
+          new_quota = llm_model.llm_quotas.find_by(group: group3)
+          expect(new_quota).to be_present
+          expect(new_quota.max_tokens).to eq(3000)
+          expect(new_quota.max_usages).to eq(30)
+          expect(new_quota.duration_seconds).to eq(86_400)
+        end
+      end
+
       it "updates the model" do
         put "/admin/plugins/discourse-ai/ai-llms/#{llm_model.id}.json",
             params: {