Merge pull request #709 from Kiln-AI/leonard/kil-128-the-add-document-screen-should-let-me-add-tags

feat: allow picking tags when uploading documents

Author: leonardmq

app/web_ui/src/lib/api_schema.d.ts

@@ -1969,6 +1969,8 @@ export interface components {
             files?: string[] | null;
             /** Names */
             names?: string[] | null;
+            /** Tags */
+            tags?: string[] | null;
         };
         /** Body_edit_tags_api_projects__project_id__documents_edit_tags_post */
         Body_edit_tags_api_projects__project_id__documents_edit_tags_post: {

app/web_ui/src/lib/stores/document_tag_store.ts

@@ -17,6 +17,7 @@ export const document_tags_errors_by_project_id = writable<
 
 export async function load_document_tags(
   project_id: string,
+  options?: { invalidate_cache: boolean },
 ): Promise<DocumentTagCounts> {
   try {
     // Return early if already loading
@@ -32,7 +33,7 @@ export async function load_document_tags(
     const existing_tag_counts = get(document_tag_store_by_project_id)[
       project_id
     ]
-    if (existing_tag_counts) {
+    if (existing_tag_counts && !options?.invalidate_cache) {
       return existing_tag_counts
     }
     const { data, error } = await client.GET(

app/web_ui/src/lib/ui/tag_dropdown.svelte

@@ -124,14 +124,16 @@
     switch (example_tag_set) {
       case "doc": {
         if (project_id !== null) {
-          current_tag_counts = document_tag_counts[project_id] || {}
+          // fallback here in case the document_tag_counts is still loading
+          current_tag_counts = document_tag_counts?.[project_id] ?? {}
         }
         default_tags = DEFAULT_DOC_TAGS
         break
       }
       case "task_run": {
         if (task_id !== null) {
-          current_tag_counts = task_tag_counts[task_id] || {}
+          // fallback here in case the task_tag_counts is still loading
+          current_tag_counts = task_tag_counts?.[task_id] ?? {}
         }
         default_tags = DEFAULT_TASK_TAGS
         break

app/web_ui/src/routes/(app)/docs/library/[project_id]/upload_file_dialog.svelte

@@ -4,12 +4,17 @@
   import Dialog from "$lib/ui/dialog.svelte"
   import TrashIcon from "$lib/ui/icons/trash_icon.svelte"
   import UploadIcon from "$lib/ui/icons/upload_icon.svelte"
+  import TagDropdown from "$lib/ui/tag_dropdown.svelte"
   import { ragProgressStore } from "$lib/stores/rag_progress_store"
+  import { load_document_tags } from "$lib/stores/document_tag_store"
   import type { BulkCreateDocumentsResponse } from "$lib/types"
   import posthog from "posthog-js"
+  import { createKilnError, KilnError } from "$lib/utils/error_handlers"
+  import FormElement from "$lib/utils/form_element.svelte"
 
   export let onUploadCompleted: () => void
 
+  let upload_error: KilnError | null = null
   let selected_files: File[] = []
   let file_input: HTMLInputElement
   let drag_over = false
@@ -97,7 +102,12 @@
   let unsupported_files_count = 0
   let show_success_dialog = false
 
+  // tags
+  let selected_tags: Set<string> = new Set()
+  let current_tag = ""
+
   async function handleUpload(): Promise<boolean> {
+    upload_error = null
     upload_in_progress = true
     upload_progress = 0
     upload_total = selected_files.length
@@ -112,6 +122,9 @@
         return false
       }
       return success
+    } catch (e) {
+      upload_error = createKilnError(e)
+      return false
     } finally {
       upload_in_progress = false
       upload_progress = 0
@@ -131,6 +144,12 @@
       formData.append(`names`, file.name)
     })
 
+    if (selected_tags.size > 0) {
+      Array.from(selected_tags).forEach((tag) => {
+        formData.append("tags", tag)
+      })
+    }
+
     const { data, error } = await client.POST(
       "/api/projects/{project_id}/documents/bulk",
       {
@@ -158,6 +177,11 @@
     selected_files = []
     onUploadCompleted()
 
+    // reload document tags for the project - because total counts have changed
+    // and we cannot know which ones due to partial upload rejection possibly
+    // happening on the backend
+    await load_document_tags(project_id, { invalidate_cache: true })
+
     ragProgressStore.run_all_rag_configs(project_id).catch((error) => {
       console.error("Error running all rag configs", error)
     })
@@ -244,6 +268,8 @@
     show_upload_result = false
     show_success_dialog = false
     unsupported_files_count = 0
+    selected_tags = new Set()
+    current_tag = ""
   }
 
   export function close() {
@@ -253,6 +279,8 @@
     show_upload_result = false
     show_success_dialog = false
     unsupported_files_count = 0
+    selected_tags = new Set()
+    current_tag = ""
     return true
   }
 
@@ -329,51 +357,98 @@
       {/if}
 
       {#if !show_success_dialog}
-        <!-- Dropzone -->
-        <div
-          class="border-2 border-dashed rounded-lg p-8 text-center transition-colors cursor-pointer {drag_over
-            ? 'border-primary bg-primary/5'
-            : 'border-gray-300 hover:border-gray-400'}"
-          on:dragover={handleDragOver}
-          on:dragleave={handleDragLeave}
-          on:drop={handleDrop}
-          on:click={openFileDialog}
-          role="button"
-          tabindex="0"
-          on:keydown={(e) => {
-            if (e.key === "Enter" || e.key === " ") {
-              e.preventDefault()
-              openFileDialog()
-            }
-          }}
-        >
-          <div class="space-y-2">
-            <div class="w-10 h-10 mx-auto text-gray-500">
-              <UploadIcon />
-            </div>
-            <div>
-              <p class="text-gray-500">Drop files here or click to select</p>
+        <div class="pb-2">
+          <!-- Dropzone -->
+          <div
+            class="border-2 border-dashed rounded-lg p-8 text-center transition-colors cursor-pointer {drag_over
+              ? 'border-primary bg-primary/5'
+              : 'border-gray-300 hover:border-gray-400'}"
+            on:dragover={handleDragOver}
+            on:dragleave={handleDragLeave}
+            on:drop={handleDrop}
+            on:click={openFileDialog}
+            role="button"
+            tabindex="0"
+            on:keydown={(e) => {
+              if (e.key === "Enter" || e.key === " ") {
+                e.preventDefault()
+                openFileDialog()
+              }
+            }}
+          >
+            <div class="space-y-2">
+              <div class="w-10 h-10 mx-auto text-gray-500">
+                <UploadIcon />
+              </div>
+              <div>
+                <p class="text-gray-500">Drop files here or click to select</p>
+              </div>
             </div>
           </div>
+
+          <!-- Hidden file input -->
+          <input
+            bind:this={file_input}
+            type="file"
+            multiple
+            class="hidden"
+            on:change={handleFileSelect}
+            accept={supported_file_types.join(",")}
+          />
+
+          {#if unsupported_files_count > 0}
+            <div class="text-error text-sm">
+              {unsupported_files_count} file{unsupported_files_count === 1
+                ? ""
+                : "s"} skipped due to unsupported format
+            </div>
+          {/if}
         </div>
 
-        <!-- Hidden file input -->
-        <input
-          bind:this={file_input}
-          type="file"
-          multiple
-          class="hidden"
-          on:change={handleFileSelect}
-          accept={supported_file_types.join(",")}
-        />
-
-        {#if unsupported_files_count > 0}
-          <div class="text-error text-sm">
-            {unsupported_files_count} file{unsupported_files_count === 1
-              ? ""
-              : "s"} skipped due to unsupported format
+        <!-- Tag selection -->
+        <div>
+          <FormElement
+            inputType="header_only"
+            label="Tags"
+            id="tags_section"
+            description="Add tags to organize your documents"
+            info_description="Any tags set here will be added to each document you add. Tags can be used to filter your document set."
+            optional={true}
+            value=""
+          />
+          {#if selected_tags.size > 0}
+            <div class="flex flex-row flex-wrap gap-2 my-2">
+              {#each Array.from(selected_tags).sort() as tag}
+                <div
+                  class="badge bg-gray-200 text-gray-500 py-3 px-3 max-w-full"
+                >
+                  <span class="truncate">{tag}</span>
+                  <button
+                    class="pl-3 font-medium shrink-0"
+                    on:click={() => {
+                      selected_tags.delete(tag)
+                      selected_tags = selected_tags
+                    }}>✕</button
+                  >
+                </div>
+              {/each}
+            </div>
+          {/if}
+          <div class="flex flex-row gap-2 items-center">
+            <TagDropdown
+              bind:tag={current_tag}
+              {project_id}
+              example_tag_set="doc"
+              on_select={(tag) => {
+                selected_tags.add(tag)
+                selected_tags = selected_tags
+                current_tag = ""
+              }}
+              on_escape={() => {}}
+              focus_on_mount={false}
+            />
           </div>
-        {/if}
+        </div>
 
         {#if show_upload_result && upload_result}
           {#if upload_result.created_documents.length > 0}
@@ -447,6 +522,12 @@
           </div>
         {/if}
       {/if}
+
+      {#if upload_error}
+        <div class="text-error text-sm">
+          {upload_error.getMessage() || "An unknown error occurred"}
+        </div>
+      {/if}
     </div>
   </div>
 </Dialog>

libs/server/kiln_server/document_api.py

@@ -598,6 +598,7 @@ async def create_documents_bulk(
         project_id: str,
         files: Annotated[List[UploadFile] | None, File()] = None,
         names: Annotated[List[str] | None, Form()] = None,
+        tags: Annotated[List[str] | None, Form()] = None,
     ) -> BulkCreateDocumentsResponse:
         project = project_from_id(project_id)
 
@@ -652,6 +653,7 @@ async def create_documents_bulk(
                     name_override=document_name,
                     description="",  # No description support in bulk upload
                     kind=kind,
+                    tags=tags if tags else [],
                     original_file=FileInfo(
                         filename=file.filename,
                         mime_type=mime_type,

libs/server/kiln_server/test_document_api.py

@@ -3871,6 +3871,101 @@ async def test_create_rag_config_invalid_tool_fields(
     assert "error_messages" in error_detail
 
 
+@pytest.mark.parametrize(
+    "tags",
+    [
+        [],
+        ["tag1"],
+        ["tag1", "tag2"],
+        ["tag1", "tag2", "tag3"],
+    ],
+)
+async def test_create_documents_bulk_with_tags_success(client, mock_project, tags):
+    """Test successful bulk upload with various tag combinations"""
+    project = mock_project
+    test_content_1 = b"test file content 1"
+    test_content_2 = b"test file content 2"
+
+    with (
+        patch("kiln_server.document_api.project_from_id") as mock_project_from_id,
+    ):
+        mock_project_from_id.return_value = project
+
+        files = [
+            ("files", ("test1.txt", io.BytesIO(test_content_1), "text/plain")),
+            ("files", ("test2.txt", io.BytesIO(test_content_2), "text/plain")),
+        ]
+        data = {"names": ["Custom Name 1", "Custom Name 2"]}
+
+        # Add tags to the data
+        for tag in tags:
+            data.setdefault("tags", []).append(tag)
+
+        response = client.post(
+            f"/api/projects/{project.id}/documents/bulk", files=files, data=data
+        )
+
+    assert response.status_code == 200, response.text
+    result = response.json()
+    assert "created_documents" in result
+    assert "failed_files" in result
+    assert len(result["created_documents"]) == 2
+    assert len(result["failed_files"]) == 0
+
+    # Check that both documents have all the tags
+    for doc in result["created_documents"]:
+        assert "tags" in doc
+        assert doc["tags"] == tags
+        assert len(doc["tags"]) == len(tags)
+        assert sorted(tags) == sorted(doc["tags"])
+
+
+@pytest.mark.parametrize(
+    "invalid_tags",
+    [
+        ["tag with spaces"],
+        ["tag with spaces", "valid_tag"],
+        ["", "valid_tag"],
+        ["   ", "valid_tag"],
+    ],
+)
+async def test_create_documents_bulk_with_invalid_tags_failure(
+    client, mock_project, invalid_tags
+):
+    """Test bulk upload failure due to invalid tags (spaces, empty strings)"""
+    project = mock_project
+    test_content = b"test file content"
+
+    with (
+        patch("kiln_server.document_api.project_from_id") as mock_project_from_id,
+    ):
+        mock_project_from_id.return_value = project
+
+        files = [
+            ("files", ("test.txt", io.BytesIO(test_content), "text/plain")),
+        ]
+        data = {"names": ["Custom Name"]}
+
+        # Add invalid tags to the data
+        for tag in invalid_tags:
+            data.setdefault("tags", []).append(tag)
+
+        response = client.post(
+            f"/api/projects/{project.id}/documents/bulk", files=files, data=data
+        )
+
+    # Should return 422 for invalid tags
+    assert response.status_code == 422, response.text
+    result = response.json()
+    assert "message" in result
+    assert "failed_files" in result["message"]
+    assert len(result["message"]["failed_files"]) == 1
+    assert (
+        "Tags cannot contain spaces" in result["message"]["failed_files"][0]
+        or "Tags cannot be empty strings" in result["message"]["failed_files"][0]
+    )
+
+
 @pytest.mark.asyncio
 async def test_delete_extraction_extractor_config_not_found(
     client, mock_project, mock_document