优化ChromaDB检索：提升PDF文档问答完整性

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

# 设置OpenAI API密钥 (如果使用OpenAIEmbeddings和OpenAI LLM)
# os.environ["OPENAI_API_KEY"] = "YOUR_OPENAI_API_KEY"

def load_documents_from_pdf_directory(directory_path: str = './static/upload/') -> list:
    """
    从指定目录加载所有PDF文档。
    """
    print(f"Loading PDF documents from: {directory_path}")
    loader = DirectoryLoader(directory_path, glob="./*.pdf", loader_cls=PyPDFLoader)
    documents = loader.load()
    print(f"Loaded {len(documents)} documents.")
    return documents

def split_documents(documents: list, chunk_size: int = 1000, chunk_overlap: int = 100) -> list:
    """
    将加载的文档切分成文本块。
    """
    print(f"Splitting documents with chunk_size={chunk_size} and chunk_overlap={chunk_overlap}")
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    texts = text_splitter.split_documents(documents)
    print(f"Split into {len(texts)} text chunks.")
    return texts

def create_and_persist_vectordb(texts: list, persist_directory: str = './ChromaDb') -> Chroma:
    """
    创建ChromaDB向量数据库并持久化。
    """
    print(f"Creating embeddings and ChromaDB at: {persist_directory}")
    # 可以选择OpenAIEmbeddings或HuggingFaceEmbeddings
    # embeddings = HuggingFaceEmbeddings(model_name="bert-base-multilingual-cased")
    embeddings = OpenAIEmbeddings()

    vectordb = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=persist_directory)
    vectordb.persist()
    print("ChromaDB created and persisted.")
    return vectordb

def setup_retrieval_qa_chain(vectordb: Chroma, llm_model_name: str = "text-davinci-003", k_retrieval: int = 4) -> RetrievalQA:
    """
    设置RetrievalQA链。
    """
    print(f"Setting up RetrievalQA chain with LLM: {llm_model_name} and retriever k={k_retrieval}")
    llm = OpenAI(temperature=0, model_name=llm_model_name)

    # 配置检索器，指定 k 参数
    retriever = vectordb.as_retriever(search_kwargs={"k": k_retrieval})

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        chain_type="stuff", # 确保所有检索到的文档被一次性送入LLM
        return_source_documents=True
    )
    print("RetrievalQA chain setup complete.")
    return qa_chain

if __name__ == "__main__":
    # 1. 加载文档
    loaded_documents = load_documents_from_pdf_directory()

    # 2. 切分文档 (调整 chunk_size 和 chunk_overlap)
    # 示例中将 chunk_overlap 增加到 100
    split_texts = split_documents(loaded_documents, chunk_size=1000, chunk_overlap=100)

    # 3. 创建并持久化向量数据库
    vector_database = create_and_persist_vectordb(split_texts)

    # 4. 设置检索QA链 (调整 k 参数)
    # 示例中将 k 增加到 4，以检索更多相关文档
    qa_pipeline = setup_retrieval_qa_chain(vector_database, k_retrieval=4)

    # 5. 执行查询
    query = "请总结这份文件中的主要内容" # 假设你的PDF有类似“书”的内容
    print(f"\nQuerying: '{query}'")
    response = qa_pipeline({"query": query})

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

    if "source_documents" in response:
        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[:200] + "...") # 打印前200字符作为示例
            print("-" * 20)