fix: remove knowledge base upload file limit

Author: Soulter

astrbot/dashboard/services/knowledge_base_service.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import asyncio
+import shutil
 import traceback
 import uuid
 from pathlib import Path
@@ -11,7 +12,7 @@
 from astrbot.core import logger
 from astrbot.core.core_lifecycle import AstrBotCoreLifecycle
 from astrbot.core.provider.provider import EmbeddingProvider, RerankProvider
-from astrbot.core.utils.astrbot_path import get_astrbot_temp_path
+from astrbot.core.utils.astrbot_path import get_astrbot_system_tmp_path
 from astrbot.dashboard.schemas import KnowledgeBaseRequest
 from astrbot.dashboard.utils import generate_tsne_visualization
 
@@ -116,17 +117,45 @@ def format_failed_doc_error(file_name: str, error: Exception) -> str:
             return message
         return f"{file_name}: {message}"
 
+    @staticmethod
+    def _cleanup_staging_dir(staging_dir: Path) -> None:
+        """Remove a knowledge base upload staging directory.
+
+        Args:
+            staging_dir: Task-specific system temporary directory to remove.
+        """
+        try:
+            shutil.rmtree(staging_dir)
+        except FileNotFoundError:
+            pass
+        except OSError as exc:
+            logger.warning(f"Failed to clean upload staging directory: {exc}")
+
     async def background_upload_task(
         self,
         task_id: str,
         kb_helper,
         files_to_upload: list[dict[str, Any]],
+        staging_dir: Path,
         chunk_size: int,
         chunk_overlap: int,
         batch_size: int,
         tasks_limit: int,
         max_retries: int,
     ) -> None:
+        """Process staged knowledge base files one at a time.
+
+        Args:
+            task_id: Identifier used to report upload progress and results.
+            kb_helper: Knowledge base helper that parses and stores documents.
+            files_to_upload: Metadata and temporary paths for staged files.
+            staging_dir: Task-specific system temporary directory.
+            chunk_size: Maximum size of each generated document chunk.
+            chunk_overlap: Number of overlapping characters between chunks.
+            batch_size: Number of chunks sent in each embedding batch.
+            tasks_limit: Maximum number of concurrent embedding tasks.
+            max_retries: Maximum retries for embedding operations.
+        """
         try:
             self.init_task(task_id, status="processing")
             self.upload_progress[task_id] = {
@@ -142,7 +171,11 @@ async def background_upload_task(
             failed_docs = []
 
             for file_idx, file_info in enumerate(files_to_upload):
+                file_content = None
                 try:
+                    temp_file_path = Path(file_info["temp_file_path"])
+                    async with aiofiles.open(temp_file_path, "rb") as file_obj:
+                        file_content = await file_obj.read()
                     self.update_progress(
                         task_id,
                         status="processing",
@@ -157,7 +190,7 @@ async def background_upload_task(
                     )
                     doc = await kb_helper.upload_document(
                         file_name=file_info["file_name"],
-                        file_content=file_info["file_content"],
+                        file_content=file_content,
                         file_type=file_info["file_type"],
                         chunk_size=chunk_size,
                         chunk_overlap=chunk_overlap,
@@ -177,6 +210,10 @@ async def background_upload_task(
                             ),
                         },
                     )
+                finally:
+                    # Release the current file before reading the next one.
+                    file_content = None
+                    Path(file_info["temp_file_path"]).unlink(missing_ok=True)
 
             self.set_task_result(
                 task_id,
@@ -194,6 +231,8 @@ async def background_upload_task(
             logger.error(f"后台上传任务 {task_id} 失败: {exc}")
             logger.error(traceback.format_exc())
             self.set_task_result(task_id, "failed", error=str(exc))
+        finally:
+            self._cleanup_staging_dir(staging_dir)
 
     async def background_import_task(
         self,
@@ -520,52 +559,62 @@ async def upload_document(
                 file_list.extend(files.getlist(key))
         if not file_list:
             raise KnowledgeBaseServiceError("缺少文件")
-        if len(file_list) > 10:
-            raise KnowledgeBaseServiceError("最多只能上传10个文件")
 
+        task_id = str(uuid.uuid4())
+        system_temp_root = Path(get_astrbot_system_tmp_path())
+        system_temp_root.mkdir(mode=0o700, parents=True, exist_ok=True)
+        staging_dir = system_temp_root / f"kb_upload_{task_id}"
+        staging_dir.mkdir(mode=0o700)
         files_to_upload = []
-        for file in file_list:
-            file_name = Path(str(file.filename or "document").replace("\\", "/")).name
-            if file_name in {"", ".", ".."}:
-                file_name = "document"
-            temp_file_path = (
-                Path(get_astrbot_temp_path()) / f"kb_upload_{uuid.uuid4()}_{file_name}"
-            )
-            await file.save(temp_file_path)
-            try:
-                async with aiofiles.open(temp_file_path, "rb") as file_obj:
-                    file_content = await file_obj.read()
+        try:
+            for file in file_list:
+                file_name = Path(
+                    str(file.filename or "document").replace("\\", "/")
+                ).name
+                if file_name in {"", ".", ".."}:
+                    file_name = "document"
+                temp_file_path = staging_dir / f"{uuid.uuid4()}_{file_name}"
                 file_type = (
                     file_name.rsplit(".", 1)[-1].lower() if "." in file_name else ""
                 )
                 files_to_upload.append(
                     {
                         "file_name": file_name,
-                        "file_content": file_content,
+                        "temp_file_path": temp_file_path,
                         "file_type": file_type,
                     },
                 )
-            finally:
-                temp_file_path.unlink(missing_ok=True)
+                await file.save(temp_file_path)
+        except Exception:
+            self._cleanup_staging_dir(staging_dir)
+            raise
 
-        kb_helper = await self.get_kb_manager().get_kb(kb_id)
-        if not kb_helper:
-            raise KnowledgeBaseServiceError("知识库不存在")
+        try:
+            kb_helper = await self.get_kb_manager().get_kb(kb_id)
+            if not kb_helper:
+                raise KnowledgeBaseServiceError("知识库不存在")
+        except Exception:
+            self._cleanup_staging_dir(staging_dir)
+            raise
 
-        task_id = str(uuid.uuid4())
-        self.init_task(task_id, status="pending")
-        asyncio.create_task(
-            self.background_upload_task(
-                task_id=task_id,
-                kb_helper=kb_helper,
-                files_to_upload=files_to_upload,
-                chunk_size=chunk_size,
-                chunk_overlap=chunk_overlap,
-                batch_size=batch_size,
-                tasks_limit=tasks_limit,
-                max_retries=max_retries,
-            ),
-        )
+        try:
+            self.init_task(task_id, status="pending")
+            asyncio.create_task(
+                self.background_upload_task(
+                    task_id=task_id,
+                    kb_helper=kb_helper,
+                    files_to_upload=files_to_upload,
+                    staging_dir=staging_dir,
+                    chunk_size=chunk_size,
+                    chunk_overlap=chunk_overlap,
+                    batch_size=batch_size,
+                    tasks_limit=tasks_limit,
+                    max_retries=max_retries,
+                ),
+            )
+        except Exception:
+            self._cleanup_staging_dir(staging_dir)
+            raise
         return {
             "task_id": task_id,
             "file_count": len(files_to_upload),

dashboard/src/views/knowledge-base/components/DocumentsTab.vue

@@ -86,7 +86,6 @@
                 <p class="mt-4 text-h6">{{ t('upload.dropzone') }}</p>
                 <p class="text-caption text-medium-emphasis mt-2">{{ t('upload.supportedFormats') }}</p>
                 <p class="text-caption text-medium-emphasis">{{ t('upload.maxSize') }}</p>
-                <p class="text-caption text-medium-emphasis">最多可上传 10 个文件</p>
                 <input ref="fileInput" type="file" multiple hidden accept=".txt,.md,.markdown,.rst,.adoc,.pdf,.docx,.epub,.xls,.xlsx"
                   @change="handleFileSelect" />
               </div>
@@ -391,13 +390,8 @@ const handleFileSelect = (event: Event) => {
   target.value = ''
 }
 
-// 添加文件（检查数量限制）
+// Add files
 const addFiles = (files: File[]) => {
-  const totalFiles = selectedFiles.value.length + files.length
-  if (totalFiles > 10) {
-    showSnackbar('最多只能选择 10 个文件', 'warning')
-    return
-  }
   selectedFiles.value.push(...files)
 }
 

tests/unit/test_knowledge_base_service_contract.py

@@ -1,3 +1,5 @@
+import asyncio
+from pathlib import Path
 from types import SimpleNamespace
 from unittest.mock import AsyncMock, MagicMock
 
@@ -8,6 +10,7 @@
 from astrbot.dashboard.api.knowledge_bases import (
     list_knowledge_bases,
 )
+from astrbot.dashboard.api.multipart import MultiDict
 from astrbot.dashboard.schemas import (
     KnowledgeBaseRequest,
 )
@@ -250,7 +253,9 @@ async def test_create_kb_raises_when_embedding_provider_is_missing():
     kb_manager = MagicMock()
     service = make_service(kb_manager)
 
-    with pytest.raises(KnowledgeBaseServiceError, match="缺少参数 embedding_provider_id"):
+    with pytest.raises(
+        KnowledgeBaseServiceError, match="缺少参数 embedding_provider_id"
+    ):
         await service.create_kb({"kb_name": "Test KB"})
 
 
@@ -264,3 +269,78 @@ async def test_create_kb_raises_when_embedding_provider_is_invalid():
         await service.create_kb(
             {"kb_name": "Test KB", "embedding_provider_id": "missing-provider"}
         )
+
+
+@pytest.mark.asyncio
+async def test_upload_document_accepts_more_than_ten_files_and_cleans_temporary_files(
+    tmp_path, monkeypatch
+):
+    """Upload files without a count limit and remove their temporary copies.
+
+    Args:
+        tmp_path: Temporary directory provided by pytest.
+        monkeypatch: Pytest fixture used to isolate staging and task scheduling.
+    """
+    kb_helper = SimpleNamespace(
+        upload_document=AsyncMock(
+            return_value=SimpleNamespace(model_dump=lambda: {"doc_id": "doc-1"})
+        )
+    )
+    kb_manager = SimpleNamespace(get_kb=AsyncMock(return_value=kb_helper))
+    service = make_service(kb_manager)
+    uploads = []
+    for index in range(11):
+        content = f"content-{index}".encode()
+        uploads.append(
+            (
+                f"file{index}",
+                SimpleNamespace(
+                    filename=f"document-{index}.txt",
+                    save=AsyncMock(
+                        side_effect=lambda destination, content=content: Path(
+                            destination
+                        ).write_bytes(content)
+                    ),
+                ),
+            )
+        )
+
+    created_tasks = []
+    create_task = asyncio.create_task
+
+    def capture_task(coroutine):
+        """Capture a scheduled background task for deterministic waiting.
+
+        Args:
+            coroutine: Upload coroutine passed to ``asyncio.create_task``.
+
+        Returns:
+            The scheduled asyncio task.
+        """
+        task = create_task(coroutine)
+        created_tasks.append(task)
+        return task
+
+    monkeypatch.setattr(
+        "astrbot.dashboard.services.knowledge_base_service.get_astrbot_system_tmp_path",
+        lambda: tmp_path,
+    )
+    monkeypatch.setattr(
+        "astrbot.dashboard.services.knowledge_base_service.asyncio.create_task",
+        capture_task,
+    )
+
+    result = await service.upload_document(
+        content_type="multipart/form-data",
+        form_data=MultiDict([("kb_id", "kb-1")]),
+        files=MultiDict(uploads),
+    )
+    await created_tasks[0]
+
+    assert result["file_count"] == 11
+    assert kb_helper.upload_document.await_count == 11
+    assert [
+        call.kwargs["file_content"]
+        for call in kb_helper.upload_document.await_args_list
+    ] == [f"content-{index}".encode() for index in range(11)]
+    assert not list(tmp_path.glob("kb_upload_*"))