From 7b4c0996730f153e134e87178ccb50ecb3c8d55e Mon Sep 17 00:00:00 2001
From: Roman Rizzi <roman@discourse.org>
Date: Tue, 6 Aug 2024 14:35:35 -0300
Subject: [PATCH] FIX: LlmModel validations. (#742)

- Validate fields to reduce the chance of breaking features by a misconfigured model.
- Fixed a bug where the URL might get deleted during an update.
- Display a warning when a model is currently in use.
---
 .../discourse_ai/admin/ai_llms_controller.rb  |  5 +-
 app/models/llm_model.rb                       |  7 +++
 app/serializers/llm_model_serializer.rb       |  8 +--
 .../components/ai-llm-editor-form.gjs         | 17 +++++++
 config/locales/client.en.yml                  |  4 ++
 spec/fabricators/llm_model_fabricator.rb      |  2 +
 .../requests/admin/ai_llms_controller_spec.rb | 50 ++++++++++++++-----
 7 files changed, 76 insertions(+), 17 deletions(-)

diff --git a/app/controllers/discourse_ai/admin/ai_llms_controller.rb b/app/controllers/discourse_ai/admin/ai_llms_controller.rb
index 4707cb6b..635a4a2f 100644
--- a/app/controllers/discourse_ai/admin/ai_llms_controller.rb
+++ b/app/controllers/discourse_ai/admin/ai_llms_controller.rb
@@ -97,6 +97,8 @@ module DiscourseAi
       private
 
       def ai_llm_params(updating: nil)
+        return {} if params[:ai_llm].blank?
+
         permitted =
           params.require(:ai_llm).permit(
             :display_name,
@@ -112,7 +114,8 @@ module DiscourseAi
         provider = updating ? updating.provider : permitted[:provider]
         permit_url = provider != LlmModel::BEDROCK_PROVIDER_NAME
 
-        permitted[:url] = params.dig(:ai_llm, :url) if permit_url
+        new_url = params.dig(:ai_llm, :url)
+        permitted[:url] = new_url if permit_url && new_url
 
         extra_field_names = LlmModel.provider_params.dig(provider&.to_sym, :fields).to_a
         received_prov_params = params.dig(:ai_llm, :provider_params)
diff --git a/app/models/llm_model.rb b/app/models/llm_model.rb
index 8e6049b4..d3e92854 100644
--- a/app/models/llm_model.rb
+++ b/app/models/llm_model.rb
@@ -6,6 +6,13 @@ class LlmModel < ActiveRecord::Base
 
   belongs_to :user
 
+  validates :display_name, presence: true, length: { maximum: 100 }
+  validates :tokenizer, presence: true, inclusion: DiscourseAi::Completions::Llm.tokenizer_names
+  validates :provider, presence: true, inclusion: DiscourseAi::Completions::Llm.provider_names
+  validates :url, presence: true, unless: -> { provider == BEDROCK_PROVIDER_NAME }
+  validates_presence_of :name, :api_key
+  validates :max_prompt_tokens, numericality: { greater_than: 0 }
+
   def self.provider_params
     {
       aws_bedrock: {
diff --git a/app/serializers/llm_model_serializer.rb b/app/serializers/llm_model_serializer.rb
index 2ef1b28b..ca852409 100644
--- a/app/serializers/llm_model_serializer.rb
+++ b/app/serializers/llm_model_serializer.rb
@@ -12,13 +12,13 @@ class LlmModelSerializer < ApplicationSerializer
              :api_key,
              :url,
              :enabled_chat_bot,
-             :shadowed_by_srv,
              :provider_params,
-             :vision_enabled
+             :vision_enabled,
+             :used_by
 
   has_one :user, serializer: BasicUserSerializer, embed: :object
 
-  def shadowed_by_srv
-    object.url.to_s.starts_with?("srv://")
+  def used_by
+    DiscourseAi::Configuration::LlmValidator.new.modules_using(object)
   end
 end
diff --git a/assets/javascripts/discourse/components/ai-llm-editor-form.gjs b/assets/javascripts/discourse/components/ai-llm-editor-form.gjs
index 48d239a5..172ba0c7 100644
--- a/assets/javascripts/discourse/components/ai-llm-editor-form.gjs
+++ b/assets/javascripts/discourse/components/ai-llm-editor-form.gjs
@@ -65,6 +65,17 @@ export default class AiLlmEditorForm extends Component {
     return this.metaProviderParams.url_editable !== false;
   }
 
+  get modulesUsingModel() {
+    return this.args.model.used_by?.join(", ");
+  }
+
+  get inUseWarning() {
+    return I18n.t("discourse_ai.llms.in_use_warning", {
+      settings: this.modulesUsingModel,
+      count: this.args.model.used_by.length,
+    });
+  }
+
   @computed("args.model.provider")
   get metaProviderParams() {
     return (
@@ -166,6 +177,12 @@ export default class AiLlmEditorForm extends Component {
   }
 
   <template>
+    {{#if this.modulesUsingModel}}
+      <div class="alert alert-info">
+        {{icon "exclamation-circle"}}
+        {{this.inUseWarning}}
+      </div>
+    {{/if}}
     <form class="form-horizontal ai-llm-editor">
       <div class="control-group">
         <label>{{i18n "discourse_ai.llms.display_name"}}</label>
diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml
index 5f4a9117..3f25599b 100644
--- a/config/locales/client.en.yml
+++ b/config/locales/client.en.yml
@@ -236,6 +236,10 @@ en:
         back: "Back"
         confirm_delete: Are you sure you want to delete this model?
         delete: Delete
+        in_use_warning: 
+          one: "This model is currently used by the %{settings} setting. If misconfigured, the feature won't work as expected."
+          other: "This model is currently used by the following settings: %{settings}. If misconfigured, features won't work as expected. "
+
 
         preconfigured_llms: "Select your LLM"
         preconfigured:
diff --git a/spec/fabricators/llm_model_fabricator.rb b/spec/fabricators/llm_model_fabricator.rb
index 6ef5863c..5dd4adcb 100644
--- a/spec/fabricators/llm_model_fabricator.rb
+++ b/spec/fabricators/llm_model_fabricator.rb
@@ -43,6 +43,8 @@ Fabricator(:fake_model, from: :llm_model) do
   provider "fake"
   tokenizer "DiscourseAi::Tokenizer::OpenAiTokenizer"
   max_prompt_tokens 32_000
+  api_key "fake"
+  url "https://fake.test/"
 end
 
 Fabricator(:gemini_model, from: :llm_model) do
diff --git a/spec/requests/admin/ai_llms_controller_spec.rb b/spec/requests/admin/ai_llms_controller_spec.rb
index 68a76d48..77d0b7ae 100644
--- a/spec/requests/admin/ai_llms_controller_spec.rb
+++ b/spec/requests/admin/ai_llms_controller_spec.rb
@@ -20,19 +20,19 @@ RSpec.describe DiscourseAi::Admin::AiLlmsController do
   end
 
   describe "POST #create" do
-    context "with valid attributes" do
-      let(:valid_attrs) do
-        {
-          display_name: "My cool LLM",
-          name: "gpt-3.5",
-          provider: "open_ai",
-          url: "https://test.test/v1/chat/completions",
-          api_key: "test",
-          tokenizer: "DiscourseAi::Tokenizer::OpenAiTokenizer",
-          max_prompt_tokens: 16_000,
-        }
-      end
+    let(:valid_attrs) do
+      {
+        display_name: "My cool LLM",
+        name: "gpt-3.5",
+        provider: "open_ai",
+        url: "https://test.test/v1/chat/completions",
+        api_key: "test",
+        tokenizer: "DiscourseAi::Tokenizer::OpenAiTokenizer",
+        max_prompt_tokens: 16_000,
+      }
+    end
 
+    context "with valid attributes" do
       it "creates a new LLM model" do
         post "/admin/plugins/discourse-ai/ai-llms.json", params: { ai_llm: valid_attrs }
 
@@ -84,6 +84,19 @@ RSpec.describe DiscourseAi::Admin::AiLlmsController do
         expect(created_model.lookup_custom_param("access_key_id")).to be_nil
       end
     end
+
+    context "with invalid attributes" do
+      it "doesn't create a model" do
+        post "/admin/plugins/discourse-ai/ai-llms.json",
+             params: {
+               ai_llm: valid_attrs.except(:url),
+             }
+
+        created_model = LlmModel.last
+
+        expect(created_model).to be_nil
+      end
+    end
   end
 
   describe "PUT #update" do
@@ -130,6 +143,19 @@ RSpec.describe DiscourseAi::Admin::AiLlmsController do
       end
     end
 
+    context "with invalid update params" do
+      it "doesn't update the model" do
+        put "/admin/plugins/discourse-ai/ai-llms/#{llm_model.id}.json",
+            params: {
+              ai_llm: {
+                url: "",
+              },
+            }
+
+        expect(response.status).to eq(422)
+      end
+    end
+
     context "with provider-specific params" do
       it "updates provider-specific config params" do
         provider_params = { organization: "Discourse" }