LangchainChatChat批量上传文档到知识库的方法

概要

使用 LangchainChatChat 时，要上传的知识库文件非常多，且分布在不同的文件夹中，使用 Web UI 前端界面进行批量选择上传一次只能全选一个文件夹，所以想到了调用接口上传的办法。

思路

• 访问 fastapi 的 api 文档
  http://LangchainChatChatIP地址:7861/docs
  
• 找到需要的接口
  
• 点击 Try it out
• 随便选择一个文件
• 填写知识库名称
• 点击 Execute 执行
• 得到curl请求

curl -X 'POST' \
  'http://ip:7861/knowledge_base/upload_docs' \
  -H 'accept: application/json' \
  -H 'Content-Type: multipart/form-data' \
  -F 'to_vector_store=true' \
  -F 'override=false' \
  -F 'not_refresh_vs_cache=false' \
  -F 'chunk_size=250' \
  -F 'chunk_overlap=50' \
  -F 'zh_title_enhance=false' \
  -F 'files=@1602298756694.txt;type=text/plain' \
  -F 'knowledge_base_name=aiops' \
  -F 'docs={
  "test.txt": [
    {
      "page_content": "custom doc",
      "metadata": {},
      "type": "Document"
    }
  ]
}'

• 把 curl 改写成 request 请求

代码实现

import os
import requests
from tqdm import tqdm

def upload_txt_files(folder_path, upload_url, knowledge_base_name):
    # 初始化一个空列表，用于存储找到的txt文件的路径
    txt_files = []
    # 初始化一个变量，用于记录找到的txt文件总数
    total_files = 0
    
    # 使用os.walk遍历指定文件夹及其子文件夹
    for root, dirs, files in os.walk(folder_path):
        # 遍历当前文件夹中的所有文件
        for file in files:
            # 如果文件是txt文件，则将其路径添加到txt_files列表中
            if file.endswith(".txt"):
                txt_files.append(os.path.join(root, file))
    # 更新找到的txt文件总数
    total_files = len(txt_files)
    # 打印找到的txt文件总数
    print("Total txt files found:", total_files)

    # 初始化一个变量，用于记录已上传的文件数量
    uploaded_count = 0
    
    with tqdm(total=total_files, desc="Uploading") as pbar:
        # 遍历所有找到的txt文件
        for txt_file in txt_files:
            # 准备上传的文件，包括文件名、文件对象和文件类型
            files = {'files': (os.path.basename(txt_file), open(txt_file, 'rb'), 'text/plain')}
            # 准备上传的数据，包括一些设置参数和知识库名称
            data = {
                'to_vector_store': 'true',
                'override': 'false',
                'not_refresh_vs_cache': 'false',
                'chunk_size': '500',
                'chunk_overlap': '50',
                'zh_title_enhance': 'true',
                'knowledge_base_name': knowledge_base_name
            }
            # 发送POST请求，上传文件
            response = requests.post(upload_url, files=files, data=data)
            # 更新已上传的文件数量
            uploaded_count += 1
            # 更新进度条
            pbar.update(1)
            # 更新进度条的后缀，显示已上传的文件数量
            pbar.set_postfix({"Uploaded": uploaded_count})
            # 打印服务器的响应
            print("Response:", response.text)

if __name__ == "__main__":
    # 设置需要上传的txt文件所在的文件夹路径
    folder_path = "知识库文件夹路径"
    # 设置LangChainChatChat的API上传地址
    upload_url = 'http://填写LangchainChatChatIP地址:7861/knowledge_base/upload_docs'
    # 设置要上传到的知识库的名称
    knowledge_base_name = "填写知识库名称"  
    # 调用函数，开始上传txt文件
    upload_txt_files(folder_path, upload_url, knowledge_base_name)

TODO 进一步优化速度

这个接口是支持批量的，所以可以每次把1000个txt作为一个请求发送