怎么快速实现一个 ChatPDF

import os
import pathlib
import tiktoken
from typing import List, Tuple
from dotenv import load_dotenv

from llama_index.core import (
    Document,
    PromptTemplate,
    Settings,
    StorageContext,
    VectorStoreIndex,
)
from llama_index.embeddings.jinaai import JinaEmbedding
from llama_index.llms.openai_like import OpenAILike
from llama_index.postprocessor.jinaai_rerank import JinaRerank
from llama_index.vector_stores.qdrant import QdrantVectorStore
from llama_index.node_parser.docling import DoclingNodeParser
from llama_index.readers.docling import DoclingReader
from qdrant_client.models import Distance, VectorParams
from qdrant_client import QdrantClient

root_path = pathlib.Path(__file__).parent
load_dotenv()

# ============================================================
# 第一步: 配置大语言模型 (LLM)
# ============================================================
# 使用 Deepseek 作为推理模型,负责根据检索到的文档生成最终答案
llm = OpenAILike(
    api_base=os.getenv("DEEPSEEK_BASE_URL"),
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    model="deepseek-chat",
    temperature=0.0,
    max_tokens=1024,
    context_window=8192,
    is_chat_model=True,
    timeout=60,
)

Settings.llm = llm
Settings.tokenizer = tiktoken.encoding_for_model("gpt-4o").encode

# ============================================================
# 第二步: 配置嵌入模型 (Embedding Model)
# ============================================================
# 文本嵌入模型: 将文档文本转换为 256 维向量,用于存储到向量数据库
text_embed_model = JinaEmbedding(
    api_key=os.getenv("JINA_API_KEY"),
    model="jina-embeddings-v3",
    embed_batch_size=32,
    dimensions=256,
    task="retrieval.passage",  # 针对文档段落优化
)

# 查询嵌入模型: 将用户问题转换为 256 维向量,用于检索相似文档
query_embed_model = JinaEmbedding(
    api_key=os.getenv("JINA_API_KEY"),
    model="jina-embeddings-v3",
    embed_batch_size=32,
    dimensions=256,
    task="retrieval.query",  # 针对查询优化
)

# ============================================================
# 第三步: 配置重排序模型 (Reranker)
# ============================================================
# 重排序模型: 对初步检索结果进行精排,提高最相关文档的排名
jina_rerank = JinaRerank(
    api_key=os.getenv("JINA_API_KEY"),
    model="jina-reranker-v2-base-multilingual",
    top_n=3,  # 重排序后保留前 3 个最相关的文档
)

# ============================================================
# 第四步: 连接向量数据库 (Qdrant)
# ============================================================
# Qdrant: 存储文档向量,支持高效的相似度搜索
qdrant = QdrantClient(host="localhost", port=6333)


# ============================================================
# 核心函数 1: 初始化向量存储
# ============================================================
def init_vector_store(collection_name: str) -> Tuple[QdrantVectorStore, StorageContext]:
    """
    初始化向量存储和存储上下文

    功能说明:
    - 创建 Qdrant 向量存储的连接
    - 配置 LlamaIndex 的存储上下文,用于后续索引创建

    Args:
        collection_name: Qdrant 集合名称

    Returns:
        Tuple[QdrantVectorStore, StorageContext]: 向量存储和存储上下文的元组
    """
    vector_store = QdrantVectorStore(client=qdrant, collection_name=collection_name)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)
    return vector_store, storage_context


# ============================================================
# 核心函数 2: 确保集合存在
# ============================================================
def ensure_collection_exists(collection_name: str) -> bool:
    """
    确保 Qdrant 集合存在,如果不存在则创建

    功能说明:
    - 检查指定的集合是否已存在
    - 如果不存在,创建新集合并配置向量参数:
      * size=256: 向量维度为 256 (与嵌入模型的 dimensions 参数一致)
      * distance=COSINE: 使用余弦相似度计算文档相似性

    Args:
        collection_name: Qdrant 集合名称

    Returns:
        bool: True 表示新创建的集合, False 表示集合已存在
    """
    if not qdrant.collection_exists(collection_name):
        qdrant.create_collection(
            collection_name=collection_name,
            vectors_config=VectorParams(size=256, distance=Distance.COSINE),
        )
        return True
    return False


# ============================================================
# 核心函数 3: PDF 文档处理
# ============================================================
def process_pdf_to_documents(file_path: str) -> List[Document]:
    """
    处理 PDF 文件并转换为 LlamaIndex Document 列表

    功能说明:
    这个函数完成三个关键步骤:
    1. 读取 PDF: 使用 Docling 解析 PDF 的文本、表格、图片等内容
    2. 分块 (Chunking): 将长文档切分为适合嵌入的小块
    3. 元数据提取: 保留页码、位置等信息,用于后续引用溯源

    Args:
        file_path: PDF 文件路径

    Returns:
        List[Document]: Document 对象列表,每个对象包含文本和元数据
    """
    # 步骤 1: 读取 PDF 文件
    # DoclingReader 可以解析复杂的 PDF 结构 (表格、图片、公式等)
    reader = DoclingReader(export_type=DoclingReader.ExportType.JSON)
    docs = reader.load_data(file_path)

    # 步骤 2: 文档分块 (Chunking)
    # DoclingNodeParser 使用 HierarchicalChunker 进行智能分块:
    # - 保持段落完整性
    # - 根据文档结构 (标题、章节) 进行分块
    # - 避免在句子中间切断
    node_parser = DoclingNodeParser()
    nodes = node_parser.get_nodes_from_documents(docs)

    # 打印分块结果，理解文档如何被切分
    print("\n" + "=" * 70)
    print("文档分块结果")
    print("=" * 70)
    print(f"总节点数: {len(nodes)} 个")
    print(f"平均每块长度: {sum(len(n.text) for n in nodes) // len(nodes)} 字符")
    print(f"最小块长度: {min(len(n.text) for n in nodes)} 字符")
    print(f"最大块长度: {max(len(n.text) for n in nodes)} 字符")

    # 展示第一个节点的详细信息
    if nodes:
        first_node = nodes[0]
        print("\n第一个节点详情:")
        print(f"   文本长度: {len(first_node.text)} 字符")
        print(f"   文本预览: {first_node.text[:150]}...")
        print(f"   元数据: {first_node.metadata}")
    print("=" * 70 + "\n")

    # 步骤 3: 转换为 LlamaIndex Document 格式
    documents = []
    for node in nodes:
        # 提取块级元数据 (用于引用溯源)
        # - page: 页码,用于在答案中标注信息来源
        # - bbox: 边界框坐标,标记文本在页面中的位置
        # - headings: 标题层级,帮助理解文本的上下文
        metadata = {
            "page": str(node.metadata["doc_items"][0]["prov"][0]["page_no"]),
            "bbox": node.metadata["doc_items"][0]["prov"][0]["bbox"],
            "headings": node.metadata.get("headings", []),
        }
        llama_document = Document(
            text=node.text,
            metadata=metadata,
            text_template="Metadata: {metadata_str}\n-----\nContent: {content}",
        )
        documents.append(llama_document)

    return documents


# ============================================================
# 核心函数 4: 创建或加载向量索引 (主流程编排)
# ============================================================
def create_or_get_index():
    """
    创建或获取现有的向量存储索引

    功能说明:
    这是整个 RAG 系统的核心流程编排函数,负责:
    1. 初始化向量存储连接
    2. 检查集合是否存在
    3. 如果是新集合:
       - 处理 PDF 文档
       - 进行向量嵌入 (关键步骤)
       - 存储到 Qdrant
    4. 如果集合已存在:
       - 直接加载现有索引

    Returns:
        VectorStoreIndex: 向量存储索引对象,用于后续查询
    """
    collection_name = "chatpdf"

    # 步骤 1: 初始化向量存储
    vector_store, storage_context = init_vector_store(collection_name)

    # 步骤 2: 检查并创建集合
    is_new_collection = ensure_collection_exists(collection_name)

    if is_new_collection:
        # 分支 A: 新集合 - 需要处理文档并创建索引
        print("检测到新集合,开始处理文档...")

        # 这里解析 A 股一家上市公司的季报作为示例
        file_path = f"{str(root_path)}/data/disu.pdf"
        documents = process_pdf_to_documents(file_path)
        print(f"文档处理完成,共 {len(documents)} 个文档块")

        # 关键步骤: 向量嵌入
        # VectorStoreIndex.from_documents() 内部会:
        # 1. 调用 text_embed_model 将每个文档转换为 256 维向量
        # 2. 将向量存储到 Qdrant 的 chatpdf 集合中
        # 3. 建立文本-向量的映射关系
        print("开始向量嵌入...")
        vector_index = VectorStoreIndex.from_documents(
            documents,
            storage_context=storage_context,
            embed_model=text_embed_model,
        )
        print("向量嵌入完成,索引已创建!")
    else:
        # 分支 B: 已存在集合 - 直接加载现有索引
        print("检测到已存在的集合,加载现有索引...")
        vector_index = VectorStoreIndex.from_vector_store(
            vector_store=vector_store,
            embed_model=query_embed_model,
        )
        print("索引加载完成!")

    return vector_index


# ============================================================
# 第五步: 创建向量索引
# ============================================================
index = create_or_get_index()

# ============================================================
# 第六步: 配置查询提示词
# ============================================================
# 通过精心设计的提示词,引导大模型:
# 1. 仅使用检索到的文档内容回答
# 2. 标注信息来源的页码
# 3. 以结构化的方式输出答案
custom_prompt = PromptTemplate(
    "你是一个乐于助人的助手。使用以下内容作为你所学习的知识的来源。请仔细分析提供的文档内容，回答以下问题。避免使用你自己的知识。\n"
    "文档内容：{context_str}\n"
    "问题：{query_str}\n"
    "请遵循以下要求：\n"
    "1. 仅使用上述文档内容回答问题。\n"
    "2. 如果信息分散在多个部分，请将相关信息整合在一起回答。\n"
    "3. 即使信息出现在表格的不同行或列中，也要能够正确关联。\n"
    "4. 如果上下文中没有直接相关信息，进行合理推算或估计，并说明是推算结果。\n"
    "5. 以项目符号（bullet points）的形式列出陈述。\n"
    "6. 每个陈述必须注明信息来源的页码，格式为 [页码]。页码信息位于每段文本的元数据中的 'page' 字段。\n"
    "7. 如果一个陈述来自多个页面，请标注所有相关页码，如 [4][5]。\n"
)

# ============================================================
# 第七步: 创建查询引擎
# ============================================================
# 查询引擎是 RAG 系统的核心,负责:
# 1. 将用户问题转换为向量
# 2. 在 Qdrant 中检索相似文档
# 3. 使用重排序模型精排结果
# 4. 将检索结果和问题一起发送给 LLM
# 5. 生成最终答案
query_engine = index.as_query_engine(
    llm=llm,  # 使用 Deepseek 生成答案
    embed_model=query_embed_model,  # 使用查询嵌入模型
    streaming=True,  # 开启流式输出,实时返回答案
    similarity_top_k=5,  # 初步检索 5 个最相似的文档
    text_qa_template=custom_prompt,  # 使用自定义提示词
    node_postprocessors=[jina_rerank],  # 重排序: 5 个文档 → 精选 3 个
    verbose=True,  # 打印检索和生成的详细过程
)


# ============================================================
# 主函数: RAG 查询流程演示
# ============================================================
def main():
    """
    RAG 查询的完整流程:
    1. 用户提问
    2. 问题向量化
    3. 向量检索 (在 Qdrant 中找相似文档)
    4. 重排序 (精选最相关的文档)
    5. 生成答案 (LLM 基于检索结果回答)
    6. 展示答案和引用来源
    """
    # 步骤 1: 用户提问
    query_str = "李军是谁？"

    # 步骤 2-5: 查询引擎自动完成向量检索、重排序、生成答案
    # 内部流程:
    # 1. query_embed_model 将问题转为 256 维向量
    # 2. Qdrant 检索出 5 个最相似的文档块
    # 3. jina_rerank 重排序,精选出 3 个最相关的
    # 4. 将 3 个文档块 + 问题 + 提示词发送给 Deepseek
    # 5. Deepseek 生成答案并流式返回
    response = query_engine.query(query_str)

    # 步骤 6: 展示答案 (流式输出)
    print("\n=== 模型回答 ===")
    response_text = ""
    for text in response.response_gen:
        print(text, end="", flush=True)
        response_text += text
    print()

    # 步骤 7: 展示引用来源 (用于验证答案的可靠性)
    print("\n=== 参考文本片段 ===")
    for i, node in enumerate(response.source_nodes, 1):
        print(f"\n文本片段 {i}")
        print("=" * 50)
        print(f"相关度分数: {node.score:.3f}")  # 重排序后的相关度分数
        print(f"元数据信息: {node.metadata}")  # 页码、位置等信息
        print("-" * 50)
        print("引用文本内容:")
        # 只显示前 200 个字符,避免输出过长
        print(node.text[:200] + "..." if len(node.text) > 200 else node.text)
        print("=" * 50)


if __name__ == "__main__":
    main()

# 使用 Docling 智能分块
node_parser = DoclingNodeParser()
nodes = node_parser.get_nodes_from_documents(docs)

# 提取元数据用于引用溯源
for node in nodes:
    metadata = {
        "page": str(node.metadata["doc_items"][0]["prov"][0]["page_no"]),  # 页码
        "bbox": node.metadata["doc_items"][0]["prov"][0]["bbox"],          # 位置
        "headings": node.metadata.get("headings", []),                     # 标题
    }

# 创建向量索引（自动完成嵌入）
vector_index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,  # 指向 Qdrant
    embed_model=text_embed_model,     # 使用 Jina 嵌入模型
)

query_engine = index.as_query_engine(
    similarity_top_k=5,               # 粗排：检索 5 个候选
    node_postprocessors=[jina_rerank], # 精排：重排序后取 Top 3
)

custom_prompt = PromptTemplate(
    "你是一个乐于助人的助手。使用以下内容作为你所学习的知识的来源...\n"
    "6. 每个陈述必须注明信息来源的页码，格式为 [页码]。\n"
)