优化ChromaDB检索，提升RAG系统响应完整性

from langchain.text_splitter import RecursiveCharacterTextSplitter

# 原始文档加载后，进行文本分块
documents = [...] # 假设这里是已加载的文档列表

# 调整 chunk_size 和 chunk_overlap
# chunk_size=1000 意味着每个块最大1000字符
# chunk_overlap=100 意味着相邻块之间有100字符的重叠
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
texts = text_splitter.split_documents(documents)

# 打印一些块以观察重叠效果
# for i, text in enumerate(texts[:3]):
#     print(f"--- Chunk {i} ---")
#     print(text.page_content[:200]) # 打印前200字符

from langchain.vectorstores import Chroma
from langchain.embeddings import HuggingFaceEmbeddings # 也可以使用 OpenAIEmbeddings

# 选择嵌入模型
# embeddings = OpenAIEmbeddings() # 如果使用OpenAI API
embeddings = HuggingFaceEmbeddings(model_name="bert-base-multilingual-cased") # 使用HuggingFace模型

persist_directory = "./ChromaDb" # 定义ChromaDB的持久化目录

# 从文本块创建ChromaDB向量存储
# 如果ChromaDb目录已存在，from_documents会加载现有数据并追加
vectordb = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=persist_directory)

# 持久化向量存储，以便下次可以直接加载而无需重新创建
vectordb.persist()

from langchain.chains import RetrievalQA
from langchain.llms import OpenAI

# 初始化LLM
llm = OpenAI(temperature=0, model_name="text-davinci-003")

# 配置检索器，可以指定检索多少个文档 (k)
# 默认k=4，可以根据需要调整，增加k值可能有助于获取更多上下文
# retriever = vectordb.as_retriever(search_kwargs={"k": 6}) 
retriever = vectordb.as_retriever()

# 创建RetrievalQA链
qa_chain = RetrievalQA.from_chain_type(
    llm=llm,
    retriever=retriever,
    chain_type="stuff", # 将所有检索到的文档填充到一个提示中
    return_source_documents=True # 返回源文档，便于调试
)

# 示例查询
query = "请总结这本书的内容"
response = qa_chain(query)

print("LLM响应：", response["result"])
if response.get("source_documents"):
    print("\n检索到的源文档：")
    for doc in response["source_documents"]:
        print(f"- {doc.page_content[:150]}...") # 打印每个源文档的前150字符

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

# --- 1. 文档加载 ---
def load_documents(directory_path='./static/upload/'):
    """加载指定目录下的PDF文档。"""
    loader = DirectoryLoader(directory_path, glob="./*.pdf", loader_cls=PyPDFLoader)
    documents = loader.load()
    return documents

# --- 2. 文本分块 ---
def split_documents(documents, chunk_size=1000, chunk_overlap=100):
    """将文档分割成带有重叠的文本块。"""
    text_splitter = RecursiveCharacterTextSplitter(chunk_size=chunk_size, chunk_overlap=chunk_overlap)
    texts = text_splitter.split_documents(documents)
    return texts

# --- 3. 创建或加载ChromaDB向量存储 ---
def create_or_load_vectordb(texts, persist_directory='./ChromaDb'):
    """创建或从持久化目录加载ChromaDB向量存储。"""
    # 选择嵌入模型
    # embeddings = OpenAIEmbeddings()
    embeddings = HuggingFaceEmbeddings(model_name="bert-base-multilingual-cased")

    if not os.path.exists(persist_directory) or not os.listdir(persist_directory):
        print(f"ChromaDB目录 {persist_directory} 不存在或为空，正在从文档创建...")
        vectordb = Chroma.from_documents(documents=texts, embedding=embeddings, persist_directory=persist_directory)
        vectordb.persist()
        print("ChromaDB创建并持久化完成。")
    else:
        print(f"ChromaDB目录 {persist_directory} 已存在，正在加载...")
        vectordb = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
        print("ChromaDB加载完成。")
    return vectordb

# --- 4. 配置并执行检索QA链 ---
def run_qa_chain(vectordb, query):
    """配置RetrievalQA链并执行查询。"""
    llm = OpenAI(temperature=0, model_name="text-davinci-003")

    # 可以通过 search_kwargs 调整检索器的参数，例如 k (检索的文档数量)
    # retriever = vectordb.as_retriever(search_kwargs={"k": 5}) 
    retriever = vectordb.as_retriever()

    qa_chain = RetrievalQA.from_chain_type(
        llm=llm,
        retriever=retriever,
        chain_type="stuff",
        return_source_documents=True
    )

    response = qa_chain(query)
    return response

# --- 主执行流程 ---
if __name__ == "__main__":
    # 确保存在一个用于测试的PDF文件，例如在 './static/upload/' 目录下放置 'sample.pdf'
    # 示例中使用了 '/tmp/'，实际应用中请根据你的文件路径修改

    # 1. 加载文档
    documents = load_documents(directory_path='./static/upload/')
    if not documents:
        print("未找到任何PDF文档，请确保 './static/upload/' 目录下有PDF文件。")
    else:
        print(f"成功加载 {len(documents)} 份文档。")

        # 2. 分割文档
        texts = split_documents(documents, chunk_size=1000, chunk_overlap=100)
        print(f"文档被分割成 {len(texts)} 个文本块。")

        # 3. 创建或加载ChromaDB
        vectordb = create_or_load_vectordb(texts, persist_directory='./ChromaDb')

        # 4. 执行查询
        user_query = "请总结这份文档的主要内容"
        print(f"\n正在查询: '{user_query}'")
        qa_response = run_qa_chain(vectordb, user_query)

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

        print("\n--- 检索到的源文档 ---")
        if qa_response.get("source_documents"):
            for i, doc in enumerate(qa_response["source_documents"]):
                print(f"文档 {i+1}:")
                print(f"  内容片段: {doc.page_content[:200]}...") # 打印前200字符
                print(f"  来源: {doc.metadata.get('source', '未知')}")
        else:
            print("未检索到源文档。")