AstrBot 4.8.0__py3-none-any.whl → 4.9.1__py3-none-any.whl

astrbot/dashboard/routes/knowledge_base.py

@@ -48,6 +48,7 @@ class KnowledgeBaseRoute(Route):
             # 文档管理
             "/kb/document/list": ("GET", self.list_documents),
             "/kb/document/upload": ("POST", self.upload_document),
+            "/kb/document/import": ("POST", self.import_documents),
             "/kb/document/upload/url": ("POST", self.upload_document_from_url),
             "/kb/document/upload/progress": ("GET", self.get_upload_progress),
             "/kb/document/get": ("GET", self.get_document),
@@ -66,6 +67,65 @@ class KnowledgeBaseRoute(Route):
     def _get_kb_manager(self):
         return self.core_lifecycle.kb_manager
 
+    def _init_task(self, task_id: str, status: str = "pending") -> None:
+        self.upload_tasks[task_id] = {
+            "status": status,
+            "result": None,
+            "error": None,
+        }
+
+    def _set_task_result(
+        self, task_id: str, status: str, result: any = None, error: str | None = None
+    ) -> None:
+        self.upload_tasks[task_id] = {
+            "status": status,
+            "result": result,
+            "error": error,
+        }
+        if task_id in self.upload_progress:
+            self.upload_progress[task_id]["status"] = status
+
+    def _update_progress(
+        self,
+        task_id: str,
+        *,
+        status: str | None = None,
+        file_index: int | None = None,
+        file_name: str | None = None,
+        stage: str | None = None,
+        current: int | None = None,
+        total: int | None = None,
+    ) -> None:
+        if task_id not in self.upload_progress:
+            return
+        p = self.upload_progress[task_id]
+        if status is not None:
+            p["status"] = status
+        if file_index is not None:
+            p["file_index"] = file_index
+        if file_name is not None:
+            p["file_name"] = file_name
+        if stage is not None:
+            p["stage"] = stage
+        if current is not None:
+            p["current"] = current
+        if total is not None:
+            p["total"] = total
+
+    def _make_progress_callback(self, task_id: str, file_idx: int, file_name: str):
+        async def _callback(stage: str, current: int, total: int):
+            self._update_progress(
+                task_id,
+                status="processing",
+                file_index=file_idx,
+                file_name=file_name,
+                stage=stage,
+                current=current,
+                total=total,
+            )
+
+        return _callback
+
     async def _background_upload_task(
         self,
         task_id: str,
@@ -80,11 +140,7 @@ class KnowledgeBaseRoute(Route):
         """后台上传任务"""
         try:
             # 初始化任务状态
-            self.upload_tasks[task_id] = {
-                "status": "processing",
-                "result": None,
-                "error": None,
-            }
+            self._init_task(task_id, status="processing")
             self.upload_progress[task_id] = {
                 "status": "processing",
                 "file_index": 0,
@@ -100,30 +156,20 @@ class KnowledgeBaseRoute(Route):
             for file_idx, file_info in enumerate(files_to_upload):
                 try:
                     # 更新整体进度
-                    self.upload_progress[task_id].update(
-                        {
-                            "status": "processing",
-                            "file_index": file_idx,
-                            "file_name": file_info["file_name"],
-                            "stage": "parsing",
-                            "current": 0,
-                            "total": 100,
-                        },
+                    self._update_progress(
+                        task_id,
+                        status="processing",
+                        file_index=file_idx,
+                        file_name=file_info["file_name"],
+                        stage="parsing",
+                        current=0,
+                        total=100,
                     )
 
                     # 创建进度回调函数
-                    async def progress_callback(stage, current, total):
-                        if task_id in self.upload_progress:
-                            self.upload_progress[task_id].update(
-                                {
-                                    "status": "processing",
-                                    "file_index": file_idx,
-                                    "file_name": file_info["file_name"],
-                                    "stage": stage,
-                                    "current": current,
-                                    "total": total,
-                                },
-                            )
+                    progress_callback = self._make_progress_callback(
+                        task_id, file_idx, file_info["file_name"]
+                    )
 
                     doc = await kb_helper.upload_document(
                         file_name=file_info["file_name"],
@@ -154,23 +200,99 @@ class KnowledgeBaseRoute(Route):
                 "failed_count": len(failed_docs),
             }
 
-            self.upload_tasks[task_id] = {
-                "status": "completed",
-                "result": result,
-                "error": None,
-            }
-            self.upload_progress[task_id]["status"] = "completed"
+            self._set_task_result(task_id, "completed", result=result)
 
         except Exception as e:
             logger.error(f"后台上传任务 {task_id} 失败: {e}")
             logger.error(traceback.format_exc())
-            self.upload_tasks[task_id] = {
-                "status": "failed",
-                "result": None,
-                "error": str(e),
+            self._set_task_result(task_id, "failed", error=str(e))
+
+    async def _background_import_task(
+        self,
+        task_id: str,
+        kb_helper,
+        documents: list,
+        batch_size: int,
+        tasks_limit: int,
+        max_retries: int,
+    ):
+        """后台导入预切片文档任务"""
+        try:
+            # 初始化任务状态
+            self._init_task(task_id, status="processing")
+            self.upload_progress[task_id] = {
+                "status": "processing",
+                "file_index": 0,
+                "file_total": len(documents),
+                "stage": "waiting",
+                "current": 0,
+                "total": 100,
+            }
+
+            uploaded_docs = []
+            failed_docs = []
+
+            for file_idx, doc_info in enumerate(documents):
+                file_name = doc_info.get("file_name", f"imported_doc_{file_idx}")
+                chunks = doc_info.get("chunks", [])
+
+                try:
+                    # 更新整体进度
+                    self._update_progress(
+                        task_id,
+                        status="processing",
+                        file_index=file_idx,
+                        file_name=file_name,
+                        stage="importing",
+                        current=0,
+                        total=100,
+                    )
+
+                    # 创建进度回调函数
+                    progress_callback = self._make_progress_callback(
+                        task_id, file_idx, file_name
+                    )
+
+                    # 调用 upload_document，传入 pre_chunked_text
+                    doc = await kb_helper.upload_document(
+                        file_name=file_name,
+                        file_content=None,  # 预切片模式下不需要原始内容
+                        file_type=doc_info.get("file_type")
+                        or (
+                            file_name.rsplit(".", 1)[-1].lower()
+                            if "." in file_name
+                            else "txt"
+                        ),
+                        batch_size=batch_size,
+                        tasks_limit=tasks_limit,
+                        max_retries=max_retries,
+                        progress_callback=progress_callback,
+                        pre_chunked_text=chunks,
+                    )
+
+                    uploaded_docs.append(doc.model_dump())
+                except Exception as e:
+                    logger.error(f"导入文档 {file_name} 失败: {e}")
+                    failed_docs.append(
+                        {"file_name": file_name, "error": str(e)},
+                    )
+
+            # 更新任务完成状态
+            result = {
+                "task_id": task_id,
+                "uploaded": uploaded_docs,
+                "failed": failed_docs,
+                "total": len(documents),
+                "success_count": len(uploaded_docs),
+                "failed_count": len(failed_docs),
             }
-            if task_id in self.upload_progress:
-                self.upload_progress[task_id]["status"] = "failed"
+
+            self._set_task_result(task_id, "completed", result=result)
+
+        except Exception as e:
+            logger.error(f"后台导入任务 {task_id} 失败: {e}")
+            logger.error(traceback.format_exc())
+            self._set_task_result(task_id, "failed", error=str(e))
 
     async def list_kbs(self):
         """获取知识库列表
@@ -614,11 +736,7 @@ class KnowledgeBaseRoute(Route):
             task_id = str(uuid.uuid4())
 
             # 初始化任务状态
-            self.upload_tasks[task_id] = {
-                "status": "pending",
-                "result": None,
-                "error": None,
-            }
+            self._init_task(task_id, status="pending")
 
             # 启动后台任务
             asyncio.create_task(
@@ -653,6 +771,93 @@ class KnowledgeBaseRoute(Route):
             logger.error(traceback.format_exc())
             return Response().error(f"上传文档失败: {e!s}").__dict__
 
+    def _validate_import_request(self, data: dict):
+        kb_id = data.get("kb_id")
+        if not kb_id:
+            raise ValueError("缺少参数 kb_id")
+
+        documents = data.get("documents")
+        if not documents or not isinstance(documents, list):
+            raise ValueError("缺少参数 documents 或格式错误")
+
+        for doc in documents:
+            if "file_name" not in doc or "chunks" not in doc:
+                raise ValueError("文档格式错误，必须包含 file_name 和 chunks")
+            if not isinstance(doc["chunks"], list):
+                raise ValueError("chunks 必须是列表")
+            if not all(
+                isinstance(chunk, str) and chunk.strip() for chunk in doc["chunks"]
+            ):
+                raise ValueError("chunks 必须是非空字符串列表")
+
+        batch_size = data.get("batch_size", 32)
+        tasks_limit = data.get("tasks_limit", 3)
+        max_retries = data.get("max_retries", 3)
+        return kb_id, documents, batch_size, tasks_limit, max_retries
+
+    async def import_documents(self):
+        """导入预切片文档
+
+        Body:
+        - kb_id: 知识库 ID (必填)
+        - documents: 文档列表 (必填)
+            - file_name: 文件名 (必填)
+            - chunks: 切片列表 (必填, list[str])
+            - file_type: 文件类型 (可选, 默认从文件名推断或为 txt)
+        - batch_size: 批处理大小 (可选, 默认32)
+        - tasks_limit: 并发任务限制 (可选, 默认3)
+        - max_retries: 最大重试次数 (可选, 默认3)
+        """
+        try:
+            kb_manager = self._get_kb_manager()
+            data = await request.json
+
+            kb_id, documents, batch_size, tasks_limit, max_retries = (
+                self._validate_import_request(data)
+            )
+
+            # 获取知识库
+            kb_helper = await kb_manager.get_kb(kb_id)
+            if not kb_helper:
+                return Response().error("知识库不存在").__dict__
+
+            # 生成任务ID
+            task_id = str(uuid.uuid4())
+
+            # 初始化任务状态
+            self._init_task(task_id, status="pending")
+
+            # 启动后台任务
+            asyncio.create_task(
+                self._background_import_task(
+                    task_id=task_id,
+                    kb_helper=kb_helper,
+                    documents=documents,
+                    batch_size=batch_size,
+                    tasks_limit=tasks_limit,
+                    max_retries=max_retries,
+                ),
+            )
+
+            return (
+                Response()
+                .ok(
+                    {
+                        "task_id": task_id,
+                        "doc_count": len(documents),
+                        "message": "import task created, processing in background",
+                    },
+                )
+                .__dict__
+            )
+
+        except ValueError as e:
+            return Response().error(str(e)).__dict__
+        except Exception as e:
+            logger.error(f"导入文档失败: {e}")
+            logger.error(traceback.format_exc())
+            return Response().error(f"导入文档失败: {e!s}").__dict__
+
     async def get_upload_progress(self):
         """获取上传进度和结果
 
@@ -960,11 +1165,7 @@ class KnowledgeBaseRoute(Route):
             task_id = str(uuid.uuid4())
 
             # 初始化任务状态
-            self.upload_tasks[task_id] = {
-                "status": "pending",
-                "result": None,
-                "error": None,
-            }
+            self._init_task(task_id, status="pending")
 
             # 启动后台任务
             asyncio.create_task(
@@ -1017,11 +1218,7 @@ class KnowledgeBaseRoute(Route):
         """后台上传URL任务"""
         try:
             # 初始化任务状态
-            self.upload_tasks[task_id] = {
-                "status": "processing",
-                "result": None,
-                "error": None,
-            }
+            self._init_task(task_id, status="processing")
             self.upload_progress[task_id] = {
                 "status": "processing",
                 "file_index": 0,
@@ -1033,18 +1230,7 @@ class KnowledgeBaseRoute(Route):
             }
 
             # 创建进度回调函数
-            async def progress_callback(stage, current, total):
-                if task_id in self.upload_progress:
-                    self.upload_progress[task_id].update(
-                        {
-                            "status": "processing",
-                            "file_index": 0,
-                            "file_name": f"URL: {url}",
-                            "stage": stage,
-                            "current": current,
-                            "total": total,
-                        },
-                    )
+            progress_callback = self._make_progress_callback(task_id, 0, f"URL: {url}")
 
             # 上传文档
             doc = await kb_helper.upload_from_url(
@@ -1069,20 +1255,9 @@ class KnowledgeBaseRoute(Route):
                 "failed_count": 0,
             }
 
-            self.upload_tasks[task_id] = {
-                "status": "completed",
-                "result": result,
-                "error": None,
-            }
-            self.upload_progress[task_id]["status"] = "completed"
+            self._set_task_result(task_id, "completed", result=result)
 
         except Exception as e:
             logger.error(f"后台上传URL任务 {task_id} 失败: {e}")
             logger.error(traceback.format_exc())
-            self.upload_tasks[task_id] = {
-                "status": "failed",
-                "result": None,
-                "error": str(e),
-            }
-            if task_id in self.upload_progress:
-                self.upload_progress[task_id]["status"] = "failed"
+            self._set_task_result(task_id, "failed", error=str(e))

astrbot/dashboard/routes/log.py

@@ -1,6 +1,8 @@
 import asyncio
 import json
+from typing import cast
 
+from quart import Response as QuartResponse
 from quart import make_response
 
 from astrbot.core import LogBroker, logger
@@ -39,14 +41,17 @@ class LogRoute(Route):
                 if queue:
                     self.log_broker.unregister(queue)
 
-        response = await make_response(
-            stream(),
-            {
-                "Content-Type": "text/event-stream",
-                "Cache-Control": "no-cache",
-                "Connection": "keep-alive",
-                "Transfer-Encoding": "chunked",
-            },
+        response = cast(
+            QuartResponse,
+            await make_response(
+                stream(),
+                {
+                    "Content-Type": "text/event-stream",
+                    "Cache-Control": "no-cache",
+                    "Connection": "keep-alive",
+                    "Transfer-Encoding": "chunked",
+                },
+            ),
         )
         response.timeout = None
         return response

astrbot/dashboard/routes/platform.py

@@ -82,7 +82,7 @@ class PlatformRoute(Route):
         """
         for platform in self.platform_manager.platform_insts:
             if platform.config.get("webhook_uuid") == webhook_uuid:
-                if platform.config.get("unified_webhook_mode", False):
+                if platform.unified_webhook():
                     return platform
         return None