知识库的创建（3） - 知识库上传入口update_docs

---

前言

langchain-chatchat用到的底层函数很多，为了大家思路更清晰，我们换个顺序，先从上传文档+更新、保存向量库的入口 update_docs 开始看起。

update_docs 是一个用于更新知识库文档的函数。这个函数能够处理多个文件，并将它们转换为知识库文档，还可以处理自定义的文档并将其向量化。以下是对该函数的详细介绍：

一、添加了注释的源码

def update_docs(
        knowledge_base_name: str = Body(..., description="知识库名称", examples=["samples"]),
        file_names: List[str] = Body(..., description="文件名称，支持多文件", examples=[["file_name1", "text.txt"]]),
        chunk_size: int = Body(CHUNK_SIZE, description="知识库中单段文本最大长度"),
        chunk_overlap: int = Body(OVERLAP_SIZE, description="知识库中相邻文本重合长度"),
        zh_title_enhance: bool = Body(ZH_TITLE_ENHANCE, description="是否开启中文标题加强"),
        override_custom_docs: bool = Body(False, description="是否覆盖之前自定义的docs"),
        docs: Json = Body({}, description="自定义的docs，需要转为json字符串",
                          examples=[{"test.txt": [Document(page_content="custom doc")]}]),
        not_refresh_vs_cache: bool = Body(False, description="暂不保存向量库（用于FAISS）"),
) -> BaseResponse:
    """
    更新知识库文档
    """
    # 验证知识库名称是否合法
    if not validate_kb_name(knowledge_base_name):
        return BaseResponse(code=403, msg="Don't attack me")

    # 获取对应的知识库服务
    kb = KBServiceFactory.get_service_by_name(knowledge_base_name)
    if kb is None:
        return BaseResponse(code=404, msg=f"未找到知识库 {knowledge_base_name}")

    failed_files = {}  # 记录加载失败的文件
    kb_files = []  # 需要处理的知识库文件列表

    # 生成需要加载文档的文件列表
    for file_name in file_names:
        file_detail = get_file_detail(kb_name=knowledge_base_name, filename=file_name)
        # 如果文件之前使用了自定义文档，根据参数决定略过或覆盖
        if file_detail.get("custom_docs") and not override_custom_docs:
            continue
        if file_name not in docs:
            try:
                kb_files.append(KnowledgeFile(filename=file_name, knowledge_base_name=knowledge_base_name))
            except Exception as e:
                # 记录加载文档时的错误信息
                msg = f"加载文档 {file_name} 时出错：{e}"
                logger.error(f'{e.__class__.__name__}: {msg}',
                             exc_info=e if log_verbose else None)
                failed_files[file_name] = msg

    # 从文件生成文档，并进行向量化
    for status, result in files2docs_in_thread(kb_files,
                                               chunk_size=chunk_size,
                                               chunk_overlap=chunk_overlap,
                                               zh_title_enhance=zh_title_enhance):
        if status:
            # 如果成功，更新知识库文档
            kb_name, file_name, new_docs = result
            kb_file = KnowledgeFile(filename=file_name,
                                    knowledge_base_name=knowledge_base_name)
            kb_file.splited_docs = new_docs
            kb.update_doc(kb_file, not_refresh_vs_cache=True)
        else:
            # 如果失败，记录错误信息
            kb_name, file_name, error = result
            failed_files[file_name] = error

    # 将自定义的文档进行向量化
    for file_name, v in docs.items():
        try:
            v = [x if isinstance(x, Document) else Document(**x) for x in v]
            kb_file = KnowledgeFile(filename=file_name, knowledge_base_name=knowledge_base_name)
            kb.update_doc(kb_file, docs=v, not_refresh_vs_cache=True)
        except Exception as e:
            # 记录自定义文档添加时的错误信息
            msg = f"为 {file_name} 添加自定义docs时出错：{e}"
            logger.error(f'{e.__class__.__name__}: {msg}',
                         exc_info=e if log_verbose else None)
            failed_files[file_name] = msg

    # 保存向量库
    if not not_refresh_vs_cache:
        kb.save_vector_store()

    # 返回处理结果
    return BaseResponse(code=200, msg=f"更新文档完成", data={"failed_files": failed_files})

二、 参数说明

1. knowledge_base_name (str):

   • 描述: 知识库的名称。
   • 示例: “samples”。
   • 用法: 必填参数，用于指定需要更新的知识库。

2. file_names (List[str]):

   • 描述: 需要更新的文件名列表，支持多文件。
   • 示例: [“file_name1”, “text.txt”]。
   • 用法: 必填参数，用于指定需要处理的文件。

3. chunk_size (int):

   • 描述: 知识库中单段文本的最大长度。
   • 默认值: CHUNK_SIZE。
   • 用法: 可选参数，用于指定文本分块的大小。

4. chunk_overlap (int):

   • 描述: 知识库中相邻文本重合的长度。
   • 默认值: OVERLAP_SIZE。
   • 用法: 可选参数，用于指定文本分块时的重叠部分长度。

5. zh_title_enhance (bool):

   • 描述: 是否开启中文标题加强功能。
   • 默认值: ZH_TITLE_ENHANCE。
   • 用法: 可选参数，用于指定是否增强中文标题。

6. override_custom_docs (bool):

   • 描述: 是否覆盖之前自定义的文档。
   • 默认值: False。
   • 用法: 可选参数，用于指定是否覆盖已经存在的自定义文档。

7. docs (Json):

   • 描述: 自定义的文档，需要转换为JSON字符串。
   • 示例: {“test.txt”: [Document(page_content=“custom doc”)]}。
   • 用法: 可选参数，用于传入自定义的文档内容。

8. not_refresh_vs_cache (bool):

   • 描述: 暂不保存向量库（用于FAISS）。
   • 默认值: False。
   • 用法: 可选参数，用于指定是否保存向量库缓存。

三、 功能说明

该函数主要完成以下任务：

1. 验证知识库名称：

   • 首先调用 validate_kb_name 函数验证知识库名称是否有效。如果无效，返回403错误。

2. 获取知识库服务：

   • 使用 KBServiceFactory.get_service_by_name 函数获取知识库服务实例。如果未找到对应的知识库，返回404错误。

3. 生成需要加载文档的文件列表：

   • 对传入的文件名列表进行遍历，生成需要加载文档的文件列表。如果文件之前使用了自定义文档，根据参数决定略过或覆盖。

4. 从文件生成文档并进行向量化：

   • 调用 files2docs_in_thread 函数将文件生成文档，并在多线程中进行处理。对于处理成功的文件，更新知识库文档；对于处理失败的文件，记录错误信息。

5. 处理自定义文档并进行向量化：

   • 对传入的自定义文档进行处理，将其转换为 Document 对象，并更新到知识库中。如果处理失败，记录错误信息。

6. 保存向量库：

   • 如果未指定 not_refresh_vs_cache 参数，则保存向量库。

四、 返回值

该函数返回一个 BaseResponse 对象，包含以下信息：

• code: 状态码，200表示成功。
• msg: 消息，更新文档完成。
• data: 数据，包含处理失败的文件信息。

总结

通过上述步骤，update_docs 函数能够有效地更新知识库中的文档，并提供处理失败的详细信息，以便于后续的调试和处理。