Move allowing or disallowing pdf/images to a site setting

SamSaffron · SamSaffron · commit 25c97cad9687 · 2025-02-12T11:32:59.000+11:00
diff --git a/admin/assets/javascripts/discourse/routes/admin-plugins-show-discourse-ai-tools-edit.js b/admin/assets/javascripts/discourse/routes/admin-plugins-show-discourse-ai-tools-edit.js
@@ -15,5 +15,6 @@ export default class DiscourseAiToolsEditRoute extends DiscourseRoute {
     controller.set("allTools", toolsModel);
     controller.set("presets", toolsModel.resultSetMeta.presets);
     controller.set("llms", toolsModel.resultSetMeta.llms);
+    controller.set("settings", toolsModel.resultSetMeta.settings);
   }
 }
diff --git a/admin/assets/javascripts/discourse/routes/admin-plugins-show-discourse-ai-tools-new.js b/admin/assets/javascripts/discourse/routes/admin-plugins-show-discourse-ai-tools-new.js
@@ -12,5 +12,6 @@ export default class DiscourseAiToolsNewRoute extends DiscourseRoute {
     controller.set("allTools", toolsModel);
     controller.set("presets", toolsModel.resultSetMeta.presets);
     controller.set("llms", toolsModel.resultSetMeta.llms);
+    controller.set("settings", toolsModel.resultSetMeta.settings);
   }
 }
diff --git a/admin/assets/javascripts/discourse/templates/admin-plugins/show/discourse-ai-tools/edit.hbs b/admin/assets/javascripts/discourse/templates/admin-plugins/show/discourse-ai-tools/edit.hbs
@@ -4,5 +4,6 @@
     @model={{this.model}}
     @presets={{this.presets}}
     @llms={{this.llms}}
+    @settings={{this.settings}}
   />
 </section>
diff --git a/admin/assets/javascripts/discourse/templates/admin-plugins/show/discourse-ai-tools/new.hbs b/admin/assets/javascripts/discourse/templates/admin-plugins/show/discourse-ai-tools/new.hbs
@@ -4,5 +4,6 @@
     @model={{this.model}}
     @presets={{this.presets}}
     @llms={{this.llms}}
+    @settings={{this.settings}}
   />
 </section>
diff --git a/app/controllers/discourse_ai/admin/ai_personas_controller.rb b/app/controllers/discourse_ai/admin/ai_personas_controller.rb
@@ -35,7 +35,16 @@ def index
           DiscourseAi::Configuration::LlmEnumerator.values_for_serialization(
             allowed_seeded_llm_ids: SiteSetting.ai_bot_allowed_seeded_models_map,
           )
-        render json: { ai_personas: ai_personas, meta: { tools: tools, llms: llms } }
+        render json: {
+                 ai_personas: ai_personas,
+                 meta: {
+                   tools: tools,
+                   llms: llms,
+                   settings: {
+                     rag_pdf_images_enabled: SiteSetting.ai_rag_pdf_images_enabled,
+                   },
+                 },
+               }
       end
 
       def new
diff --git a/app/controllers/discourse_ai/admin/rag_document_fragments_controller.rb b/app/controllers/discourse_ai/admin/rag_document_fragments_controller.rb
@@ -48,7 +48,8 @@ def upload_file
 
       def validate_extension!(filename)
         extension = File.extname(filename)[1..-1] || ""
-        authorized_extensions = %w[txt md pdf png jpg jpeg]
+        authorized_extensions = %w[txt md]
+        authorized_extensions.concat(%w[pdf png jpg jpeg]) if SiteSetting.ai_rag_pdf_images_enabled
         if !authorized_extensions.include?(extension)
           raise Discourse::InvalidParameters.new(
                   I18n.t(
diff --git a/app/jobs/regular/digest_rag_upload.rb b/app/jobs/regular/digest_rag_upload.rb
@@ -164,6 +164,11 @@ def first_chunk(text, chunk_tokens:, tokenizer:, splitters: ["\n\n", "\n", ".",
     end
 
     def get_uploaded_file(upload:, target:)
+      if %w[pdf png jpg jpeg].include?(upload.extension) && !SiteSetting.ai_rag_pdf_images_enabled
+        raise Discourse::InvalidAccess.new(
+                "The setting ai_rag_pdf_images_enabled is false, can not index images and pdfs.",
+              )
+      end
       if upload.extension == "pdf"
         pages =
           DiscourseAi::Utils::PdfToImages.new(
diff --git a/app/serializers/ai_custom_tool_list_serializer.rb b/app/serializers/ai_custom_tool_list_serializer.rb
@@ -9,6 +9,9 @@ def meta
     {
       presets: AiTool.presets,
       llms: DiscourseAi::Configuration::LlmEnumerator.values_for_serialization,
+      settings: {
+        rag_pdf_images_enabled: SiteSetting.ai_rag_pdf_images_enabled,
+      },
     }
   end
 
diff --git a/assets/javascripts/discourse/components/ai-persona-editor.gjs b/assets/javascripts/discourse/components/ai-persona-editor.gjs
@@ -596,11 +596,13 @@ export default class PersonaEditor extends Component {
             @target={{this.editingModel}}
             @updateUploads={{this.updateUploads}}
             @onRemove={{this.removeUpload}}
+            @allowPdfsAndImages={{@personas.resultSetMeta.settings.rag_pdf_images_enabled}}
           />
         </div>
         <RagOptions
           @model={{this.editingModel}}
           @llms={{@personas.resultSetMeta.llms}}
+          @allowPdfsAndImages={{@personas.resultSetMeta.settings.rag_pdf_images_enabled}}
         >
           <div class="control-group">
             <label>{{i18n
diff --git a/assets/javascripts/discourse/components/ai-tool-editor.gjs b/assets/javascripts/discourse/components/ai-tool-editor.gjs
@@ -245,9 +245,14 @@ export default class AiToolEditor extends Component {
               @target={{this.editingModel}}
               @updateUploads={{this.updateUploads}}
               @onRemove={{this.removeUpload}}
+              @allowPdfsAndImages={{@settings.rag_pdf_images_enabled}}
             />
           </div>
-          <RagOptions @model={{this.editingModel}} @llms={{this.args.llms}} />
+          <RagOptions
+            @model={{this.editingModel}}
+            @llms={{this.args.llms}}
+            @allowPdfsAndImages={{@settings.rag_pdf_images_enabled}}
+          />
         {{/if}}
 
         <div class="control-group ai-tool-editor__action_panel">
diff --git a/assets/javascripts/discourse/components/rag-options.gjs b/assets/javascripts/discourse/components/rag-options.gjs
@@ -81,18 +81,20 @@ export default class RagOptions extends Component {
           }}
         />
       </div>
-      <div class="control-group">
-        <label>{{i18n "discourse_ai.rag.options.rag_llm_model"}}</label>
-        <AiLlmSelector
-          class="ai-persona-editor__llms"
-          @value={{this.visionLlmId}}
-          @llms={{this.visionLlms}}
-        />
-        <DTooltip
-          @icon="circle-question"
-          @content={{i18n "discourse_ai.rag.options.rag_llm_model_help"}}
-        />
-      </div>
+      {{#if @allowPdfsAndImages}}
+        <div class="control-group">
+          <label>{{i18n "discourse_ai.rag.options.rag_llm_model"}}</label>
+          <AiLlmSelector
+            class="ai-persona-editor__llms"
+            @value={{this.visionLlmId}}
+            @llms={{this.visionLlms}}
+          />
+          <DTooltip
+            @icon="circle-question"
+            @content={{i18n "discourse_ai.rag.options.rag_llm_model_help"}}
+          />
+        </div>
+      {{/if}}
       {{yield}}
     {{/if}}
   </template>
diff --git a/assets/javascripts/discourse/components/rag-uploader.gjs b/assets/javascripts/discourse/components/rag-uploader.gjs
@@ -77,6 +77,14 @@ export default class RagUploader extends Component {
     this.updateUploads(this.ragUploads);
   }
 
+  get acceptedFileTypes() {
+    if (this.args?.allowPdfsAndImages) {
+      return ".txt,.md,.pdf,.png,.jpg,.jpeg";
+    } else {
+      return ".txt,.md";
+    }
+  }
+
   @action
   submitFiles() {
     this.uppyUpload.openPicker();
@@ -119,7 +127,11 @@ export default class RagUploader extends Component {
   <template>
     <div class="rag-uploader">
       <h3>{{i18n "discourse_ai.rag.uploads.title"}}</h3>
-      <p>{{i18n "discourse_ai.rag.uploads.description"}}</p>
+      {{#if @allowPdfsAndImages}}
+        <p>{{i18n "discourse_ai.rag.uploads.description_with_pdfs"}}</p>
+      {{else}}
+        <p>{{i18n "discourse_ai.rag.uploads.description"}}</p>
+      {{/if}}
 
       {{#if this.ragUploads}}
         <div class="rag-uploader__search-input-container">
@@ -187,7 +199,7 @@ export default class RagUploader extends Component {
         disabled={{this.uploading}}
         type="file"
         multiple="multiple"
-        accept=".txt,.md,.pdf,.png,.jpg,.jpeg"
+        accept={{this.acceptedFileTypes}}
       />
       <DButton
         @label="discourse_ai.rag.uploads.button"
diff --git a/config/locales/client.en.yml b/config/locales/client.en.yml
@@ -280,7 +280,8 @@ en:
           hide_indexing_options: "Hide upload options"
         uploads:
           title: "Uploads"
-          description: "Plaintext (.txt), markdown (.md), PDF (.pdf) or image (.png, .jpeg)."
+          description: "Plaintext (.txt) or markdown (.md)."
+          description_with_pdfs: "Plaintext (.txt), markdown (.md), PDF (.pdf) or image (.png, .jpeg)."
           button: "Add files"
           filter: "Filter uploads"
           indexed: "Indexed"
diff --git a/config/settings.yml b/config/settings.yml
@@ -347,12 +347,14 @@ discourse_ai:
   ai_spam_detection_enabled:
     default: false
     hidden: true
-
   ai_spam_detection_user_id:
     default: ""
     hidden: true
-
   ai_spam_detection_model_allowed_seeded_models:
     default: ""
     hidden: true
     type: list
+
+  ai_rag_pdf_images_enabled:
+    default: false
+    hidden: true
diff --git a/spec/jobs/regular/digest_rag_upload_spec.rb b/spec/jobs/regular/digest_rag_upload_spec.rb
@@ -3,6 +3,7 @@
 RSpec.describe Jobs::DigestRagUpload do
   fab!(:persona) { Fabricate(:ai_persona) }
   fab!(:upload) { Fabricate(:upload, extension: "txt") }
+  fab!(:pdf_upload) { Fabricate(:upload, extension: "pdf") }
   let(:document_file) { StringIO.new("some text" * 200) }
 
   fab!(:cloudflare_embedding_def)
@@ -30,6 +31,19 @@
   end
 
   describe "#execute" do
+    context "when processing a PDF upload" do
+      it "will reject the indexing if the site setting is not enabled" do
+        SiteSetting.ai_rag_pdf_images_enabled = false
+
+        expect {
+          described_class.new.execute(
+            upload_id: pdf_upload.id,
+            target_id: persona.id,
+            target_type: persona.class.to_s,
+          )
+        }.to raise_error(Discourse::InvalidAccess)
+      end
+    end
     context "when processing an upload containing metadata" do
       it "correctly splits on metadata boundary" do
         # be explicit here about chunking strategy
diff --git a/spec/requests/admin/rag_document_fragments_controller_spec.rb b/spec/requests/admin/rag_document_fragments_controller_spec.rb
@@ -12,6 +12,8 @@
     SiteSetting.ai_embeddings_enabled = true
   end
 
+  after { @cleanup_files&.each(&:unlink) }
+
   describe "GET #indexing_status_check" do
     it "works for AiPersona" do
       get "/admin/plugins/discourse-ai/rag-document-fragments/files/status.json?target_type=AiPersona&target_id=#{ai_persona.id}"
@@ -22,6 +24,15 @@
   end
 
   describe "POST #upload_file" do
+    let :fake_pdf do
+      @cleanup_files ||= []
+      tempfile = Tempfile.new(%w[test .pdf])
+      tempfile.write("fake pdf")
+      tempfile.rewind
+      @cleanup_files << tempfile
+      tempfile
+    end
+
     it "works" do
       post "/admin/plugins/discourse-ai/rag-document-fragments/files/upload.json",
            params: {
@@ -33,5 +44,28 @@
       upload = Upload.last
       expect(upload.original_filename).to end_with("spec.txt")
     end
+
+    it "rejects PDF files if site setting is not enabled" do
+      SiteSetting.ai_rag_pdf_images_enabled = false
+
+      post "/admin/plugins/discourse-ai/rag-document-fragments/files/upload.json",
+           params: {
+             file: Rack::Test::UploadedFile.new(fake_pdf),
+           }
+
+      expect(response.status).to eq(400)
+    end
+
+    it "allows PDF files if site setting is enabled" do
+      SiteSetting.ai_rag_pdf_images_enabled = true
+
+      post "/admin/plugins/discourse-ai/rag-document-fragments/files/upload.json",
+           params: {
+             file: Rack::Test::UploadedFile.new(fake_pdf),
+           }
+
+      upload = Upload.last
+      expect(upload.original_filename).to end_with(".pdf")
+    end
   end
 end

Original file line number	Diff line number	Diff line change
`@@ -15,5 +15,6 @@ export default class DiscourseAiToolsEditRoute extends DiscourseRoute {`
`15`	`15`	`controller.set("allTools", toolsModel);`
`16`	`16`	`controller.set("presets", toolsModel.resultSetMeta.presets);`
`17`	`17`	`controller.set("llms", toolsModel.resultSetMeta.llms);`
	`18`	`+ controller.set("settings", toolsModel.resultSetMeta.settings);`
`18`	`19`	`}`
`19`	`20`	`}`
Original file line number	Diff line number	Diff line change
`@@ -12,5 +12,6 @@ export default class DiscourseAiToolsNewRoute extends DiscourseRoute {`
`12`	`12`	`controller.set("allTools", toolsModel);`
`13`	`13`	`controller.set("presets", toolsModel.resultSetMeta.presets);`
`14`	`14`	`controller.set("llms", toolsModel.resultSetMeta.llms);`
	`15`	`+ controller.set("settings", toolsModel.resultSetMeta.settings);`
`15`	`16`	`}`
`16`	`17`	`}`
Original file line number	Diff line number	Diff line change
`@@ -9,6 +9,9 @@ def meta`
`9`	`9`	`{`
`10`	`10`	`presets: AiTool.presets,`
`11`	`11`	`llms: DiscourseAi::Configuration::LlmEnumerator.values_for_serialization,`
	`12`	`+ settings: {`
	`13`	`+ rag_pdf_images_enabled: SiteSetting.ai_rag_pdf_images_enabled,`
	`14`	`+ },`
`12`	`15`	`}`
`13`	`16`	`end`
`14`	`17`