FIX: Move the bot user toggling to the controller. (#688)

Having this as a callback prevents deploys of sites with a vLLM SRV configured and pending migrations. Additionally, this fixes a bug where we didn't delete/deactivate the companion user after deleting an LLM.
This commit is contained in:
Roman Rizzi 2024-06-25 12:45:19 -03:00 committed by GitHub
parent b6f0ad157c
commit e39e0bdb4a
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 51 additions and 6 deletions

View File

@ -35,6 +35,7 @@ module DiscourseAi
def create
llm_model = LlmModel.new(ai_llm_params)
if llm_model.save
llm_model.toggle_companion_user
render json: { ai_persona: llm_model }, status: :created
else
render_json_error llm_model
@ -70,6 +71,10 @@ module DiscourseAi
)
end
# Clean up companion users
llm_model.enabled_chat_bot = false
llm_model.toggle_companion_user
if llm_model.destroy
head :no_content
else

View File

@ -8,7 +8,6 @@ class LlmModel < ActiveRecord::Base
belongs_to :user
validates :url, exclusion: { in: [RESERVED_VLLM_SRV_URL] }
before_save :toggle_companion_user_before_save
def self.enable_or_disable_srv_llm!
srv_model = find_by(url: RESERVED_VLLM_SRV_URL)
@ -42,10 +41,6 @@ class LlmModel < ActiveRecord::Base
}
end
def toggle_companion_user_before_save
toggle_companion_user if enabled_chat_bot_changed? || new_record?
end
def toggle_companion_user
return if name == "fake" && Rails.env.production?

View File

@ -3,7 +3,10 @@
RSpec.describe DiscourseAi::Admin::AiLlmsController do
fab!(:admin)
before { sign_in(admin) }
before do
sign_in(admin)
SiteSetting.ai_bot_enabled = true
end
describe "GET #index" do
it "includes all available providers metadata" do
@ -42,6 +45,17 @@ RSpec.describe DiscourseAi::Admin::AiLlmsController do
expect(created_model.max_prompt_tokens).to eq(valid_attrs[:max_prompt_tokens])
end
it "creates a companion user" do
post "/admin/plugins/discourse-ai/ai-llms.json",
params: {
ai_llm: valid_attrs.merge(enabled_chat_bot: true),
}
created_model = LlmModel.last
expect(created_model.user_id).to be_present
end
it "stores provider-specific config params" do
provider_params = { organization: "Discourse" }
@ -93,6 +107,27 @@ RSpec.describe DiscourseAi::Admin::AiLlmsController do
expect(response.status).to eq(404)
end
it "creates a companion user" do
put "/admin/plugins/discourse-ai/ai-llms/#{llm_model.id}.json",
params: {
ai_llm: update_attrs.merge(enabled_chat_bot: true),
}
expect(llm_model.reload.user_id).to be_present
end
it "removes the companion user when desabling the chat bot option" do
llm_model.update!(enabled_chat_bot: true)
llm_model.toggle_companion_user
put "/admin/plugins/discourse-ai/ai-llms/#{llm_model.id}.json",
params: {
ai_llm: update_attrs.merge(enabled_chat_bot: false),
}
expect(llm_model.reload.user_id).to be_nil
end
end
context "with provider-specific params" do
@ -190,5 +225,15 @@ RSpec.describe DiscourseAi::Admin::AiLlmsController do
expect(response.status).to eq(409)
expect(fake_llm.reload).to eq(fake_llm)
end
it "cleans up companion users before deleting the model" do
llm_model.update!(enabled_chat_bot: true)
llm_model.toggle_companion_user
companion_user = llm_model.user
delete "/admin/plugins/discourse-ai/ai-llms/#{llm_model.id}.json"
expect { companion_user.reload }.to raise_error(ActiveRecord::RecordNotFound)
end
end
end