FIX: RAG uploader must support multi-file indexing. (#592)

Updating the editing model's rag_uploads in the editor component broke multi-file uploading. Instead, we'll keep the uploads in the uploader and update the model when we finish. This PR also fast-tracks the initial update so we can show feedback to the user quickly, and allows uploading MD files. Bug reported on https://meta.discourse.org/t/discourse-ai-persona-upload-support/304049/11
discourse · Apr 25, 2024 · 283445c · 283445c
1 parent 0c4069a
commit 283445c
Show file tree

Hide file tree

Showing 5 changed files with 56 additions and 22 deletions.
diff --git a/app/controllers/discourse_ai/admin/ai_personas_controller.rb b/app/controllers/discourse_ai/admin/ai_personas_controller.rb
@@ -152,10 +152,13 @@ def permit_commands(commands)
 
       def validate_extension!(filename)
         extension = File.extname(filename)[1..-1] || ""
-        authorized_extension = "txt"
-        if extension != authorized_extension
+        authorized_extensions = %w[txt md]
+        if !authorized_extensions.include?(extension)
           raise Discourse::InvalidParameters.new(
-                  I18n.t("upload.unauthorized", authorized_extensions: authorized_extension),
+                  I18n.t(
+                    "upload.unauthorized",
+                    authorized_extensions: authorized_extensions.join(" "),
+                  ),
                 )
         end
       end

diff --git a/app/jobs/regular/digest_rag_upload.rb b/app/jobs/regular/digest_rag_upload.rb
@@ -26,6 +26,8 @@ def execute(args)
         document = get_uploaded_file(upload)
         return if document.nil?
 
+        RagDocumentFragment.publish_status(upload, { total: 0, indexed: 0, left: 0 })
+
         fragment_ids = []
         idx = 0
 
@@ -54,11 +56,6 @@ def execute(args)
         end
       end
 
-      RagDocumentFragment.publish_status(
-        upload,
-        { total: fragment_ids.size, indexed: 0, left: fragment_ids.size },
-      )
-
       fragment_ids.each_slice(50) do |slice|
         Jobs.enqueue(:generate_rag_embeddings, fragment_ids: slice)
       end

diff --git a/assets/javascripts/discourse/components/ai-persona-editor.gjs b/assets/javascripts/discourse/components/ai-persona-editor.gjs
@@ -207,11 +207,8 @@ export default class PersonaEditor extends Component {
   }
 
   @action
-  addUpload(upload) {
-    const newUpload = upload;
-    newUpload.status = "uploaded";
-    newUpload.statusText = I18n.t("discourse_ai.ai_persona.uploads.uploaded");
-    this.editingModel.rag_uploads.addObject(newUpload);
+  updateUploads(uploads) {
+    this.editingModel.rag_uploads = uploads;
   }
 
   @action
@@ -460,8 +457,7 @@ export default class PersonaEditor extends Component {
         <div class="control-group">
           <PersonaRagUploader
             @persona={{this.editingModel}}
-            @ragUploads={{this.editingModel.rag_uploads}}
-            @onAdd={{this.addUpload}}
+            @updateUploads={{this.updateUploads}}
             @onRemove={{this.removeUpload}}
           />
           <a

diff --git a/assets/javascripts/discourse/components/persona-rag-uploader.gjs b/assets/javascripts/discourse/components/persona-rag-uploader.gjs
@@ -3,6 +3,7 @@ import Component, { Input } from "@ember/component";
 import { fn } from "@ember/helper";
 import { on } from "@ember/modifier";
 import { action } from "@ember/object";
+import willDestroy from "@ember/render-modifiers/modifiers/will-destroy";
 import { inject as service } from "@ember/service";
 import DButton from "discourse/components/d-button";
 import { ajax } from "discourse/lib/ajax";
@@ -20,6 +21,7 @@ export default class PersonaRagUploader extends Component.extend(
   @tracked term = null;
   @tracked filteredUploads = null;
   @tracked ragIndexingStatuses = null;
+  @tracked ragUploads = null;
   id = "discourse-ai-persona-rag-uploader";
   maxFiles = 20;
   uploadUrl = "/admin/plugins/discourse-ai/ai-personas/files/upload";
@@ -32,7 +34,8 @@ export default class PersonaRagUploader extends Component.extend(
       this._uppyInstance?.cancelAll();
     }
 
-    this.filteredUploads = this.ragUploads || [];
+    this.ragUploads = this.persona?.rag_uploads || [];
+    this.filteredUploads = this.ragUploads;
 
     if (this.ragUploads?.length && this.persona?.id) {
       ajax(
@@ -41,10 +44,27 @@ export default class PersonaRagUploader extends Component.extend(
         this.set("ragIndexingStatuses", statuses);
       });
     }
+
+    this.appEvents.on(
+      `upload-mixin:${this.id}:all-uploads-complete`,
+      this,
+      "_updatePersonaWithUploads"
+    );
+  }
+
+  removeListener() {
+    this.appEvents.off(`upload-mixin:${this.id}:all-uploads-complete`);
+  }
+
+  _updatePersonaWithUploads() {
+    this.updateUploads(this.ragUploads);
   }
 
   uploadDone(uploadedFile) {
-    this.onAdd(uploadedFile.upload);
+    const newUpload = uploadedFile.upload;
+    newUpload.status = "uploaded";
+    newUpload.statusText = I18n.t("discourse_ai.ai_persona.uploads.uploaded");
+    this.ragUploads.pushObject(newUpload);
     this.debouncedSearch();
   }
 
@@ -79,8 +99,16 @@ export default class PersonaRagUploader extends Component.extend(
     discourseDebounce(this, this.search, 100);
   }
 
+  @action
+  removeUpload(upload) {
+    this.ragUploads.removeObject(upload);
+    this.onRemove(upload);
+
+    this.debouncedSearch();
+  }
+
   <template>
-    <div class="persona-rag-uploader">
+    <div class="persona-rag-uploader" {{willDestroy this.removeListener}}>
       <h3>{{I18n.t "discourse_ai.ai_persona.uploads.title"}}</h3>
       <p>{{I18n.t "discourse_ai.ai_persona.uploads.description"}}</p>
       <p>{{I18n.t "discourse_ai.ai_persona.uploads.hint"}}</p>
@@ -118,7 +146,7 @@ export default class PersonaRagUploader extends Component.extend(
                 <DButton
                   @icon="times"
                   @title="discourse_ai.ai_persona.uploads.remove"
-                  @action={{fn @onRemove upload}}
+                  @action={{fn this.removeUpload upload}}
                   @class="btn-flat"
                 />
               </td>
@@ -153,7 +181,7 @@ export default class PersonaRagUploader extends Component.extend(
         disabled={{this.uploading}}
         type="file"
         multiple="multiple"
-        accept=".txt"
+        accept=".txt,.md"
       />
       <DButton
         @label="discourse_ai.ai_persona.uploads.button"

diff --git a/assets/javascripts/discourse/components/rag-upload-progress.gjs b/assets/javascripts/discourse/components/rag-upload-progress.gjs
@@ -30,17 +30,27 @@ export default class RagUploadProgress extends Component {
   @bind
   onIndexingUpdate(data) {
     // Order not guaranteed. Discard old updates.
-    if (!this.updatedProgress || this.updatedProgress.left > data.left) {
+    if (
+      !this.updatedProgress ||
+      data.total === 0 ||
+      this.updatedProgress.left > data.left
+    ) {
       this.updatedProgress = data;
     }
   }
 
   get calculateProgress() {
+    if (this.progress.total === 0) {
+      return 0;
+    }
+
     return Math.ceil((this.progress.indexed * 100) / this.progress.total);
   }
 
   get fullyIndexed() {
-    return this.progress && this.progress.left === 0;
+    return (
+      this.progress && this.progress.total !== 0 && this.progress.left === 0
+    );
   }
 
   get progress() {