LlamaIndex RAG框架

from llama_index.core import SimpleDirectoryReader

# 基本用法：从目录加载所有支持文件
reader = SimpleDirectoryReader(input_dir="./data")
documents = reader.load_data()
print(f"加载了 {len(documents)} 个文档")

# 指定文件列表
reader = SimpleDirectoryReader(input_files=["./data/report.pdf", "./data/notes.txt"])
documents = reader.load_data()

# 递归读取子目录
reader = SimpleDirectoryReader(input_dir="./data", recursive=True, filename_as_id=True)
documents = reader.load_data()

# 排除特定文件
reader = SimpleDirectoryReader(
    input_dir="./data",
    exclude=["*.tmp", "*.log"],  # 排除模式
    required_exts=[".pdf", ".txt", ".md"],  # 只包含指定扩展名
)
documents = reader.load_data()

# PDF 读取器
from llama_index.readers.file import PDFReader

# Web 页面读取器
from llama_index.readers.web import SimpleWebPageReader
reader = SimpleWebPageReader()
docs = reader.load_data(["https://example.com/article"])

# 数据库读取器（需安装 llama-index-readers-database）
# from llama_index.readers.database import DatabaseReader
# reader = DatabaseReader(uri="sqlite:///mydb.sqlite")
# docs = reader.load_data(query="SELECT content FROM articles")

from llama_index.core import Document

# 创建文档
doc = Document(
    text="这是一段示例文本，用于演示Document的创建。",
    metadata={
        "filename": "example.txt",
        "author": "张三",
        "category": "技术文档",
    },
    metadata_template="{key}: {value}",  # 元数据格式化模板
    excluded_llm_metadata_keys=["author"],  # LLM不看到的元数据字段
    excluded_embed_metadata_keys=["filename"],  # 嵌入不包含的元数据字段
)

# 文档属性
print(doc.text)         # 文本内容
print(doc.metadata)     # 元数据字典
print(doc.doc_id)       # 文档唯一ID
print(doc.hash)         # 内容哈希

from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import Document

# 使用文本切分器将 Document 转为 Node
doc = Document(text="这是第一句话。这是第二句话。这是第三句话。", metadata={"source": "demo"})

parser = SentenceSplitter(
    chunk_size=1024,      # 每个Node最大字符数
    chunk_overlap=200,    # 相邻Node重叠字符数
)
nodes = parser.get_nodes_from_documents([doc])

for node in nodes:
    print(f"Node ID: {node.node_id}")
    print(f"文本: {node.text}")
    print(f"元数据: {node.metadata}")
    print(f"关系: {node.relationships}")  # 与其他Node的关系

from llama_index.core.node_parser import (
    SentenceSplitter,       # 按句子切分
    TokenTextSplitter,      # 按Token切分
    HierarchicalNodeParser, # 层级切分（父子关系）
    SentenceWindowNodeParser,  # 句子窗口切分
    SemanticSplitterNodeParser, # 语义切分
)

# 语义切分器（需嵌入模型）
from llama_index.core.node_parser import SemanticSplitterNodeParser
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-zh-v1.5")
semantic_parser = SemanticSplitterNodeParser(
    buffer_size=1,           # 缓冲句子数
    breakpoint_percentile_threshold=95,  # 语义断点阈值
    embed_model=embed_model,
)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, StorageContext
from llama_index.vector_stores.chroma import ChromaVectorStore
import chromadb

# 方式1：从文档直接构建
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)

# 方式2：从Node构建
nodes = parser.get_nodes_from_documents(documents)
index = VectorStoreIndex(nodes)

# 方式3：使用持久化向量存储
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_or_create_collection("my_collection")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(documents, storage_context=storage_context)

# 方式4：增量插入
index.insert(doc=Document(text="新文档内容"))

from llama_index.core import SummaryIndex

index = SummaryIndex.from_documents(documents)

from llama_index.core import KeywordTableIndex

index = KeywordTableIndex.from_documents(documents)

# 使用检索
retriever = index.as_retriever(retriever_mode="simple")
# retriever_mode: "simple"（精确匹配）, "rake"（RAKE算法提取关键词）

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)

# 基本检索器
retriever = index.as_retriever(
    similarity_top_k=5,  # 返回最相似的top-k个节点
)

# 执行检索
nodes = retriever.retrieve("什么是RAG？")
for node_with_score in nodes:
    print(f"分数: {node_with_score.score:.4f}")
    print(f"文本: {node_with_score.node.text[:100]}")
    print(f"元数据: {node_with_score.node.metadata}")

from llama_index.core import VectorStoreIndex
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.retrievers import QueryFusionRetriever

# 向量检索器
vector_retriever = index.as_retriever(similarity_top_k=5)

# 关键词检索器（需支持BM25的向量存储）
# 以 Qdrant 为例：
# from llama_index.vector_stores.qdrant import QdrantVectorStore
# hybrid_retriever = index.as_retriever(
#     vector_store_query_mode="hybrid",
#     similarity_top_k=5,
#     sparse_top_k=5,
#     alpha=0.5,  # 向量与稀疏检索的权重，0=纯稀疏，1=纯向量
# )

# QueryFusion：融合多个检索器结果
fusion_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever],  # 可添加多个检索器
    num_queries=1,        # 生成额外查询的数量（Query Expansion）
    similarity_top_k=5,
    mode="reciprocal_rerank",  # 融合模式：reciprocal_rerank 或 dist_based
)

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)

# 创建查询引擎
query_engine = index.as_query_engine(
    similarity_top_k=3,              # 检索top-k节点
    response_mode="tree_summarize",   # 响应合成模式
    streaming=False,                  # 是否流式输出
)

# 执行查询
response = query_engine.query("什么是RAG？")
print(response)                    # 响应文本
print(response.source_nodes)       # 来源节点
print(response.metadata)           # 元数据

query_engine = index.as_query_engine(streaming=True, similarity_top_k=3)
response = query_engine.query("解释LlamaIndex的架构")

# 方式1：逐token打印
response.print_response_stream()

# 方式2：迭代处理
for text in response.response_gen:
    print(text, end="", flush=True)

from llama_index.core import get_response_synthesizer
from llama_index.core import VectorStoreIndex, PromptTemplate

index = VectorStoreIndex.from_documents(documents)

# 创建响应合成器
synthesizer = get_response_synthesizer(
    response_mode="compact",       # 合成模式
    streaming=True,                # 流式输出
    # 自定义QA提示模板
    text_qa_template=PromptTemplate(
        "上下文信息如下：\n"
        "---------------------\n"
        "{context_str}\n"
        "---------------------\n"
        "根据上下文信息（不依赖先验知识），回答以下问题。\n"
        "问题: {query_str}\n"
        "回答: "
    ),
)

# 与查询引擎配合
query_engine = index.as_query_engine(response_synthesizer=synthesizer)

from llama_index.core.agent import ReActAgent
from llama_index.core.tools import FunctionTool
from llama_index.llms.openai import OpenAI

# 定义工具函数
def multiply(a: int, b: int) -> int:
    """将两个整数相乘并返回结果。"""
    return a * b

def add(a: int, b: int) -> int:
    """将两个整数相加并返回结果。"""
    return a + b

# 包装为 LlamaIndex 工具
multiply_tool = FunctionTool.from_defaults(fn=multiply)
add_tool = FunctionTool.from_defaults(fn=add)

# 创建 Agent
llm = OpenAI(model="gpt-4o")
agent = ReActAgent.from_tools(
    [multiply_tool, add_tool],
    llm=llm,
    verbose=True,  # 打印推理过程
)

# 执行任务
response = agent.chat("3加5等于多少？然后把结果乘以2")
print(response)

from llama_index.core.agent import FunctionCallingAgent

agent = FunctionCallingAgent.from_tools(
    [multiply_tool, add_tool],
    llm=llm,
    verbose=True,
)

response = agent.chat("计算 (10 + 20) * 3")

from llama_index.core.tools import QueryEngineTool

# 将查询引擎包装为工具
query_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine,
    name="knowledge_base",
    description="包含公司内部技术文档的知识库，可用于查询技术相关问题和公司政策。",
)

agent = ReActAgent.from_tools([query_tool, add_tool], llm=llm, verbose=True)
response = agent.chat("查阅公司文档中关于API限流的规定，然后计算限流阈值的2倍是多少")

from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings

# 使用本地 HuggingFace 模型
embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",  # 中文嵌入模型
    max_length=512,                        # 最大序列长度
    cache_folder="./model_cache",          # 模型缓存目录
    embed_batch_size=10,                   # 批量嵌入大小
)

# 设置为全局默认
Settings.embed_model = embed_model

# 直接使用
embeddings = embed_model.get_text_embedding("Hello, world!")
print(f"嵌入维度: {len(embeddings)}")  # 通常为 384 或 768

# 批量嵌入
embeddings = embed_model.get_text_embedding_batch(
    ["第一句话", "第二句话", "第三句话"]
)

from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core import Settings

embed_model = OpenAIEmbedding(
    model="text-embedding-3-small",  # 嵌入模型名称
    dimensions=1536,                 # 输出维度（text-embedding-3 支持自定义）
    api_key="sk-...",               # API Key，也可通过环境变量设置
)

Settings.embed_model = embed_model

from llama_index.core import Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 全局配置 LLM
Settings.llm = OpenAI(
    model="gpt-4o",
    temperature=0.1,      # 低温度使输出更确定性
    max_tokens=1024,      # 最大生成token数
    system_prompt="你是一个专业的技术助手，用中文回答问题。",  # 系统提示
)

# 全局配置嵌入模型
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-zh-v1.5")

# 全局配置切分参数
Settings.chunk_size = 1024       # Node 最大大小
Settings.chunk_overlap = 200     # Node 重叠大小

# 全局配置回调
from llama_index.core.callbacks import TokenCountingHandler
token_counter = TokenCountingHandler()
Settings.callback_manager.add_handler(token_counter)

# 查看token使用量
# print(token_counter.total_llm_token_count)

import os
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings, PromptTemplate
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 1. 配置模型
Settings.llm = OpenAI(model="gpt-4o", temperature=0)
Settings.embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-small-zh-v1.5")

# 2. 加载文档
documents = SimpleDirectoryReader("./knowledge_base").load_data()
print(f"加载了 {len(documents)} 个文档")

# 3. 构建索引
index = VectorStoreIndex.from_documents(documents)

# 4. 创建查询引擎（自定义提示词）
qa_template = PromptTemplate(
    "上下文信息如下：\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "请仅根据上下文信息回答问题。如果上下文中没有相关信息，请回答"我不知道"。\n"
    "问题: {query_str}\n"
    "回答: "
)

query_engine = index.as_query_engine(
    similarity_top_k=3,
    text_qa_template=qa_template,
    response_mode="compact",
)

# 5. 查询
response = query_engine.query("公司的API限流策略是什么？")
print(response)

# 6. 查看来源
for node in response.source_nodes:
    print(f"[分数: {node.score:.4f}] 来源: {node.node.metadata.get('filename', 'unknown')}")

from llama_index.core import SummaryIndex, Document
from llama_index.core.query_engine import SubQuestionQueryEngine
from llama_index.core.tools import QueryEngineTool, ToolMetadata

# 为每个文档构建独立的索引和查询引擎
doc1 = Document(text="产品A的特性：轻量级、易部署...", metadata={"product": "A"})
doc2 = Document(text="产品B的特性：高性能、可扩展...", metadata={"product": "B"})

index_a = VectorStoreIndex.from_documents([doc1])
index_b = VectorStoreIndex.from_documents([doc2])

engine_a = index_a.as_query_engine(similarity_top_k=3)
engine_b = index_b.as_query_engine(similarity_top_k=3)

# 将查询引擎包装为工具
tool_a = QueryEngineTool.from_defaults(
    query_engine=engine_a,
    metadata=ToolMetadata(
        name="product_a",
        description="包含产品A的详细信息，包括特性、价格、部署方式等",
    ),
)
tool_b = QueryEngineTool.from_defaults(
    query_engine=engine_b,
    metadata=ToolMetadata(
        name="product_b",
        description="包含产品B的详细信息，包括特性、价格、部署方式等",
    ),
)

# 创建子问题查询引擎（自动分解问题）
sub_question_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=[tool_a, tool_b]
)

# 提出对比问题
response = sub_question_engine.query("产品A和产品B在部署方式上有什么区别？")
print(response)

from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core import VectorStoreIndex

memory = ChatMemoryBuffer.from_defaults(token_limit=3900)  # 对话记忆token上限

index = VectorStoreIndex.from_documents(documents)

chat_engine = index.as_chat_engine(
    chat_mode="condense_question",  # 对话模式
    memory=memory,
    similarity_top_k=3,
    verbose=True,
)

# 多轮对话
response1 = chat_engine.chat("公司的年假政策是什么？")
print(response1)

response2 = chat_engine.chat("那病假呢？")  # 会自动结合上下文理解"那病假"
print(response2)

# 重置对话
chat_engine.reset()

from llama_index.core.query_pipeline import QueryPipeline
from llama_index.core import PromptTemplate

# 定义管道步骤
prompt_str = "请根据以下上下文回答问题：\n上下文: {context}\n问题: {query}"
prompt_tmpl = PromptTemplate(prompt_str)

# 组装管道
p = QueryPipeline()
p.add_modules({
    "retriever": retriever,
    "prompt": prompt_tmpl,
    "llm": llm,
})
p.add_link("retriever", "prompt", dest_key="context")
p.add_link("prompt", "llm")

# 执行
response = p.run(query="什么是RAG？")

from llama_index.core.vector_stores import MetadataFilters, MetadataFilter

filters = MetadataFilters(
    filters=[MetadataFilter(key="category", value="技术文档")]
)
retriever = index.as_retriever(filters=filters, similarity_top_k=5)

# 保存索引
index.storage_context.persist(persist_dir="./storage")

# 加载索引
from llama_index.core import load_index_from_storage, StorageContext
storage_context = StorageContext.from_defaults(persist_dir="./storage")
index = load_index_from_storage(storage_context)

from llama_index.core.postprocessor import SentenceTransformerRerank

reranker = SentenceTransformerRerank(
    model="cross-encoder/ms-marco-MiniLM-L-2-v2",
    top_n=3,  # 重排后保留的节点数
)

query_engine = index.as_query_engine(
    similarity_top_k=10,         # 先多检索
    node_postprocessors=[reranker],  # 再精排
)

from llama_index.core.evaluation import RetrieverEvaluator

evaluator = RetrieverEvaluator.from_metric_names(
    ["mrr", "hit_rate", "precision", "recall", "ap", "ndcg"],
)
# 评估结果可根据指标调整 similarity_top_k 和检索策略