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

# Document: 原始文档对象
doc = Document(
    text="LlamaIndex 是一个 RAG 框架...",
    metadata={"file_name": "intro.md", "category": "tech"}
)

# Node: 从 Document 切出的可检索单元
splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50)
nodes = splitter.get_nodes_from_documents([doc])
for n in nodes:
    print(n.text, n.metadata, n.node_id)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)

# Retriever: 只负责找资料，返回 NodeWithScore 列表
retriever = index.as_retriever(similarity_top_k=3)
nodes = retriever.retrieve("什么是 RAG？")
for n in nodes:
    print(f"score={n.score:.3f} | {n.node.text[:80]}")

# Query Engine: 完整问答引擎，找资料+重排+生成
query_engine = index.as_query_engine(similarity_top_k=3)
response = query_engine.query("什么是 RAG？")
print(response)
print(response.source_nodes)  # 溯源节点

# v0.10+ 拆分式安装
pip install llama-index-core               # 核心
pip install llama-index-llms-openai        # OpenAI LLM
pip install llama-index-embeddings-openai  # OpenAI Embedding
pip install llama-index-vector-stores-chroma  # 向量库
pip install llama-index-readers-file       # 文件 Reader
# 一键全装
pip install llama-index[all]

from llama_index.core import SimpleDirectoryReader

# 1. 加载整个目录（自动识别格式）
docs = SimpleDirectoryReader(input_dir="./data").load_data()

# 2. 递归加载子目录
docs = SimpleDirectoryReader(input_dir="./data", recursive=True).load_data()

# 3. 加载指定文件
docs = SimpleDirectoryReader(
    input_files=["a.pdf", "b.docx", "c.md"]
).load_data()

# 4. 限制扩展名 + 文件数
docs = SimpleDirectoryReader(
    input_dir="./data",
    required_exts=[".pdf", ".docx"],
    num_files_limit=100
).load_data()

# 5. 并行加载（Linux/Mac 才有收益）
docs = SimpleDirectoryReader(
    input_dir="./data",
    recursive=True
).load_data(num_workers=4)

# 方案1: PyPDFReader (轻量，文本型 PDF)
from llama_index.readers.file import PyPDFReader
reader = PyPDFReader()
docs = reader.load_data(file="./paper.pdf")

# 方案2: PyMuPDFReader (速度快，支持扫描件) — 推荐
from llama_index.readers.file import PyMuPDFReader
reader = PyMuPDFReader()
docs = reader.load_data(file="./paper.pdf")

# 方案3: 在 SimpleDirectoryReader 中按需指定
from llama_index.core import SimpleDirectoryReader
from llama_index.readers.file import PDFReader
file_extractor = {".pdf": PDFReader()}
docs = SimpleDirectoryReader(
    "./data", file_extractor=file_extractor
).load_data()

from llama_index.core.node_parser import SentenceSplitter, TokenTextSplitter

# 1. SentenceSplitter (推荐): 按句子边界切，再按 token 切
splitter = SentenceSplitter(
    chunk_size=512,
    chunk_overlap=50,
    paragraph_separator="\n\n\n",
    tokenizer=tiktoken.encoding_for_model("gpt-4").encode
)

# 2. TokenTextSplitter: 严格按 token 切，可能切到句子中间
splitter = TokenTextSplitter(
    chunk_size=512,
    chunk_overlap=50,
    separator=" "
)

doc = Document(
    text=content,
    metadata={
        # 1. 来源追溯
        "source": "合同库",
        "file_name": "服务协议_2024.pdf",
        "file_path": "/data/contracts/...",
        "page": 5,
        # 2. 业务分类
        "category": "服务合同",
        "department": "商务部",
        # 3. 权限控制
        "permission_level": "L2",  # 后续做 metadata filter
        # 4. 时效性
        "created_at": "2024-01-15",
        "effective_date": "2024-02-01",
        "expire_at": "2025-12-31",
        "version": "v2.1",
        # 5. 文档结构
        "heading_path": "第三章 违约责任 > 第二节 违约金",
        # 6. 业务标签
        "tags": ["SLA", "违约金", "高优先级"]
    }
)

from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.extractors import (
    TitleExtractor,
    SummaryExtractor,
    QuestionsAnsweredExtractor,
)
from llama_index.core import Settings

# 1. 定义 Pipeline
pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=512, chunk_overlap=50),  # 切分
        TitleExtractor(nodes=5),                              # 抽取小标题
        SummaryExtractor(summaries=["self"]),                 # 摘要
        QuestionsAnsweredExtractor(questions=3),              # 抽取假设问题
    ]
)

# 2. 跑 Pipeline
nodes = pipeline.run(documents=documents)

# 3. 带缓存的 Pipeline（生产必备）
pipeline_with_cache = IngestionPipeline(
    transformations=[...],
    cache=True,  # 避免重复跑
)

# 示例：加载 Notion
from llama_index.readers.notion import NotionPageReader
reader = NotionPageReader(integration_token="...")
docs = reader.load_data(page_ids=["..."])

# LlamaParse 示例
from llama_parse import LlamaParse
parser = LlamaParse(
    api_key="llx-...",
    result_type="markdown",  # 或 "text"
    parsing_instruction="保留表格结构和数学公式"
)
docs = parser.load_data("./complex.pdf")

from llama_index.core.ingestion import IngestionPipeline, DocstoreStrategy
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core import StorageContext, SimpleDocumentStore

# 1. 初始化带 cache 的 pipeline（默认基于输入 hash）
pipeline = IngestionPipeline(
    transformations=[SentenceSplitter(chunk_size=512)],
    cache=True  # 默认 SimpleCache
)

# 2. 用 docstore 配合去重（生产推荐）
docstore = SimpleDocumentStore()
pipeline = IngestionPipeline(
    transformations=[SentenceSplitter()],
    docstore=docstore,
    docstore_strategy=DocstoreStrategy.UPSERTS,  # 增量更新
)

# 3. 跑一次，后续相同文件跳过
nodes = pipeline.run(documents=docs)
nodes2 = pipeline.run(documents=docs)  # 命中缓存，不重跑

# 完整流程
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter

# 1. 加载 + 切分
docs = SimpleDirectoryReader("./data").load_data()
nodes = SentenceSplitter(chunk_size=512).get_nodes_from_documents(docs)

# 2. 构建索引（内部调 embed_model）
index = VectorStoreIndex(nodes)

# 3. 查询（内部用向量相似度）
retriever = index.as_retriever(similarity_top_k=3)
results = retriever.retrieve("RAG 的核心思想是什么？")

from llama_index.core import SummaryIndex
index = SummaryIndex.from_documents(docs)
query_engine = index.as_retriever(retriever_mode="embedding")  # 嵌入检索
# 或
query_engine = index.as_query_engine(response_mode="tree_summarize")

from llama_index.core import TreeIndex
index = TreeIndex.from_documents(docs, num_children=10)
query_engine = index.as_query_engine(child_branch_factor=2)

from llama_index.core import PropertyGraphIndex
from llama_index.core.indices.property_graph import (
    SimpleLLMPathExtractor, ImplicitPathExtractor
)

index = PropertyGraphIndex.from_documents(
    docs,
    kg_extractors=[
        SimpleLLMPathExtractor(llm=llm, max_paths_per_chunk=10),
        ImplicitPathExtractor()
    ]
)
retriever = index.as_retriever(include_text=True, similarity_top_k=2)

from llama_index.core import StorageContext
from llama_index.core.vector_stores import SimpleVectorStore
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core.storage.index_store import SimpleIndexStore
from llama_index.core.graph_stores import SimpleGraphStore

storage_context = StorageContext.from_defaults(
    vector_store=SimpleVectorStore(),       # 向量
    docstore=SimpleDocumentStore(),         # 原始 Node
    index_store=SimpleIndexStore(),         # 索引元数据
    graph_store=SimpleGraphStore(),         # 图结构
)

index = VectorStoreIndex(nodes, storage_context=storage_context)
index.storage_context.persist(persist_dir="./storage")  # 持久化

from llama_index.core.node_parser import SentenceWindowNodeParser
from llama_index.core.postprocessor import MetadataReplacementPostProcessor

# 1. 切分（每个 Node 是单个句子，但带 window 上下文）
node_parser = SentenceWindowNodeParser.from_defaults(
    window_size=3,  # 包含前后各3个句子
)
nodes = node_parser.get_nodes_from_documents(docs)

# 2. 建索引
index = VectorStoreIndex(nodes)

# 3. 检索（按句子相似度）
retriever = index.as_retriever(similarity_top_k=5)

# 4. 关键：检索后用 MetadataReplacementPostProcessor 把句子替换为窗口
postproc = MetadataReplacementPostProcessor(
    target_metadata_key="window"
)
query_engine = index.as_query_engine(
    similarity_top_k=5,
    node_postprocessors=[postproc]
)
response = query_engine.query("...")

from llama_index.core.node_parser import HierarchicalNodeParser, get_leaf_nodes
from llama_index.core.retrievers import AutoMergingRetriever

# 1. 层级切分
node_parser = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128]  # 父→中→子
)
nodes = node_parser.get_nodes_from_documents(docs)

# 2. 只对叶子节点建索引（细粒度）
index = VectorStoreIndex(get_leaf_nodes(nodes))
index.storage_context.docstore.add_documents(nodes)  # 父节点放 docstore

# 3. 检索（自动合并）
retriever = AutoMergingRetriever(
    index.as_retriever(similarity_top_k=6),
    storage_context=index.storage_context
)

# 混合检索：BM25 + 向量 + Rerank
from llama_index.core.retrievers import VectorIndexRetriever
from llama_index.retrievers.bm25 import BM25Retriever

vector_retriever = index.as_retriever(similarity_top_k=10)
bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=10)

# 融合（用 QueryFusionRetriever）
from llama_index.core.retrievers import QueryFusionRetriever
fusion_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=5,
    num_queries=4,  # 自动生成多个 query
    mode="reciprocal_rerank"  # RRF 融合
)

from llama_index.core.postprocessor import (
    SimilarityPostprocessor,           # 相似度阈值过滤
    KeywordNodePostprocessor,         # 必含/排除关键词
    MetadataReplacementPostProcessor,  # 窗口替换
    SentenceEmbeddingOptimizer,       # 句子级精排
    CohereRerank,                     # Cohere Rerank
    LongContextReorder,               # 长上下文重排
    FixedRecencyPostprocessor,        # 时效性优先
)

# 1. 相似度过滤
postproc = SimilarityPostprocessor(similarity_cutoff=0.7)

# 2. Rerank（强烈推荐）
from llama_index.postprocessor.cohere_rerank import CohereRerank
rerank = CohereRerank(api_key=cohere_api_key, top_n=5)

# 3. BGE Rerank（开源免费）
from llama_index.postprocessor.bge_rerank import BGERerank
bge = BGERerank(top_n=5, model="BAAI/bge-reranker-large")

# 4. 链式：先过滤 + 再 Rerank
query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[
        SimilarityPostprocessor(similarity_cutoff=0.6),
        bge,
    ]
)

# 1. BGE Rerank（开源、本地、免费）— 推荐生产
from llama_index.postprocessor.bge_rerank import BGERerank
bge = BGERerank(
    model="BAAI/bge-reranker-large",  # 或 bge-reranker-base
    top_n=5
)

# 2. Cohere Rerank（云端、要钱）
from llama_index.postprocessor.cohere_rerank import CohereRerank
cohere = CohereRerank(api_key="...", top_n=5)

# 3. 接入 Query Engine
query_engine = index.as_query_engine(
    similarity_top_k=20,  # 粗筛 20
    node_postprocessors=[bge]  # 精排到 5
)

# 4. 配合元数据过滤
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
sentence_window = MetadataReplacementPostProcessor(target_metadata_key="window")
query_engine = index.as_query_engine(
    node_postprocessors=[bge, sentence_window]
)

# 方案1: QueryFusionRetriever（推荐）
from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.retrievers.bm25 import BM25Retriever

vector_retriever = index.as_retriever(similarity_top_k=10)
bm25_retriever = BM25Retriever.from_defaults(nodes=nodes, similarity_top_k=10)

fusion_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=5,
    num_queries=4,  # 用 LLM 生成多 query 变体
    mode="reciprocal_rerank",  # RRF 融合算法
    use_async=True
)

# 方案2: 自定义融合（更灵活）
def hybrid_search(query, k=5):
    vector_nodes = vector_retriever.retrieve(query)
    bm25_nodes = bm25_retriever.retrieve(query)
    # RRF 融合
    fused = reciprocal_rerank_fusion([vector_nodes, bm25_nodes])
    return fused[:k]

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

# 1. 检索器
retriever = VectorIndexRetriever(index=index, similarity_top_k=3)

# 2. 合成器
synthesizer = get_response_synthesizer(response_mode="compact")

# 3. 引擎
query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=synthesizer
)

# 切换 response_mode
qe = index.as_query_engine(response_mode="tree_summarize")
qe = index.as_query_engine(response_mode="refine")
qe = index.as_query_engine(response_mode="compact")

# 1. 启用流式（需要 LLM 支持 streaming）
from llama_index.core import VectorStoreIndex
query_engine = index.as_query_engine(
    streaming=True,
    similarity_top_k=3
)

# 2. 调用 query 返回 StreamingResponse
streaming_response = query_engine.query("什么是 RAG？")

# 3. 迭代 token
for text in streaming_response.response_gen:
    print(text, end="", flush=True)

# 4. 完整响应（流结束后可用）
print(streaming_response.get_response())

# 5. 或直接 print
streaming_response.print_response_stream()

from pydantic import BaseModel, Field
from typing import List
from llama_index.core import get_response_synthesizer
from llama_index.core.query_engine import RetrieverQueryEngine

class PhoneInfo(BaseModel):
    name: str = Field(description="手机型号")
    description: str = Field(description="型号描述")
    features: List[str] = Field(description="主要功能")

# 方案1: 通过 output_cls
synthesizer = get_response_synthesizer(
    response_mode="tree_summarize",
    output_cls=PhoneInfo
)
query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=synthesizer
)
response = query_engine.query("介绍 iPhone 16 Pro")
# response 是结构化对象
print(response.name, response.features)

# 方案2: 通过 Pydantic 风格的 program
from llama_index.core.program import LLMTextCompletionProgram
prompt = """
根据上下文信息提取手机参数:
{context}
"""
program = LLMTextCompletionProgram.from_defaults(
    output_cls=PhoneInfo,
    prompt_template_str=prompt,
    llm=llm
)
result = program(context=context_str)

query_engine = index.as_query_engine(
    # 检索参数
    similarity_top_k=3,                # 向量检索返回 top-k
    vector_store_query_mode="default",  # default / sparse / hybrid
    alpha=0.7,                          # hybrid 模式下向量权重
    filters=MetadataFilters(...),       # 元数据过滤
    
    # 响应参数
    response_mode="compact",            # compact / refine / tree_summarize
    text_qa_template=None,              # 自定义 QA prompt
    refine_template=None,               # 自定义 refine prompt
    
    # LLM
    llm=llm,                            # 自定义 LLM
    
    # 流式
    streaming=False,
    
    # 调试
    verbose=True,                       # 打印检索过程
    
    # 后处理
    node_postprocessors=[bge_rerank],   # 节点后处理器
    
    # 回调
    callback_manager=CallbackManager([handler])
)

# 方案1: 设置相似度阈值过滤
from llama_index.core.postprocessor import SimilarityPostprocessor
qe = index.as_query_engine(
    similarity_top_k=5,
    node_postprocessors=[
        SimilarityPostprocessor(similarity_cutoff=0.7)
    ]
)

# 方案2: 自定义 Prompt 强制 LLM 答"不知道"
from llama_index.core import PromptTemplate
qa_prompt = PromptTemplate(
    "你是一个严谨的助手。基于以下上下文回答问题。\n"
    "如果上下文中没有答案，直接回答'我不知道'，不要编造。\n"
    "上下文：\n{context_str}\n"
    "问题：{query_str}\n"
    "答案："
)
qe = index.as_query_engine(text_qa_template=qa_prompt)

# 方案3: 配合 Faithfulness 评估检测幻觉

from llama_index.core.query_engine import SubQuestionQueryEngine
from llama_index.llms.openai import OpenAI

sub_qe = SubQuestionQueryEngine.from_args(
    index=index,
    llm=OpenAI(model="gpt-4"),
    use_async=True,  # 并行子问题
)
response = sub_qe.query("比较 LlamaIndex 和 LangChain 在 RAG 实现上的差异")

from llama_index.core.query_engine import MultiStepQueryEngine
multi_qe = MultiStepQueryEngine(
    query_engine=base_qe,
    query_transform=step_decompose_transform
)

from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector, LLMMultiSelector
from llama_index.core.tools import QueryEngineTool

# 1. 定义子引擎
vector_tool = QueryEngineTool.from_defaults(
    query_engine=vector_qe,
    name="vector_search",
    description="用于语义相似度检索，例如事实查询、概念解释"
)
summary_tool = QueryEngineTool.from_defaults(
    query_engine=summary_qe,
    name="summary_search",
    description="用于文档总结、概览类问题"
)
kg_tool = QueryEngineTool.from_defaults(
    query_engine=kg_qe,
    name="kg_search",
    description="用于实体关系推理，例如'X 和 Y 的关系'"
)

# 2. 路由引擎
router_qe = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(llm=llm),
    query_engine_tools=[vector_tool, summary_tool, kg_tool]
)
response = router_qe.query("总结今年的产品趋势")

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

# 1. 为每个数据源建索引
contract_index = VectorStoreIndex.from_documents(contract_docs)
hr_index = VectorStoreIndex.from_documents(hr_docs)
product_index = VectorStoreIndex.from_documents(product_docs)

# 2. 各自建 QueryEngine
contract_qe = contract_index.as_query_engine(similarity_top_k=3)
hr_qe = hr_index.as_query_engine(similarity_top_k=3)
product_qe = product_index.as_query_engine(similarity_top_k=3)

# 3. 包装成 Tool
contract_tool = QueryEngineTool(
    query_engine=contract_qe,
    metadata=ToolMetadata(
        name="合同库",
        description="查询合同条款、违约责任、SLA 条款"
    )
)
hr_tool = QueryEngineTool(
    query_engine=hr_qe,
    metadata=ToolMetadata(
        name="HR制度库",
        description="查询员工手册、年假、考勤制度"
    )
)
product_tool = QueryEngineTool(
    query_engine=product_qe,
    metadata=ToolMetadata(
        name="产品手册库",
        description="查询产品参数、功能、使用方法"
    )
)

# 4. RouterQueryEngine 整合
multi_doc_qe = RouterQueryEngine.from_defaults(
    query_engine_tools=[contract_tool, hr_tool, product_tool]
)

# 5. 复杂查询
response = multi_doc_qe.query("合同里关于员工培训的违约条款是什么？")

# Query Engine
query_engine = index.as_query_engine()
response = query_engine.query("什么是 RAG？")

# Chat Engine（带记忆）
chat_engine = index.as_chat_engine()
response = chat_engine.chat("什么是 RAG？")
response = chat_engine.chat("它和微调有什么区别？")  # 记得上文

# 1. CondenseQuestion（默认）
from llama_index.core.chat_engine import CondenseQuestionChatEngine
chat_engine = CondenseQuestionChatEngine.from_defaults(
    query_engine=query_engine,
    llm=llm
)

# 2. Context 模式（带 chat history 作为 context）
from llama_index.core.chat_engine import ContextChatEngine
chat_engine = ContextChatEngine.from_defaults(
    retriever=retriever,
    llm=llm,
    chat_memory=memory
)

# 3. ReAct Agent Chat
chat_engine = index.as_chat_engine(chat_mode="react", verbose=True)

from llama_index.core.memory import (
    ChatMemoryBuffer,
    VectorMemory,
    SimpleComposableMemory
)

# 1. 简单 buffer 记忆
memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,  # 控制 token 上限
)

# 2. 注入 ChatEngine
chat_engine = index.as_chat_engine(
    chat_mode="condense_question",
    memory=memory,
    system_prompt="你是一个专业的助手..."
)

# 3. 手动管理
memory.put("user", "你好")
memory.put("assistant", "你好，有什么可以帮你？")
chat_memory = memory.get()  # 获取全部

# 4. 自定义 Memory（多个来源组合）
custom_memory = SimpleComposableMemory.from_defaults(
    primary_memory=ChatMemoryBuffer.from_defaults(token_limit=2000),
    secondary_memory=[VectorMemory.from_defaults(...)]
)

# 方案1: 序列化到文件
import json
from llama_index.core.memory import ChatMemoryBuffer

# 保存
memory = ChatMemoryBuffer.from_defaults(token_limit=3000)
memory.put("user", "你好")
memory_data = {"chat_history": memory.chat_store.get_all()}  # 持久化

# 加载
from llama_index.core.llms import ChatMessage
restored = ChatMemoryBuffer.from_defaults(token_limit=3000)
for msg in memory_data["chat_history"]:
    restored.put(msg.role, msg.content)

# 方案2: 用外部存储（Redis / DB）
import redis
r = redis.Redis()
r.set(f"chat:{user_id}", json.dumps(memory.chat_store.get_all()))

from llama_index.core import PromptTemplate

# 1. 直接传字符串
chat_engine = index.as_chat_engine(
    system_prompt="""
你是一个专业的金融知识助手，名叫小金。
规则：
1. 只基于上下文回答，不要编造。
2. 如果不知道，答"未找到相关信息"。
3. 回答不超过 200 字，使用专业语气。
4. 引用时附上文件名。
    """
)

# 2. 用 PromptTemplate 动态注入
prompt_tmpl = PromptTemplate("""
你是 {role}，名叫 {name}。
当前用户：{user_name}
上下文：{context_str}
问题：{query_str}
""")
query_engine.update_prompts({"response_synthesizer:text_qa_template": prompt_tmpl})

# 3. 不同模式用不同模板
text_qa_template = PromptTemplate("基于上下文：{context_str}\n问题：{query_str}\n答案：")
refine_template = PromptTemplate("现有答案：{existing_answer}\n新上下文：{context_msg}\n问题：{query_str}\n优化答案：")
query_engine = index.as_query_engine(
    text_qa_template=text_qa_template,
    refine_template=refine_template
)

# 1. 启用 streaming
chat_engine = index.as_chat_engine(streaming=True)

# 2. 流式调用
streaming_response = chat_engine.stream_chat("什么是 RAG？")
for token in streaming_response.response_gen:
    print(token, end="", flush=True)

# 3. 同步 API
response = chat_engine.chat("什么是 RAG？")

import asyncio
async def stream_chat():
    chat_engine = index.as_chat_engine(streaming=True)
    streaming_response = await chat_engine.astream_chat("什么是 RAG？")
    async for token in streaming_response.async_response_gen():
        print(token, end="", flush=True)
asyncio.run(stream_chat())

# ReActAgent 示例
from llama_index.core.agent.workflow import ReActAgent
from llama_index.core.tools import FunctionTool

def multiply(a: int, b: int) -> int:
    """Multiply two numbers"""
    return a * b

multiply_tool = FunctionTool.from_defaults(fn=multiply)
agent = ReActAgent(tools=[multiply_tool], llm=llm)
response = await agent.run(query="3 * 5 = ?")

Thought: 用户问 3*5，我需要计算
Action: multiply(3, 5)
Action Input: {"a": 3, "b": 5}
Observation: 15
Thought: 拿到结果了，可以回答
Final Answer: 3 * 5 = 15

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

def search(query: str) -> str:
    """搜索工具"""
    return f"搜索结果: {query}"

agent = ReActAgent(
    tools=[FunctionTool.from_defaults(fn=search)],
    llm=OpenAI(model="gpt-4o"),
    verbose=True
)

response = await agent.run("北京今天天气如何？")
print(response)

# FunctionCallingAgent
from llama_index.core.agent.workflow import FunctionCallingAgent
agent = FunctionCallingAgent(
    tools=tools,
    llm=OpenAI(model="gpt-4o"),
    verbose=True
)

# ReActAgent
from llama_index.core.agent.workflow import ReActAgent
agent = ReActAgent(tools=tools, llm=llm, verbose=True)

from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core import VectorStoreIndex

# 1. 假设有两个索引
contract_index = VectorStoreIndex.from_documents(contract_docs)
hr_index = VectorStoreIndex.from_documents(hr_docs)

# 2. 把每个索引变成 Tool
contract_tool = QueryEngineTool(
    query_engine=contract_index.as_query_engine(),
    metadata=ToolMetadata(
        name="contract_search",
        description="查询合同条款、违约责任、SLA 条款相关问题"
    )
)
hr_tool = QueryEngineTool(
    query_engine=hr_index.as_query_engine(),
    metadata=ToolMetadata(
        name="hr_search",
        description="查询员工手册、年假、考勤制度等 HR 问题"
    )
)

# 3. 给 Agent 使用
from llama_index.core.agent.workflow import FunctionCallingAgent
agent = FunctionCallingAgent(
    tools=[contract_tool, hr_tool],
    llm=llm
)
response = await agent.run("合同里关于年假的规定是什么？")

from llama_index.core.tools import FunctionTool
from typing import Optional, List
from pydantic import BaseModel, Field

# 1. 简单函数
def search_weather(city: str) -> str:
    """查询指定城市的天气"""
    return f"{city} 今天晴，25度"

tool = FunctionTool.from_defaults(fn=search_weather)

# 2. 复杂参数（用 Pydantic）
class SearchQuery(BaseModel):
    keyword: str = Field(description="搜索关键词")
    max_results: int = Field(default=10, description="返回结果数量")
    language: str = Field(default="zh", description="语言")

def advanced_search(query: SearchQuery) -> List[str]:
    """高级搜索工具"""
    return [f"结果 {i}: {query.keyword}" for i in range(query.max_results)]

tool = FunctionTool.from_defaults(fn=advanced_search)

# 3. 异步函数
async def async_fetch(url: str) -> str:
    """异步抓取网页"""
    import aiohttp
    async with aiohttp.ClientSession() as session:
        async with session.get(url) as resp:
            return await resp.text()

async_tool = FunctionTool.from_defaults(fn=async_fetch)

from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)
from llama_index.llms.openai import OpenAI

# 1. 定义事件
class JokeEvent(Event):
    joke: str

class AnalysisEvent(Event):
    analysis: str

# 2. 定义 Workflow
class JokeFlow(Workflow):
    llm = OpenAI(model="gpt-4o-mini")

    @step
    async def generate_joke(self, ev: StartEvent) -> JokeEvent:
        topic = ev.topic
        prompt = f"讲一个关于{topic}的笑话"
        response = await self.llm.acomplete(prompt)
        return JokeEvent(joke=str(response))

    @step
    async def critique_joke(self, ev: JokeEvent) -> StopEvent:
        prompt = f"分析笑话笑点: {ev.joke}"
        response = await self.llm.acomplete(prompt)
        return StopEvent(result=str(response))

# 3. 运行
w = JokeFlow(timeout=60, verbose=True)
result = await w.run(topic="程序员")
print(result)

from llama_index.core.workflow import Context

class MyWorkflow(Workflow):
    @step
    async def step1(self, ctx: Context, ev: StartEvent) -> SomeEvent:
        # 写 Context
        await ctx.set("user_name", ev.user_name)
        await ctx.set("query", ev.query)
        return SomeEvent(...)

    @step
    async def step2(self, ctx: Context, ev: SomeEvent) -> StopEvent:
        # 读 Context
        user_name = await ctx.get("user_name")
        query = await ctx.get("query")
        return StopEvent(result=...)

# 1. Context 存 LLM 调用的中间结果
@step
async def retrieve_step(self, ctx: Context, ev: StartEvent) -> NextEvent:
    nodes = retriever.retrieve(ev.query)
    await ctx.set("nodes", nodes)
    return NextEvent(...)

# 2. 多步骤共享数据
@step
async def process_step(self, ctx: Context, ev: NextEvent) -> StopEvent:
    nodes = await ctx.get("nodes")
    # ...

# 3. 流式事件（context 在多 step 间共享）
async for event in wf.stream(topic="..."):
    print(event)

from llama_index.core.workflow import (
    InputRequiredEvent, HumanResponseEvent, Context
)

class DangerousTool:
    async def execute(self, ctx: Context) -> str:
        question = "确认要删除这条记录吗？(yes/no)"
        # 发出中断事件
        response = await ctx.wait_for_event(
            HumanResponseEvent,
            waiter_id=question,
            waiter_event=InputRequiredEvent(question)
        )
        if response.user_response.lower().startswith("y"):
            return "已删除"
        else:
            return "操作取消"

# 配合外部系统
# 前端监听 InputRequiredEvent → 展示问题
# 用户回答后 → 触发 HumanResponseEvent → 继续 Workflow

# 1. 启动 Workflow，序列化状态
ctx_dict = ctx.to_dict()
db.save(f"workflow:{user_id}", ctx_dict)

# 2. 用户回答后，从 DB 恢复
ctx = Context.from_dict(workflow, ctx_dict)
ctx.send_event(HumanResponseEvent(user_response="yes"))

# 3. 继续运行
result = await wf.run(ctx=ctx)

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

# 1. 定义多个 Agent
def search_docs(query: str) -> str:
    """搜索文档"""
    return f"找到 {query} 相关文档"

def write_report(topic: str) -> str:
    """写报告"""
    return f"已生成 {topic} 报告"

research_agent = ReActAgent(
    name="research_agent",
    description="负责检索资料",
    tools=[FunctionTool.from_defaults(fn=search_docs)],
    llm=OpenAI(model="gpt-4o"),
    system_prompt="你是研究助理"
)

write_agent = ReActAgent(
    name="write_agent",
    description="负责写报告",
    tools=[FunctionTool.from_defaults(fn=write_report)],
    llm=OpenAI(model="gpt-4o"),
    system_prompt="你是写作助理"
)

# 2. 编排多 Agent
workflow = AgentWorkflow(
    agents=[research_agent, write_agent],
    root_agent="research_agent",  # 入口
)

# 3. 运行（agent 间自动 handoff）
response = await workflow.run("研究 AI Agent 最新进展并写报告")
print(response)

# 安装工具
# pip install llama-index-utils-workflow

from llama_index.utils.workflow import (
    draw_all_possible_flows,        # 静态流程图
    draw_most_recent_execution      # 动态执行轨迹
)

# 1. 静态流程图（基于代码分析）
draw_all_possible_flows(MyWorkflow, filename="workflow_structure.html")

# 2. 动态执行图（最近一次）
workflow = MyWorkflow()
await workflow.run(topic="程序员")
draw_most_recent_execution(workflow, filename="recent_execution.html")

# 3. 流式事件调试
async for event in workflow.stream(topic="..."):
    print(f"事件: {type(event).__name__}, 数据: {event}")

from llama_index.core import VectorStoreIndex, StorageContext
from llama_index.core.vector_stores import ...

# 1. Chroma
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_store = ChromaVectorStore(chroma_collection=chroma_client.get_or_create_collection("docs"))

# 2. Milvus
from llama_index.vector_stores.milvus import MilvusVectorStore
milvus_store = MilvusVectorStore(
    uri="./milvus.db",
    collection_name="docs",
    dim=1536,
    overwrite=True
)

# 3. Qdrant
from llama_index.vector_stores.qdrant import QdrantVectorStore
import qdrant_client
qclient = qdrant_client.QdrantClient(path="./qdrant_db")
qdrant_store = QdrantVectorStore(client=qclient, collection_name="docs")

# 4. FAISS
from llama_index.vector_stores.faiss import FaissVectorStore
import faiss
faiss_index = faiss.IndexFlatL2(1536)
faiss_store = FaissVectorStore(faiss_index=faiss_index)

storage_context = StorageContext.from_defaults(vector_store=<your_store>)
index = VectorStoreIndex(nodes, storage_context=storage_context)

# 1. OpenAI（默认）
from llama_index.embeddings.openai import OpenAIEmbedding
embed_model = OpenAIEmbedding(model="text-embedding-3-small", api_key="...")

# 2. BGE（开源、中文强）
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
embed_model = HuggingFaceEmbedding(model_name="BAAI/bge-large-zh-v1.5")

# 3. Cohere
from llama_index.embeddings.cohere import CohereEmbedding
embed_model = CohereEmbedding(api_key="...", model="embed-english-v3.0")

# 4. 本地 Sentence-Transformers
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
embed_model = HuggingFaceEmbedding(model_name="sentence-transformers/all-MiniLM-L6-v2")

# 5. 全局设置（v0.10+ 推荐）
from llama_index.core import Settings
Settings.embed_model = embed_model

from llama_index.core.embeddings import BaseEmbedding
from typing import List
import numpy as np

class MyCustomEmbedding(BaseEmbedding):
    def __init__(self, model_path: str, **kwargs):
        super().__init__(**kwargs)
        # 加载自定义模型
        from sentence_transformers import SentenceTransformer
        self._model = SentenceTransformer(model_path)

    def _get_query_embedding(self, query: str) -> List[float]:
        emb = self._model.encode(query)
        return emb.tolist()

    def _get_text_embedding(self, text: str) -> List[float]:
        return self._model.encode(text).tolist()

    async def _aget_query_embedding(self, query: str) -> List[float]:
        return self._get_query_embedding(query)

    async def _aget_text_embedding(self, text: str) -> List[float]:
        return self._get_text_embedding(text)

# 使用
embed_model = MyCustomEmbedding(model_path="BAAI/bge-large-zh-v1.5")

# 1. 持久化
index = VectorStoreIndex.from_documents(docs)
index.storage_context.persist(persist_dir="./storage")

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

# 3. 继续用
query_engine = index.as_query_engine()
response = query_engine.query("...")

# 方案1: 用支持 sparse+dense 的向量库（如 Qdrant / Milvus）
from qdrant_client import models
# 上传时同时给 dense 和 sparse vector
qclient.upsert(
    collection_name="docs",
    points=[
        models.PointStruct(
            id=1,
            vector={
                "dense": dense_emb,           # BGE
                "sparse": sparse_emb,         # BM25
            },
            payload={"text": "..."}
        )
    ]
)

# 方案2: LlamaIndex 内置 hybrid（QueryFusionRetriever）
# 见 Q30

# 1. 准备数据（query, positive_doc, negative_doc）
train_data = [
    {
        "query": "如何申请年假？",
        "positive": "年假申请流程：登录 OA → 提交申请",
        "negative": "病假申请流程：..."
    },
    # ...
]

# 2. 用 sentence-transformers 微调
from sentence_transformers import InputExample, losses
from sentence_transformers import SentenceTransformer

model = SentenceTransformer("BAAI/bge-small-zh")
train_examples = [InputExample(texts=[d["query"], d["positive"], d["negative"]]) for d in train_data]
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
train_loss = losses.MultipleNegativesRankingLoss(model)

model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    output_path="./fine_tuned_bge"
)

# 3. 接入 LlamaIndex
from llama_index.embeddings.huggingface import HuggingFaceEmbedding
embed_model = HuggingFaceEmbedding(model_name="./fine_tuned_bge")

# 查询重写示例
from llama_index.core.query_engine import TransformQueryEngine
from llama_index.core.indices.query.query_transform import HyDEQueryTransform

hyde = HyDEQueryTransform(llm=llm, include_original=True)
transformed_qe = TransformQueryEngine(qe, query_transform=hyde)
response = transformed_qe.query("用户口语化的问题")

from llama_index.core import (
    VectorStoreIndex, SimpleDirectoryReader, Settings
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
from llama_index.core.postprocessor import SimilarityPostprocessor
from llama_index.postprocessor.bge_rerank import BGERerank

# 1. 全局配置
Settings.llm = OpenAI(model="gpt-4o-mini", temperature=0)
Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
Settings.node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=50)

# 2. 加载
docs = SimpleDirectoryReader("./data", recursive=True).load_data()

# 3. 注入 metadata
for doc in docs:
    doc.metadata["source"] = "公司制度"
    doc.metadata["version"] = "v2.1"

# 4. 建索引
index = VectorStoreIndex.from_documents(docs)

# 5. 建 Query Engine（带 Rerank + 阈值）
bge = BGERerank(top_n=5, model="BAAI/bge-reranker-large")
sim = SimilarityPostprocessor(similarity_cutoff=0.7)
query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[sim, bge],
    response_mode="compact",
    streaming=True
)

# 6. 查询
response = query_engine.query("公司年假政策是什么？")
print(response)
for node in response.source_nodes:
    print(f"  score={node.score:.3f} | {node.node.text[:100]}")

from llama_index.core import PropertyGraphIndex, SimpleDirectoryReader
from llama_index.core.indices.property_graph import (
    SimpleLLMPathExtractor, ImplicitPathExtractor
)

# 1. 加载
docs = SimpleDirectoryReader("./data").load_data()

# 2. 建图索引
pg_index = PropertyGraphIndex.from_documents(
    docs,
    kg_extractors=[
        SimpleLLMPathExtractor(llm=llm, max_paths_per_chunk=10),
        ImplicitPathExtractor()
    ]
)

# 3. 查询
query_engine = pg_index.as_query_engine(
    include_text=True,         # 返回 source text
    similarity_top_k=2
)
response = query_engine.query("苹果公司的 CEO 是谁？")
# 自动找 实体(苹果公司) - 关系(CEO_of) - 实体(Tim Cook)

# 4. 自定义子检索器
from llama_index.core.indices.property_graph import (
    VectorContextRetriever, LLMSynonymRetriever
)

retriever = pg_index.as_retriever(
    sub_retrievers=[
        VectorContextRetriever(pg_index.property_graph_store, embed_model=embed_model),
        LLMSynonymRetriever(llm=llm)
    ]
)

# 方案1: LlamaParse + Multi-Modal Index（最简单）
from llama_parse import LlamaParse
from llama_index.core import SimpleDirectoryReader

parser = LlamaParse(api_key="llx-...", result_type="markdown")
file_extractor = {".pdf": parser}
docs = SimpleDirectoryReader("./pdfs", file_extractor=file_extractor).load_data()

# 启用多模态检索
from llama_index.indices.managed.llama_cloud import LlamaParseIndex
index = LlamaParseIndex(name="my_index", project_name="default", api_key="llx-...")
retriever = index.as_retriever(retrieve_image_nodes=True)

# 方案2: 自定义多模态查询引擎
from llama_index.core.query_engine import CustomQueryEngine
from llama_index.multi_modal_llms.openai import OpenAIMultiModal

class MultimodalQueryEngine(CustomQueryEngine):
    qa_prompt: PromptTemplate
    retriever: BaseRetriever
    multi_modal_llm: OpenAIMultiModal

    def custom_query(self, query_str: str):
        nodes = self.retriever.retrieve(query_str)
        img_nodes = [n for n in nodes if isinstance(n.node, ImageNode)]
        text_nodes = [n for n in nodes if isinstance(n.node, TextNode)]
        context_str = "\n\n".join([r.get_content() for r in nodes])
        prompt = self.qa_prompt.format(context_str=context_str, query_str=query_str)
        llm_response = self.multi_modal_llm.complete(
            prompt=prompt,
            image_documents=[n.node for n in img_nodes]
        )
        return Response(response=str(llm_response), source_nodes=nodes)

from llama_index.core.query_engine import NLSQLTableQueryEngine, SQLTableRetrieverQueryEngine
from llama_index.core import SQLDatabase
from sqlalchemy import create_engine

# 1. 连接数据库
engine = create_engine("sqlite:///example.db")
sql_db = SQLDatabase(engine, include_tables=["users", "orders"])

# 2. NLSQLTableQueryEngine（单表）
query_engine = NLSQLTableQueryEngine(
    sql_database=sql_db,
    tables=["users", "orders"]
)
response = query_engine.query("查询 2024 年注册用户数")

# 3. SQLTableRetrieverQueryEngine（多表 + 表检索）
from llama_index.core.indices.struct_store import SQLTableRetrieverQueryEngine
from llama_index.core.objects import SQLTableNodeMapping, ObjectIndex, SQLTableSchema
from llama_index.core import VectorStoreIndex

table_node_mapping = SQLTableNodeMapping(sql_db)
table_schema_objs = [
    SQLTableSchema(table_name="users"),
    SQLTableSchema(table_name="orders")
]
obj_index = ObjectIndex.from_objects(
    table_schema_objs,
    table_node_mapping,
    VectorStoreIndex,
)
table_retriever = obj_index.as_retriever(similarity_top_k=1)

query_engine = SQLTableRetrieverQueryEngine(
    sql_database=sql_db,
    table_retriever=table_retriever
)
response = query_engine.query("查 2024 年订单量最高的用户")

from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.agent.workflow import FunctionCallingAgent

# 1. 多种检索策略建多个 Tool
vector_tool = QueryEngineTool(
    query_engine=vector_qe,
    metadata=ToolMetadata(name="vector_search", description="语义检索")
)
keyword_tool = QueryEngineTool(
    query_engine=bm25_qe,
    metadata=ToolMetadata(name="keyword_search", description="关键词检索")
)
sql_tool = QueryEngineTool(
    query_engine=sql_qe,
    metadata=ToolMetadata(name="sql_search", description="查结构化数据库")
)

# 2. Agent 决策
agent = FunctionCallingAgent(
    tools=[vector_tool, keyword_tool, sql_tool],
    llm=llm
)
response = await agent.run("查 2024 年 Q3 销售额最高的产品的客户评价")
# Agent 自动选：先 SQL 查产品，再向量查评价

from llama_index.core.indices.query.query_transform import HyDEQueryTransform
from llama_index.core.query_engine import TransformQueryEngine

hyde = HyDEQueryTransform(llm=llm, include_original=True)
transformed_qe = TransformQueryEngine(query_engine, query_transform=hyde)

response = transformed_qe.query("RAG 的核心思想")
# 1. LLM 生成"假设答案"：RAG 是检索增强生成...
# 2. 用假设答案做向量检索
# 3. 用检索到的真实 Node 生成最终答案

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

# 场景：doc 里有 [REF:其他doc_id] 这种引用
recursive_retriever = RecursiveRetriever(
    root_id="root",
    retriever_dict={
        "root": vector_retriever,
        "other_doc": other_retriever,  # 引用目标
    }
)

# 用法
nodes = recursive_retriever.retrieve("查找所有相关引用")

from fastapi import FastAPI
from pydantic import BaseModel

app = FastAPI()
query_engine = index.as_query_engine(streaming=True)

class Query(BaseModel):
    question: str

@app.post("/query")
async def query(q: Query):
    response = await query_engine.aquery(q.question)
    return {"answer": str(response), "sources": [n.node.text for n in response.source_nodes]}

from llama_index.server import LlamaIndexServer
server = LlamaIndexServer(index=index)
server.run(host="0.0.0.0", port=8000)

from llama_index.core.evaluation import (
    FaithfulnessEvaluator,         # 答案是否忠实于上下文
    AnswerRelevancyEvaluator,      # 答案与问题相关性
    ContextRelevancyEvaluator,     # 上下文与问题相关性
    CorrectnessEvaluator,          # 答案正确性（需 ground_truth）
    PairwiseComparisonEvaluator,   # 两个引擎对比
)

# 1. Faithfulness
evaluator = FaithfulnessEvaluator(llm=llm)
result = evaluator.evaluate_response(
    query="什么是 RAG？",
    response=response
)
print(f"Faithfulness: {result.score}, passing: {result.passing}")

# 2. 批量评估
from llama_index.core.evaluation import BatchEvalRunner
runner = BatchEvalRunner(
    {
        "faithfulness": FaithfulnessEvaluator(llm=llm),
        "relevancy": AnswerRelevancyEvaluator(llm=llm),
    },
    workers=4
)
results = await runner.aevaluate_queries(
    query_engine=query_engine,
    queries=["q1", "q2", ...]
)

from ragas.metrics import (
    faithfulness,
    answer_relevancy,
    context_precision,
    context_recall,
    context_relevancy,
)
from ragas.llama_index import evaluate

# 1. 准备数据
questions = ["什么是 RAG？", "RAG 和微调的区别？"]
ground_truth = ["RAG 是检索增强生成...", "RAG 实时检索，微调改模型参数..."]

# 2. 评估
metrics = [faithfulness, answer_relevancy, context_precision, context_recall]
result = evaluate(query_engine, metrics, questions, ground_truth)
print(result)
# {'faithfulness': 0.9, 'answer_relevancy': 0.89, ...}

# 3. 转 DataFrame 分析
df = result.to_pandas()
df.to_csv("eval_results.csv", index=False)

# Faithfulness
from llama_index.core.evaluation import FaithfulnessEvaluator
f_eval = FaithfulnessEvaluator(llm=llm)
result = f_eval.evaluate(query="X是什么？", response=response)
# score 接近 1 = 答案都来自上下文
# score 接近 0 = LLM 编造内容

# Correctness
from llama_index.core.evaluation import CorrectnessEvaluator
c_eval = CorrectnessEvaluator(llm=llm)
result = c_eval.evaluate(
    query="X是什么？",
    response=response,
    reference="X 的标准答案..."
)
# score 0-5，5 = 完全正确

from deepeval.metrics import (
    AnswerRelevancyMetric,
    FaithfulnessMetric,
    ContextualPrecisionMetric,
)
from deepeval.test_case import LLMTestCase

# 1. 定义指标
answer_relevancy = AnswerRelevancyMetric()
faithfulness = FaithfulnessMetric()
contextual_precision = ContextualPrecisionMetric()

# 2. 运行 RAG 应用拿到 response
rag_application = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=synthesizer,
    node_postprocessors=[SimilarityPostprocessor(similarity_cutoff=0.7)]
)
response = rag_application.query("What is LlamaIndex?")

# 3. 创建 TestCase
test_case = LLMTestCase(
    input="What is LlamaIndex?",
    actual_output=str(response),
    retrieval_context=[n.node.get_content() for n in response.source_nodes]
)

# 4. 评估
answer_relevancy.measure(test_case)
faithfulness.measure(test_case)
contextual_precision.measure(test_case)
print(f"Answer Relevancy: {answer_relevancy.score}")
print(f"Faithfulness: {faithfulness.score}")
print(f"Contextual Precision: {contextual_precision.score}")

from llama_index.core.callbacks import CallbackManager, TokenCountingHandler
import tiktoken

# 1. Token 计数 handler
token_counter = TokenCountingHandler(
    tokenizer=tiktoken.encoding_for_model("gpt-4").encode,
    verbose=True
)

# 2. 注册到 Settings
Settings.callback_manager = CallbackManager([token_counter])

# 3. 跑 query
response = query_engine.query("什么是 RAG？")

# 4. 查看 token 用量
print(f"Total LLM tokens: {token_counter.total_llm_token_count}")
print(f"Prompt tokens: {token_counter.prompt_llm_token_count}")
print(f"Completion tokens: {token_counter.completion_llm_token_count}")
print(f"Embedding tokens: {token_counter.total_embedding_token_count}")

# 5. 重置
token_counter.reset_counts()

from llama_index.core.callbacks import BaseCallbackHandler
from llama_index.core.callbacks.schema import CBEventType

class MyObservabilityHandler(BaseCallbackHandler):
    def __init__(self):
        super().__init__()
        self.llm_calls = []
        self.retrieve_calls = []

    def on_event_start(
        self,
        event_type: CBEventType,
        payload: Optional[Dict] = None,
        event_id: str = "",
        parent_id: str = "",
        **kwargs
    ) -> str:
        return event_id

    def on_event_end(
        self,
        event_type: CBEventType,
        payload: Optional[Dict] = None,
        event_id: str = "",
        **kwargs
    ) -> None:
        if event_type == CBEventType.LLM:
            self.llm_calls.append({
                "timestamp": time.time(),
                "payload": payload,
            })
        elif event_type == CBEventType.RETRIEVE:
            self.retrieve_calls.append({
                "timestamp": time.time(),
                "nodes": len(payload.get("nodes", [])),
            })

# 注册
handler = MyObservabilityHandler()
Settings.callback_manager = CallbackManager([handler])

from llama_index.core import set_global_handler
set_global_handler("arize_phoenix")  # 推送到 Phoenix
# 或
set_global_handler("langfuse")        # 推送到 Langfuse

# RAGAs Testset Generator
from ragas.testset.generator import TestsetGenerator
from ragas.testset.evolutions import simple, reasoning, multi_context
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# 1. 加载文档
documents = SimpleDirectoryReader("./data").load_data()

# 2. 初始化生成器
generator_llm = OpenAI(model="gpt-3.5-turbo-16k")
critic_llm = OpenAI(model="gpt-4")
embeddings = OpenAIEmbedding()
generator = TestsetGenerator.from_llama_index(
    generator_llm=generator_llm,
    critic_llm=critic_llm,
    embeddings=embeddings,
)

# 3. 生成测试集
testset = generator.generate_with_llamaindex_docs(
    documents,
    test_size=50,
    distributions={simple: 0.4, reasoning: 0.3, multi_context: 0.3}
)

from llama_index.core.evaluation import PairwiseComparisonEvaluator

# 比较两个 query engine 的输出
evaluator = PairwiseComparisonEvaluator(llm=llm)
result = evaluator.evaluate(
    query="什么是 RAG？",
    response_a=engine_a.query("什么是 RAG？"),
    response_b=engine_b.query("什么是 RAG？"),
)
# 返回哪个更好
print(f"Winner: {result.choice}")  # "A" / "B" / "Tie"
print(f"Score: {result.score}")

# Hit Rate: 至少有一个相关 doc 在 top-k
hit_rate = 1 if any(g in retrieved[:k] for g in ground_truth) else 0

# MRR (Mean Reciprocal Rank): 第一个相关 doc 排名的倒数
rr = 1 / rank_of_first_relevant_doc
mrr = mean(rr for all queries)

# NDCG (Normalized Discounted Cumulative Gain)
import numpy as np
def dcg(relevances, k):
    relevances = relevances[:k]
    return sum(rel / np.log2(i + 2) for i, rel in enumerate(relevances))
def ndcg(retrieved, ground_truth, k=10):
    relevances = [1 if d in ground_truth else 0 for d in retrieved[:k]]
    ideal = sorted(relevances, reverse=True)
    return dcg(relevances, k) / (dcg(ideal, k) + 1e-10)

import asyncio

# 1. 异步查询
async def async_query():
    response = await query_engine.aquery("什么是 RAG？")
    return response

# 2. 异步流式
async def stream_query():
    streaming_response = await query_engine.astream_query("什么是 RAG？")
    async for token in streaming_response.async_response_gen():
        print(token, end="", flush=True)

# 3. 并发多个 query
async def batch_queries():
    queries = ["q1", "q2", "q3"]
    tasks = [query_engine.aquery(q) for q in queries]
    responses = await asyncio.gather(*tasks)
    return responses

# 4. 异步 + 多个索引
async def multi_index_query():
    tasks = [
        index_a.aquery(q) for q in queries
    ] + [
        index_b.aquery(q) for q in queries
    ]
    responses = await asyncio.gather(*tasks)
    return responses

asyncio.run(async_query())

# 方案1: LangChain Cache（如果有 LangChain 集成）
from langchain.cache import InMemoryCache
import langchain
langchain.llm_cache = InMemoryCache()

# 方案2: LlamaIndex 自带缓存（基于 query hash）
from llama_index.core.callbacks import CallbackManager
from llama_index.core import Settings

# 方案3: Redis 分布式缓存
import redis
r = redis.Redis()

def cached_query(q):
    key = f"rag:{hash(q)}"
    cached = r.get(key)
    if cached:
        return cached.decode()
    response = query_engine.query(q)
    r.setex(key, 3600, str(response))  # TTL 1h
    return str(response)

from llama_index.core.postprocessor import LongContextReorder
from llama_index.core.response_synthesizers import CompactAndRefine

# 上下文压缩（用小 LLM 提取关键信息）
from llama_index.core.postprocessor import SentenceEmbeddingOptimizer
optimizer = SentenceEmbeddingOptimizer(percentile_cutoff=0.5)

# 1. 删除旧索引
import shutil
shutil.rmtree("./storage")

# 2. 重新构建
docs = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(docs)
index.storage_context.persist(persist_dir="./storage")

# 1. 加载旧索引
storage_context = StorageContext.from_defaults(persist_dir="./storage")
index = load_index_from_storage(storage_context)

# 2. 增量插入新文档
new_docs = SimpleDirectoryReader("./new_data").load_data()
for doc in new_docs:
    index.insert(doc)

# 3. 删除过期文档
index.delete_ref_doc(doc_id="old_doc_id")

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

# 1. 加载阶段
reader = SimpleDirectoryReader(
    "./data",
    recursive=True,
    required_exts=[".pdf", ".docx"]
)
docs = reader.load_data(num_workers=8)  # 并行

# 2. 切分 + 嵌入（分批处理避免 OOM）
from llama_index.core.node_parser import SentenceSplitter
splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50)

# 3. 写入生产级向量库
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import VectorStoreIndex, StorageContext

chroma_client = chromadb.PersistentClient(path="./chroma_prod")
collection = chroma_client.get_or_create_collection(
    "docs",
    metadata={"hnsw:space": "cosine"}
)
chroma_store = ChromaVectorStore(chroma_collection=collection)

storage_context = StorageContext.from_defaults(vector_store=chroma_store)

# 4. 分批构建
batch_size = 1000
all_nodes = splitter.get_nodes_from_documents(docs)
for i in range(0, len(all_nodes), batch_size):
    batch = all_nodes[i:i+batch_size]
    index = VectorStoreIndex(batch, storage_context=storage_context)
    print(f"Processed {i+len(batch)}/{len(all_nodes)} nodes")

# 1. Phoenix (Arize)
import phoenix as px
from llama_index.core import set_global_handler
set_global_handler("arize_phoenix")

# 2. Langfuse
set_global_handler("langfuse")

# 3. OpenLLMetry (Traceloop)
set_global_handler("traceloop")

# 4. HoneyHive
set_global_handler("honeyhive")

# 5. Weights & Biases
set_global_handler("wandb")

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

# 场景：小 node（精确检索）→ 大 node（完整上下文）
recursive_retriever = RecursiveRetriever(
    root_id="root",
    retriever_dict={
        "root": small_node_retriever,    # 小 node 检索
        # 自动跳转到对应父节点
    },
    node_dict=node_dict,  # id → Node 映射
)

# 旧
from llama_index.core import KnowledgeGraphIndex
kg_index = KnowledgeGraphIndex.from_documents(docs)

# 新
from llama_index.core import PropertyGraphIndex
pg_index = PropertyGraphIndex.from_documents(
    docs,
    kg_extractors=[SimpleLLMPathExtractor(llm=llm)]
)

from pydantic import BaseModel, Field
from llama_index.core.program import LLMTextCompletionProgram

class Invoice(BaseModel):
    """发票信息"""
    invoice_no: str = Field(description="发票号码")
    amount: float = Field(description="金额")
    date: str = Field(description="开票日期")
    items: list[str] = Field(description="项目明细")

# 1. Pydantic 风格抽取
prompt = """
从以下文本中抽取发票信息:
{text}
"""
program = LLMTextCompletionProgram.from_defaults(
    output_cls=Invoice,
    prompt_template_str=prompt,
    llm=llm
)
invoice = program(text="发票号 12345，金额 1000 元...")
print(invoice.invoice_no, invoice.amount)

# 2. 批量抽取
from llama_index.core.extractors import PydanticExtractor
extractor = PydanticExtractor(
    output_cls=Invoice,
    llm=llm
)
extracted_nodes = extractor.extract(nodes)