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
from chromadb.config import Settings
from langchain.vectorstores import Chroma # 假设 embeddings 已定义并传入

DOC_LOADERS_MAPPING = {
    ".txt": (TextLoader, {"encoding": "utf8"}),
    # 可以根据需要添加更多文档加载器，例如：
    # ".pdf": (PyPDFLoader, {}),
    # ".md": (UnstructuredMarkdownLoader, {}),
}

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:
            # 使用 glob 查找目录及其子目录中的所有匹配文件
            all_files.extend(
                glob.glob(os.path.join(path, f"**/*{ext}"), recursive=True)
                )

        # 使用 load_document 函数加载所有找到的文件
        return [load_document(file_path) for file_path in all_files]
    except Exception as exception:
        raise RuntimeError(f"加载目录 '{path}' 中的文件时出错: {exception}")

# 假设您已经定义了 embeddings 对象，例如：
# from langchain.embeddings import OpenAIEmbeddings
# embeddings = OpenAIEmbeddings()

# 示例：使用一个虚拟的 embeddings 对象，实际应用中请替换为您的具体实现
class MockEmbeddings:
    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        # 模拟生成向量，每个文本生成一个1536维的随机向量
        import random
        return [[random.random() for _ in range(1536)] for _ in texts]

    def embed_query(self, text: str) -> List[float]:
        import random
        return [random.random() for _ in range(1536)]

embeddings = MockEmbeddings() # 实际使用时请替换为您的 Embedding 模型实例

# 1. 定义存储ChromaDB的目录
chromaDirectory = "./folder/chroma_db"

# 2. 从指定目录加载所有文档
print(f"正在从 '{chromaDirectory.replace('/chroma_db', '')}' 目录加载文档...")
documents = load_documents_from_dir(chromaDirectory.replace('/chroma_db', ''))
print(f"已加载 {len(documents)} 个文档。")

# 3. 初始化 RecursiveCharacterTextSplitter
# chunk_size 定义每个文本块的最大长度
# chunk_overlap 定义相邻文本块之间的重叠部分，有助于保持上下文连贯性
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=300,
    chunk_overlap=50
)
print("正在切分文档...")
texts = text_splitter.split_documents(documents)
print(f"文档已切分为 {len(texts)} 个文本块。")

# 4. 初始化 ChromaDB 并从切分后的文本创建向量存储
print("正在创建或加载 ChromaDB 向量存储...")
chroma_db = Chroma.from_documents(
    texts,
    embeddings,
    persist_directory=chromaDirectory,
    client_settings= Settings(
            persist_directory=chromaDirectory,
            chroma_db_impl="duckdb+parquet", # 指定 ChromaDB 的实现方式
            anonymized_telemetry=False, # 关闭匿名遥测
        ),    
)

# 5. 持久化 ChromaDB
# 这一步非常重要，确保数据被写入磁盘
chroma_db.persist()
print(f"ChromaDB 已持久化到 '{chromaDirectory}'。")

# 释放 ChromaDB 实例以确保所有数据写入完成（可选，但推荐）
chroma_db = None
print("ChromaDB 处理完成。")