From ba44f9117b618a49867418b1ccfe8834a26c9ee2 Mon Sep 17 00:00:00 2001 From: Soulter <905617992@qq.com> Date: Fri, 24 Oct 2025 16:37:37 +0800 Subject: [PATCH] feat: enhance document upload process with batch settings and improved chunk handling --- astrbot/core/db/vec_db/base.py | 15 + .../db/vec_db/faiss_impl/embedding_storage.py | 16 ++ astrbot/core/db/vec_db/faiss_impl/vec_db.py | 51 ++++ astrbot/core/knowledge_base/chunking/base.py | 2 +- .../knowledge_base/chunking/fixed_size.py | 11 +- astrbot/core/knowledge_base/kb_helper.py | 31 ++- astrbot/core/provider/provider.py | 59 ++++ astrbot/dashboard/routes/knowledge_base.py | 168 +++++++---- .../en-US/features/knowledge-base/detail.json | 9 +- .../zh-CN/features/knowledge-base/detail.json | 9 +- .../components/DocumentsTab.vue | 262 ++++++++++-------- 11 files changed, 444 insertions(+), 189 deletions(-) diff --git a/astrbot/core/db/vec_db/base.py b/astrbot/core/db/vec_db/base.py index d71cf6bf8..d100aa71e 100644 --- a/astrbot/core/db/vec_db/base.py +++ b/astrbot/core/db/vec_db/base.py @@ -24,6 +24,21 @@ class BaseVecDB: """ ... + @abc.abstractmethod + async def insert_batch( + self, + contents: list[str], + metadatas: list[dict] | None = None, + ids: list[str] | None = None, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + ) -> int: + """ + 批量插入文本和其对应向量，自动生成 ID 并保持一致性。 + """ + ... + @abc.abstractmethod async def retrieve( self, diff --git a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py index f29084148..581032aa7 100644 --- a/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py +++ b/astrbot/core/db/vec_db/faiss_impl/embedding_storage.py @@ -36,6 +36,22 @@ class EmbeddingStorage: self.index.add_with_ids(vector.reshape(1, -1), np.array([id])) await self.save_index() + async def insert_batch(self, vectors: np.ndarray, ids: list[int]): + """批量插入向量 + + Args: + vectors (np.ndarray): 要插入的向量数组 + ids (list[int]): 向量的ID列表 + Raises: + ValueError: 如果向量的维度与存储的维度不匹配 + """ + assert self.index is not None, "FAISS index is not initialized." + if vectors.shape[1] != self.dimension: + raise ValueError( + f"向量维度不匹配, 期望: {self.dimension}, 实际: {vectors.shape[1]}" + ) + self.index.add_with_ids(vectors, np.array(ids)) + async def search(self, vector: np.ndarray, k: int) -> tuple: """搜索最相似的向量 diff --git a/astrbot/core/db/vec_db/faiss_impl/vec_db.py b/astrbot/core/db/vec_db/faiss_impl/vec_db.py index dfbfa3f31..77a1e605a 100644 --- a/astrbot/core/db/vec_db/faiss_impl/vec_db.py +++ b/astrbot/core/db/vec_db/faiss_impl/vec_db.py @@ -1,11 +1,13 @@ import uuid import json +import time import numpy as np from .document_storage import DocumentStorage from .embedding_storage import EmbeddingStorage from ..base import Result, BaseVecDB from astrbot.core.provider.provider import EmbeddingProvider from astrbot.core.provider.provider import RerankProvider +from astrbot import logger class FaissVecDB(BaseVecDB): @@ -60,6 +62,55 @@ class FaissVecDB(BaseVecDB): await self.embedding_storage.insert(vector, int_id) return int_id + async def insert_batch( + self, + contents: list[str], + metadatas: list[dict] | None = None, + ids: list[str] | None = None, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, + ) -> list[int]: + """ + 批量插入文本和其对应向量，自动生成 ID 并保持一致性。 + """ + assert self.document_storage.connection is not None, ( + "Database connection is not initialized." + ) + metadatas = metadatas or [{} for _ in contents] + ids = ids or [str(uuid.uuid4()) for _ in contents] + + start = time.time() + logger.debug(f"Generating embeddings for {len(contents)} contents...") + vectors = await self.embedding_provider.get_embeddings_batch( + contents, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + ) + end = time.time() + logger.debug( + f"Generated embeddings for {len(contents)} contents in {end - start:.2f} seconds." + ) + + int_ids = [] + async with self.document_storage.connection.cursor() as cursor: + for str_id, content, metadata in zip(ids, contents, metadatas): + await cursor.execute( + "INSERT INTO documents (doc_id, text, metadata) VALUES (?, ?, ?)", + (str_id, content, json.dumps(metadata)), + ) + await self.document_storage.connection.commit() + + for str_id in ids: + result = await self.document_storage.get_document_by_doc_id(str_id) + int_ids.append(result["id"]) + + # 批量插入向量到 FAISS + vectors_array = np.array(vectors).astype("float32") + await self.embedding_storage.insert_batch(vectors_array, int_ids) + return int_ids + async def retrieve( self, query: str, diff --git a/astrbot/core/knowledge_base/chunking/base.py b/astrbot/core/knowledge_base/chunking/base.py index bcc29a5cf..5aaf84ba1 100644 --- a/astrbot/core/knowledge_base/chunking/base.py +++ b/astrbot/core/knowledge_base/chunking/base.py @@ -13,7 +13,7 @@ class BaseChunker(ABC): """ @abstractmethod - async def chunk(self, text: str) -> list[str]: + async def chunk(self, text: str, **kwargs) -> list[str]: """将文本分块 Args: diff --git a/astrbot/core/knowledge_base/chunking/fixed_size.py b/astrbot/core/knowledge_base/chunking/fixed_size.py index 4d1a1b280..850c5adf7 100644 --- a/astrbot/core/knowledge_base/chunking/fixed_size.py +++ b/astrbot/core/knowledge_base/chunking/fixed_size.py @@ -22,7 +22,7 @@ class FixedSizeChunker(BaseChunker): self.chunk_size = chunk_size self.chunk_overlap = chunk_overlap - async def chunk(self, text: str) -> list[str]: + async def chunk(self, text: str, **kwargs) -> list[str]: """固定大小分块 Args: @@ -31,22 +31,25 @@ class FixedSizeChunker(BaseChunker): Returns: list[str]: 分块后的文本列表 """ + chunk_size = kwargs.get("chunk_size", self.chunk_size) + chunk_overlap = kwargs.get("chunk_overlap", self.chunk_overlap) + chunks = [] start = 0 text_len = len(text) while start < text_len: - end = start + self.chunk_size + end = start + chunk_size chunk = text[start:end] if chunk: chunks.append(chunk) # 移动窗口,保留重叠部分 - start = end - self.chunk_overlap + start = end - chunk_overlap # 防止无限循环: 如果重叠过大,直接移到end - if start >= end or self.chunk_overlap >= self.chunk_size: + if start >= end or chunk_overlap >= chunk_size: start = end return chunks diff --git a/astrbot/core/knowledge_base/kb_helper.py b/astrbot/core/knowledge_base/kb_helper.py index 32730f4c8..389d2cb5f 100644 --- a/astrbot/core/knowledge_base/kb_helper.py +++ b/astrbot/core/knowledge_base/kb_helper.py @@ -101,13 +101,18 @@ class KBHelper: file_name: str, file_content: bytes, file_type: str, + chunk_size: int = 512, + chunk_overlap: int = 50, + batch_size: int = 32, + tasks_limit: int = 3, + max_retries: int = 3, ) -> KBDocument: """上传并处理文档（带原子性保证和失败清理） 流程: 1. 保存原始文件 2. 解析文档内容 - 3. 提取多媒体资源 (TODO) + 3. 提取多媒体资源 4. 分块处理 5. 生成向量并存储 6. 保存元数据（事务） @@ -116,7 +121,6 @@ class KBHelper: await self._ensure_vec_db() doc_id = str(uuid.uuid4()) media_paths: list[Path] = [] - vec_doc_ids = [] # file_path = self.kb_files_dir / f"{doc_id}.{file_type}" # async with aiofiles.open(file_path, "wb") as f: @@ -144,18 +148,27 @@ class KBHelper: media_paths.append(Path(media.file_path)) # 分块并生成向量 - # saved_chunks = [] - chunks_text = await self.chunker.chunk(text_content) + chunks_text = await self.chunker.chunk( + text_content, chunk_size=chunk_size, chunk_overlap=chunk_overlap + ) + contents = [] + metadatas = [] for idx, chunk_text in enumerate(chunks_text): - vec_doc_id = await self.vec_db.insert( - content=chunk_text, - metadata={ + contents.append(chunk_text) + metadatas.append( + { "kb_id": self.kb.kb_id, "doc_id": doc_id, "chunk_index": idx, - }, + } ) - vec_doc_ids.append(str(vec_doc_id)) + await self.vec_db.insert_batch( + contents=contents, + metadatas=metadatas, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + ) # 保存文档的元数据 doc = KBDocument( diff --git a/astrbot/core/provider/provider.py b/astrbot/core/provider/provider.py index 901036b11..99ec9443b 100644 --- a/astrbot/core/provider/provider.py +++ b/astrbot/core/provider/provider.py @@ -1,4 +1,5 @@ import abc +import asyncio from typing import List from typing import AsyncGenerator from astrbot.core.agent.tool import ToolSet @@ -203,6 +204,64 @@ class EmbeddingProvider(AbstractProvider): """获取向量的维度""" ... + async def get_embeddings_batch( + self, + texts: list[str], + batch_size: int = 16, + tasks_limit: int = 3, + max_retries: int = 3, + ) -> list[list[float]]: + """批量获取文本的向量，分批处理以节省内存 + + Args: + texts: 文本列表 + batch_size: 每批处理的文本数量 + tasks_limit: 并发任务数量限制 + max_retries: 失败时的最大重试次数 + + Returns: + 向量列表 + """ + semaphore = asyncio.Semaphore(tasks_limit) + all_embeddings: list[list[float]] = [] + failed_batches: list[tuple[int, list[str]]] = [] + + async def process_batch(batch_idx: int, batch_texts: list[str]): + async with semaphore: + for attempt in range(max_retries): + try: + batch_embeddings = await self.get_embeddings(batch_texts) + all_embeddings.extend(batch_embeddings) + return + except Exception as e: + if attempt == max_retries - 1: + # 最后一次重试失败，记录失败的批次 + failed_batches.append((batch_idx, batch_texts)) + raise Exception( + f"批次 {batch_idx} 处理失败，已重试 {max_retries} 次: {str(e)}" + ) + # 等待一段时间后重试，使用指数退避 + await asyncio.sleep(2**attempt) + + tasks = [] + for i in range(0, len(texts), batch_size): + batch_texts = texts[i : i + batch_size] + batch_idx = i // batch_size + tasks.append(process_batch(batch_idx, batch_texts)) + + # 收集所有任务的结果，包括失败的任务 + results = await asyncio.gather(*tasks, return_exceptions=True) + + # 检查是否有失败的任务 + errors = [r for r in results if isinstance(r, Exception)] + if errors: + error_msg = ( + f"有 {len(errors)} 个批次处理失败: {'; '.join(str(e) for e in errors)}" + ) + raise Exception(error_msg) + + return all_embeddings + class RerankProvider(AbstractProvider): def __init__(self, provider_config: dict, provider_settings: dict) -> None: diff --git a/astrbot/dashboard/routes/knowledge_base.py b/astrbot/dashboard/routes/knowledge_base.py index a22110260..9bd71b0e6 100644 --- a/astrbot/dashboard/routes/knowledge_base.py +++ b/astrbot/dashboard/routes/knowledge_base.py @@ -428,17 +428,18 @@ class KnowledgeBaseRoute(Route): """上传文档 支持两种方式: - 1. multipart/form-data 文件上传 - 2. JSON 格式 base64 编码上传 + 1. multipart/form-data 文件上传（支持多文件，最多10个） + 2. JSON 格式 base64 编码上传（支持多文件，最多10个） Form Data (multipart/form-data): - kb_id: 知识库 ID (必填) - - file: 文件对象 (必填) + - file: 文件对象 (必填，可多个，字段名为 file, file1, file2, ... 或 files[]) JSON Body (application/json): - kb_id: 知识库 ID (必填) - - file_name: 文件名 (必填) - - file_content: base64 编码的文件内容 (必填) + - files: 文件数组 (必填) + - file_name: 文件名 (必填) + - file_content: base64 编码的文件内容 (必填) """ try: kb_manager = self._get_kb_manager() @@ -446,20 +447,46 @@ class KnowledgeBaseRoute(Route): # 检查 Content-Type content_type = request.content_type kb_id = None + chunk_size = None + chunk_overlap = None + batch_size = 32 + tasks_limit = 3 + max_retries = 3 + files_to_upload = [] # 存储待上传的文件信息列表 - if content_type and "multipart/form-data" in content_type: - # 方式 1: multipart/form-data - form_data = await request.form - files = await request.files + if content_type and "multipart/form-data" not in content_type: + return ( + Response().error("Content-Type 须为 multipart/form-data").__dict__ + ) + form_data = await request.form + files = await request.files - kb_id = form_data.get("kb_id") - if not kb_id: - return Response().error("缺少参数 kb_id").__dict__ + kb_id = form_data.get("kb_id") + chunk_size = int(form_data.get("chunk_size", 512)) + chunk_overlap = int(form_data.get("chunk_overlap", 50)) + batch_size = int(form_data.get("batch_size", 32)) + tasks_limit = int(form_data.get("tasks_limit", 3)) + max_retries = int(form_data.get("max_retries", 3)) + if not kb_id: + return Response().error("缺少参数 kb_id").__dict__ - if "file" not in files: - return Response().error("缺少文件").__dict__ + # 收集所有文件 + file_list = [] + # 支持 file, file1, file2, ... 或 files[] 格式 + for key in files.keys(): + if key == "file" or key.startswith("file") or key == "files[]": + file_items = files.getlist(key) + file_list.extend(file_items) - file = files["file"] + if not file_list: + return Response().error("缺少文件").__dict__ + + # 限制文件数量 + if len(file_list) > 10: + return Response().error("最多只能上传10个文件").__dict__ + + # 处理每个文件 + for file in file_list: file_name = file.filename # 保存到临时文件 @@ -470,64 +497,83 @@ class KnowledgeBaseRoute(Route): # 异步读取文件内容 async with aiofiles.open(temp_file_path, "rb") as f: file_content = await f.read() + + # 提取文件类型 + 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, + "file_type": file_type, + } + ) finally: # 清理临时文件 if os.path.exists(temp_file_path): os.remove(temp_file_path) - else: - # 方式 2: JSON base64 - import base64 - - data = await request.json - - kb_id = data.get("kb_id") - file_name = data.get("file_name") - file_content_b64 = data.get("file_content") - - if not kb_id or not file_name or not file_content_b64: - return ( - Response() - .error("缺少参数 kb_id, file_name 或 file_content") - .__dict__ - ) - - try: - file_content = base64.b64decode(file_content_b64) - except Exception: - return ( - Response() - .error("file_content 必须是有效的 base64 编码") - .__dict__ - ) - - # 提取文件类型 - file_type = file_name.rsplit(".", 1)[-1].lower() if "." in file_name else "" - + # 获取知识库 kb_helper = await kb_manager.get_kb(kb_id) if not kb_helper: return Response().error("知识库不存在").__dict__ - # 上传文档 - doc = await kb_helper.upload_document( - file_name=file_name, - file_content=file_content, - file_type=file_type, - ) + # 上传所有文档 + uploaded_docs = [] + failed_docs = [] - doc_dict = { - "doc_id": doc.doc_id, - "kb_id": doc.kb_id, - "doc_name": doc.doc_name, - "file_type": doc.file_type, - "file_size": doc.file_size, - "chunk_count": doc.chunk_count, - "media_count": doc.media_count, - "created_at": doc.created_at.isoformat(), - "updated_at": doc.updated_at.isoformat(), + for file_info in files_to_upload: + try: + doc = await kb_helper.upload_document( + file_name=file_info["file_name"], + file_content=file_info["file_content"], + file_type=file_info["file_type"], + chunk_size=chunk_size, + chunk_overlap=chunk_overlap, + batch_size=batch_size, + tasks_limit=tasks_limit, + max_retries=max_retries, + ) + + doc_dict = { + "doc_id": doc.doc_id, + "kb_id": doc.kb_id, + "doc_name": doc.doc_name, + "file_type": doc.file_type, + "file_size": doc.file_size, + "chunk_count": doc.chunk_count, + "media_count": doc.media_count, + "created_at": doc.created_at.isoformat(), + "updated_at": doc.updated_at.isoformat(), + } + uploaded_docs.append(doc_dict) + except Exception as e: + logger.error(f"上传文档 {file_info['file_name']} 失败: {e}") + failed_docs.append( + {"file_name": file_info["file_name"], "error": str(e)} + ) + + # 返回结果 + result = { + "uploaded": uploaded_docs, + "failed": failed_docs, + "total": len(files_to_upload), + "success_count": len(uploaded_docs), + "failed_count": len(failed_docs), } - return Response().ok(doc_dict, "上传文档成功").__dict__ + if failed_docs: + message = ( + f"部分文档上传成功 ({len(uploaded_docs)}/{len(files_to_upload)})" + ) + else: + message = ( + f"所有文档上传成功 ({len(uploaded_docs)}/{len(files_to_upload)})" + ) + + return Response().ok(result, message).__dict__ except ValueError as e: return Response().error(str(e)).__dict__ diff --git a/dashboard/src/i18n/locales/en-US/features/knowledge-base/detail.json b/dashboard/src/i18n/locales/en-US/features/knowledge-base/detail.json index d1a9daeee..9af7ef593 100644 --- a/dashboard/src/i18n/locales/en-US/features/knowledge-base/detail.json +++ b/dashboard/src/i18n/locales/en-US/features/knowledge-base/detail.json @@ -46,12 +46,19 @@ "selectFile": "Select File", "dropzone": "Drop files here or click to select", "supportedFormats": "Supported formats: TXT, PDF, Markdown", - "maxSize": "Max file size: 50MB", + "maxSize": "Max file size: 128MB", "chunkSettings": "Chunk Settings", + "batchSettings": "Batch Settings", "chunkSize": "Chunk Size", "chunkSizeHint": "Number of characters per chunk (default: 512)", "chunkOverlap": "Chunk Overlap", "chunkOverlapHint": "Overlapping characters between chunks (default: 50)", + "batchSize": "Batch Size", + "batchSizeHint": "Number of chunks to process in each batch (default: 32)", + "tasksLimit": "Concurrent Tasks Limit", + "tasksLimitHint": "Maximum number of concurrent upload tasks (default: 3)", + "maxRetries": "Max Retries", + "maxRetriesHint": "Number of times to retry a failed upload task (default: 3)", "cancel": "Cancel", "submit": "Upload", "fileRequired": "Please select a file to upload" diff --git a/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/detail.json b/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/detail.json index 59dc27c16..3dd9f751e 100644 --- a/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/detail.json +++ b/dashboard/src/i18n/locales/zh-CN/features/knowledge-base/detail.json @@ -47,12 +47,19 @@ "selectFile": "选择文件", "dropzone": "拖放文件到这里或点击选择", "supportedFormats": "支持的格式: TXT, PDF, Markdown", - "maxSize": "最大文件大小: 50MB", + "maxSize": "最大文件大小: 128MB", "chunkSettings": "分块设置", + "batchSettings": "批处理设置", "chunkSize": "分块大小", "chunkSizeHint": "每个文本块的字符数 (默认: 512)", "chunkOverlap": "分块重叠", "chunkOverlapHint": "相邻文本块之间的重叠字符数 (默认: 50)", + "batchSize": "批处理大小", + "batchSizeHint": "每批处理的文本块数量 (默认: 32)", + "tasksLimit": "并发任务限制", + "tasksLimitHint": "最大并发上传任务数 (默认: 3)", + "maxRetries": "最大重试次数", + "maxRetriesHint": "上传失败任务的重试次数 (默认: 3)", "cancel": "取消", "submit": "上传", "fileRequired": "请选择要上传的文件" diff --git a/dashboard/src/views/knowledge-base/components/DocumentsTab.vue b/dashboard/src/views/knowledge-base/components/DocumentsTab.vue index 814443415..4dd2ff622 100644 --- a/dashboard/src/views/knowledge-base/components/DocumentsTab.vue +++ b/dashboard/src/views/knowledge-base/components/DocumentsTab.vue @@ -2,35 +2,16 @@
- + {{ t('documents.upload') }} - +
- +