优化LangChain与ChromaDB：提升RAG响应完整性与准确性

from langchain.document_loaders import DirectoryLoader, PyPDFLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter

def load_and_split_documents(directory_path: str, chunk_size: int = 1000, chunk_overlap: int = 100):
    """
    加载指定目录下的PDF文档，并进行递归字符分块。

    Args:
        directory_path (str): 包含PDF文档的目录路径。
        chunk_size (int): 每个文本块的最大字符数。
        chunk_overlap (int): 相邻文本块之间的重叠字符数。

    Returns:
        list: 分割后的文档文本块列表。
    """
    # 使用PyPDFLoader加载PDF文档
    loader = DirectoryLoader(directory_path, glob="./*.pdf", loader_cls=PyPDFLoader)
    documents = loader.load()

    # 初始化递归字符文本分块器
    # RecursiveCharacterTextSplitter会尝试根据不同分隔符（如换行符、空格）智能地分割文本
    text_splitter = RecursiveCharacterTextTextSplitter(
        chunk_size=chunk_size,
        chunk_overlap=chunk_overlap
    )
    texts = text_splitter.split_documents(documents)
    print(f"原始文档数量: {len(documents)}, 分割后文本块数量: {len(texts)}")
    return texts

# 示例使用
# documents_to_process = load_and_split_documents('./static/upload/', chunk_size=1000, chunk_overlap=100)

import os
from langchain.embeddings import OpenAIEmbeddings, HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI
from langchain.prompts import PromptTemplate

# 假设已经通过 load_and_split_documents 函数获取了 texts
# documents_to_process = load_and_split_documents('./static/upload/', chunk_size=1000, chunk_overlap=100)

def create_and_query_vectordb(texts: list, persist_directory: str = './ChromaDb', k_retrieved_docs: int = 6):
    """
    创建ChromaDB向量数据库，并使用RetrievalQA链进行查询。

    Args:
        texts (list): 经过分块处理的文档文本块列表。
        persist_directory (str): ChromaDB持久化存储的目录。
        k_retrieved_docs (int): 检索器返回的最相关文档块数量。

    Returns:
        RetrievalQA: 配置好的检索问答链。
    """
    # 确保OpenAI API Key已设置
    # os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"

    # 选择嵌入模型
    # 可以使用OpenAIEmbeddings，或者本地的HuggingFaceEmbeddings
    # embeddings = OpenAIEmbeddings()
    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2") # 示例使用本地模型

    # 从文档创建ChromaDB，并持久化
    vectordb = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=persist_directory)
    vectordb.persist()
    print(f"ChromaDB已创建并持久化到: {persist_directory}")

    # 定义LLM
    llm = OpenAI(temperature=0, model_name="text-davinci-003") # 或者 gpt-3.5-turbo 等

    # 自定义提示模板，以更好地引导LLM利用检索到的上下文
    # template = """使用以下上下文来回答最后的问题。如果你不知道答案，就说你不知道，不要试图编造答案。
    # 
    # {context}
    # 
    # 问题: {question}
    # 有用的回答:"""
    # QA_CHAIN_PROMPT = PromptTemplate(input_variables=["context", "question"], template=template)

    # 配置RetrievalQA链
    # 关键参数：
    #   retriever=vectordb.as_retriever(search_kwargs={"k": k_retrieved_docs})：设置检索器，指定检索k个文档
    #   chain_type="stuff"：将所有检索到的文档块合并成一个大字符串作为LLM的输入
    #   return_source_documents=True：返回检索到的源文档，便于调试和验证
    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=vectordb.as_retriever(search_kwargs={"k": k_retrieved_docs}),
        chain_type="stuff",
        # chain_type_kwargs={"prompt": QA_CHAIN_PROMPT}, # 如果使用自定义提示模板
        return_source_documents=True
    )
    return qa_chain

# 完整流程示例
if __name__ == "__main__":
    # 1. 加载并分块文档
    # 确保 './static/upload/' 目录下有PDF文件
    # 调整 chunk_size 和 chunk_overlap 以优化上下文保留
    print("加载并分块文档...")
    documents_to_process = load_and_split_documents('./static/upload/', chunk_size=1500, chunk_overlap=150)

    # 2. 创建ChromaDB并配置检索链
    # 调整 k_retrieved_docs 以确保检索到足够的信息
    print("创建ChromaDB并配置检索链...")
    qa_retrieval_chain = create_and_query_vectordb(documents_to_process, k_retrieved_docs=8) # 增加k值

    # 3. 执行查询
    print("执行查询...")
    query = "请总结这本书的主要内容。" # 假设PDF是一本书
    response = qa_retrieval_chain({"query": query})

    print("\n--- LLM 响应 ---")
    print(response["result"])

    print("\n--- 检索到的源文档 ---")
    for i, doc in enumerate(response["source_documents"]):
        print(f"文档 {i+1}:")
        print(f"  内容: {doc.page_content[:200]}...") # 打印前200字符
        print(f"  来源: {doc.metadata}")
        print("-" * 20)