LlamaIndex框架¶

from llama_index.core import Document

# 手动创建Document
doc = Document(
    text="LlamaIndex是一个数据框架...",
    metadata={
        "source": "tutorial.md",
        "author": "AI Team",
        "created_at": "2025-01-01"
    }
)

print(f"文档ID: {doc.doc_id}")
print(f"文档内容: {doc.text[:100]}")
print(f"元数据: {doc.metadata}")

from llama_index.core.node_parser import SentenceSplitter

# 创建节点解析器
parser = SentenceSplitter(
    chunk_size=1024,
    chunk_overlap=200
)

# 将Document转换为Nodes
nodes = parser.get_nodes_from_documents([doc])

for i, node in enumerate(nodes):  # enumerate同时获取索引和元素
    print(f"Node {i}: {node.text[:80]}...")
    print(f"  - Node ID: {node.node_id}")
    print(f"  - 元数据: {node.metadata}")
    print(f"  - 关系: {node.relationships}")

from llama_index.core import VectorStoreIndex

# 从Nodes创建索引
index = VectorStoreIndex(nodes)

# 也可以直接从Documents创建（内部自动分块）
index = VectorStoreIndex.from_documents(documents)

# 创建查询引擎
query_engine = index.as_query_engine()

# 执行查询
response = query_engine.query("LlamaIndex是什么？")
print(response)

import os
from llama_index.core import Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# 设置API密钥
os.environ["OPENAI_API_KEY"] = "your-openai-api-key"

# 全局设置（推荐方式）
Settings.llm = OpenAI(model="gpt-4o", temperature=0.1)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
Settings.chunk_size = 1024
Settings.chunk_overlap = 200

print("LlamaIndex配置完成！")
print(f"LLM: {Settings.llm.model}")
print(f"Embed Model: {Settings.embed_model.model_name}")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# Step 1: 加载数据
documents = SimpleDirectoryReader("./data").load_data()
print(f"加载了 {len(documents)} 个文档")

# Step 2: 构建索引
index = VectorStoreIndex.from_documents(documents)

# Step 3: 创建查询引擎
query_engine = index.as_query_engine()

# Step 4: 进行查询
response = query_engine.query("请总结这些文档的主要内容")
print(response)

# 查看引用的源文档
for node in response.source_nodes:
    print(f"\n来源: {node.metadata.get('file_name', 'Unknown')}")
    print(f"相关度: {node.score:.4f}")
    print(f"内容: {node.text[:200]}...")

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

# 使用本地HuggingFace模型
Settings.llm = HuggingFaceLLM(
    model_name="Qwen/Qwen2.5-7B-Instruct",
    tokenizer_name="Qwen/Qwen2.5-7B-Instruct",
    context_window=4096,
    max_new_tokens=512,
    device_map="auto",
    model_kwargs={"torch_dtype": "float16"}
)

# 使用本地Embedding模型
Settings.embed_model = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5"
)

print("本地模型配置完成！")

from llama_index.llms.ollama import Ollama
from llama_index.embeddings.ollama import OllamaEmbedding
from llama_index.core import Settings

# 使用Ollama运行的本地模型
Settings.llm = Ollama(
    model="qwen2.5:7b",
    base_url="http://localhost:11434",
    request_timeout=120.0
)

Settings.embed_model = OllamaEmbedding(
    model_name="nomic-embed-text",
    base_url="http://localhost:11434"
)

from llama_index.core import SimpleDirectoryReader

# 基础用法 - 加载目录下所有支持格式的文件
documents = SimpleDirectoryReader("./data").load_data()

# 指定文件类型
documents = SimpleDirectoryReader(
    input_dir="./data",
    required_exts=[".pdf", ".txt", ".md", ".docx"],
    recursive=True,  # 递归加载子目录
    filename_as_id=True  # 使用文件名作为文档ID
).load_data()

# 加载单个文件
documents = SimpleDirectoryReader(
    input_files=["./data/report.pdf", "./data/readme.md"]
).load_data()

# 排除特定文件
documents = SimpleDirectoryReader(
    input_dir="./data",
    exclude=["*.tmp", "*.log"],
    recursive=True
).load_data()

for doc in documents:
    print(f"文件: {doc.metadata['file_name']}")
    print(f"大小: {len(doc.text)} 字符")
    print(f"类型: {doc.metadata.get('file_type', 'Unknown')}")
    print("---")

from llama_index.readers.file import PDFReader

# 基础PDF加载
pdf_reader = PDFReader()
documents = pdf_reader.load_data(file_path="./data/technical_doc.pdf")

# 使用SimpleDirectoryReader加载PDF（自动识别）
from llama_index.core import SimpleDirectoryReader
documents = SimpleDirectoryReader(
    input_files=["./data/paper.pdf"]
).load_data()

# 使用pymupdf加速PDF解析
# pip install llama-index-readers-file pymupdf
from llama_index.readers.file import PyMuPDFReader

reader = PyMuPDFReader()
documents = reader.load_data(file_path="./data/paper.pdf")

for doc in documents:
    print(f"页码: {doc.metadata.get('page_label', 'N/A')}")
    print(f"内容预览: {doc.text[:200]}")
    print("---")

from llama_index.readers.web import SimpleWebPageReader, BeautifulSoupWebReader

# 简单网页加载
web_reader = SimpleWebPageReader()
documents = web_reader.load_data(
    urls=["https://docs.llamaindex.ai/en/stable/"]
)

# 使用BeautifulSoup解析
bs_reader = BeautifulSoupWebReader()
documents = bs_reader.load_data(
    urls=["https://example.com/article1", "https://example.com/article2"]
)

for doc in documents:
    print(f"URL: {doc.metadata.get('url', 'N/A')}")
    print(f"内容长度: {len(doc.text)}")

from llama_index.readers.json import JSONReader

# 加载JSON数据
json_reader = JSONReader(
    levels_back=0,
    collapse_length=None,
    ensure_ascii=False
)
documents = json_reader.load_data(input_file="./data/api_data.json")

# 从JSON行文件加载
import json
from llama_index.core import Document

def load_jsonl(file_path: str) -> list:
    """加载JSONL文件"""
    documents = []
    with open(file_path, 'r', encoding='utf-8') as f:  # with自动管理文件关闭
        for line in f:
            data = json.loads(line.strip())  # json.loads将JSON字符串→Python对象
            doc = Document(
                text=data.get("content", ""),
                metadata={
                    "title": data.get("title", ""),
                    "source": data.get("source", ""),
                    "date": data.get("date", "")
                }
            )
            documents.append(doc)
    return documents

documents = load_jsonl("./data/articles.jsonl")

from llama_index.readers.database import DatabaseReader

# MySQL加载器
db_reader = DatabaseReader(
    uri="mysql+pymysql://user:password@localhost:3306/mydb"
)

# 使用SQL查询加载数据
documents = db_reader.load_data(
    query="SELECT id, title, content FROM articles WHERE status='published'"
)

# PostgreSQL加载器
pg_reader = DatabaseReader(
    uri="postgresql://user:password@localhost:5432/mydb"
)
documents = pg_reader.load_data(
    query="SELECT * FROM knowledge_base"
)

# 自定义数据库加载函数
from sqlalchemy import create_engine, text
from llama_index.core import Document

def load_from_database(connection_string: str, query: str) -> list:
    """从数据库加载数据为Document列表"""
    engine = create_engine(connection_string)
    documents = []

    with engine.connect() as conn:
        result = conn.execute(text(query))
        columns = result.keys()

        for row in result:
            row_dict = dict(zip(columns, row))
            content = row_dict.pop("content", str(row_dict))
            doc = Document(
                text=content,
                metadata=row_dict
            )
            documents.append(doc)

    return documents

# Notion加载器
# pip install llama-index-readers-notion
from llama_index.readers.notion import NotionPageReader

notion_reader = NotionPageReader(integration_token="your-notion-token")
documents = notion_reader.load_data(page_ids=["page_id_1", "page_id_2"])

# Slack加载器
# pip install llama-index-readers-slack
from llama_index.readers.slack import SlackReader

slack_reader = SlackReader(slack_token="your-slack-token")
documents = slack_reader.load_data(channel_ids=["C01234ABCDE"])

# GitHub加载器
# pip install llama-index-readers-github
from llama_index.readers.github import GithubRepositoryReader

github_reader = GithubRepositoryReader(
    github_token="your-github-token",
    owner="owner",
    repo="repo",
    filter_file_extensions=[".py", ".md", ".txt"]
)
documents = github_reader.load_data(branch="main")

# CSV加载器
# pip install llama-index-readers-file
from llama_index.readers.file import PandasCSVReader

csv_reader = PandasCSVReader()
documents = csv_reader.load_data(file_path="./data/dataset.csv")

from llama_index.core.readers.base import BaseReader
from llama_index.core import Document

class CustomAPIReader(BaseReader):
    """自定义API数据加载器"""

    def __init__(self, api_base_url: str, api_key: str):
        self.api_base_url = api_base_url
        self.api_key = api_key

    def load_data(
        self,
        endpoint: str = "/articles",
        params: dict | None = None
    ) -> list[Document]:
        """从API加载数据"""
        import requests

        headers = {"Authorization": f"Bearer {self.api_key}"}
        response = requests.get(
            f"{self.api_base_url}{endpoint}",
            headers=headers,
            params=params or {}
        )
        response.raise_for_status()

        documents = []
        for item in response.json().get("data", []):
            doc = Document(
                text=item.get("content", ""),
                metadata={
                    "title": item.get("title", ""),
                    "id": item.get("id", ""),
                    "created_at": item.get("created_at", ""),
                    "source": f"{self.api_base_url}{endpoint}"
                }
            )
            documents.append(doc)

        return documents

# 使用自定义加载器
reader = CustomAPIReader(
    api_base_url="https://api.example.com",
    api_key="your-api-key"
)
documents = reader.load_data(endpoint="/articles", params={"limit": 100})

from llama_index.core.node_parser import SentenceSplitter

# 基础句子分块
splitter = SentenceSplitter(
    chunk_size=1024,      # 每块最大token数
    chunk_overlap=200,    # 块间重叠token数
    separator=" ",        # 分隔符
    paragraph_separator="\n\n\n"  # 段落分隔符
)

nodes = splitter.get_nodes_from_documents(documents)
print(f"生成了 {len(nodes)} 个节点")

for i, node in enumerate(nodes[:3]):
    print(f"\nNode {i}:")
    print(f"  文本长度: {len(node.text)}")
    print(f"  预览: {node.text[:100]}...")

from llama_index.core.node_parser import SentenceWindowNodeParser

# 句子窗口解析器 - 用于高级RAG
window_parser = SentenceWindowNodeParser.from_defaults(
    window_size=3,               # 窗口大小（前后各3句）
    window_metadata_key="window",
    original_text_metadata_key="original_text"
)

nodes = window_parser.get_nodes_from_documents(documents)

# 查看节点和窗口
for node in nodes[:2]:
    print(f"原始句子: {node.text[:80]}...")
    print(f"窗口文本: {node.metadata.get('window', '')[:200]}...")
    print("---")

from llama_index.core.node_parser import HierarchicalNodeParser, get_leaf_nodes

# 层级解析器 - 用于自动合并检索
hierarchical_parser = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128]  # 三层层级：大块→中块→小块
)

nodes = hierarchical_parser.get_nodes_from_documents(documents)
leaf_nodes = get_leaf_nodes(nodes)

print(f"总节点数: {len(nodes)}")
print(f"叶子节点数: {len(leaf_nodes)}")

# 查看层级关系
for node in nodes[:5]:
    parent = node.relationships.get("parent", None)
    children = node.relationships.get("children", [])
    print(f"Node: {node.node_id[:8]}...")
    print(f"  层级: 父={parent is not None}, 子={len(children)}")
    print(f"  文本长度: {len(node.text)}")

from llama_index.core.node_parser import MarkdownNodeParser

# Markdown感知解析器
md_parser = MarkdownNodeParser()

nodes = md_parser.get_nodes_from_documents(documents)

for node in nodes[:5]:
    print(f"标题: {node.metadata.get('header_path', 'N/A')}")
    print(f"内容: {node.text[:100]}...")
    print("---")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# 加载文档并构建向量索引
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)

# 查询（默认返回top-2最相似的节点）
query_engine = index.as_query_engine(similarity_top_k=3)
response = query_engine.query("什么是深度学习？")
print(response)

# 查看检索到的源节点
for node in response.source_nodes:
    print(f"相关度分数: {node.score:.4f}")
    print(f"内容: {node.text[:200]}")
    print("---")

from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.core.node_parser import SentenceSplitter

# 自定义分块
parser = SentenceSplitter(chunk_size=512, chunk_overlap=100)
nodes = parser.get_nodes_from_documents(documents)

# 从节点创建索引（更多控制）
index = VectorStoreIndex(
    nodes=nodes,
    show_progress=True  # 显示进度
)

# 增量添加文档
new_docs = SimpleDirectoryReader("./new_data").load_data()
for doc in new_docs:
    index.insert(doc)

# 删除文档
index.delete_ref_doc("doc_id_to_delete")

from llama_index.core import SummaryIndex

# 创建摘要索引
summary_index = SummaryIndex.from_documents(documents)

# 查询（默认使用所有节点）
query_engine = summary_index.as_query_engine(
    response_mode="tree_summarize"
)
response = query_engine.query("请对所有文档进行总结")
print(response)

# 适用场景：
# - 需要全文摘要时
# - 文档量较小时
# - 不需要精确定位具体段落时

from llama_index.core import TreeIndex

# 创建树形索引
tree_index = TreeIndex.from_documents(
    documents,
    num_children=10,  # 每个节点的子节点数
    show_progress=True
)

# 查询
query_engine = tree_index.as_query_engine(
    child_branch_factor=2  # 每次遍历的分支数
)
response = query_engine.query("这些文档的主要主题是什么？")
print(response)

# 适用场景：
# - 层级化的文档结构
# - 需要从概览到细节的多层次查询
# - 大规模文档集合的摘要

from llama_index.core import KeywordTableIndex

# 创建关键词索引
keyword_index = KeywordTableIndex.from_documents(documents)

# 查询
query_engine = keyword_index.as_query_engine()
response = query_engine.query("Transformer架构")
print(response)

# 查看提取的关键词
print("\n关键词表:")
for keyword, node_ids in keyword_index.index_struct.table.items():
    print(f"  {keyword}: {len(node_ids)} 个节点")

# 适用场景：
# - 基于精确关键词的检索
# - 技术文档搜索
# - 与向量检索互补使用

from llama_index.core import KnowledgeGraphIndex
from llama_index.core import StorageContext
from llama_index.graph_stores.nebula import NebulaGraphStore

# 基础知识图谱索引（内存存储）
kg_index = KnowledgeGraphIndex.from_documents(
    documents,
    max_triplets_per_chunk=10,
    include_embeddings=True,
    show_progress=True
)

# 查询
query_engine = kg_index.as_query_engine(
    include_text=True,
    response_mode="tree_summarize"
)
response = query_engine.query("LlamaIndex和LangChain有什么关系？")
print(response)

# 使用NebulaGraph存储
# graph_store = NebulaGraphStore(
#     space_name="knowledge_graph",
#     edge_types=["关系"],
#     rel_prop_names=["name"],
#     tags=["entity"]
# )
# storage_context = StorageContext.from_defaults(graph_store=graph_store)
# kg_index = KnowledgeGraphIndex.from_documents(
#     documents,
#     storage_context=storage_context,
# )

# 适用场景：
# - 实体关系问答
# - 多跳推理
# - 结构化知识检索

from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    StorageContext,
    load_index_from_storage
)

# 构建索引
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)

# 持久化到本地
index.storage_context.persist(persist_dir="./storage")
print("索引已保存到 ./storage")

# 从本地加载
storage_context = StorageContext.from_defaults(persist_dir="./storage")
loaded_index = load_index_from_storage(storage_context)
print("索引已从 ./storage 加载")

# 使用加载的索引进行查询
query_engine = loaded_index.as_query_engine()
response = query_engine.query("请回答问题")
print(response)

# pip install llama-index-vector-stores-chroma chromadb
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import VectorStoreIndex, StorageContext

# 创建Chroma客户端
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)

# 构建索引（数据自动存储到Chroma）
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context
)

# 后续加载（无需重新构建索引）
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_collection("my_collection")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)
index = VectorStoreIndex.from_vector_store(vector_store)

query_engine = index.as_query_engine()
response = query_engine.query("你的问题")

# pip install llama-index-vector-stores-pinecone pinecone
from pinecone import Pinecone, ServerlessSpec
from llama_index.vector_stores.pinecone import PineconeVectorStore
from llama_index.core import VectorStoreIndex, StorageContext

# 初始化Pinecone
pc = Pinecone(api_key="your-pinecone-api-key")

# 创建索引（如果不存在）
index_name = "llama-index-demo"
if index_name not in pc.list_indexes().names():
    pc.create_index(
        name=index_name,
        dimension=1536,  # OpenAI text-embedding-3-small的维度
        metric="cosine",
        spec=ServerlessSpec(cloud="aws", region="us-east-1")
    )

pinecone_index = pc.Index(index_name)

# 配置向量存储
vector_store = PineconeVectorStore(pinecone_index=pinecone_index)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

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

# 查询
query_engine = index.as_query_engine()
response = query_engine.query("你的问题")

# pip install llama-index-vector-stores-weaviate weaviate-client
import weaviate
from llama_index.vector_stores.weaviate import WeaviateVectorStore
from llama_index.core import VectorStoreIndex, StorageContext

# 连接Weaviate
client = weaviate.connect_to_local()  # 或 weaviate.connect_to_wcs(...)

# 配置向量存储
vector_store = WeaviateVectorStore(
    weaviate_client=client,
    index_name="LlamaIndexDemo"
)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

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

# 查询
query_engine = index.as_query_engine()
response = query_engine.query("你的问题")

# 记得关闭连接
client.close()

# pip install llama-index-vector-stores-milvus pymilvus
from llama_index.vector_stores.milvus import MilvusVectorStore
from llama_index.core import VectorStoreIndex, StorageContext

# 连接Milvus
vector_store = MilvusVectorStore(
    uri="http://localhost:19530",  # Milvus服务地址
    collection_name="llama_index_demo",
    dim=1536,  # 嵌入维度
    overwrite=True  # 覆盖已有集合
)

storage_context = StorageContext.from_defaults(vector_store=vector_store)

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

# 使用Milvus Lite（轻量级，无需服务器）
vector_store_lite = MilvusVectorStore(
    uri="./milvus_demo.db",  # 本地文件
    collection_name="demo",
    dim=1536
)

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)

# === QueryEngine: 一问一答模式 ===
query_engine = index.as_query_engine(
    similarity_top_k=3,
    response_mode="compact"
)
response = query_engine.query("什么是RAG？")
print("QueryEngine:", response)

# === ChatEngine: 多轮对话模式 ===
chat_engine = index.as_chat_engine(
    chat_mode="condense_plus_context",
    similarity_top_k=3,
    verbose=True
)

# 多轮对话
response1 = chat_engine.chat("什么是RAG？")
print("回合1:", response1)

response2 = chat_engine.chat("它有什么优势？")  # 自动携带上下文
print("回合2:", response2)

response3 = chat_engine.chat("如何实现？")
print("回合3:", response3)

# 重置对话
chat_engine.reset()

# QueryEngine流式输出
query_engine = index.as_query_engine(streaming=True)
streaming_response = query_engine.query("详细解释Transformer架构")

# 逐token输出
for text in streaming_response.response_gen:
    print(text, end="", flush=True)
print()  # 换行

# ChatEngine流式输出
chat_engine = index.as_chat_engine(
    chat_mode="condense_plus_context",
    streaming=True
)
streaming_response = chat_engine.stream_chat("介绍RAG的工作原理")

for token in streaming_response.response_gen:
    print(token, end="", flush=True)
print()

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

# 假设有多个索引（不同数据源）
tech_index = VectorStoreIndex.from_documents(tech_docs)
business_index = VectorStoreIndex.from_documents(business_docs)

# 创建查询引擎工具
query_engine_tools = [
    QueryEngineTool(
        query_engine=tech_index.as_query_engine(),
        metadata=ToolMetadata(
            name="tech_knowledge",
            description="包含技术文档的知识库，用于回答技术相关问题"
        )
    ),
    QueryEngineTool(
        query_engine=business_index.as_query_engine(),
        metadata=ToolMetadata(
            name="business_knowledge",
            description="包含商业文档的知识库，用于回答商业相关问题"
        )
    ),
]

# 创建子问题查询引擎
sub_question_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=query_engine_tools,
    verbose=True
)

# 复杂查询会被自动分解
response = sub_question_engine.query(
    "从技术和商业角度分析AI在医疗领域的应用前景"
)
print(response)

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

# 创建路由查询引擎
router_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=query_engine_tools,
    verbose=True
)

# LLM自动选择最合适的引擎
response = router_engine.query("最新的AI技术趋势是什么？")
print(response)

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)

# 基础向量检索
retriever = index.as_retriever(
    similarity_top_k=5
)

nodes = retriever.retrieve("什么是深度学习？")
for node in nodes:
    print(f"分数: {node.score:.4f}")
    print(f"内容: {node.text[:150]}...")
    print("---")

from llama_index.core import KeywordTableIndex

keyword_index = KeywordTableIndex.from_documents(documents)

retriever = keyword_index.as_retriever(
    retriever_mode="default"  # 或 "simple"
)

nodes = retriever.retrieve("Transformer注意力机制")
for node in nodes:
    print(f"内容: {node.text[:150]}...")

from llama_index.core.retrievers import QueryFusionRetriever

# 创建多个检索器
vector_retriever = vector_index.as_retriever(similarity_top_k=5)
keyword_retriever = keyword_index.as_retriever()

# 查询融合检索器（混合检索 + 倒数排名融合）
fusion_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, keyword_retriever],
    similarity_top_k=5,
    num_queries=4,  # 生成查询变体数量
    mode="reciprocal_rerank",  # 倒数排名融合
    use_async=True,
    verbose=True
)

# 检索
nodes = fusion_retriever.retrieve("解释注意力机制的工作原理")
for node in nodes:
    print(f"分数: {node.score:.4f}")
    print(f"内容: {node.text[:150]}...")

from llama_index.core.retrievers import VectorIndexAutoRetriever
from llama_index.core.vector_stores.types import MetadataInfo, VectorStoreInfo

# 定义元数据信息
vector_store_info = VectorStoreInfo(
    content_info="技术文档",
    metadata_info=[
        MetadataInfo(
            name="category",
            type="str",
            description="文档类别，如：AI、Web、Cloud"
        ),
        MetadataInfo(
            name="date",
            type="str",
            description="文档发布日期，格式：YYYY-MM-DD"
        ),
        MetadataInfo(
            name="author",
            type="str",
            description="文档作者"
        ),
    ]
)

# 创建自动检索器
auto_retriever = VectorIndexAutoRetriever(
    index=vector_index,
    vector_store_info=vector_store_info,
    similarity_top_k=5,
    verbose=True
)

# 查询时自动添加过滤条件
nodes = auto_retriever.retrieve(
    "2025年关于AI的最新文档"
)
# LLM会自动生成: category="AI", date>="2025-01-01" 的过滤条件

from llama_index.core.node_parser import SentenceWindowNodeParser
from llama_index.core.indices.postprocessor import MetadataReplacementPostProcessor
from llama_index.core import VectorStoreIndex, Settings

# Step 1: 使用句子窗口解析器
node_parser = SentenceWindowNodeParser.from_defaults(
    window_size=3,  # 每侧包含3句
    window_metadata_key="window",
    original_text_metadata_key="original_text"
)

# Step 2: 构建索引（嵌入基于单句）
nodes = node_parser.get_nodes_from_documents(documents)
sentence_index = VectorStoreIndex(nodes)

# Step 3: 查询时用窗口文本替换原始句子
postprocessor = MetadataReplacementPostProcessor(
    target_metadata_key="window"
)

# Step 4: 创建查询引擎
query_engine = sentence_index.as_query_engine(
    similarity_top_k=3,
    node_postprocessors=[postprocessor]
)

response = query_engine.query("什么是Attention机制？")
print(response)

# 优势：
# - 嵌入匹配更精确（基于单句）
# - 上下文更丰富（返回窗口内多句）
# - 减少信息截断问题

from llama_index.core.node_parser import (
    HierarchicalNodeParser,
    get_leaf_nodes,
    get_root_nodes,
)
from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core import VectorStoreIndex, StorageContext

# Step 1: 构建层级节点
node_parser = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128]  # 大→中→小 三层
)

nodes = node_parser.get_nodes_from_documents(documents)
leaf_nodes = get_leaf_nodes(nodes)

print(f"总节点: {len(nodes)}, 叶节点: {len(leaf_nodes)}")

# Step 2: 存储所有节点（包括父节点）
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(nodes)

# Step 3: 只索引叶子节点的嵌入
index = VectorStoreIndex(
    leaf_nodes,
    storage_context=storage_context
)

# Step 4: 使用自动合并检索器
base_retriever = index.as_retriever(similarity_top_k=12)
auto_merging_retriever = AutoMergingRetriever(
    base_retriever,
    storage_context=storage_context,
    simple_ratio_thresh=0.5  # 超过50%的子节点被检索到就合并
)

# 查询
nodes = auto_merging_retriever.retrieve("深度学习的基本原理是什么？")
for node in nodes:
    print(f"文本长度: {len(node.text)}")
    print(f"内容: {node.text[:200]}...")
    print("---")

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

# 递归检索：先检索摘要，再深入检索相关细节
# 适用于多文档、多层级的复杂知识库

# Step 1: 创建文档摘要索引
from llama_index.core import SummaryIndex
summary_index = SummaryIndex.from_documents(documents)

# Step 2: 创建详细向量索引
vector_index = VectorStoreIndex.from_documents(documents)

# Step 3: 使用递归检索
# 先通过摘要定位相关文档，再通过向量检索获取细节
from llama_index.core.schema import IndexNode

# 为每个文档创建索引节点
index_nodes = []
for doc in documents:
    # 创建指向详细索引的索引节点
    index_node = IndexNode(
        text=f"关于 {doc.metadata.get('file_name', 'unknown')} 的摘要",
        index_id=doc.doc_id
    )
    index_nodes.append(index_node)

# 构建顶层索引
top_index = VectorStoreIndex(index_nodes)

# 创建递归检索器
recursive_retriever = RecursiveRetriever(
    "vector",
    retriever_dict={
        "vector": top_index.as_retriever(similarity_top_k=3)
    },
    query_engine_dict={
        doc.doc_id: vector_index.as_query_engine()
        for doc in documents
    },
    verbose=True
)

# 创建查询引擎
query_engine = RetrieverQueryEngine.from_args(recursive_retriever)
response = query_engine.query("请详细介绍相关技术")

from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.core.schema import IndexNode

# 小块用于精确匹配,大块用于上下文
small_parser = SentenceSplitter(chunk_size=128, chunk_overlap=20)
large_parser = SentenceSplitter(chunk_size=1024, chunk_overlap=200)

# 创建大块节点
large_nodes = large_parser.get_nodes_from_documents(documents)

# 创建小块节点,并链接到大块节点
all_nodes = []
for large_node in large_nodes:
    small_nodes = small_parser.get_nodes_from_documents(
        [Document(text=large_node.text)]
    )
    for small_node in small_nodes:
        # 创建引用大块的索引节点
        index_node = IndexNode(
            text=small_node.text,
            index_id=large_node.node_id
        )
        all_nodes.append(index_node)

# 存储大块节点和索引节点
storage_context = StorageContext.from_defaults()
storage_context.docstore.add_documents(large_nodes)

# 用小块建索引
index = VectorStoreIndex(all_nodes, storage_context=storage_context)

# 检索时返回大块上下文
query_engine = index.as_query_engine(similarity_top_k=3)
response = query_engine.query("你的问题")

# pip install llama-index-postprocessor-cohere-rerank
from llama_index.postprocessor.cohere_rerank import CohereRerank

# 创建Cohere重排序器
cohere_reranker = CohereRerank(
    api_key="your-cohere-api-key",
    model="rerank-multilingual-v3.0",
    top_n=3  # 重排后保留top-3
)

# 在查询引擎中使用
query_engine = index.as_query_engine(
    similarity_top_k=10,  # 先检索10个
    node_postprocessors=[cohere_reranker]  # 重排后保留3个
)

response = query_engine.query("深度学习与机器学习的区别")
print(response)

from llama_index.core.postprocessor import SentenceTransformerRerank

# 使用CrossEncoder模型重排
reranker = SentenceTransformerRerank(
    model="cross-encoder/ms-marco-MiniLM-L-6-v2",
    top_n=3
)

# 在查询引擎中使用
query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[reranker]
)

response = query_engine.query("你的问题")

from llama_index.core.postprocessor import LLMRerank

# 使用LLM进行重排序
llm_reranker = LLMRerank(
    top_n=3,
    choice_batch_size=5  # 每批处理5个节点
)

query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[llm_reranker]
)

response = query_engine.query("你的问题")

from llama_index.core import get_response_synthesizer
from llama_index.core.query_engine import RetrieverQueryEngine

# === Refine模式 ===
# 逐个节点精化答案（默认模式）
refine_synthesizer = get_response_synthesizer(response_mode="refine")

# === Compact模式 ===
# 将所有节点压缩到一个prompt中（减少LLM调用）
compact_synthesizer = get_response_synthesizer(response_mode="compact")

# === Tree Summarize模式 ===
# 自底向上树形摘要（适合大量节点）
tree_synthesizer = get_response_synthesizer(response_mode="tree_summarize")

# === Simple Summarize模式 ===
# 简单截断拼接后一次性摘要
simple_synthesizer = get_response_synthesizer(response_mode="simple_summarize")

# === Accumulate模式 ===
# 对每个节点单独生成答案，然后拼接
accumulate_synthesizer = get_response_synthesizer(response_mode="accumulate")

# === No Text模式 ===
# 只返回检索到的节点，不生成答案
no_text_synthesizer = get_response_synthesizer(response_mode="no_text")

# 在查询引擎中使用
retriever = index.as_retriever(similarity_top_k=5)
query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=compact_synthesizer
)

response = query_engine.query("综合分析深度学习的发展历程")
print(response)

from llama_index.core import PromptTemplate

# 自定义QA Prompt
qa_prompt_tmpl = PromptTemplate(
    "上下文信息如下：\n"
    "---------------------\n"
    "{context_str}\n"
    "---------------------\n"
    "请基于以上上下文信息（而非先验知识）回答以下问题。\n"
    "如果上下文信息不足以回答问题，请明确说明'根据提供的信息无法回答此问题'。\n"
    "问题: {query_str}\n"
    "答案: "
)

# 自定义Refine Prompt
refine_prompt_tmpl = PromptTemplate(
    "原始问题: {query_str}\n"
    "已有答案: {existing_answer}\n"
    "新的上下文信息:\n"
    "---------------------\n"
    "{context_msg}\n"
    "---------------------\n"
    "请根据新的上下文信息优化已有答案。如果新信息无用，返回原始答案。\n"
    "优化后的答案: "
)

# 创建使用自定义prompt的合成器
synthesizer = get_response_synthesizer(
    response_mode="compact",
    text_qa_template=qa_prompt_tmpl,
    refine_template=refine_prompt_tmpl
)

query_engine = index.as_query_engine(
    response_synthesizer=synthesizer,
    similarity_top_k=5
)

response = query_engine.query("你的问题")

from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.agent import ReActAgent
from llama_index.core import VectorStoreIndex, SummaryIndex

# 为每个文档创建独立的索引和工具
doc_tools = []

for i, doc in enumerate(documents):
    # 向量索引（细粒度检索）
    vector_index = VectorStoreIndex.from_documents([doc])
    vector_query_engine = vector_index.as_query_engine(similarity_top_k=3)

    # 摘要索引（全文摘要）
    summary_index = SummaryIndex.from_documents([doc])
    summary_query_engine = summary_index.as_query_engine(
        response_mode="tree_summarize"
    )

    doc_name = doc.metadata.get("file_name", f"doc_{i}")

    # 创建向量搜索工具
    vector_tool = QueryEngineTool(
        query_engine=vector_query_engine,
        metadata=ToolMetadata(
            name=f"vector_search_{i}",
            description=f"用于在 {doc_name} 中搜索具体信息"
        )
    )

    # 创建摘要工具
    summary_tool = QueryEngineTool(
        query_engine=summary_query_engine,
        metadata=ToolMetadata(
            name=f"summary_{i}",
            description=f"用于获取 {doc_name} 的整体摘要"
        )
    )

    doc_tools.extend([vector_tool, summary_tool])

# 创建多文档Agent
agent = ReActAgent.from_tools(
    tools=doc_tools,
    verbose=True,
    max_iterations=10
)

# 跨文档查询
response = agent.chat("比较不同文档中提到的技术方案，给出推荐")
print(response)

from llama_index.core.evaluation import FaithfulnessEvaluator

# 评估答案是否忠实于检索到的上下文
faithfulness_evaluator = FaithfulnessEvaluator()

# 获取查询结果
query_engine = index.as_query_engine()
response = query_engine.query("什么是Transformer？")

# 评估
eval_result = faithfulness_evaluator.evaluate_response(response=response)

print(f"忠实度评估 - 通过: {eval_result.passing}")
print(f"分数: {eval_result.score}")
print(f"反馈: {eval_result.feedback}")

from llama_index.core.evaluation import RelevancyEvaluator

# 评估答案与查询的相关性
relevancy_evaluator = RelevancyEvaluator()

query = "深度学习有哪些应用？"
response = query_engine.query(query)

eval_result = relevancy_evaluator.evaluate_response(
    query=query,
    response=response
)

print(f"相关性评估 - 通过: {eval_result.passing}")
print(f"分数: {eval_result.score}")
print(f"反馈: {eval_result.feedback}")

from llama_index.core.evaluation import CorrectnessEvaluator

# 评估答案的正确性（需要参考答案）
correctness_evaluator = CorrectnessEvaluator()

query = "PyTorch的创建者是谁？"
response = query_engine.query(query)
reference = "PyTorch由Meta AI（前Facebook AI Research）团队创建"

eval_result = correctness_evaluator.evaluate(
    query=query,
    response=str(response),
    reference=reference
)

print(f"正确性评估 - 分数: {eval_result.score}")
print(f"反馈: {eval_result.feedback}")