提升LangChain与ChromaDB检索问答系统响应完整度的策略

from langchain.document_loaders import PyPDFLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain.embeddings import OpenAIEmbeddings # 或 HuggingFaceEmbeddings
from langchain.vectorstores import Chroma
from langchain.chains import RetrievalQA
from langchain.llms import OpenAI
import os

# 确保设置了OpenAI API Key
# os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"

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: 分割后的文档文本块列表。
    """
    print(f"Loading documents from: {directory_path}")
    loader = DirectoryLoader(directory_path, glob="./*.pdf", loader_cls=PyPDFLoader)
    documents = loader.load()

    print(f"Splitting documents into chunks (size={chunk_size}, overlap={chunk_overlap})...")
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    texts = text_splitter.split_documents(documents)
    print(f"Total {len(texts)} chunks generated.")
    return texts

def create_or_load_vectordb(texts: list, persist_directory: str = './ChromaDb'):
    """
    创建或加载ChromaDB向量数据库。

    Args:
        texts (list): 待嵌入的文本块列表。
        persist_directory (str): 向量数据库的持久化目录。

    Returns:
        Chroma: ChromaDB向量数据库实例。
    """
    embeddings = OpenAIEmbeddings() # 可以替换为HuggingFaceEmbeddings

    # 检查持久化目录是否存在，如果存在则加载，否则创建
    if os.path.exists(persist_directory) and os.listdir(persist_directory):
        print(f"Loading existing ChromaDB from: {persist_directory}")
        vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
    else:
        print(f"Creating new ChromaDB at: {persist_directory}")
        vectordb = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=persist_directory)
        vectordb.persist()
        print("ChromaDB created and persisted.")
    return vectordb

def setup_qa_chain(vectordb: Chroma, k_value: int = 8):
    """
    设置RetrievalQA链。

    Args:
        vectordb (Chroma): ChromaDB向量数据库实例。
        k_value (int): 检索器返回的最相关文档块数量。

    Returns:
        RetrievalQA: 配置好的RetrievalQA链。
    """
    llm = OpenAI(temperature=0, model_name="text-davinci-003") # 也可以使用gpt-3.5-turbo等

    # 配置检索器，增加返回的文档数量
    retriever = vectordb.as_retriever(search_kwargs={"k": k_value})

    # 注意：如果chain_type_kwargs包含与retriever冲突的参数，可能会导致问题
    # 确保chain_type_kwargs是适合stuff链的
    # 例如，可以用于设置prompt模板等
    # chain_type_kwargs = {"prompt": YOUR_CUSTOM_PROMPT_TEMPLATE} 

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        chain_type="stuff", # "stuff"会将所有检索到的文档拼接起来
        # chain_type_kwargs=chain_type_kwargs, # 如果有自定义prompt，可以在这里传入
        return_source_documents=True # 方便调试，查看哪些文档被检索到
    )
    print(f"RetrievalQA chain setup with k={k_value} for retriever.")
    return qa_chain

# 主运行逻辑
if __name__ == "__main__":
    UPLOAD_DIR = './static/upload/' # 确保此目录存在并包含PDF文件
    PERSIST_DIR = './ChromaDb'

    # 1. 加载和分割文档
    # 尝试更大的chunk_size和chunk_overlap以捕获更多上下文
    texts = load_and_split_documents(UPLOAD_DIR, chunk_size=1500, chunk_overlap=150) 

    # 2. 创建或加载向量数据库
    vectordb = create_or_load_vectordb(texts, PERSIST_DIR)

    # 3. 设置QA链，并增加检索的k值
    # k值应根据期望的响应长度和LLM的上下文窗口进行调整
    # 例如，如果每个chunk 1500字，k=8，总上下文将是12000字，需要LLM支持
    qa_chain = setup_qa_chain(vectordb, k_value=8) 

    # 4. 提问并获取响应
    query = "请总结这篇文档的主要内容" # 替换为你的具体查询
    print(f"\nQuerying: '{query}'")
    response = qa_chain({"query": query})

    print("\n--- Full Response ---")
    print(response["result"])

    print("\n--- Source Documents ---")
    for i, doc in enumerate(response["source_documents"]):
        print(f"Document {i+1} (Page {doc.metadata.get('page', 'N/A')}):")
        print(doc.page_content[:500] + "...") # 打印前500字
        print("-" * 20)