优化Langchain文档加载与文本切分：解决ChromaDB多文档处理难题

import os
import glob
from typing import List

from langchain.docstore.document import Document
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.document_loaders import TextLoader, PyPDFLoader # 示例：可扩展支持更多文件类型
from langchain_community.vectorstores import Chroma
from langchain_community.embeddings import OpenAIEmbeddings # 或者其他你选择的嵌入模型
from chromadb.config import Settings

# 定义支持的文档加载器映射
DOC_LOADERS_MAPPING = {
    ".txt": (TextLoader, {"encoding": "utf8"}),
    ".pdf": (PyPDFLoader, {}), # 示例：添加PDF加载器
    # 可以根据需要添加更多文件类型，例如 .md, .docx 等
}

def load_document(path: str) -> Document:
    """
    根据文件路径加载单个文档。
    支持DOC_LOADERS_MAPPING中定义的文件类型。
    """
    try:
        # 提取文件扩展名
        ext = "." + path.rsplit(".", 1)[-1]
        if ext in DOC_LOADERS_MAPPING:
            loader_class, loader_args = DOC_LOADERS_MAPPING[ext]
            loader = loader_class(path, **loader_args)
            # load()方法通常返回一个Document列表，这里我们取第一个
            return loader.load()[0]

        raise ValueError(f"不支持的文件扩展名: {ext}")
    except Exception as exception:
        raise ValueError(f"加载文档时出错 '{path}': {exception}")

def load_documents_from_dir(path: str) -> List[Document]:
    """
    从指定目录及其子目录加载所有支持的文档。
    """
    try:
        all_files = []
        # 遍历所有支持的文件扩展名，收集匹配的文件路径
        for ext in DOC_LOADERS_MAPPING:
            all_files.extend(
                glob.glob(os.path.join(path, f"**/*{ext}"), recursive=True)
            )

        # 使用列表推导式加载所有文档
        return [load_document(file_path) for file_path in all_files]
    except Exception as exception:
        raise RuntimeError(f"加载文件时出错: {exception}")

# 假设你的嵌入模型已经初始化
# 例如：
embeddings = OpenAIEmbeddings() 

# 你的文档存放目录
document_folder = "./folder/"
chroma_db_directory = "./folder/chroma_db"

# 1. 加载所有文档
print(f"正在从目录 '{document_folder}' 加载文档...")
documents = load_documents_from_dir(document_folder)
print(f"已加载 {len(documents)} 个文档。")

# 2. 初始化 RecursiveCharacterTextSplitter
# chunk_size: 每个文本块的最大字符数
# chunk_overlap: 相邻文本块之间的重叠字符数，有助于保持上下文连贯性
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=300,
    chunk_overlap=50
)

# 3. 切分文档
print("正在切分文档...")
texts = text_splitter.split_documents(documents)
print(f"文档已被切分为 {len(texts)} 个文本块。")
# 打印第一个文本块的长度作为示例，验证切分效果
if texts:
    print(f"第一个文本块的长度: {len(texts[0].page_content)} 字符。")

# 4. 初始化ChromaDB并持久化文本块
print(f"正在将文本块存储到 ChromaDB '{chroma_db_directory}'...")
chroma_db = Chroma.from_documents(
    texts,
    embeddings,
    persist_directory=chroma_db_directory,
    client_settings= Settings(
        persist_directory=chroma_db_directory,
        chroma_db_impl="duckdb+parquet", # 指定ChromaDB的实现方式
        anonymized_telemetry=False, # 关闭匿名遥测
    ),    
)

# 5. 强制持久化到磁盘
# 调用persist()方法确保所有数据写入磁盘
chroma_db.persist()
print("ChromaDB 数据已成功持久化。")

# 6. 清理内存中的ChromaDB实例（可选）
# 如果不再需要当前会话中的ChromaDB实例，可以将其设置为None
chroma_db = None

# 后续可以重新加载ChromaDB进行检索
# reloaded_chroma_db = Chroma(
#     persist_directory=chroma_db_directory,
#     embedding_function=embeddings,
#     client_settings= Settings(
#         persist_directory=chroma_db_directory,
#         chroma_db_impl="duckdb+parquet",
#         anonymized_telemetry=False,
#     ),
# )
# print(f"重新加载的 ChromaDB 中包含 {reloaded_chroma_db._collection.count()} 个文档。")