Python:        3.12 或更高版本（推荐 3.12+）
操作系统:      Windows / macOS / Linux
内存:          建议 8GB+（运行本地模型时建议 16GB+）
IDE:           推荐使用 VS Code 或 PyCharm
包管理工具:    pip 或 poetry / uv（推荐）

# 核心框架
pip install llama-index>=0.14.0

# LLM 提供商（按需选择）
pip install llama-index-llms-openai
pip install llama-index-llms-anthropic
pip install llama-index-llms-deepseek

# Embedding 模型
pip install llama-index-embeddings-openai
pip install llama-index-embeddings-huggingface

# 向量数据库（按需选择）
pip install llama-index-vector-stores-chroma
pip install llama-index-vector-stores-qdrant
pip install llama-index-vector-stores-milvus

# 文档解析
pip install llama-index-readers-file
pip install llama-parse

---
title: RAG (检索增强生成) 全生命周期流水线
---
graph LR
    %% ==================== 样式与色彩矩阵 (大框透明，小框带底色+纯白字) ====================
    classDef startEndStyle fill:#666666,stroke:#444444,stroke-width:2px,color:#ffffff;
    classDef ingestStyle fill:#0078d4,stroke:#005a9e,stroke-width:2px,color:#ffffff;
    classDef queryStyle fill:#107c41,stroke:#0b5930,stroke-width:2px,color:#ffffff;

    %% ==================== 1. 离线数据管道 (Data Ingestion Pipeline) ====================
    subgraph Pipeline_Ingest["📥 离线数据管道 / Data Ingestion Pipeline"]
        direction LR
        P_Start["数据源"]:::startEndStyle
        P1["数据连接器"]:::ingestStyle
        P2["文档预处理"]:::ingestStyle
        P3["文本分块"]:::ingestStyle
        P4["向量化"]:::ingestStyle
        P_End["向量存储"]:::ingestStyle

        %% 内部完全横向一字长蛇阵
        P_Start ===> P1 ===> P2 ===> P3 ===> P4 ===> P_End
    end

    %% ==================== 2. 在线检索生成 (Retrieval & Generation) ====================
    subgraph Pipeline_Query["⚡ 在线检索生成 / Retrieval & Generation"]
        direction LR
        Q_Start["用户提问"]:::startEndStyle
        Q1["问题向量化"]:::queryStyle
        Q2["语义检索"]:::queryStyle
        Q3["重排序"]:::queryStyle
        Q4["上下文组装"]:::queryStyle
        Q5["LLM生成"]:::queryStyle
        Q_End["最终回答"]:::startEndStyle

        %% 内部完全横向一字长蛇阵
        Q_Start ===> Q1 ===> Q2 ===> Q3 ===> Q4 ===> Q5 ===> Q_End
    end

    %% ==================== 上下两层管道的逻辑关联 ====================
    P_End -.-> |"提供索引支撑"| Q2

    %% ==================== 容器边框与色彩自适应 (大框强制 fill:none 无背景) ====================
    style Pipeline_Ingest fill:none,stroke:#0078d4,stroke-width:2px,stroke-dasharray: 5 5
    style Pipeline_Query fill:none,stroke:#107c41,stroke-width:2px,stroke-dasharray: 5 5

    %% 全局连线颜色自适应
    linkStyle default stroke:#888888,stroke-width:2px;

选择 LlamaIndex 的场景：
├── 企业知识库问答系统（RAG 为核心）
├── 文档智能解析与检索
├── 知识图谱构建与多跳推理
├── 大规模文档处理管道
└── 需要高质量语义检索的应用

选择 LangChain 的场景：
├── 复杂的多步骤任务自动化
├── 需要大量工具调用的 Agent 系统
├── 多 LLM 协作的链式推理
├── 复杂的 Prompt 管理和模板系统
└── 需要高度自定义的 LLM 交互流程

联合使用的场景：
├── 使用 LlamaIndex 构建 RAG 数据管道
│   └── 将 LlamaIndex 的 QueryEngine 作为 LangChain 的 Retriever
├── 使用 LangChain 编排复杂 Agent 逻辑
│   └── Agent 的"知识检索工具"由 LlamaIndex 提供
└── LlamaIndex 处理数据索引，LangChain 处理对话管理

# LlamaIndex 与 LangChain 的互操作
# LlamaIndex 的 QueryEngine 可以作为 LangChain 的检索工具使用

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from langchain.agents import AgentExecutor, create_openai_tools_agent

# 用 LlamaIndex 构建 RAG 引擎
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()

# 将 LlamaIndex 的查询能力封装为 LangChain 工具
from llama_index.core.tools import QueryEngineTool

rag_tool = QueryEngineTool.from_defaults(
    query_engine=query_engine,
    name="knowledge_base",
    description="用于查询企业内部知识库的工具",
)

# 在 LangChain Agent 中使用该工具
# （具体集成方式详见后续章节）

用户提问："北京今天的天气如何？如果适合户外活动，推荐几个公园"

Agent 内部循环：
┌─────────────────────────────────────────────────────────┐
│ 思考 (Thought): 我需要先查询北京的天气                    │
│ 行动 (Action):  调用天气查询工具                          │
│ 观察 (Observe): 北京今天晴，25°C，空气质量优              │
│                                                         │
│ 思考 (Thought): 天气很好，适合户外活动，我需要推荐公园     │
│ 行动 (Action):  调用公园推荐工具                          │
│ 观察 (Observe): 推荐：颐和园、奥森公园、朝阳公园           │
│                                                         │
│ 思考 (Thought): 我已经有了所有需要的信息，可以回答了       │
│ 回答 (Answer):  北京今天天气很好，推荐以下公园...          │
└─────────────────────────────────────────────────────────┘

# 创建虚拟环境（推荐）
python -m venv llamaindex-env
# Windows:
llamaindex-env\Scripts\activate
# macOS/Linux:
# source llamaindex-env/bin/activate

# 安装 LlamaIndex 核心框架
pip install llama-index

# 安装 OpenAI LLM 和 Embedding 适配器
pip install llama-index-llms-openai
pip install llama-index-embeddings-openai

# 安装文件读取器
pip install llama-index-readers-file

import os

# 方式一：直接设置（仅用于开发环境，不要提交到代码仓库）
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

# 方式二：从 .env 文件加载（推荐用于生产环境）
# 首先安装 python-dotenv: pip install python-dotenv
from dotenv import load_dotenv
load_dotenv()  # 自动加载 .env 文件中的环境变量

# .env 文件示例：
# OPENAI_API_KEY=sk-your-api-key-here

mkdir data

LlamaIndex 是一个用于构建 LLM 应用的数据框架。
它提供了数据连接器，可以从各种数据源（PDF、Word、数据库、API等）加载数据。
LlamaIndex 的核心功能是 RAG（检索增强生成），它将文档分割成小块，转化为向量，存储在向量数据库中。
当用户提问时，系统会检索最相关的文档片段，将其作为上下文提供给 LLM，从而生成准确的回答。
LlamaIndex 支持多种索引类型，包括向量索引、知识图谱索引和摘要索引。
向量索引适合语义搜索场景，知识图谱索引适合多跳推理，摘要索引适合全局概览。
LlamaIndex 还提供了 Agent 框架，允许 LLM 自主调用工具完成复杂任务。

# quickstart.py - 你的第一个 LlamaIndex 应用
import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

# 第 1 步：加载 data 目录下的所有文档
documents = SimpleDirectoryReader("data").load_data()

# 第 2 步：构建向量索引（自动完成：文本分块 → Embedding 向量化 → 存入内存向量库）
index = VectorStoreIndex.from_documents(documents)

# 第 3 步：创建查询引擎
query_engine = index.as_query_engine()

# 第 4 步：向你的知识库提问！
response = query_engine.query("LlamaIndex 支持哪些索引类型？")
print(response)

python quickstart.py

LlamaIndex 支持多种索引类型，包括向量索引、知识图谱索引和摘要索引。
向量索引适合语义搜索场景，知识图谱索引适合多跳推理，摘要索引适合全局概览。

# ===== 第 1 步：文档加载 =====
documents = SimpleDirectoryReader("data").load_data()

# 这一行做了什么：
# 1. SimpleDirectoryReader 扫描 "data" 目录下的所有文件
# 2. 根据文件扩展名自动选择对应的解析器（.txt, .pdf, .docx, .csv 等）
# 3. 将每个文件解析为一个 Document 对象
# 4. Document 对象包含：文本内容（text）、元数据（metadata，如文件名、创建时间）
# 
# documents 是一个列表，每个元素是一个 Document 对象
# 例如：documents[0].text = "LlamaIndex 是一个用于构建 LLM 应用的数据框架..."
#       documents[0].metadata = {"file_name": "sample.txt", "file_size": 512, ...}


# ===== 第 2 步：构建向量索引 =====
index = VectorStoreIndex.from_documents(documents)

# 这一行背后发生了很多事情：
# 1. 文本分块（Chunking）：将每个 Document 按句子边界分割成多个 Node
#    - 默认使用 SentenceSplitter
#    - 默认 chunk_size=1024 tokens, chunk_overlap=20 tokens
#    - 例如原始文档被拆分成 3-4 个 Node
#
# 2. 向量化（Embedding）：将每个 Node 的文本转化为高维向量
#    - 默认使用 OpenAI 的 text-embedding-ada-002 模型
#    - 每个文本被转化为 1536 维的浮点数向量
#    - 语义相近的文本，其向量在空间中的距离也更近
#
# 3. 存储（Storage）：将向量和原始文本存入向量数据库
#    - 默认使用内存向量库（SimpleVectorStore）
#    - 生产环境中通常会换成 Chroma、Qdrant、Milvus 等持久化向量库


# ===== 第 3 步：创建查询引擎 =====
query_engine = index.as_query_engine()

# 这一行做了什么：
# 1. 创建一个默认的查询引擎，它内部包含完整的 RAG 管道
# 2. 查询引擎默认配置：
#    - 检索器（Retriever）：从索引中检索 top-k 最相关的 Node（默认 k=2）
#    - 响应合成器（Response Synthesizer）：将检索到的 Node 作为上下文，
#      结合用户问题，构造 Prompt 发送给 LLM 生成最终回答
#
# 内部流程：用户问题 → 问题向量化 → 在向量库中相似度检索 → 
#           取回最相关的文本块 → 组装 Prompt → 发送给 GPT → 返回回答


# ===== 第 4 步：执行查询 =====
response = query_engine.query("LlamaIndex 支持哪些索引类型？")

# 这一行做了什么：
# 1. 将问题 "LlamaIndex 支持哪些索引类型？" 转化为向量
# 2. 在向量库中计算余弦相似度，找到最相关的 2 个文本块
# 3. 构造 Prompt：
#    "根据以下上下文信息回答问题：
#     上下文：LlamaIndex 支持多种索引类型，包括向量索引、知识图谱索引和摘要索引...
#     问题：LlamaIndex 支持哪些索引类型？"
# 4. 将 Prompt 发送给 GPT（默认使用 gpt-4o-mini）
# 5. 返回 LLM 生成的回答


print(response)
# response 是一个 Response 对象，包含：
# - response.response: 回答文本
# - response.source_nodes: 引用的源文档节点列表
# - response.metadata: 其他元数据

# advanced_quickstart.py - 查看 RAG 过程的详细信息
import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

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

# 查看索引中的节点信息
print(f"=== 索引中共有 {len(index.docstore.docs)} 个文档节点 ===\n")
for i, (node_id, node) in enumerate(index.docstore.docs.items()):
    print(f"节点 {i + 1}:")
    print(f"  ID: {node_id}")
    print(f"  内容前100字: {node.get_content()[:100]}...")
    print(f"  元数据: {node.metadata}")
    print()

# 使用 retriever 查看检索结果（不经过 LLM）
retriever = index.as_retriever(similarity_top_k=3)
retrieved_nodes = retriever.retrieve("LlamaIndex 的核心功能是什么？")

print("=== 检索结果（未经 LLM 处理） ===\n")
for i, node_with_score in enumerate(retrieved_nodes):
    print(f"检索结果 {i + 1}（相似度得分: {node_with_score.score:.4f}）:")
    print(f"  内容: {node_with_score.node.get_content()[:150]}...")
    print()

# 使用查询引擎获取完整回答
query_engine = index.as_query_engine()
response = query_engine.query("LlamaIndex 的核心功能是什么？")

print("=== LLM 生成的回答 ===")
print(response.response)

print("\n=== 引用的源节点 ===")
for i, source_node in enumerate(response.source_nodes):
    print(f"来源 {i + 1}（得分: {source_node.score:.4f}）:")
    print(f"  {source_node.node.get_content()[:100]}...")

from llama_index.core import Document

# 手动创建一个 Document 对象
doc = Document(
    text="LlamaIndex 是一个强大的 RAG 框架，它帮助开发者将私有数据与 LLM 结合使用。",
    metadata={
        "source": "官方文档",
        "author": "技术团队",
        "category": "技术介绍",
        "version": "2025",
    },
    # doc_id 如果不指定，会自动生成一个 UUID
    doc_id="llamaindex-intro-001",
)

# Document 的核心属性
print(f"文档ID: {doc.doc_id}")
print(f"文本内容: {doc.text}")
print(f"元数据: {doc.metadata}")

# 修改元数据（这在后续处理中很有用）
doc.metadata["updated_at"] = "2026-06-01"
print(f"更新后元数据: {doc.metadata}")

# 获取文档内容（等同于 .text）
print(f"获取内容: {doc.get_content()}")

from llama_index.core import Document

# 合理配置元数据的排除策略
doc = Document(
    text="向量数据库是 RAG 系统的核心存储组件。",
    metadata={
        "file_path": "/data/technical/vector_db.md",  # 对语义无意义
        "file_size": 2048,                              # 对语义无意义
        "author": "张三",                                # 可能对检索有用
        "category": "向量数据库",                         # 对检索很有用
        "keywords": "向量,数据库,RAG,存储",              # 对检索非常有用
    },
)

# 排除对 Embedding 无意义的元数据
doc.excluded_embed_metadata_keys = ["file_path", "file_size"]
# 这样在生成文本的向量表示时，不会把文件路径和大小信息编入向量

# 排除对 LLM 生成回答无意义的元数据
doc.excluded_llm_metadata_keys = ["file_size"]
# 这样在将文本发送给 LLM 时，仍然会包含文件路径（方便引用来源）
# 但不会包含文件大小（对回答问题没有帮助）

from llama_index.core.schema import TextNode, ImageNode, NodeRelationship, RelatedNodeInfo

# ---- TextNode：最常用的节点类型 ----
node1 = TextNode(
    text="LlamaIndex 提供了向量索引，用于语义搜索场景。",
    metadata={"source": "索引文档", "section": "向量索引"},
    id_="node-vector-index-001",  # 节点ID
)

node2 = TextNode(
    text="知识图谱索引适用于需要多跳推理的复杂查询场景。",
    metadata={"source": "索引文档", "section": "知识图谱"},
    id_="node-kg-index-002",
)

# ---- 建立节点之间的关系 ----
# NodeRelationship 定义了节点之间的关系类型：
# - SOURCE: 该节点来自的源文档
# - NEXT: 下一个节点（顺序关系）
# - PREVIOUS: 上一个节点（顺序关系）
# - PARENT: 父节点
# - CHILD: 子节点

# 设置 node1 和 node2 的前后顺序关系
node1.relationships[NodeRelationship.NEXT] = RelatedNodeInfo(
    node_id=node2.node_id
)
node2.relationships[NodeRelationship.PREVIOUS] = RelatedNodeInfo(
    node_id=node1.node_id
)

print(f"node1 的下一个节点: {node1.relationships[NodeRelationship.NEXT]}")
print(f"node2 的上一个节点: {node2.relationships[NodeRelationship.PREVIOUS]}")

# ---- ImageNode：图片类型的节点 ----
image_node = ImageNode(
    image_path="./diagrams/architecture.png",  # 图片路径
    text="LlamaIndex 架构图",  # 图片的文字描述（用于检索）
    metadata={"type": "架构图", "format": "png"},
    id_="node-image-001",
)

print(f"图片节点路径: {image_node.image_path}")
print(f"图片节点描述: {image_node.text}")

一个 Document（完整文档）
    │
    ├──→ Node 1（第 1-3 段）     ← 索引和检索的最小单元
    ├──→ Node 2（第 4-6 段）
    ├──→ Node 3（第 7-9 段）
    └──→ Node 4（第 10-12 段）

关系：
- Document 是 "原始数据"，Node 是 "处理后的数据块"
- 一个 Document 会被拆分为多个 Node
- Node 保留了指向源 Document 的 SOURCE 关系
- Node 之间可以建立 NEXT/PREVIOUS 顺序关系
- Node 之间也可以建立 PARENT/CHILD 层级关系

from llama_index.core import Document
from llama_index.core.schema import (
    TextNode,
    NodeRelationship,
    RelatedNodeInfo,
)

# 创建源文档
source_doc = Document(
    text="""Python 是一种高级编程语言，由 Guido van Rossum 于 1991 年首次发布。
Python 的设计哲学强调代码的可读性和简洁性。
Python 支持多种编程范式，包括面向对象、命令式、函数式编程。
Python 拥有丰富的标准库和活跃的社区，是目前最流行的编程语言之一。
Python 广泛应用于 Web 开发、数据科学、人工智能、自动化等领域。""",
    metadata={"language": "Chinese", "topic": "Python"},
    doc_id="python-intro",
)

# 将文档手动拆分为节点
node_a = TextNode(
    text="Python 是一种高级编程语言，由 Guido van Rossum 于 1991 年首次发布。Python 的设计哲学强调代码的可读性和简洁性。",
    metadata={"section": "简介"},
    id_="python-node-1",
)

node_b = TextNode(
    text="Python 支持多种编程范式，包括面向对象、命令式、函数式编程。",
    metadata={"section": "编程范式"},
    id_="python-node-2",
)

node_c = TextNode(
    text="Python 拥有丰富的标准库和活跃的社区，是目前最流行的编程语言之一。Python 广泛应用于 Web 开发、数据科学、人工智能、自动化等领域。",
    metadata={"section": "应用"},
    id_="python-node-3",
)

# 建立节点与源文档的关系
for node in [node_a, node_b, node_c]:
    node.relationships[NodeRelationship.SOURCE] = RelatedNodeInfo(
        node_id=source_doc.doc_id
    )

# 建立节点之间的顺序关系
node_a.relationships[NodeRelationship.NEXT] = RelatedNodeInfo(node_id=node_b.node_id)
node_b.relationships[NodeRelationship.PREVIOUS] = RelatedNodeInfo(node_id=node_a.node_id)
node_b.relationships[NodeRelationship.NEXT] = RelatedNodeInfo(node_id=node_c.node_id)
node_c.relationships[NodeRelationship.PREVIOUS] = RelatedNodeInfo(node_id=node_b.node_id)

# 验证关系
print(f"节点A来源: {node_a.relationships[NodeRelationship.SOURCE].node_id}")
print(f"节点A下一个: {node_a.relationships[NodeRelationship.NEXT].node_id}")
print(f"节点B上一个: {node_b.relationships[NodeRelationship.PREVIOUS].node_id}")
print(f"节点B下一个: {node_b.relationships[NodeRelationship.NEXT].node_id}")
print(f"节点C上一个: {node_c.relationships[NodeRelationship.PREVIOUS].node_id}")

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.md"])
documents = reader.load_data()

# ===== 递归加载子目录 =====
reader = SimpleDirectoryReader(
    input_dir="./data",
    recursive=True,           # 递归扫描子目录
    required_exts=[".pdf", ".md", ".txt"],  # 只加载指定扩展名的文件
)
documents = reader.load_data()

# ===== 排除特定文件 =====
reader = SimpleDirectoryReader(
    input_dir="./data",
    exclude=["*.tmp", "draft_*"],  # 排除临时文件和草稿文件
)
documents = reader.load_data()

# ===== 查看加载文档的元数据 =====
for doc in documents:
    print(f"文件: {doc.metadata.get('file_name', 'unknown')}")
    print(f"  路径: {doc.metadata.get('file_path', 'unknown')}")
    print(f"  大小: {doc.metadata.get('file_size', 0)} bytes")
    print(f"  类型: {doc.metadata.get('file_type', 'unknown')}")
    print(f"  文本长度: {len(doc.text)} 字符")
    print()

# 安装 LlamaParse
# pip install llama-parse

# 获取 API Key：访问 https://cloud.llamaindex.ai 注册获取
import os
os.environ["LLAMA_CLOUD_API_KEY"] = "llx-your-llama-parse-api-key"

from llama_parse import LlamaParse

# ===== 基本用法 =====
parser = LlamaParse(
    result_type="markdown",  # 输出格式：markdown 或 text
    language="chinese",       # 文档语言
    verbose=True,             # 打印解析过程
)

# 解析 PDF 文件
documents = parser.load_data("./data/complex_report.pdf")

# ===== 高级配置 =====
parser = LlamaParse(
    result_type="markdown",
    language="chinese",
    # 使用 GPT-4o 进行视觉解析（处理图片和复杂表格）
    use_vendor_multimodal_model=True,
    vendor_multimodal_model_name="openai-gpt4o",
    # 处理多页文档时合并相关内容
    page_separator="\n---\n",
    # 提取图片
    take_screenshot=True,
)

documents = parser.load_data("./data/financial_report.pdf")

# LlamaParse 对比 SimpleDirectoryReader 的优势
# SimpleDirectoryReader 解析 PDF 时：
# - 只能提取纯文本，表格结构丢失
# - 多栏布局可能导致文本顺序混乱
# - 图片和图表信息完全丢失
#
# LlamaParse 解析 PDF 时：
# - 保留表格结构（转为 Markdown 表格）
# - 正确处理多栏布局
# - 使用多模态模型理解图片和图表
# - 输出结构化的 Markdown 格式

# ===== 数据库连接器示例 =====

# --- PostgreSQL 连接器 ---
# pip install llama-index-readers-database
from llama_index.readers.database import DatabaseReader

db_reader = DatabaseReader(uri="postgresql://user:password@localhost:5432/mydb")

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

# 每个查询结果行会被转化为一个 Document
for doc in documents:
    print(f"文章: {doc.text[:100]}...")


# --- MongoDB 连接器 ---
# pip install llama-index-readers-mongodb
from llama_index.readers.mongodb import SimpleMongoReader

mongo_reader = SimpleMongoReader(
    host="localhost",
    port=27017,
)

documents = mongo_reader.load_data(
    db_name="knowledge_base",
    collection_name="articles",
    field_names=["title", "content"],  # 要提取的字段
    query_dict={"status": "published"},  # 查询过滤条件
)


# ===== API 连接器示例 =====

# --- Notion 连接器 ---
# pip install llama-index-readers-notion
from llama_index.readers.notion import NotionPageReader

notion_reader = NotionPageReader(
    integration_token="ntn_your_notion_integration_token"
)

# 加载指定 Notion 页面
documents = notion_reader.load_data(
    page_ids=["your-page-id-1", "your-page-id-2"]
)

# 加载整个 Notion 数据库
documents = notion_reader.load_data(
    database_id="your-database-id"
)


# --- Google Drive 连接器 ---
# pip install llama-index-readers-google
from llama_index.readers.google import GoogleDocsReader

gdoc_reader = GoogleDocsReader()

documents = gdoc_reader.load_data(
    document_ids=["your-google-doc-id-1", "your-google-doc-id-2"]
)


# ===== Web 连接器示例 =====

# --- 网页爬虫连接器 ---
# pip install llama-index-readers-web
from llama_index.readers.web import SimpleWebPageReader

web_reader = SimpleWebPageReader(html_to_text=True)

documents = web_reader.load_data([
    "https://docs.llamaindex.ai/en/stable/",
    "https://docs.llamaindex.ai/en/stable/getting_started/",
])

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

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

# 创建示例文档
doc = Document(text="""
LlamaIndex 是一个用于构建 LLM 应用的数据框架。它提供了丰富的数据连接器，
可以从 PDF、Word、数据库、API 等多种数据源加载数据。框架的核心功能是 RAG
（检索增强生成），它将文档分割成小块，转化为向量表示，存储在向量数据库中。

当用户提出问题时，系统首先将问题转化为向量，然后在向量数据库中执行相似度搜索，
找到与问题最相关的文档片段。这些片段被作为上下文信息，与用户问题一起发送给 LLM，
由 LLM 生成最终的回答。这种方式既保证了回答的准确性，又解决了 LLM 上下文
窗口有限的限制。

LlamaIndex 支持多种索引类型。向量索引适合通用语义搜索场景。知识图谱索引适合
需要多跳推理的复杂查询。摘要索引适合需要对大量文档生成全局概览的场景。
""")

# 使用 SentenceSplitter 进行分割
splitter = SentenceSplitter(
    chunk_size=256,         # 每个块的目标大小（token 数），默认 1024
    chunk_overlap=32,       # 相邻块之间的重叠 token 数，默认 20
    # chunk_size 和 chunk_overlap 的平衡：
    # - chunk_size 越大，每个块包含的信息越多，但检索精度可能下降
    # - chunk_overlap 越大，块之间的上下文连续性越好，但会产生更多冗余
)

nodes = splitter.get_nodes_from_documents([doc])

print(f"原文档被分割为 {len(nodes)} 个节点：\n")
for i, node in enumerate(nodes):
    print(f"--- 节点 {i + 1} (长度: {len(node.text)} 字符) ---")
    print(node.text.strip())
    print()

from llama_index.core.node_parser import TokenTextSplitter

# TokenTextSplitter 基于 tiktoken（OpenAI 的 tokenizer）计算 token 数
# 适合需要精确控制 token 数量的场景（如适配 LLM 的上下文窗口）
token_splitter = TokenTextSplitter(
    chunk_size=512,              # 每个块的目标 token 数
    chunk_overlap=50,            # 重叠 token 数
    separator=" ",               # 分割分隔符
    # 默认使用 OpenAI 的 cl100k_base tokenizer
    # 可以指定其他 tokenizer: tokenizer_name="gpt2"
)

nodes = token_splitter.get_nodes_from_documents([doc])

for i, node in enumerate(nodes):
    print(f"节点 {i + 1}: {len(node.text)} 字符")

from llama_index.core.node_parser import SentenceWindowNodeParser

# SentenceWindowNodeParser 的特殊之处：
# 每个 Node 不仅包含自身的文本（window），还包含周围句子的上下文（window_metadata）
# 这解决了"块太小导致上下文丢失"的问题

window_parser = SentenceWindowNodeParser(
    window_size=3,           # 每个节点包含前后各 3 个句子作为上下文窗口
    # 例如：Node B 的实际内容只有 B 的文本，
    # 但 metadata 中包含了 A、B、C 的完整文本（窗口）
)

nodes = window_parser.get_nodes_from_documents([doc])

print(f"共 {len(nodes)} 个节点\n")
for i, node in enumerate(nodes):
    print(f"--- 节点 {i + 1} ---")
    print(f"  核心文本: {node.text[:80]}...")
    # 窗口文本存储在 metadata 中
    window_text = node.metadata.get("window", "")
    print(f"  窗口文本: {window_text[:120]}...")
    print()

# 使用场景：
# 检索时只匹配核心文本（精确度高）
# 生成时使用窗口文本（上下文丰富）

from llama_index.core.node_parser import HierarchicalNodeParser

# HierarchicalNodeParser 创建多层级的节点：
# 第 1 层：大块（父节点），包含完整段落
# 第 2 层：中块（子节点），包含句子组
# 第 3 层：小块（孙节点），包含单个句子

hierarchical_parser = HierarchicalNodeParser(
    chunk_sizes=[2048, 512, 128],  # 三个层级的块大小
    # 2048 tokens → 第 1 层（粗略但全面）
    # 512 tokens  → 第 2 层（平衡）
    # 128 tokens  → 第 3 层（精确但片段化）
)

nodes = hierarchical_parser.get_nodes_from_documents([doc])

# 统计各层级的节点数
level_counts = {}
for node in nodes:
    level = node.metadata.get("chunk_size", "unknown")
    level_counts[level] = level_counts.get(level, 0) + 1

print("层级节点分布：")
for level, count in sorted(level_counts.items(), reverse=True):
    print(f"  块大小 {level} tokens: {count} 个节点")

# 使用场景：
# 配合 HierarchicalRetriever 实现"先粗后精"的检索策略
# 先在大块中定位相关区域，再在小块中精确匹配

from llama_index.core.node_parser import MarkdownNodeParser

# MarkdownNodeParser 按照 Markdown 的标题层级进行分割
# 特别适合处理有良好结构的 Markdown 文档

markdown_doc = Document(text="""
# LlamaIndex 使用指南

## 1. 安装

使用 pip 安装 LlamaIndex：
```bash
pip install llama-index

**CodeSplitter：按代码语法分割**

```python
from llama_index.core.node_parser import CodeSplitter

# CodeSplitter 基于 tree-sitter 理解代码语法，按函数/类等语义单元分割
# 支持 Python、JavaScript、TypeScript、Java 等多种编程语言

code_doc = Document(text="""
import os
from typing import List

class DataProcessor:
    def __init__(self, config: dict):
        self.config = config
        self.data = []

    def load_data(self, path: str) -> None:
        with open(path, 'r') as f:
            self.data = f.readlines()

    def process(self) -> List[str]:
        results = []
        for line in self.data:
            cleaned = line.strip().lower()
            results.append(cleaned)
        return results

    def save(self, output_path: str) -> None:
        processed = self.process()
        with open(output_path, 'w') as f:
            f.writelines(processed)

def main():
    processor = DataProcessor({"mode": "default"})
    processor.load_data("input.txt")
    processor.save("output.txt")

if __name__ == "__main__":
    main()
""")

# 按 Python 语法分割
python_splitter = CodeSplitter(
    language="python",      # 编程语言
    chunk_lines=30,         # 每个块的目标行数
    chunk_lines_overlap=5,  # 重叠行数
    max_chars=1000,         # 每个块的最大字符数
)

nodes = python_splitter.get_nodes_from_documents([code_doc])

print(f"代码被分割为 {len(nodes)} 个节点：\n")
for i, node in enumerate(nodes):
    print(f"--- 代码块 {i + 1} ---")
    print(node.text)
    print()

from llama_index.core import Document
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import TransformComponent

# ===== 基本用法：构建简单的摄取管道 =====

# 准备示例文档
documents = [
    Document(text="LlamaIndex 是一个数据框架，专注于 RAG 应用的构建。",
             metadata={"source": "doc1"}),
    Document(text="向量索引是 LlamaIndex 中最常用的索引类型。",
             metadata={"source": "doc2"}),
    Document(text="Agent 可以自主使用工具来完成复杂的任务。",
             metadata={"source": "doc3"}),
]

# 创建摄取管道
pipeline = IngestionPipeline(
    transformations=[
        # 第 1 步：文本分割
        SentenceSplitter(chunk_size=256, chunk_overlap=20),
        # 可以在这里添加更多 Transformation（如 Embedding、元数据提取等）
    ],
)

# 执行管道
nodes = pipeline.run(documents=documents)

print(f"管道处理完成，共生成 {len(nodes)} 个节点")
for i, node in enumerate(nodes):
    print(f"节点 {i + 1}: {node.text[:60]}...")

import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import Document
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core.schema import MetadataMode

# 准备较长的文档
long_doc = Document(text="""
LlamaIndex 框架的设计理念是围绕数据展开的。它的核心目标是让开发者能够方便地将
各种格式的外部数据与大语言模型（LLM）进行连接。在当今的 AI 应用开发中，数据是
最宝贵的资产，而如何高效地组织、索引和检索这些数据，是构建高质量 AI 应用的关键。

RAG（Retrieval-Augmented Generation，检索增强生成）是 LlamaIndex 支持的核心
技术范式。RAG 的基本流程是：首先将文档进行预处理和分块，然后通过 Embedding 模型
将每个文本块转化为高维向量，存储在向量数据库中。当用户提出问题时，系统将问题也
转化为向量，在向量数据库中执行相似度搜索，找到最相关的文本块，最后将这些文本块
作为上下文与问题一起发送给 LLM 生成回答。

向量数据库的选择对 RAG 系统的性能有重要影响。常见的向量数据库包括 Chroma、
Qdrant、Milvus、Pinecone 等。Chroma 轻量级，适合开发和原型阶段；Qdrant 性能
优秀，适合中等规模的生产环境；Milvus 支持大规模分布式部署，适合海量数据场景。
""", metadata={"category": "技术文档"})

# 构建完整的摄取管道
pipeline = IngestionPipeline(
    transformations=[
        # 第 1 步：文本分割
        SentenceSplitter(chunk_size=256, chunk_overlap=30),
        # 第 2 步：生成 Embedding 向量
        OpenAIEmbedding(
            model="text-embedding-3-small",  # 使用 OpenAI 最新的小模型
            # model="text-embedding-3-large"  # 更高质量，但更贵更慢
        ),
    ],
)

# 执行管道
nodes = pipeline.run(documents=[long_doc])

print(f"共生成 {len(nodes)} 个节点\n")
for i, node in enumerate(nodes):
    print(f"节点 {i + 1}:")
    print(f"  文本: {node.text[:80]}...")
    print(f"  是否有 Embedding: {node.embedding is not None}")
    if node.embedding:
        print(f"  Embedding 维度: {len(node.embedding)}")
    print()

import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import Document, VectorStoreIndex
from llama_index.core.ingestion import IngestionPipeline, IngestionCache
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core.node_parser import SentenceSplitter

# 使用 docstore 实现增量摄取和去重
# 核心思想：记录已处理过的文档，只处理新增或变化的文档

docstore = SimpleDocumentStore()

pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=256, chunk_overlap=20),
    ],
    docstore=docstore,           # 文档存储（用于去重）
    cache=IngestionCache(),      # 缓存（加速重复处理）
)

# 第一次运行：处理所有文档
documents_v1 = [
    Document(text="Python 是一种流行的编程语言。", doc_id="doc-1"),
    Document(text="JavaScript 主要用于 Web 开发。", doc_id="doc-2"),
]
nodes_v1 = pipeline.run(documents=documents_v1)
print(f"第一次运行：处理了 {len(nodes_v1)} 个节点")

# 第二次运行：只处理新增/变化的文档
documents_v2 = [
    Document(text="Python 是一种流行的编程语言。", doc_id="doc-1"),       # 未变化
    Document(text="JavaScript 主要用于 Web 开发。", doc_id="doc-2"),       # 未变化
    Document(text="Rust 是一种注重安全和性能的系统编程语言。", doc_id="doc-3"),  # 新增
]
nodes_v2 = pipeline.run(documents=documents_v2)
print(f"第二次运行：只处理了 {len(nodes_v2)} 个节点（仅新增的 doc-3）")

# 持久化 docstore（重启后仍可保持去重状态）
docstore.persist(persist_path="./docstore.json")

# 从持久化文件恢复
restored_docstore = SimpleDocumentStore.from_persist_path("./docstore.json")

import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import Document
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.extractors import (
    TitleExtractor,
    SummaryExtractor,
    QuestionsAnsweredExtractor,
    KeywordExtractor,
)
from llama_index.llms.openai import OpenAI

# 准备示例文档
document = Document(text="""
大语言模型（LLM）是近年来人工智能领域最重要的技术突破之一。LLM 基于 Transformer
架构，通过在大规模文本语料上进行预训练，学会了理解和生成自然语言。GPT 系列、
Claude 系列、Gemini 系列等都是知名的 LLM 产品。

LLM 的一个关键特点是其"涌现能力"——当模型参数量达到一定规模后，模型会展现出
训练时未明确教授的能力，如逻辑推理、代码生成、多语言翻译等。这种能力使得 LLM
成为了通用型 AI 助手的基础。

然而，LLM 也面临一些挑战。首先是"幻觉问题"，即模型可能会生成看似合理但实际上
不正确的信息。其次是"知识时效性"，模型的训练数据有截止日期，无法获取最新信息。
最后是"上下文窗口限制"，即使是最新的模型，其能处理的最大文本长度也是有限的。

RAG（检索增强生成）技术是解决上述问题的重要手段。通过将外部知识库与 LLM 结合，
RAG 可以为模型提供实时、准确的信息，有效减少幻觉，并突破知识时效性的限制。
""")

llm = OpenAI(model="gpt-4o-mini", temperature=0.1)

# ===== TitleExtractor：自动提取标题 =====
title_extractor = TitleExtractor(
    llm=llm,
    nodes=5,          # 从前 5 个节点中提取标题
    # 它会让 LLM 阅读文档内容，然后生成一个简洁的标题
)

# ===== SummaryExtractor：自动提取摘要 =====
summary_extractor = SummaryExtractor(
    llm=llm,
    summaries=["self"],  # 为当前节点生成摘要
    # 也可以设置 "prev"、"next" 来为相邻节点生成摘要
    prompt_template="请用一句话概括以下文本的核心内容：\n{text}",
)

# ===== QuestionsAnsweredExtractor：提取该文本能回答的问题 =====
qa_extractor = QuestionsAnsweredExtractor(
    llm=llm,
    questions=3,       # 为每个节点生成 3 个它能回答的问题
    # 这些问题会作为元数据，在检索时提供额外的匹配维度
)

# ===== KeywordExtractor：提取关键词 =====
keyword_extractor = KeywordExtractor(
    llm=llm,
    keywords=5,        # 提取 5 个关键词
    # 关键词可以用于后续的关键词过滤检索
)

# ===== 构建完整的元数据提取管道 =====
pipeline = IngestionPipeline(
    transformations=[
        # 第 1 步：文本分割
        SentenceSplitter(chunk_size=256, chunk_overlap=30),
        # 第 2-5 步：元数据提取（这些会自动调用 LLM）
        title_extractor,
        summary_extractor,
        qa_extractor,
        keyword_extractor,
    ],
)

# 执行管道（注意：元数据提取会调用 LLM API，需要消耗 token）
nodes = pipeline.run(documents=[document])

# 查看提取的元数据
print(f"共生成 {len(nodes)} 个节点\n")
for i, node in enumerate(nodes):
    print(f"=== 节点 {i + 1} ===")
    print(f"文本: {node.text[:100]}...")
    print(f"元数据:")
    for key, value in node.metadata.items():
        print(f"  {key}: {value}")
    print()

# 元数据的用途：
# 1. title（标题）：在检索结果中展示，帮助用户理解来源
# 2. section_summary（摘要）：可用于摘要检索（先匹配摘要，再获取全文）
# 3. questions_this_excerpt_can_answer（可回答问题）：
#    将问题与原文一起索引，当用户提出类似问题时可以提高检索命中率
# 4. excerpt_keywords（关键词）：可用于关键词过滤，缩小检索范围

from llama_index.core.schema import TransformComponent
from llama_index.core.schema import BaseNode
from typing import Sequence

# 方式一：基于规则的自定义提取器（不调用 LLM，零成本）
class CustomMetadataExtractor(TransformComponent):
    """自定义元数据提取器：基于规则提取元数据"""

    def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]:
        for node in nodes:
            text = node.text

            # 提取文本长度
            node.metadata["text_length"] = len(text)

            # 提取是否包含代码块
            node.metadata["has_code"] = "```" in text or "def " in text

            # 提取是否包含数字/统计数据
            import re
            numbers = re.findall(r'\d+\.?\d*', text)
            node.metadata["has_numbers"] = len(numbers) > 0
            node.metadata["number_count"] = len(numbers)

            # 提取段落数
            paragraphs = [p for p in text.split('\n\n') if p.strip()]
            node.metadata["paragraph_count"] = len(paragraphs)

            # 提取语言标记（简单的启发式方法）
            if any('\u4e00' <= c <= '\u9fff' for c in text):
                node.metadata["language"] = "chinese"
            else:
                node.metadata["language"] = "english"

        return nodes


# 使用自定义提取器
extractor = CustomMetadataExtractor()

pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=256),
        extractor,  # 在分割后应用自定义提取器
    ],
)

nodes = pipeline.run(documents=[document])

for i, node in enumerate(nodes):
    print(f"节点 {i + 1} 自定义元数据:")
    for key in ["text_length", "has_code", "has_numbers", "number_count",
                 "paragraph_count", "language"]:
        print(f"  {key}: {node.metadata.get(key)}")
    print()

from llama_index.core.schema import TransformComponent, BaseNode
from typing import Sequence

class QualityFilter(TransformComponent):
    """低质量内容过滤器"""

    def __init__(
        self,
        min_text_length: int = 50,        # 最小文本长度
        max_text_length: int = 10000,     # 最大文本长度
        min_unique_ratio: float = 0.3,    # 最小唯一字符比例
        max_special_char_ratio: float = 0.5,  # 最大特殊字符比例
    ):
        self.min_text_length = min_text_length
        self.max_text_length = max_text_length
        self.min_unique_ratio = min_unique_ratio
        self.max_special_char_ratio = max_special_char_ratio

    def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]:
        filtered_nodes = []

        for node in nodes:
            text = node.text

            # 过滤 1：文本太短（可能是标题、目录等无意义内容）
            if len(text) < self.min_text_length:
                print(f"[过滤] 文本过短 ({len(text)} 字符): {text[:30]}...")
                continue

            # 过滤 2：文本太长（可能是整本书未分割）
            if len(text) > self.max_text_length:
                print(f"[过滤] 文本过长 ({len(text)} 字符): {text[:30]}...")
                continue

            # 过滤 3：重复内容过多（如大量重复的分隔符）
            unique_chars = len(set(text))
            unique_ratio = unique_chars / max(len(text), 1)
            if unique_ratio < self.min_unique_ratio:
                print(f"[过滤] 重复度过高 (唯一字符比例: {unique_ratio:.2f})")
                continue

            # 过滤 4：特殊字符过多（可能是乱码或格式错误）
            import string
            special_chars = sum(1 for c in text if c not in string.printable and c not in ' \n\t')
            special_ratio = special_chars / max(len(text), 1)
            if special_ratio > self.max_special_char_ratio:
                print(f"[过滤] 特殊字符过多 (比例: {special_ratio:.2f})")
                continue

            # 通过所有过滤器
            filtered_nodes.append(node)

        print(f"\n过滤结果: {len(nodes)} → {len(filtered_nodes)} 个节点")
        return filtered_nodes


# 在管道中使用质量过滤器
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter

pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=256, chunk_overlap=20),
        QualityFilter(min_text_length=30),  # 过滤掉太短的片段
    ],
)

nodes = pipeline.run(documents=[document])
print(f"最终保留了 {len(nodes)} 个高质量节点")

graph LR
    %% 全局连线中性化样式
    linkStyle default stroke:#888888,stroke-width:2px;

    %% 外层方向：整体流转并列
    direction LR

    %% ==========================================
    %% 上路：离线数据灌库流水线 (Ingestion Flow)
    %% ==========================================
    subgraph Flow_Ingestion [离线数据写入流水线]
        direction LR
        I_Txt[文本数据] --> I_Mdl[Embedding 模型]
        I_Mdl --> I_Vec[高维向量]
    end

    %% ==========================================
    %% 中枢：向量数据库集群 (Vector Database)
    %% ==========================================
    subgraph Cluster_DB [向量数据库数据中心]
        direction TB
        DB_Store[(向量数据库存储)] --> DB_Search[近似最近邻搜索 ANN]
    end

    %% ==========================================
    %% 下路：在线查询检索流水线 (Query Search Flow)
    %% ==========================================
    subgraph Flow_Query [在线查询检索流水线]
        direction LR
        Q_Txt[查询文本] --> Q_Mdl[Embedding 模型]
        Q_Mdl --> Q_Vec[查询向量]
        Q_Search[执行最近邻搜索] --> Q_Res[返回最相似文本块]
    end

    %% ==========================================
    %% 核心跨轴数据交互连线 (精准还原您的向下箭头分支)
    %% ==========================================
    %% 1. 上路离线向量存入数据库
    I_Vec --> DB_Store

    %% 2. 下路实时查询向量送入数据库执行 ANN 搜索
    Q_Vec --> DB_Search

    %% 3. 数据库搜索结果路由回下路的响应节点
    DB_Search --> Q_Search

    %% ==========================================
    %% 视觉与排版美学硬编码样式定义 (Dual-Mode Adaptive)
    %% ==========================================
    %% 大框全透明、虚线边框样式
    classDef transparentBox fill:none,stroke:#888888,stroke-width:1.5px,stroke-dasharray: 5 5;
    class Flow_Ingestion,Cluster_DB,Flow_Query transparentBox;

    %% 流程起点与终点样式（高饱和度对比色）
    style I_Txt fill:#0078d4,stroke:#0078d4,color:#ffffff,font-weight:bold
    style Q_Txt fill:#0078d4,stroke:#0078d4,color:#ffffff,font-weight:bold
    style Q_Res fill:#107c41,stroke

import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Document
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI

# ===== 方式一：从 Document 列表创建 =====
documents = [
    Document(text="Python 是一种解释型、面向对象的高级编程语言。它由 Guido van Rossum 创建，首次发布于 1991 年。Python 以其简洁的语法和强大的生态系统而闻名。"),
    Document(text="JavaScript 是一种动态类型的脚本语言，主要用于 Web 开发。它可以在浏览器端运行（前端），也可以在服务器端运行（通过 Node.js）。JavaScript 是目前世界上最流行的编程语言之一。"),
    Document(text="Rust 是一种系统编程语言，由 Mozilla 主导开发。Rust 的核心设计理念是内存安全、并发安全和高性能。它通过所有权（Ownership）系统在编译时防止内存相关的错误。"),
    Document(text="Go（也称为 Golang）是 Google 开发的一种静态类型编译语言。Go 的设计目标是简洁、高效、并支持高并发编程。Go 在云原生和微服务领域非常流行。"),
    Document(text="TypeScript 是 JavaScript 的超集，由 Microsoft 开发和维护。TypeScript 添加了静态类型系统和面向对象编程特性。大型前端项目通常使用 TypeScript 来提高代码的可维护性和开发体验。"),
]

# 配置 Embedding 模型
embed_model = OpenAIEmbedding(
    model="text-embedding-3-small",   # OpenAI 最新 Embedding 模型
    # model="text-embedding-3-large"  # 更高质量（3072 维）
    # model="text-embedding-ada-002"  # 经典模型（1536 维）
)

# 配置 LLM
llm = OpenAI(model="gpt-4o-mini", temperature=0.1)

# 创建向量索引
# from_documents() 会自动完成：文本分块 → Embedding 向量化 → 存入向量库
index = VectorStoreIndex.from_documents(
    documents,
    embed_model=embed_model,  # 指定 Embedding 模型
)

# ===== 基本查询 =====
query_engine = index.as_query_engine(
    llm=llm,                    # 指定 LLM
    similarity_top_k=2,         # 检索最相关的 2 个文档块
)

response = query_engine.query("哪种编程语言最适合系统级开发？")
print(f"回答: {response.response}")
print(f"\n引用来源:")
for source in response.source_nodes:
    print(f"  - {source.node.text[:60]}... (相似度: {source.score:.4f})")


# ===== 方式二：从已有的 Node 列表创建 =====
from llama_index.core.node_parser import SentenceSplitter

# 先手动创建节点
splitter = SentenceSplitter(chunk_size=256)
nodes = splitter.get_nodes_from_documents(documents)

# 从节点创建索引
index_v2 = VectorStoreIndex(
    nodes=nodes,
    embed_model=embed_model,
)

query_engine_v2 = index_v2.as_query_engine(llm=llm)
response_v2 = query_engine_v2.query("Python 和 JavaScript 有什么共同点？")
print(f"\n回答: {response_v2.response}")


# ===== 方式三：空的索引逐步添加文档 =====
index_v3 = VectorStoreIndex(
    nodes=[],                   # 先创建空索引
    embed_model=embed_model,
)

# 逐个添加文档
for doc in documents:
    nodes = splitter.get_nodes_from_documents([doc])
    index_v3.insert_nodes(nodes)

query_engine_v3 = index_v3.as_query_engine(llm=llm)
response_v3 = query_engine_v3.query("Go 语言的优势是什么？")
print(f"\n回答: {response_v3.response}")

# ===== 使用 HuggingFace 本地 Embedding 模型（免费、离线可用）=====
# pip install llama-index-embeddings-huggingface

from llama_index.embeddings.huggingface import HuggingFaceEmbedding

# 使用开源的 Embedding 模型（无需 API Key）
local_embed = HuggingFaceEmbedding(
    model_name="BAAI/bge-small-zh-v1.5",   # 中文优化的小模型
    # model_name="BAAI/bge-large-zh-v1.5"  # 中文优化的高质量模型
    # model_name="sentence-transformers/all-MiniLM-L6-v2"  # 英文通用模型
)

index_local = VectorStoreIndex.from_documents(
    documents,
    embed_model=local_embed,
)

query_engine_local = index_local.as_query_engine(llm=llm)
response_local = query_engine_local.query("Rust 的核心特点是什么？")
print(response_local.response)

# Embedding 模型选择建议：
# | 模型 | 维度 | 适用语言 | 特点 |
# |------|------|---------|------|
# | text-embedding-3-small | 1536 | 多语言 | OpenAI 出品，性价比高 |
# | text-embedding-3-large | 3072 | 多语言 | OpenAI 最高质量 |
# | bge-small-zh-v1.5 | 512 | 中文 | 开源免费，中文优化 |
# | bge-large-zh-v1.5 | 1024 | 中文 | 开源免费，中文高质量 |
# | e5-mistral-7b-instruct | 4096 | 英文 | 基于 LLM 的 Embedding |

# ===== 使用 Chroma 作为向量数据库 =====
# 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(
    name="my_knowledge_base",
    metadata={"hnsw:space": "cosine"},  # 使用余弦相似度
)

# 将 Chroma 集合包装为 LlamaIndex 的 VectorStore
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,
    embed_model=embed_model,
)

# 查询（从 Chroma 中检索）
query_engine = index.as_query_engine(llm=llm)
response = query_engine.query("TypeScript 相比 JavaScript 有什么优势？")
print(response.response)


# ===== 从已有的 Chroma 集合加载索引（不需要重新构建）=====
from llama_index.core import VectorStoreIndex

# 重新连接已有的 Chroma 数据库
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_collection("my_knowledge_base")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

# 从已有的向量存储中加载索引（不会重新 Embedding）
index = VectorStoreIndex.from_vector_store(
    vector_store=vector_store,
    embed_model=embed_model,
)

query_engine = index.as_query_engine(llm=llm)
response = query_engine.query("Python 的特点是什么？")
print(response.response)

import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import PropertyGraphIndex, Document
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding

# 准备包含实体关系的文档
documents = [
    Document(text="""
    LlamaIndex 是由 Jerry Liu 创建的开源框架。LlamaIndex 的核心功能是 RAG。
    RAG 的全称是 Retrieval-Augmented Generation。RAG 需要三个关键组件：
    数据连接器、向量数据库和大语言模型。常用的向量数据库包括 Chroma、Qdrant 和 Milvus。
    Chroma 是一个轻量级的向量数据库。Qdrant 使用 Rust 语言编写，性能优秀。
    Milvus 支持大规模分布式部署。
    """),
    Document(text="""
    LangChain 是另一个流行的 LLM 应用框架。LangChain 的创建者是 Harrison Chase。
    LangChain 和 LlamaIndex 可以互操作。LlamaIndex 的 QueryEngine 可以作为
    LangChain 的 Retriever 使用。两者都支持 OpenAI、Anthropic 等多种 LLM。
    """),
]

llm = OpenAI(model="gpt-4o-mini", temperature=0.1)
embed_model = OpenAIEmbedding(model="text-embedding-3-small")

# ===== 创建 PropertyGraphIndex =====
# PropertyGraphIndex 会自动从文本中提取实体（节点）和关系（边）
pg_index = PropertyGraphIndex.from_documents(
    documents,
    llm=llm,
    embed_model=embed_model,
    show_progress=True,  # 显示构建进度
)

# 查看提取的知识图谱
print("=== 知识图谱中的三元组 ===\n")
for triplet in pg_index.property_graph.get_triplets():
    subject, relation, obj = triplet
    print(f"  {subject} --[{relation}]--> {obj}")

# ===== 查询知识图谱 =====
# PropertyGraphIndex 支持多跳推理：
# 例如："Chroma 和 Qdrant 有什么关系？" → 它们都是 RAG 使用的向量数据库
query_engine = pg_index.as_query_engine(
    llm=llm,
    include_text=True,       # 同时使用原始文本（混合检索）
    similarity_top_k=2,
)

response = query_engine.query("RAG 需要哪些组件？各有什么特点？")
print(f"\n=== 查询结果 ===")
print(response.response)

# 多跳推理查询
response2 = query_engine.query("LlamaIndex 和 LangChain 之间有什么关系？")
print(f"\n=== 多跳推理结果 ===")
print(response2.response)

# pip install llama-index-graph-stores-neo4j

# 使用 Neo4j 作为图数据库后端（适合大规模知识图谱）
# 需要先部署 Neo4j 数据库（可通过 Docker）：
# docker run -p 7474:7474 -p 7687:7687 \
#   -e NEO4J_AUTH=neo4j/password \
#   neo4j:latest

from llama_index.graph_stores.neo4j import Neo4jPropertyGraphStore

# 连接 Neo4j
graph_store = Neo4jPropertyGraphStore(
    username="neo4j",
    password="password",
    url="bolt://localhost:7687",
)

from llama_index.core import StorageContext

# 使用 Neo4j 作为后端存储构建知识图谱
storage_context = StorageContext.from_defaults(graph_store=graph_store)

pg_index = PropertyGraphIndex.from_documents(
    documents,
    storage_context=storage_context,
    llm=llm,
    embed_model=embed_model,
)

# 查询（数据存储在 Neo4j 中，支持 Cypher 查询）
query_engine = pg_index.as_query_engine(llm=llm)
response = query_engine.query("有哪些向量数据库？")
print(response.response)

import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import SummaryIndex, Document
from llama_index.llms.openai import OpenAI

# 准备文档
documents = [
    Document(text="Python 由 Guido van Rossum 创建，是一种高级编程语言。Python 支持面向对象和函数式编程。"),
    Document(text="JavaScript 是 Web 开发的核心语言。JavaScript 可以在浏览器和服务器端（Node.js）运行。"),
    Document(text="Rust 由 Mozilla 主导开发，注重内存安全和并发安全。Rust 使用所有权系统管理内存。"),
    Document(text="Go 语言由 Google 开发，适合云原生和微服务开发。Go 的并发模型基于 goroutine。"),
]

llm = OpenAI(model="gpt-4o-mini", temperature=0.1)

# ===== 创建 SummaryIndex =====
summary_index = SummaryIndex.from_documents(documents)

# SummaryIndex 的查询模式
query_engine = summary_index.as_query_engine(
    llm=llm,
    # response_mode 控制如何将节点内容传给 LLM：
    # "tree_summarize"（默认）：递归摘要，适合大量文档
    # "compact"：压缩后一次性传入
    # "refine"：逐节点精化回答
    response_mode="tree_summarize",
)

# 全局概览查询（SummaryIndex 的强项）
response = query_engine.query("请对比这四种编程语言的共同点和差异。")
print(f"=== 全局对比 ===")
print(response.response)

# SummaryIndex 也可以用于普通问答
response2 = query_engine.query("哪些语言支持并发编程？")
print(f"\n=== 并发编程 ===")
print(response2.response)

SummaryIndex 适合的场景：
├── 需要对全部文档进行全局概览/总结
├── 文档数量较少（总 token 数在 LLM 上下文窗口内）
├── 需要对比分析（"A 和 B 有什么不同"）
├── 生成报告、摘要、目录
└── 文档内容高度相关，不需要精确检索

VectorStoreIndex 适合的场景：
├── 大规模文档库（数百到数百万文档）
├── 需要精确的语义检索
├── RAG 问答系统
├── 文档内容多样化，需要按相关性检索
└── 生产环境的标准选择

# ===== TreeIndex：树形结构索引 =====
# TreeIndex 将文档构建成一棵自底向上的树形结构
# 叶子节点是原始文本块，中间节点是子节点的摘要
# 查询时从根节点向下遍历，逐步缩小范围

from llama_index.core import TreeIndex, Document

documents = [
    Document(text="Python 的 Web 框架包括 Django、Flask 和 FastAPI。Django 是全功能框架，Flask 是微框架，FastAPI 专注于高性能 API。"),
    Document(text="Python 的数据科学库包括 NumPy、Pandas 和 Matplotlib。NumPy 提供数值计算，Pandas 提供数据处理，Matplotlib 提供可视化。"),
    Document(text="JavaScript 的前端框架包括 React、Vue 和 Angular。React 由 Meta 维护，Vue 是社区驱动，Angular 由 Google 维护。"),
    Document(text="JavaScript 的后端框架包括 Express、Koa 和 Fastify。Express 最流行，Koa 最轻量，Fastify 性能最好。"),
]

# TreeIndex 构建过程（概念说明）：
# 第 1 层（叶子）：4 个原始文档
# 第 2 层（中间）：2 个摘要节点（分别摘要 Python 和 JavaScript）
# 第 3 层（根）：1 个总摘要节点

tree_index = TreeIndex.from_documents(documents)

query_engine = tree_index.as_query_engine()
response = query_engine.query("Python 有哪些数据科学相关的库？")
print(response.response)

# TreeIndex 适合场景：
# - 层次化的文档结构
# - 需要逐步缩小范围的查询
# - 注意：TreeIndex 的使用频率远低于 VectorStoreIndex

你的核心需求是什么？
│
├── "我需要在大量文档中精确搜索" ──→ VectorStoreIndex ✓
│   （95% 的 RAG 应用场景）
│
├── "我需要总结/概览全部文档" ──→ SummaryIndex ✓
│   （文档数量少，需要全局理解）
│
├── "我需要理解实体之间的关系" ──→ PropertyGraphIndex ✓
│   （知识图谱、多跳推理）
│
├── "我需要层次化地浏览文档" ──→ TreeIndex ✓
│   （较少使用，特定场景）
│
└── "我不确定" ──→ 先用 VectorStoreIndex ✓
    （最通用，覆盖大多数场景）

import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import Document, VectorStoreIndex
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.schema import TransformComponent, BaseNode
from llama_index.embeddings.openai import OpenAIEmbedding
from typing import Sequence

class MetadataEnricher(TransformComponent):
    """动态元数据增强器：为每个节点添加丰富的检索元数据"""

    def __init__(self, category: str = "general"):
        self.category = category

    def __call__(self, nodes: Sequence[BaseNode], **kwargs) -> Sequence[BaseNode]:
        for node in nodes:
            text = node.text.lower()

            # 添加内容分类
            node.metadata["category"] = self.category

            # 基于关键词的自动标签
            tags = []
            if any(kw in text for kw in ["python", "django", "flask", "fastapi"]):
                tags.append("python")
            if any(kw in text for kw in ["javascript", "react", "vue", "node"]):
                tags.append("javascript")
            if any(kw in text for kw in ["rust", "ownership", "cargo"]):
                tags.append("rust")
            if any(kw in text for kw in ["向量", "索引", "embedding", "检索"]):
                tags.append("rag-core")
            node.metadata["tags"] = ", ".join(tags) if tags else "general"

            # 内容质量评分（简单启发式）
            score = 0
            if len(text) > 100:
                score += 1
            if any(kw in text for kw in ["因此", "所以", "总结", "关键", "核心"]):
                score += 1
            if "?" in text or "？" in text:
                score += 1  # 包含问题，可能有价值
            node.metadata["quality_score"] = score

        return nodes


# 构建带元数据增强的完整管道
documents = [
    Document(text="Python 的 FastAPI 框架专注于构建高性能 REST API。FastAPI 基于类型提示自动生成 API 文档，并支持异步处理。"),
    Document(text="向量索引是 RAG 系统的核心组件。它通过 Embedding 模型将文本转化为向量，并使用近似最近邻算法进行高效检索。"),
    Document(text="React 是 Meta 开发的前端 UI 库。React 使用虚拟 DOM 和组件化架构，使 UI 开发更加高效和可维护。"),
]

pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=256, chunk_overlap=30),
        MetadataEnricher(category="tech-docs"),
        OpenAIEmbedding(model="text-embedding-3-small"),
    ],
)

nodes = pipeline.run(documents=documents)

# 查看增强后的元数据
for i, node in enumerate(nodes):
    print(f"节点 {i + 1}:")
    print(f"  文本: {node.text[:60]}...")
    print(f"  分类: {node.metadata.get('category')}")
    print(f"  标签: {node.metadata.get('tags')}")
    print(f"  质量分: {node.metadata.get('quality_score')}")
    print()

# 利用元数据进行过滤检索
from llama_index.core.vector_stores import MetadataFilters, ExactMatchFilter

# 只检索 Python 相关的内容
filters = MetadataFilters(
    filters=[
        ExactMatchFilter(key="tags", value="python"),
    ]
)

index = VectorStoreIndex(nodes=nodes)
query_engine = index.as_query_engine(filters=filters)
response = query_engine.query("FastAPI 有什么特点？")
print(f"过滤检索结果: {response.response}")

import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import VectorStoreIndex, Document, StorageContext
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.openai import OpenAIEmbedding

embed_model = OpenAIEmbedding(model="text-embedding-3-small")
splitter = SentenceSplitter(chunk_size=256)

# ===== 1. 创建索引 =====
initial_docs = [
    Document(text="Python 3.12 引入了更强大的类型系统改进。", doc_id="py-001"),
    Document(text="JavaScript ES2024 新增了 Array.groupBy 方法。", doc_id="js-001"),
]

nodes = splitter.get_nodes_from_documents(initial_docs)
index = VectorStoreIndex(nodes=nodes, embed_model=embed_model)
print(f"初始索引: {len(index.docstore.docs)} 个文档节点")


# ===== 2. 增量添加文档（insert）=====
new_doc = Document(
    text="Rust 2024 Edition 带来了改进的异步编程支持和更好的错误处理。",
    doc_id="rust-001",
)
new_nodes = splitter.get_nodes_from_documents([new_doc])
index.insert_nodes(new_nodes)
print(f"添加后: {len(index.docstore.docs)} 个文档节点")


# ===== 3. 智能增量更新（refresh）=====
# refresh() 会比较文档内容，只更新发生变化的文档
# 未变化的文档会被跳过，节省 Embedding API 调用成本

updated_docs = [
    Document(text="Python 3.12 引入了更强大的类型系统改进。", doc_id="py-001"),       # 未变
    Document(text="JavaScript ES2025 新增了多个实用 API 和语法特性。", doc_id="js-001"),   # 已更新
    Document(text="Go 1.22 版本增强了泛型支持和循环语义。", doc_id="go-001"),              # 新增
]

refreshed = index.refresh_ref_docs(updated_docs)
print(f"刷新后: {len(index.docstore.docs)} 个文档节点")
print(f"更新状态: {refreshed}")
# refreshed 是一个布尔列表，表示每个文档是否被更新
# [False, True, True] → py-001 未变，js-001 更新了，go-001 新增了


# ===== 4. 删除文档 =====
# 通过 doc_id 删除
index.delete_ref_doc("rust-001", delete_from_docstore=True)
print(f"删除后: {len(index.docstore.docs)} 个文档节点")


# ===== 5. 索引持久化与加载 =====
# 保存到磁盘
index.storage_context.persist(persist_dir="./index_storage")
print("索引已保存到磁盘")

# 从磁盘加载（避免重新 Embedding）
from llama_index.core import load_index_from_storage

storage_context = StorageContext.from_defaults(persist_dir="./index_storage")
loaded_index = load_index_from_storage(
    storage_context,
    embed_model=embed_model,  # 加载时仍需要 Embedding 模型来向量化查询
)
print(f"从磁盘加载了 {len(loaded_index.docstore.docs)} 个文档节点")

# 验证加载后的索引可用
query_engine = loaded_index.as_query_engine()
response = query_engine.query("JavaScript 有什么新特性？")
print(f"查询结果: {response.response}")

"""
当文档量非常大时（如数十万份文档），可以采用以下策略：

1. 使用支持分布式的外部向量数据库（Milvus、Qdrant Cloud）
2. 分批次构建索引，避免内存溢出
3. 使用 IngestionPipeline 的缓存和去重机制
"""
import os
os.environ["OPENAI_API_KEY"] = "sk-your-api-key-here"

from llama_index.core import VectorStoreIndex, Document, StorageContext
from llama_index.core.ingestion import IngestionPipeline
from llama_index.core.node_parser import SentenceSplitter
from llama_index.embeddings.openai import OpenAIEmbedding

embed_model = OpenAIEmbedding(model="text-embedding-3-small")

# 模拟大批量文档
all_documents = []
for i in range(100):
    all_documents.append(
        Document(
            text=f"这是第 {i + 1} 份技术文档的内容。它包含了关于分布式系统设计的重要知识，包括一致性哈希、数据分片、副本策略等核心概念。",
            doc_id=f"tech-doc-{i + 1:04d}",
            metadata={"batch": i // 20, "category": "distributed-systems"},
        )
    )

# 分批处理策略
BATCH_SIZE = 20

pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=256, chunk_overlap=20),
        embed_model,
    ],
)

# 创建空索引
index = VectorStoreIndex(nodes=[], embed_model=embed_model)

# 分批处理和添加
for batch_start in range(0, len(all_documents), BATCH_SIZE):
    batch = all_documents[batch_start:batch_start + BATCH_SIZE]
    print(f"处理批次 {batch_start // BATCH_SIZE + 1}: "
          f"文档 {batch_start + 1}-{batch_start + len(batch)}")

    # 使用管道处理当前批次
    nodes = pipeline.run(documents=batch)

    # 添加到索引
    index.insert_nodes(nodes)

    print(f"  完成，当前索引共 {len(index.docstore.docs)} 个节点")

print(f"\n索引构建完成，共 {len(index.docstore.docs)} 个节点")

# 验证索引
query_engine = index.as_query_engine(similarity_top_k=3)
response = query_engine.query("分布式系统的一致性哈希是什么？")
print(f"\n查询结果: {response.response}")

graph LR
    %% 全局连线中性化样式
    linkStyle default stroke:#888888,stroke-width:2px;

    %% 外层方向：整体容器保持水平横向排布
    direction LR

    %% ==========================================
    %% STORAGE CONTEXT CONTAINER (外层全透明虚线大框)
    %% ==========================================
    subgraph StorageContext [StorageContext 统一存储上下文管理器]
        direction LR

        %% ==========================================
        %% 第一列：核心数据资产存储 (Vector & Doc)
        %% ==========================================
        subgraph Col_1 [数据结构层]
            direction TB
            VectorStore[VectorStore<br>向量存储]
            DocStore[DocStore<br>文档存储]
        end

        %% ==========================================
        %% 第二列：拓扑与对话状态存储 (Graph & Chat)
        %% ==========================================
        subgraph Col_2 [关系与状态层]
            direction TB
            GraphStore[GraphStore<br>图存储]
            ChatStore[ChatStore<br>聊天存储]
        end

        %% ==========================================
        %% 第三列：元数据路由存储 (Index)
        %% ==========================================
        subgraph Col_3 [索引控制层]
            direction TB
            IndexStore[IndexStore<br>索引存储]
        end

    end

    %% ==========================================
    %% 视觉与排版美学硬编码样式定义 (Dual-Mode Adaptive)
    %% ==========================================
    %% 外部大框：优雅的透明度与无底色虚线框
    classDef transparentBox fill:none,stroke:#888888,stroke-width:1.5px,stroke-dasharray: 5 5;
    class StorageContext,Col_1,Col_2,Col_3 transparentBox;

    %% 内部实体节点：移除冗长描述，采用对称的高饱和度极客色块（文字强制纯白）
    style VectorStore fill:#0078d4,stroke:#0078d4,color:#ffffff,font-weight:bold;
    style DocStore fill:#7a24db,stroke:#7a24db,color:#ffffff,font-weight:bold;

    style GraphStore fill:#d83b01,stroke:#d83b01,color:#ffffff,font-weight:bold;
    style ChatStore fill:#666666,stroke:#666666,color:#ffffff,font-weight:bold;

    style IndexStore fill:#107c41,stroke:#107c41,color:#ffffff,font-weight:bold;

pip install llama-index-vector-stores-chroma chromadb

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

# ---- 第一步：初始化 Chroma 客户端 ----
# chromadb.PersistentClient 会将数据持久化到指定目录
# 你也可以用 chromadb.EphemeralClient() 创建纯内存实例（关闭后数据丢失）
chroma_client = chromadb.PersistentClient(path="./chroma_db")

# ---- 第二步：创建或获取 Collection ----
# Collection 类似于数据库中的"表"，用于存放一组向量数据
# get_or_create_collection 表示如果不存在就创建，已存在就复用
chroma_collection = chroma_client.get_or_create_collection(
    name="my_documents",
    metadata={"description": "教程文档集合"}  # 可选的集合描述
)

# ---- 第三步：将 Chroma 适配为 LlamaIndex 的 VectorStore ----
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

# ---- 第四步：构建 StorageContext ----
# StorageContext 是 LlamaIndex 的统一存储上下文，
# 它将各种存储层组合在一起，供 Index 使用
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# ---- 第五步：加载文档并构建索引 ----
# 加载数据（假设 ./data 目录下有一些文本文件）
documents = SimpleDirectoryReader("./data").load_data()

# 构建 VectorStoreIndex 时，LlamaIndex 会自动：
# 1. 将文档切分为 Node
# 2. 为每个 Node 生成 Embedding 向量
# 3. 将向量存入 Chroma
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True  # 显示进度条
)

# ---- 第六步：使用索引进行查询 ----
query_engine = index.as_query_engine()
response = query_engine.query("LlamaIndex 支持哪些存储后端？")
print(response)

# ---- 第七步：从已有的 Chroma 数据恢复索引（下次运行时） ----
# 下次运行时不需要重新构建索引，直接从持久化的 Chroma 中恢复
chroma_client_2 = chromadb.PersistentClient(path="./chroma_db")
chroma_collection_2 = chroma_client_2.get_collection("my_documents")
vector_store_2 = ChromaVectorStore(chroma_collection=chroma_collection_2)
index_restored = VectorStoreIndex.from_vector_store(vector_store_2)
query_engine_2 = index_restored.as_query_engine()
response_2 = query_engine_2.query("LlamaIndex 支持哪些存储后端？")
print(response_2)

chroma_collection = chroma_client.create_collection(
    name="my_docs_cosine",
    metadata={"hnsw:space": "cosine"}  # 使用余弦距离
    # 可选值："l2"（默认）、"cosine"、"ip"（内积）
)

pip install llama-index-vector-stores-milvus pymilvus

from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader
from llama_index.vector_stores.milvus import MilvusVectorStore

# ---- 第一步：初始化 MilvusVectorStore ----
# uri 参数：
#   - "./milvus_demo.db"  ：使用 Milvus Lite（本地文件模式，适合开发）
#   - "http://localhost:19530" ：连接本地部署的 Milvus 服务
#   - "https://xxx.cloud.zilliz.com:19530" ：连接 Zilliz Cloud（托管服务）
#
# token 参数：
#   - Milvus Lite 不需要 token
#   - 生产环境需要提供认证 token
#
# dim 参数：
#   - 必须与你的 Embedding 模型输出维度一致
#   - OpenAI text-embedding-3-small: 1536 维
#   - BAAI/bge-small-zh-v1.5: 512 维
vector_store = MilvusVectorStore(
    uri="./milvus_demo.db",       # Milvus Lite 本地模式
    token="",                     # 本地模式不需要 token
    collection_name="llama_docs", # 集合名称
    dim=1536,                     # Embedding 维度（与模型匹配）
    overwrite=False,              # False = 复用已有数据；True = 清空重建
    # 高级索引参数（可选）
    index_config={"index_type": "AUTOINDEX"},  # 自动选择索引类型
    search_config={"nprobe": 10},              # 搜索时探测的聚类数（越大越精确，越慢）
)

# ---- 第二步：构建存储上下文并创建索引 ----
storage_context = StorageContext.from_defaults(vector_store=vector_store)

documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True
)

# ---- 第三步：查询 ----
query_engine = index.as_query_engine(similarity_top_k=5)
response = query_engine.query("Milvus 和 Chroma 有什么区别？")
print(response)

# ---- 附加：生产环境连接示例（Zilliz Cloud） ----
# from llama_index.vector_stores.milvus import MilvusVectorStore
# production_store = MilvusVectorStore(
#     uri="https://your-instance.cloud.zilliz.com:19530",
#     token="your-api-token-here",
#     collection_name="production_docs",
#     dim=1536,
#     overwrite=False,
# )

pip install llama-index-vector-stores-pinecone pinecone-client

import os
from pinecone import Pinecone, ServerlessSpec
from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader
from llama_index.vector_stores.pinecone import PineconeVectorStore

# ---- 第一步：初始化 Pinecone 客户端 ----
# 从环境变量读取 API Key（安全实践，不要硬编码）
# 设置方式：export PINECONE_API_KEY="your-api-key"（Linux/Mac）
# 或：set PINECONE_API_KEY=your-api-key（Windows）
pc = Pinecone(api_key=os.environ.get("PINECONE_API_KEY"))

# ---- 第二步：创建索引（如果不存在） ----
INDEX_NAME = "llama-tutorial"
EMBEDDING_DIM = 1536  # 与你的 Embedding 模型维度一致

if INDEX_NAME not in [idx.name for idx in pc.list_indexes()]:
    pc.create_index(
        name=INDEX_NAME,
        dimension=EMBEDDING_DIM,
        metric="cosine",  # 距离度量：cosine / euclidean / dotproduct
        spec=ServerlessSpec(
            cloud="aws",       # 云服务商：aws / gcp / azure
            region="us-west-2" # 区域
        )
    )

# ---- 第三步：获取索引引用并适配为 LlamaIndex VectorStore ----
pinecone_index = pc.Index(INDEX_NAME)
vector_store = PineconeVectorStore(
    pinecone_index=pinecone_index,
    namespace="tutorial-ns"  # 命名空间，用于在同一个索引中隔离不同数据集
)

# ---- 第四步：构建索引 ----
storage_context = StorageContext.from_defaults(vector_store=vector_store)
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True
)

# ---- 第五步：查询 ----
query_engine = index.as_query_engine()
response = query_engine.query("Pinecone 的优势是什么？")
print(response)

import faiss
from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader
from llama_index.vector_stores.faiss import FaissVectorStore

# 创建 FAISS 索引
# faiss.IndexFlatL2 是精确 L2 距离搜索（暴力搜索，适合小数据集）
# 对于大数据集，可以使用 faiss.IndexIVFFlat 等近似索引
d = 1536  # Embedding 维度
faiss_index = faiss.IndexFlatL2(d)

vector_store = FaissVectorStore(faiss_index=faiss_index)
storage_context = StorageContext.from_defaults(vector_store=vector_store)

documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True
)

# 查询
query_engine = index.as_query_engine()
response = query_engine.query("FAISS 的性能特点是什么？")
print(response)

# FAISS 索引持久化（FAISS 自身的方法）
faiss.write_index(faiss_index, "faiss_index.bin")

# 下次加载
loaded_faiss_index = faiss.read_index("faiss_index.bin")
loaded_vector_store = FaissVectorStore(faiss_index=loaded_faiss_index)

import qdrant_client
from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader
from llama_index.vector_stores.qdrant import QdrantVectorStore

# 本地模式（数据存在磁盘上）
client = qdrant_client.QdrantClient(path="./qdrant_data")

# 或者连接远程 Qdrant 服务
# client = qdrant_client.QdrantClient(
#     url="http://localhost:6333",
#     api_key="your-api-key"  # 如果需要认证
# )

vector_store = QdrantVectorStore(
    client=client,
    collection_name="my_collection"
)

storage_context = StorageContext.from_defaults(vector_store=vector_store)
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True
)

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

# 加载并切分文档
documents = SimpleDirectoryReader("./data").load_data()
splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50)
nodes = splitter.get_nodes_from_documents(documents)

# 使用 insert_nodes 批量写入（内部会自动分批处理）
# 对于 Pinecone/Milvus 等外部服务，LlamaIndex 会自动使用批量 API
index = VectorStoreIndex(
    nodes=nodes,
    storage_context=storage_context,
    insert_batch_size=100,  # 每批写入 100 个 Node（部分 VectorStore 支持此参数）
    show_progress=True
)

from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core import StorageContext, VectorStoreIndex, SimpleDirectoryReader

# 创建内存文档存储
docstore = SimpleDocumentStore()

# 也可以通过本地目录持久化
# 数据会以 JSON 文件的形式保存到指定目录
docstore = SimpleDocumentStore.from_persist_dir("./storage")

# 手动向 DocStore 中添加文档
documents = SimpleDirectoryReader("./data").load_data()
docstore.add_documents(documents)

# 构建包含 DocStore 的存储上下文
storage_context = StorageContext.from_defaults(docstore=docstore)

# 构建索引（此时 Node 的原始文本会存储在 DocStore 中）
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context
)

# 持久化整个存储上下文
storage_context.persist(persist_dir="./storage")

pip install llama-index-storage-docstore-mongodb

from llama_index.storage.docstore.mongodb import MongoDocumentStore
from llama_index.core import StorageContext, VectorStoreIndex, SimpleDirectoryReader

# 连接 MongoDB（请替换为你自己的连接字符串）
docstore = MongoDocumentStore.from_uri(
    uri="mongodb://localhost:27017",
    db_name="llama_index_db",
    namespace="my_app_docs"
)

storage_context = StorageContext.from_defaults(docstore=docstore)
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context
)

pip install llama-index-storage-docstore-redis

from llama_index.storage.docstore.redis import RedisDocumentStore
from llama_index.core import StorageContext, VectorStoreIndex, SimpleDirectoryReader

docstore = RedisDocumentStore.from_host_and_port(
    host="localhost",
    port=6379,
    namespace="llama_docs"
)

storage_context = StorageContext.from_defaults(docstore=docstore)
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context
)

from llama_index.core.storage.index_store import SimpleIndexStore
from llama_index.core import StorageContext

# 创建内存索引存储
index_store = SimpleIndexStore()

# 从持久化目录加载
index_store = SimpleIndexStore.from_persist_dir("./storage")

# 集成到 StorageContext
storage_context = StorageContext.from_defaults(
    index_store=index_store
)

# 持久化
storage_context.persist(persist_dir="./storage")

pip install llama-index-graph-stores-nebula

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

# 连接 NebulaGraph 集群
# NebulaGraph 是一个开源的分布式图数据库
graph_store = NebulaGraphStore(
    space_name="llama_graph",           # 图空间名称
    edge_types=["relationship"],         # 边类型
    rel_prop_names=["name"],             # 关系属性名
    tags=["entity"],                     # 节点标签
)

storage_context = StorageContext.from_defaults(graph_store=graph_store)

documents = SimpleDirectoryReader("./data").load_data()

# 构建知识图谱索引
# LLM 会自动从文档中提取实体和关系
kg_index = KnowledgeGraphIndex.from_documents(
    documents,
    storage_context=storage_context,
    max_triplets_per_chunk=10,  # 每个文本块最多提取的三元组数量
    include_embeddings=True,    # 同时生成向量（支持向量+图混合检索）
    show_progress=True
)

# 查询知识图谱
query_engine = kg_index.as_query_engine()
response = query_engine.query("张三和李四之间有什么关系？")
print(response)

from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.core.memory import ChatMemoryBuffer

# 创建内存聊天存储
chat_store = SimpleChatStore()

# 从持久化目录加载
# chat_store = SimpleChatStore.from_persist_path("./chat_store.json")

# 将 ChatStore 与 ChatMemoryBuffer 结合使用
memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,       # 记忆窗口大小（以 token 计）
    chat_store=chat_store,  # 指定底层的 ChatStore
    chat_store_key="user_1" # 用户标识（用于多用户场景）
)

# 使用 ChatEngine 时传入 memory
chat_engine = index.as_chat_engine(memory=memory)

# 多轮对话
response1 = chat_engine.chat("LlamaIndex 是什么？")
print(f"AI: {response1}")

response2 = chat_engine.chat("它支持哪些向量数据库？")  # 会自动携带上文
print(f"AI: {response2}")

# 持久化聊天记录
chat_store.persist(persist_path="./chat_store.json")

import chromadb
from llama_index.core import (
    VectorStoreIndex, StorageContext, SimpleDirectoryReader,
    Settings
)
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core.storage.index_store import SimpleIndexStore
from llama_index.core.storage.chat_store import SimpleChatStore
from llama_index.vector_stores.chroma import ChromaVectorStore

# ---- 分别创建各个存储层 ----

# 1. VectorStore：使用 Chroma
chroma_client = chromadb.PersistentClient(path="./chroma_db")
chroma_collection = chroma_client.get_or_create_collection("my_docs")
vector_store = ChromaVectorStore(chroma_collection=chroma_collection)

# 2. DocStore：使用本地持久化的 SimpleDocumentStore
docstore = SimpleDocumentStore()

# 3. IndexStore：使用本地持久化的 SimpleIndexStore
index_store = SimpleIndexStore()

# 4. ChatStore：使用 SimpleChatStore
chat_store = SimpleChatStore()

# ---- 组装 StorageContext ----
storage_context = StorageContext.from_defaults(
    vector_store=vector_store,
    docstore=docstore,
    index_store=index_store,
    chat_store=chat_store,
)

# ---- 构建索引 ----
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True
)

# ---- 一次性持久化所有存储层 ----
storage_context.persist(persist_dir="./full_storage")
print("所有存储层已持久化到 ./full_storage 目录")

# ---- 下次运行时从持久化目录恢复 ----
# from llama_index.core import load_index_from_storage
# restored_storage_context = StorageContext.from_defaults(
#     docstore=SimpleDocumentStore.from_persist_dir("./full_storage"),
#     vector_store=ChromaVectorStore(
#         chroma_collection=chromadb.PersistentClient(path="./chroma_db")
#             .get_collection("my_docs")
#     ),
#     index_store=SimpleIndexStore.from_persist_dir("./full_storage"),
# )
# restored_index = load_index_from_storage(restored_storage_context)

from typing import Any, Dict, List, Optional
from llama_index.core.vector_stores.types import (
    BasePydanticVectorStore,
    VectorStoreQuery,
    VectorStoreQueryResult,
)
from llama_index.core.schema import TextNode, BaseNode
import numpy as np


class MyCustomVectorStore(BasePydanticVectorStore):
    """
    自定义向量存储示例。
    这里用一个简单的内存字典来模拟，实际项目中你可以替换为任何后端。
    """

    # stores_text 表明此 VectorStore 是否同时存储原始文本
    # True = 不需要额外的 DocStore；False = 需要配合 DocStore 使用
    stores_text: bool = True

    # 内部数据结构
    _data: Dict[str, dict] = {}  # node_id -> {embedding, text, metadata}

    def __init__(self, **kwargs: Any) -> None:
        """初始化自定义 VectorStore"""
        super().__init__(**kwargs)
        self._data = {}

    @classmethod
    def class_name(cls) -> str:
        """返回类名，用于序列化"""
        return "MyCustomVectorStore"

    def add(
        self,
        nodes: List[BaseNode],
        **add_kwargs: Any,
    ) -> List[str]:
        """
        将一批 Node 添加到向量存储中。

        参数：
            nodes: 要添加的 Node 列表（每个 Node 包含 text、embedding、metadata 等）
            add_kwargs: 额外参数

        返回：
            添加的 Node ID 列表
        """
        ids = []
        for node in nodes:
            node_id = node.node_id
            self._data[node_id] = {
                "embedding": node.embedding,       # Embedding 向量（list of float）
                "text": node.get_content(),         # 原始文本内容
                "metadata": node.metadata,          # 元数据字典
                "ref_doc_id": node.ref_doc_id,      # 关联的原始文档 ID
            }
            ids.append(node_id)
        return ids

    def delete(self, ref_doc_id: str, **delete_kwargs: Any) -> None:
        """
        删除指定文档关联的所有 Node。

        参数：
            ref_doc_id: 原始文档的 ID（删除该文档产生的所有 Node）
        """
        ids_to_delete = [
            node_id for node_id, data in self._data.items()
            if data.get("ref_doc_id") == ref_doc_id
        ]
        for node_id in ids_to_delete:
            del self._data[node_id]

    def query(
        self,
        query: VectorStoreQuery,
        **kwargs: Any,
    ) -> VectorStoreQueryResult:
        """
        执行向量相似度查询。

        参数：
            query: 查询对象，包含：
                - query_embedding: 查询向量
                - similarity_top_k: 返回最相似的 top_k 个结果
                - filters: 元数据过滤条件（可选）
            kwargs: 额外参数

        返回：
            VectorStoreQueryResult，包含 nodes、similarities、ids
        """
        query_embedding = np.array(query.query_embedding)
        top_k = query.similarity_top_k or 5

        # 计算查询向量与所有存储向量的余弦相似度
        similarities = []
        ids = []
        nodes = []

        for node_id, data in self._data.items():
            if data["embedding"] is None:
                continue

            # 余弦相似度计算
            stored_embedding = np.array(data["embedding"])
            dot_product = np.dot(query_embedding, stored_embedding)
            norm_product = np.linalg.norm(query_embedding) * np.linalg.norm(stored_embedding)
            similarity = float(dot_product / norm_product) if norm_product > 0 else 0.0

            similarities.append(similarity)
            ids.append(node_id)

            # 构建返回的 Node 对象
            node = TextNode(
                id_=node_id,
                text=data["text"],
                metadata=data["metadata"],
                embedding=data["embedding"],
            )
            nodes.append(node)

        # 按相似度降序排列，取 top_k
        if similarities:
            sorted_indices = np.argsort(similarities)[::-1][:top_k]
            similarities = [similarities[i] for i in sorted_indices]
            ids = [ids[i] for i in sorted_indices]
            nodes = [nodes[i] for i in sorted_indices]

        return VectorStoreQueryResult(
            nodes=nodes,
            similarities=similarities,
            ids=ids,
        )


# ---- 使用自定义 VectorStore ----
from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader

custom_vector_store = MyCustomVectorStore()
storage_context = StorageContext.from_defaults(vector_store=custom_vector_store)

documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context,
    show_progress=True
)

query_engine = index.as_query_engine()
response = query_engine.query("测试自定义向量存储")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.schema import TextNode
from llama_index.core.vector_stores.types import MetadataFilter, MetadataFilters

# 假设你已经有了一个全局索引
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents)

# ---- 写入时：为每个 Node 添加 tenant_id 元数据 ----
# 在实际应用中，你可以在数据接入阶段就为文档打上租户标签
nodes = []
for doc in documents:
    doc.metadata["tenant_id"] = "tenant_001"  # 标记所属租户
    node = TextNode(text=doc.text, metadata=doc.metadata)
    nodes.append(node)

# ---- 查询时：通过 MetadataFilters 过滤租户数据 ----
# 这样只会检索属于 tenant_001 的数据
filters = MetadataFilters(
    filters=[
        MetadataFilter(
            key="tenant_id",       # 元数据字段名
            value="tenant_001",    # 目标租户 ID
            operator="eq"          # 操作符：eq（等于）、ne（不等于）、in（包含于列表）
        )
    ]
)

# 在查询引擎中应用过滤条件
query_engine = index.as_query_engine(
    filters=filters,
    similarity_top_k=5
)
response = query_engine.query("查找租户相关的文档")
print(response)

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

# 为每个租户创建独立的 Collection
chroma_client = chromadb.PersistentClient(path="./chroma_db")

def get_tenant_index(tenant_id: str, documents=None):
    """
    获取指定租户的索引。
    每个租户使用独立的 Chroma Collection，实现物理隔离。

    参数：
        tenant_id: 租户唯一标识
        documents: 要索引的文档列表（仅在首次创建时需要）
    返回：
        VectorStoreIndex 实例
    """
    collection_name = f"tenant_{tenant_id}"
    collection = chroma_client.get_or_create_collection(collection_name)
    vector_store = ChromaVectorStore(chroma_collection=collection)
    storage_context = StorageContext.from_defaults(vector_store=vector_store)

    if documents:
        index = VectorStoreIndex.from_documents(
            documents,
            storage_context=storage_context,
            show_progress=True
        )
    else:
        index = VectorStoreIndex.from_vector_store(vector_store)

    return index


# 租户 A 的索引
tenant_a_docs = SimpleDirectoryReader("./tenant_a_data").load_data()
tenant_a_index = get_tenant_index("A", tenant_a_docs)

# 租户 B 的索引
tenant_b_docs = SimpleDirectoryReader("./tenant_b_data").load_data()
tenant_b_index = get_tenant_index("B", tenant_b_docs)

# 各自查询互不干扰
response_a = tenant_a_index.as_query_engine().query("租户 A 的问题")
response_b = tenant_b_index.as_query_engine().query("租户 B 的问题")

import os
import hashlib
from typing import Optional
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.schema import TextNode


class AccessControlManager:
    """
    访问控制管理器示例。
    在生产环境中，你应该使用成熟的权限管理框架（如 Casbin、OPA）。
    """

    def __init__(self):
        # 模拟用户权限映射：user_id -> 允许访问的 doc_categories
        self.user_permissions = {
            "user_admin": ["public", "internal", "confidential"],
            "user_employee": ["public", "internal"],
            "user_guest": ["public"],
        }

    def get_allowed_categories(self, user_id: str) -> list:
        """获取用户允许访问的文档类别"""
        return self.user_permissions.get(user_id, [])

    def can_access(self, user_id: str, doc_category: str) -> bool:
        """检查用户是否有权访问某个文档类别"""
        allowed = self.get_allowed_categories(user_id)
        return doc_category in allowed


# 使用元数据标记文档的安全级别
documents = SimpleDirectoryReader("./data").load_data()

# 为文档添加安全级别标签
for doc in documents:
    # 在实际应用中，安全级别可以从文档属性、文件路径等推断
    doc.metadata["security_level"] = "internal"  # public / internal / confidential
    doc.metadata["department"] = "engineering"

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

# 查询时根据用户权限过滤
acl = AccessControlManager()
user_id = "user_employee"
allowed_levels = acl.get_allowed_categories(user_id)

from llama_index.core.vector_stores.types import MetadataFilter, MetadataFilters

# 只允许用户访问其权限范围内的文档
security_filters = MetadataFilters(
    filters=[
        MetadataFilter(
            key="security_level",
            value=allowed_levels,
            operator="in"  # security_level 在允许列表中
        )
    ]
)

query_engine = index.as_query_engine(filters=security_filters)
response = query_engine.query("公司内部技术规范是什么？")
print(response)

from llama_index.core import (
    VectorStoreIndex, StorageContext, SimpleDirectoryReader,
    load_index_from_storage, Settings
)
from llama_index.core.storage.docstore import SimpleDocumentStore
from llama_index.core.storage.index_store import SimpleIndexStore
from llama_index.embeddings.openai import OpenAIEmbedding

PERSIST_DIR = "./my_index_storage"

# ==========================================
# 第一次运行：构建索引并持久化
# ==========================================
def build_and_persist():
    """构建索引并保存到磁盘"""
    documents = SimpleDirectoryReader("./data").load_data()

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

    # 持久化到指定目录
    # 这会保存 DocStore、IndexStore 和 VectorStore（如果使用默认的 InMemoryVectorStore）
    index.storage_context.persist(persist_dir=PERSIST_DIR)
    print(f"索引已持久化到 {PERSIST_DIR}")
    return index


# ==========================================
# 后续运行：从磁盘恢复索引
# ==========================================
def load_from_disk():
    """从持久化目录恢复索引"""
    # 恢复存储上下文
    storage_context = StorageContext.from_defaults(
        persist_dir=PERSIST_DIR
    )

    # 从存储上下文恢复索引
    index = load_index_from_storage(storage_context)
    print(f"索引已从 {PERSIST_DIR} 恢复")
    return index


# 运行逻辑
if not os.path.exists(PERSIST_DIR):
    index = build_and_persist()
else:
    index = load_from_disk()

# 使用恢复的索引查询
query_engine = index.as_query_engine()
response = query_engine.query("测试恢复后的索引")
print(response)

pip install fsspec s3fs

import os
from llama_index.core import VectorStoreIndex, StorageContext, SimpleDirectoryReader

# ---- 方法一：先本地持久化，再上传到 S3 ----
PERSIST_DIR = "./local_storage"
S3_BUCKET = "s3://my-llamaindex-bucket/index_storage"

# 1. 构建并本地持久化
documents = SimpleDirectoryReader("./data").load_data()
index = VectorStoreIndex.from_documents(documents, show_progress=True)
index.storage_context.persist(persist_dir=PERSIST_DIR)

# 2. 上传到 S3（使用 AWS CLI 或 boto3）
import subprocess
subprocess.run([
    "aws", "s3", "sync", PERSIST_DIR, S3_BUCKET
], check=True)
print(f"索引已上传到 {S3_BUCKET}")

# ---- 方法二：从 S3 下载并恢复 ----
# 1. 从 S3 下载到本地临时目录
LOCAL_RESTORE_DIR = "./restored_storage"
os.makedirs(LOCAL_RESTORE_DIR, exist_ok=True)
subprocess.run([
    "aws", "s3", "sync", S3_BUCKET, LOCAL_RESTORE_DIR
], check=True)

# 2. 恢复索引
from llama_index.core import load_index_from_storage
storage_context = StorageContext.from_defaults(persist_dir=LOCAL_RESTORE_DIR)
restored_index = load_index_from_storage(storage_context)

query_engine = restored_index.as_query_engine()
response = query_engine.query("从 S3 恢复的索引测试")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

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

# 创建查询引擎（最简单的用法）
query_engine = index.as_query_engine()

# 执行查询
response = query_engine.query("LlamaIndex 的核心组件有哪些？")
print(f"回答：{response}")

# 获取响应中的源文档信息
print("\n来源文档：")
for node_with_score in response.source_nodes:
    print(f"  - 相似度: {node_with_score.score:.4f}")
    print(f"    文本: {node_with_score.node.get_content()[:100]}...")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.query_engine import RetrieverQueryEngine
from llama_index.core.response_synthesizers import get_response_synthesizer

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

# 显式创建 RetrieverQueryEngine
# 第一步：创建 Retriever
retriever = index.as_retriever(similarity_top_k=5)

# 第二步：创建 ResponseSynthesizer
response_synthesizer = get_response_synthesizer(
    response_mode="compact"  # 使用紧凑模式
)

# 第三步：组装 QueryEngine
query_engine = RetrieverQueryEngine(
    retriever=retriever,
    response_synthesizer=response_synthesizer,
    # node_postprocessors=[...],  # 可选的后处理器列表
)

response = query_engine.query("默认查询引擎的工作流程是什么？")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.query_engine import TransformQueryEngine
from llama_index.core.indices.query.query_transform.base import StepDecomposeQueryTransform

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

# 先创建基础查询引擎
base_query_engine = index.as_query_engine()

# 使用 StepDecomposeQueryTransform 将复杂问题分解为子问题
# 例如："LlamaIndex 的存储和查询有什么区别？"
# 会被分解为：
#   子问题1："LlamaIndex 的存储系统是什么？"
#   子问题2："LlamaIndex 的查询系统是什么？"
query_transform = StepDecomposeQueryTransform()

transform_query_engine = TransformQueryEngine(
    query_engine=base_query_engine,
    query_transform=query_transform,
)

response = transform_query_engine.query(
    "LlamaIndex 的存储系统和查询系统各自的设计目标是什么？它们是如何协同工作的？"
)
print(response)

pip install llama-index-tools-llama-hub

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, SummaryIndex
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

# ---- 创建不同主题的索引 ----
# 假设你有两组不同主题的文档
tech_docs = SimpleDirectoryReader("./tech_docs").load_data()
business_docs = SimpleDirectoryReader("./business_docs").load_data()

tech_index = VectorStoreIndex.from_documents(tech_docs)
business_index = VectorStoreIndex.from_documents(business_docs)

# ---- 为每个索引创建查询工具 ----
tech_engine = tech_index.as_query_engine()
business_engine = business_index.as_query_engine()

query_tools = [
    QueryEngineTool(
        query_engine=tech_engine,
        metadata=ToolMetadata(
            name="tech_docs_engine",
            description="用于查询技术文档，包括架构设计、API 文档、技术规范等",
        ),
    ),
    QueryEngineTool(
        query_engine=business_engine,
        metadata=ToolMetadata(
            name="business_docs_engine",
            description="用于查询商业文档，包括市场分析、财务报告、商业计划等",
        ),
    ),
]

# ---- 创建 RouterQueryEngine ----
# LLMSingleSelector：使用 LLM 来选择最合适的工具
router_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=query_tools,
    verbose=True  # 打印路由决策过程
)

# LLM 会自动判断问题类型并路由到对应的索引
response = router_engine.query("我们的系统采用了哪些核心技术？")
# → 路由到 tech_docs_engine

response = router_engine.query("第三季度的营收增长了多少？")
# → 路由到 business_docs_engine

print(response)

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

# ---- 创建多个索引和查询引擎 ----
tech_docs = SimpleDirectoryReader("./tech_docs").load_data()
business_docs = SimpleDirectoryReader("./business_docs").load_data()

tech_index = VectorStoreIndex.from_documents(tech_docs)
business_index = VectorStoreIndex.from_documents(business_docs)

tech_engine = tech_index.as_query_engine()
business_engine = business_index.as_query_engine()

query_tools = [
    QueryEngineTool(
        query_engine=tech_engine,
        metadata=ToolMetadata(
            name="tech_engine",
            description="技术文档查询引擎",
        ),
    ),
    QueryEngineTool(
        query_engine=business_engine,
        metadata=ToolMetadata(
            name="business_engine",
            description="商业文档查询引擎",
        ),
    ),
]

# ---- 创建 SubQuestionQueryEngine ----
sub_qa_engine = SubQuestionQueryEngine.from_defaults(
    query_engine_tools=query_tools,
    verbose=True  # 打印子问题分解过程
)

# 这个复杂问题会被自动拆分为多个子问题：
# 子问题1（tech_engine）："系统使用了哪些核心技术？"
# 子问题2（business_engine）："这些技术的商业价值是什么？"
response = sub_qa_engine.query(
    "我们的系统采用了哪些核心技术？这些技术带来了怎样的商业价值？"
)
print(response)

pip install llama-index-readers-database sqlalchemy

import sqlalchemy
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.query_engine import SQLAutoVectorQueryEngine, NLSQLTableQueryEngine
from llama_index.core.tools import QueryEngineTool, ToolMetadata
from llama_index.core.objects import ObjectIndex, SQLTableNodeMapping, SQLTableSchema

# ---- 连接数据库 ----
engine = sqlalchemy.create_engine("sqlite:///./my_database.db")

# ---- 创建 SQL 查询引擎 ----
# NLSQLTableQueryEngine 可以将自然语言转换为 SQL 查询
sql_engine = NLSQLTableQueryEngine(
    sql_database=engine,
    tables=["employees", "departments"],  # 指定可查询的表
    verbose=True
)

# ---- 创建向量查询引擎（用于非结构化文档） ----
docs = SimpleDirectoryReader("./text_docs").load_data()
vector_index = VectorStoreIndex.from_documents(docs)
vector_engine = vector_index.as_query_engine()

# ---- 创建 SQLAutoVectorQueryEngine ----
# 它会根据问题类型自动选择 SQL 查询还是向量检索
sql_tool = QueryEngineTool(
    query_engine=sql_engine,
    metadata=ToolMetadata(
        name="sql_engine",
        description="用于查询结构化数据库中的员工和部门信息",
    ),
)
vector_tool = QueryEngineTool(
    query_engine=vector_engine,
    metadata=ToolMetadata(
        name="vector_engine",
        description="用于查询非结构化的公司文档和知识库",
    ),
)

hybrid_engine = SQLAutoVectorQueryEngine(
    sql_query_tool=sql_tool,
    vector_query_tool=vector_tool,
    verbose=True
)

# 自动路由到 SQL 查询
response1 = hybrid_engine.query("工程部门有多少员工？")

# 自动路由到向量检索
response2 = hybrid_engine.query("公司的技术愿景是什么？")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

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

# 启用流式输出：设置 streaming=True
query_engine = index.as_query_engine(streaming=True)

# query() 返回的不再是简单的字符串，而是一个 StreamingResponse 对象
streaming_response = query_engine.query("请详细介绍 LlamaIndex 的架构设计")

# 逐块获取生成的文本
print("回答：", end="", flush=True)
for text_chunk in streaming_response.response_gen:
    print(text_chunk, end="", flush=True)  # 逐块输出
print()  # 换行

# 也可以获取来源信息
print("\n来源文档：")
for source_node in streaming_response.source_nodes:
    print(f"  - [{source_node.score:.3f}] {source_node.node.get_content()[:80]}...")

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

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

# 创建聊天记忆缓冲区
# token_limit 控制记忆的窗口大小：
# 超过这个限制的早期对话会被自动丢弃或摘要化
memory = ChatMemoryBuffer.from_defaults(token_limit=4000)

# 创建聊天引擎
chat_engine = index.as_chat_engine(
    memory=memory,
    verbose=True  # 打印内部处理过程
)

# 多轮对话示例
print("=== 多轮对话演示 ===\n")

# 第一轮
response1 = chat_engine.chat("LlamaIndex 是什么框架？")
print(f"用户: LlamaIndex 是什么框架？")
print(f"AI: {response1}\n")

# 第二轮 —— 使用"它"来指代 LlamaIndex
response2 = chat_engine.chat("它支持哪些存储后端？")
print(f"用户: 它支持哪些存储后端？")
print(f"AI: {response2}\n")

# 第三轮 —— 追问细节
response3 = chat_engine.chat("Chroma 和 Milvus 分别适合什么场景？")
print(f"用户: Chroma 和 Milvus 分别适合什么场景？")
print(f"AI: {response3}\n")

# 第四轮 —— 跨轮次引用
response4 = chat_engine.chat("请帮我总结一下上面讨论的所有内容")
print(f"用户: 请帮我总结一下上面讨论的所有内容")
print(f"AI: {response4}\n")

# 查看对话历史
print("=== 对话历史 ===")
for msg in memory.get_all():
    role = "用户" if msg.role == "user" else "AI"
    print(f"[{role}]: {msg.content[:80]}...")

# 重置对话（清空记忆）
chat_engine.reset()
print("\n对话已重置")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.chat_engine import CondenseQuestionChatEngine
from llama_index.core.memory import ChatMemoryBuffer

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

memory = ChatMemoryBuffer.from_defaults(token_limit=4000)

# 创建 CondenseQuestionChatEngine
# 它会先将 "它的最新版本是什么？" 压缩为 "LlamaIndex 的最新版本是什么？"
# 然后再执行检索和生成
chat_engine = CondenseQuestionChatEngine.from_defaults(
    query_engine=index.as_query_engine(),
    memory=memory,
    verbose=True
)

# 对话示例
chat_engine.chat("LlamaIndex 是什么？")
chat_engine.chat("它的最新版本是什么？")  # 会被压缩为独立查询
response = chat_engine.chat("和 LangChain 比有什么优势？")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.chat_engine import ContextChatEngine
from llama_index.core.memory import ChatMemoryBuffer

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

memory = ChatMemoryBuffer.from_defaults(token_limit=4000)

# ContextChatEngine 在每次对话前都会执行检索
# 适合需要频繁引用文档知识的场景
chat_engine = ContextChatEngine.from_defaults(
    retriever=index.as_retriever(similarity_top_k=3),
    memory=memory,
    system_prompt=(
        "你是一个专业的技术助手。请基于提供的文档内容来回答用户的问题。"
        "如果文档中没有相关信息，请如实告知。"
    ),
    verbose=True
)

response = chat_engine.chat("请介绍一下 LlamaIndex 的存储系统设计")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.chat_engine import CondensePlusContextChatEngine
from llama_index.core.memory import ChatMemoryBuffer

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

memory = ChatMemoryBuffer.from_defaults(token_limit=4000)

# CondensePlusContextChatEngine 是功能最全面的聊天引擎
# 工作流程：
# 1. 将用户问题 + 对话历史压缩为独立查询
# 2. 使用压缩后的查询检索相关文档
# 3. 将检索到的文档上下文 + 用户消息发给 LLM
chat_engine = CondensePlusContextChatEngine.from_defaults(
    retriever=index.as_retriever(similarity_top_k=3),
    memory=memory,
    system_prompt="你是一个技术文档助手，请基于文档内容准确回答问题。",
    verbose=True
)

# 多轮对话
chat_engine.chat("LlamaIndex 有哪些核心组件？")
chat_engine.chat("存储组件具体支持哪些后端？")
response = chat_engine.chat("帮我对比一下这些后端的优缺点")
print(response)

from llama_index.core.chat_engine import SimpleChatEngine
from llama_index.core.memory import ChatMemoryBuffer

memory = ChatMemoryBuffer.from_defaults(token_limit=4000)

# SimpleChatEngine 不依赖任何索引
# 直接使用 LLM 进行对话
chat_engine = SimpleChatEngine.from_defaults(
    memory=memory,
    system_prompt="你是一个友好的 AI 助手，擅长解答编程问题。",
)

response = chat_engine.chat("Python 中列表和元组有什么区别？")
print(response)

from llama_index.core.memory import ChatMemoryBuffer
from llama_index.core.llms import ChatMessage, MessageRole

# ---- 基本配置 ----
memory = ChatMemoryBuffer.from_defaults(
    token_limit=3000,          # 记忆窗口大小（以 token 计）
    # 超过 token_limit 的早期消息会被自动丢弃
    # 建议值：3000~6000（取决于你的 LLM 上下文窗口大小）
)

# ---- 手动添加消息 ----
memory.put(ChatMessage(role=MessageRole.USER, content="我叫张三"))
memory.put(ChatMessage(role=MessageRole.ASSISTANT, content="你好，张三！有什么可以帮你的？"))

# ---- 获取所有消息 ----
all_messages = memory.get_all()
for msg in all_messages:
    print(f"[{msg.role}]: {msg.content}")

# ---- 获取指定数量的最近消息 ----
recent = memory.get(last_n=5)  # 获取最近 5 条消息

# ---- 重置记忆 ----
memory.reset()

# ---- 与持久化 ChatStore 结合使用 ----
from llama_index.core.storage.chat_store import SimpleChatStore

chat_store = SimpleChatStore()

# 为不同用户创建独立的记忆
user_memories = {}
for user_id in ["user_001", "user_002"]:
    user_memories[user_id] = ChatMemoryBuffer.from_defaults(
        token_limit=3000,
        chat_store=chat_store,
        chat_store_key=user_id  # 每个用户独立的存储 key
    )

# 用户 A 的对话
user_memories["user_001"].put(
    ChatMessage(role=MessageRole.USER, content="我是用户 A")
)

# 用户 B 的对话（完全独立）
user_memories["user_002"].put(
    ChatMessage(role=MessageRole.USER, content="我是用户 B")
)

# 持久化聊天存储
chat_store.persist(persist_path="./chat_history.json")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

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

# 方式一：从索引创建（最常用）
retriever = index.as_retriever(
    similarity_top_k=5,   # 返回最相似的 5 个 Node
    # 可选参数：
    # filters=MetadataFilters(...)  # 元数据过滤
)

# 执行检索
query_str = "LlamaIndex 的存储架构"
retrieved_nodes = retriever.retrieve(query_str)

# 处理检索结果
print(f"查询: {query_str}")
print(f"检索到 {len(retrieved_nodes)} 个相关 Node：\n")
for i, node_with_score in enumerate(retrieved_nodes):
    print(f"--- Node {i+1} (相似度: {node_with_score.score:.4f}) ---")
    print(f"{node_with_score.node.get_content()[:200]}")
    print(f"元数据: {node_with_score.node.metadata}")
    print()

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader

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

# 显式创建 VectorIndexRetriever
from llama_index.core.indices.vector_store.retrievers import VectorIndexRetriever

retriever = VectorIndexRetriever(
    index=index,
    similarity_top_k=5,      # 返回最相似的 5 个结果
    # 高级参数：
    # vector_store_info: 可选的 VectorStoreInfo（用于自动过滤）
    # filters: MetadataFilters  # 元数据过滤条件
)

nodes = retriever.retrieve("向量检索的工作原理")
for node in nodes:
    print(f"[{node.score:.3f}] {node.node.get_content()[:100]}")

from llama_index.core import SummaryIndex, SimpleDirectoryReader

documents = SimpleDirectoryReader("./data").load_data()
summary_index = SummaryIndex.from_documents(documents)

# 默认返回所有 Node
retriever = summary_index.as_retriever()
all_nodes = retriever.retrieve("总结所有文档内容")

print(f"总共检索到 {len(all_nodes)} 个 Node")

# SummaryIndexRetriever 常用于"全文总结"场景
# 配合 tree_summarize 响应模式效果最佳
query_engine = summary_index.as_query_engine(
    response_mode="tree_summarize"
)
response = query_engine.query("请对所有文档内容做一个综合总结")
print(response)

pip install llama-index-retrievers-bm25

from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from llama_index.retrievers.bm25 import BM25Retriever

# 加载文档并切分为 Node
documents = SimpleDirectoryReader("./data").load_data()
splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50)
nodes = splitter.get_nodes_from_documents(documents)

# 创建 BM25 检索器
bm25_retriever = BM25Retriever.from_defaults(
    nodes=nodes,
    similarity_top_k=5,
    # language="chinese"  # 如果是中文文档，可以指定语言以使用正确的分词器
)

# 关键词检索
results = bm25_retriever.retrieve("LlamaIndex VectorStore Chroma")
print("BM25 关键词检索结果：")
for i, node in enumerate(results):
    print(f"\n--- 结果 {i+1} (分数: {node.score:.3f}) ---")
    print(node.node.get_content()[:200])

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from llama_index.retrievers.bm25 import BM25Retriever
from llama_index.core.retrievers import QueryFusionRetriever

# 加载和切分文档
documents = SimpleDirectoryReader("./data").load_data()
splitter = SentenceSplitter(chunk_size=512, chunk_overlap=50)
nodes = splitter.get_nodes_from_documents(documents)

# 创建向量索引和检索器
vector_index = VectorStoreIndex(nodes)
vector_retriever = vector_index.as_retriever(similarity_top_k=10)

# 创建 BM25 检索器
bm25_retriever = BM25Retriever.from_defaults(
    nodes=nodes,
    similarity_top_k=10,
)

# 创建混合检索器
fusion_retriever = QueryFusionRetriever(
    retrievers=[
        vector_retriever,   # 第一路：向量语义检索
        bm25_retriever,     # 第二路：BM25 关键词检索
    ],
    retriever_weights=[0.6, 0.4],   # 两路的权重（向量 60%，关键词 40%）
    similarity_top_k=5,             # 最终返回的 top_k
    num_queries=1,                  # 查询改写次数（>1 时会改写查询做多路检索）
    mode="reciprocal_rerank",       # 融合策略：
                                    # "reciprocal_rerank" - RRF 算法（推荐）
                                    # "relative_score" - 相对分数融合
                                    # "dist_based_score" - 距离分数融合
                                    # "simple" - 简单去重合并
    verbose=True,
)

# 执行混合检索
results = fusion_retriever.retrieve("LlamaIndex 的存储系统")
print("混合检索结果：")
for i, node in enumerate(results):
    print(f"\n--- 结果 {i+1} (分数: {node.score:.4f}) ---")
    print(node.node.get_content()[:200])

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.retrievers import RecursiveRetriever
from llama_index.core.schema import IndexNode

# ---- 模拟文档层级结构 ----
# 假设有 3 篇长文档，每篇文档有一个摘要 Node 和若干详情 Node

documents = SimpleDirectoryReader("./data").load_data()

# 切分为小粒度 Node（详情 Node）
splitter = SentenceSplitter(chunk_size=256, chunk_overlap=30)
detail_nodes = splitter.get_nodes_from_documents(documents)

# 为每篇文档创建摘要 Node
summary_nodes = []
for i, doc in enumerate(documents):
    summary_text = f"文档 {i+1} 摘要：{doc.text[:200]}"
    summary_node = IndexNode(
        text=summary_text,
        metadata=doc.metadata,
        index_id=f"doc_summary_{i}",  # 用于关联详情 Node
    )
    summary_nodes.append(summary_node)

# 构建摘要索引（第一层：粗略检索）
summary_index = VectorStoreIndex(nodes=summary_nodes)
summary_retriever = summary_index.as_retriever(similarity_top_k=2)

# 构建详情索引（第二层：精确检索）
detail_index = VectorStoreIndex(nodes=detail_nodes)
detail_retriever = detail_index.as_retriever(similarity_top_k=3)

# ---- 创建 RecursiveRetriever ----
# retriever_dict 定义了递归检索的路径
# 当摘要检索器返回一个 IndexNode 时，会自动跳转到对应的详情检索器
recursive_retriever = RecursiveRetriever(
    root_id="summary",  # 根检索器 ID
    retriever_dict={
        "summary": summary_retriever,  # 第一层：摘要检索
        "detail": detail_retriever,    # 第二层：详情检索
    },
    node_dict={
        # IndexNode 的 index_id 到检索器 ID 的映射
        # 当检索到 index_id 为 "doc_summary_X" 的节点时，
        # 会自动使用 "detail" 检索器进行第二层检索
    },
    verbose=True,
)

# 执行递归检索
results = recursive_retriever.retrieve("查找特定技术细节")
for node in results:
    print(f"[{node.score:.3f}] {node.node.get_content()[:150]}")

from llama_index.core import SimpleDirectoryReader
from llama_index.core.node_parser import HierarchicalNodeParser, SentenceSplitter
from llama_index.core.retrievers import AutoMergingRetriever
from llama_index.core import VectorStoreIndex
from llama_index.core.storage.docstore import SimpleDocumentStore

# ---- 第一步：加载文档 ----
documents = SimpleDirectoryReader("./data").load_data()

# ---- 第二步：使用 HierarchicalNodeParser 创建层级节点 ----
# HierarchicalNodeParser 会自动创建多层粒度的 Node：
# 第一层：大粒度（chunk_size ~2048）- 父节点
# 第二层：中粒度（chunk_size ~512） - 子节点
# 第三层：小粒度（chunk_size ~128） - 孙节点
hierarchical_parser = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128],  # 三层粒度配置
    chunk_overlap=50,
)
hierarchical_nodes = hierarchical_parser.get_nodes_from_documents(documents)

# ---- 第三步：将最细粒度的 Node 构建索引 ----
# leaf_nodes 是最小粒度的节点（最后一层）
from llama_index.core.node_parser import get_leaf_nodes, get_root_nodes

leaf_nodes = get_leaf_nodes(hierarchical_nodes)
root_nodes = get_root_nodes(hierarchical_nodes)

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

# ---- 第四步：构建 DocStore（存储所有层级的 Node） ----
docstore = SimpleDocumentStore()
docstore.add_documents(hierarchical_nodes)

# ---- 第五步：用叶子节点构建向量索引 ----
vector_index = VectorStoreIndex(leaf_nodes)
base_retriever = vector_index.as_retriever(similarity_top_k=10)

# ---- 第六步：创建 AutoMergingRetriever ----
auto_merging_retriever = AutoMergingRetriever(
    vector_retriever=base_retriever,
    docstore=docstore,
    synthetic_ratio=0.5,  # 合并比例阈值：
                          # 当超过 50% 的子节点被检索到时，
                          # 自动合并为父节点返回
    verbose=True,
)

# ---- 第七步：使用 AutoMergingRetriever 检索 ----
results = auto_merging_retriever.retrieve("请详细解释存储架构的设计原理")
print(f"\n自动合并检索返回了 {len(results)} 个结果")
for i, node in enumerate(results):
    content = node.node.get_content()
    print(f"\n--- 结果 {i+1} ---")
    print(f"文本长度: {len(content)} 字符")
    print(f"内容: {content[:200]}...")

# 实际调优时，可以这样测试不同 top_k 的效果
for top_k in [3, 5, 8, 10]:
    retriever = index.as_retriever(similarity_top_k=top_k)
    results = retriever.retrieve("测试查询")

    # 分析检索结果
    scores = [r.score for r in results]
    print(f"top_k={top_k}: 平均分={sum(scores)/len(scores):.4f}, "
          f"最低分={min(scores):.4f}, 最高分={max(scores):.4f}")

Retriever 结果 (top_k=10)
    │
    ▼
┌──────────────────┐
│  Rerank 重排序    │ ← 用更精确的模型重新评分
└──────┬───────────┘
       │
       ▼
┌──────────────────┐
│  关键词过滤       │ ← 过滤掉不包含特定关键词的结果
└──────┬───────────┘
       │
       ▼
┌──────────────────┐
│  相似度过滤       │ ← 去掉相似度太低的结果
└──────┬───────────┘
       │
       ▼
最终结果 (top_n=5)

pip install sentence-transformers

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import SentenceTransformerRerank

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

# 创建重排序后处理器
reranker = SentenceTransformerRerank(
    model="cross-encoder/ms-marco-MiniLM-L-6-v2",  # 重排序模型
    # 常用的重排序模型：
    # - "cross-encoder/ms-marco-MiniLM-L-6-v2"  （英文，轻量）
    # - "cross-encoder/ms-marco-electra-base"    （英文，更精确）
    # - "BAAI/bge-reranker-v2-m3"                （多语言，支持中文）
    top_n=3,  # 重排序后只保留前 3 个结果
)

# 使用重排序器创建查询引擎
query_engine = index.as_query_engine(
    similarity_top_k=10,   # 先检索 10 个
    node_postprocessors=[reranker],  # 再重排为 3 个
)

response = query_engine.query("LlamaIndex 支持哪些后处理器？")
print(response)

# 也可以单独使用重排序器
retriever = index.as_retriever(similarity_top_k=10)
retrieved_nodes = retriever.retrieve("LlamaIndex 后处理器")

# 手动执行重排序
from llama_index.core.schema import QueryBundle
query_bundle = QueryBundle(query_str="LlamaIndex 后处理器")
reranked_nodes = reranker.postprocess_nodes(
    retrieved_nodes,
    query_bundle=query_bundle
)

print(f"重排序前: {len(retrieved_nodes)} 个结果")
print(f"重排序后: {len(reranked_nodes)} 个结果")
for node in reranked_nodes:
    print(f"  [{node.score:.4f}] {node.node.get_content()[:100]}")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import LLMRerank

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

# 使用 LLM 进行重排序
llm_reranker = LLMRerank(
    top_n=3,           # 重排序后保留的结果数
    # choice_batch_size: 每次让 LLM 评估多少个 Node
    # 值越小，LLM 判断越准确，但调用次数越多
    choice_batch_size=5,
)

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

response = query_engine.query("使用 LLM 重排序的效果如何？")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import KeywordNodePostprocessor

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

# 关键词过滤后处理器
keyword_filter = KeywordNodePostprocessor(
    required_keywords=["LlamaIndex", "存储"],  # 结果中必须包含这些关键词之一
    exclude_keywords=["deprecated", "已废弃"],  # 排除包含这些关键词的结果
    # lang="chinese"  # 中文分词支持
)

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

response = query_engine.query("LlamaIndex 的存储方案有哪些？")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import SimilarityPostprocessor

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

# 相似度过滤器
similarity_filter = SimilarityPostprocessor(
    similarity_cutoff=0.7  # 相似度低于 0.7 的结果会被丢弃
    # 阈值设置建议：
    # 0.5 - 宽松，保留较多结果
    # 0.7 - 适中（推荐起始值）
    # 0.85 - 严格，只保留高度相关的结果
)

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

response = query_engine.query("高度相关的存储方案")
print(response)
print(f"通过过滤的源文档数: {len(response.source_nodes)}")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
from llama_index.core.node_parser import SentenceWindowNodeParser

documents = SimpleDirectoryReader("./data").load_data()

# SentenceWindowNodeParser 会为每个 Node 创建一个"窗口"
# Node 本身的文本是当前句子（用于精确检索）
# metadata["window_text"] 包含前后若干句的完整上下文（用于生成回答）
window_parser = SentenceWindowNodeParser.from_defaults(
    window_size=3,         # 窗口大小：前后各 3 句
    window_metadata_key="window_text",    # 窗口文本存储的元数据 key
    original_text_metadata_key="original_text",  # 原始文本存储的元数据 key
)
nodes = window_parser.get_nodes_from_documents(documents)

index = VectorStoreIndex(nodes)

# MetadataReplacementPostProcessor 会在检索后将 Node 文本替换为窗口文本
# 这样检索用的是精确的短句，而 LLM 看到的是更完整的上下文
metadata_replacer = MetadataReplacementPostProcessor(
    target_metadata_key="window_text"  # 替换为窗口文本
)

query_engine = index.as_query_engine(
    similarity_top_k=5,
    node_postprocessors=[metadata_replacer],
)

response = query_engine.query("详细描述存储架构的设计")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import LongContextReorder

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

long_context_reorder = LongContextReorder()

query_engine = index.as_query_engine(
    similarity_top_k=8,
    node_postprocessors=[long_context_reorder],
)

response = query_engine.query("LlamaIndex 的所有核心特性")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.postprocessor import (
    SentenceTransformerRerank,
    SimilarityPostprocessor,
    KeywordNodePostprocessor,
    LongContextReorder,
)

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

# 构建后处理管道（按顺序执行）
postprocessors = [
    # 第一步：相似度过滤（粗筛，去掉明显不相关的）
    SimilarityPostprocessor(similarity_cutoff=0.6),

    # 第二步：重排序（用 Cross-Encoder 精确评分）
    SentenceTransformerRerank(
        model="cross-encoder/ms-marco-MiniLM-L-6-v2",
        top_n=5,
    ),

    # 第三步：关键词过滤（确保结果包含相关术语）
    KeywordNodePostprocessor(
        required_keywords=["LlamaIndex"],
    ),

    # 第四步：长上下文重排（优化给 LLM 的上下文排列顺序）
    LongContextReorder(),
]

# 将后处理管道应用到查询引擎
query_engine = index.as_query_engine(
    similarity_top_k=15,        # 先多检索一些
    node_postprocessors=postprocessors,  # 再层层过滤和精排
)

response = query_engine.query("LlamaIndex 的存储系统设计")
print(f"最终使用的源文档数: {len(response.source_nodes)}")
print(response)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.response_synthesizers import get_response_synthesizer

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

# =============================
# 1. Compact 模式（默认，推荐）
# =============================
compact_synthesizer = get_response_synthesizer(
    response_mode="compact",
    # text_qa_template: 可选，自定义 QA 提示词模板
    # refine_template: 可选，自定义精化提示词模板
)

compact_engine = index.as_query_engine(
    response_synthesizer=compact_synthesizer,
    similarity_top_k=5,
)
response = compact_engine.query("LlamaIndex 的存储架构是什么？")
print(f"[Compact] {response}\n")

# =============================
# 2. Refine 模式（迭代精化）
# =============================
# 工作流程：
# 第 1 轮：用 Node1 + 查询 → 生成初始回答
# 第 2 轮：用 Node2 + 初始回答 + 查询 → 精化回答
# 第 3 轮：用 Node3 + 精化回答 + 查询 → 进一步精化
# ...以此类推
refine_synthesizer = get_response_synthesizer(
    response_mode="refine",
    # 可以自定义精化模板
    # refine_template=PromptTemplate(
    #     "基于以下新信息，优化你的回答。\n"
    #     "新问题: {query_str}\n"
    #     "现有回答: {existing_answer}\n"
    #     "新信息: {context_msg}\n"
    #     "优化后的回答: "
    # ),
)

refine_engine = index.as_query_engine(
    response_synthesizer=refine_synthesizer,
    similarity_top_k=5,
)
response = refine_engine.query("请详细介绍 LlamaIndex 的各种存储后端及其适用场景")
print(f"[Refine] {response}\n")

# =============================
# 3. Tree Summarize 模式（树形摘要）
# =============================
# 工作流程：
# 将所有 Node 分组 → 每组生成摘要 → 合并摘要再摘要 → 最终回答
# 类似于 MapReduce 的思想
tree_synthesizer = get_response_synthesizer(
    response_mode="tree_summarize",
    summary_template=None,  # 可选：自定义摘要模板
    # max_tokens: 每次 LLM 调用的最大 token 数
    # num_children: 树结构中每个父节点包含的子节点数（默认 10）
)

tree_engine = index.as_query_engine(
    response_synthesizer=tree_synthesizer,
    similarity_top_k=20,  # 树形摘要适合处理大量 Node
)
response = tree_engine.query("请对所有文档内容做一个全面的总结")
print(f"[Tree Summarize] {response}\n")

# =============================
# 4. Simple Summarize 模式
# =============================
simple_synthesizer = get_response_synthesizer(
    response_mode="simple_summarize",
)

simple_engine = index.as_query_engine(
    response_synthesizer=simple_synthesizer,
    similarity_top_k=5,
)
response = simple_engine.query("简要总结 LlamaIndex 的存储系统")
print(f"[Simple Summarize] {response}\n")

# =============================
# 5. Generation 模式（仅生成，不用检索结果）
# =============================
gen_synthesizer = get_response_synthesizer(
    response_mode="generation",
)

gen_engine = index.as_query_engine(
    response_synthesizer=gen_synthesizer,
)
# 注意：这个模式下 LLM 不会看到任何检索结果
# 完全依赖 LLM 自身的知识来回答
response = gen_engine.query("Python 有哪些优点？")
print(f"[Generation] {response}\n")

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.response_synthesizers import get_response_synthesizer

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

# 流式响应合成器
streaming_synthesizer = get_response_synthesizer(
    response_mode="compact",
    streaming=True,  # 启用流式输出
)

streaming_engine = index.as_query_engine(
    response_synthesizer=streaming_synthesizer,
)

# 获取流式响应
streaming_response = streaming_engine.query("请详细介绍 LlamaIndex")

# 逐块输出
print("流式输出：", end="", flush=True)
for text_chunk in streaming_response.response_gen:
    print(text_chunk, end="", flush=True)
print()

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
from llama_index.core.prompts import PromptTemplate
from llama_index.core.response_synthesizers import get_response_synthesizer

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

# 自定义 QA 模板
# 可用的变量：
#   {context_str} - 检索到的上下文文本
#   {query_str} - 用户的查询
custom_qa_template = PromptTemplate(
    "你是一个专业的技术文档助手。\n"
    "请基于以下参考文档来回答用户的问题。\n"
    "如果参考文档中没有足够的信息，请明确说明。\n"
    "回答时请引用具体的文档来源。\n\n"
    "参考文档：\n{context_str}\n\n"
    "用户问题：{query_str}\n\n"
    "请给出详细且准确的回答："
)

# 自定义 Refine 模板（仅 refine 模式需要）
custom_refine_template = PromptTemplate(
    "你是一个专业的技术文档助手。\n"
    "基于以下新信息和已有的回答，优化你的回答。\n"
    "如果新信息没有提供额外价值，保持原回答不变。\n\n"
    "已有回答：{existing_answer}\n\n"
    "新的参考文档：{context_msg}\n\n"
    "用户问题：{query_str}\n\n"
    "优化后的回答："
)

# 使用自定义模板创建响应合成器
synthesizer = get_response_synthesizer(
    response_mode="refine",
    text_qa_template=custom_qa_template,
    refine_template=custom_refine_template,
)

engine = index.as_query_engine(response_synthesizer=synthesizer)
response = engine.query("LlamaIndex 的存储架构有什么特点？")
print(response)

from llama_index.core.workflow import Event

# ============================================================
# 自定义 Event 类
# ============================================================

class QueryEvent(Event):
    """查询事件：携带用户问题"""
    query: str                      # 用户原始查询
    session_id: str = "default"     # 会话 ID（带默认值）

class RetrievalEvent(Event):
    """检索事件：携带检索到的文档"""
    documents: list                 # 检索到的文档列表
    relevance_scores: list[float]   # 相关性分数列表

class GenerationEvent(Event):
    """生成事件：携带 LLM 生成的结果"""
    generated_text: str             # 生成的文本
    source_documents: list          # 来源文档
    confidence: float               # 置信度分数

class ErrorEvent(Event):
    """错误事件：携带错误信息"""
    error_message: str              # 错误描述
    step_name: str                  # 发生错误的 Step 名称
    retry_count: int = 0            # 重试次数

from llama_index.core.workflow import (
    Event,
    StartEvent,
    StopEvent,
    InputRequiredEvent,
    HumanResponseEvent,
)

# StartEvent 的用法：
# 调用 workflow.run(query="什么是机器学习?") 后，
# 在接收 StartEvent 的 Step 中可以通过 ev.query 获取参数

# StopEvent 的用法：
# return StopEvent(result="这是最终结果")
# workflow.run() 的返回值就是 "这是最终结果"

from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent

# ============================================================
# @step 装饰器详解
# ============================================================
#
# 基本语法：
#   @step
#   async def step_name(self, ev: InputEventType) -> OutputEventType:
#       ...
#       return OutputEventType(...)
#
# 关键规则：
# 1. 必须用 async def 定义（异步函数）
# 2. 第一个参数（self 之后）是事件参数，类型注解决定了该 Step 接收哪种事件
# 3. 返回值类型注解决定了该 Step 输出哪种事件
# 4. 框架通过类型注解自动完成事件路由，无需手动连接

class SimpleWorkflow(Workflow):
    @step
    async def greet(self, ev: StartEvent) -> StopEvent:
        """
        最简单的 Step：接收 StartEvent，返回 StopEvent
        - ev.name 来自 workflow.run(name="Alice") 的调用参数
        - 返回 StopEvent，工作流结束
        """
        name = ev.get("name", "World")  # 安全获取参数，提供默认值
        greeting = f"Hello, {name}!"
        return StopEvent(result=greeting)


# 运行工作流
import asyncio

async def main():
    w = SimpleWorkflow(timeout=30, verbose=False)
    result = await w.run(name="Alice")
    print(result)  # 输出: Hello, Alice!

asyncio.run(main())

# 标准签名：(self, ev: InputEvent) -> OutputEvent
@step
async def my_step(self, ev: MyEvent) -> ResultEvent:
    ...

# 带 Context 参数：Context 可以放在任意位置（通常在 ev 之前）
@step
async def my_step(self, ctx: Context, ev: MyEvent) -> ResultEvent:
    ...

# 一个 Step 可以发出多种事件（用 Union 或 | 语法）
from typing import Union

@step
async def classifier(self, ev: StartEvent) -> Union[TextEvent, ImageEvent]:
    """根据输入类型返回不同的事件"""
    if ev.get("input_type") == "text":
        return TextEvent(content=ev.get("content"))
    else:
        return ImageEvent(image_path=ev.get("content"))

"""
最简工作流示例：文本翻译器
流程：StartEvent -> translate_step -> TranslatedEvent -> format_step -> StopEvent
"""
import asyncio
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent


# 第一步：定义自定义事件
class TranslatedEvent(Event):
    """翻译完成事件"""
    original_text: str   # 原始文本
    translated_text: str # 翻译后的文本
    source_lang: str     # 源语言
    target_lang: str     # 目标语言


# 第二步：定义工作流
class TranslationWorkflow(Workflow):
    """翻译工作流：接收文本 -> 翻译 -> 格式化输出"""

    @step
    async def translate_step(self, ev: StartEvent) -> TranslatedEvent:
        """
        翻译步骤：接收 StartEvent，执行翻译，返回 TranslatedEvent
        注意：实际项目中这里应该调用 LLM 或翻译 API
        """
        text = ev.get("text", "")
        source_lang = ev.get("source_lang", "英文")
        target_lang = ev.get("target_lang", "中文")

        # 模拟翻译逻辑（实际项目中替换为 LLM 调用）
        # 这里用一个简单的字典模拟翻译结果
        mock_translations = {
            "Hello World": "你好，世界",
            "Machine Learning": "机器学习",
            "Event-Driven Architecture": "事件驱动架构",
        }
        translated = mock_translations.get(text, f"[已翻译] {text}")

        print(f"[翻译步骤] 正在将 {source_lang} 翻译为 {target_lang}...")
        print(f"  原文: {text}")
        print(f"  译文: {translated}")

        return TranslatedEvent(
            original_text=text,
            translated_text=translated,
            source_lang=source_lang,
            target_lang=target_lang,
        )

    @step
    async def format_step(self, ev: TranslatedEvent) -> StopEvent:
        """
        格式化步骤：接收 TranslatedEvent，生成格式化输出，返回 StopEvent
        """
        formatted = (
            f"=== 翻译结果 ===\n"
            f"源语言: {ev.source_lang}\n"
            f"目标语言: {ev.target_lang}\n"
            f"原文: {ev.original_text}\n"
            f"译文: {ev.translated_text}\n"
            f"================"
        )

        print(f"[格式化步骤] 生成最终输出")
        return StopEvent(result=formatted)


# 第三步：运行工作流
async def main():
    # 创建工作流实例
    # timeout: 超时时间（秒），超时后工作流会抛出异常
    # verbose: 是否打印详细的运行日志
    workflow = TranslationWorkflow(timeout=60, verbose=False)

    # 运行工作流，传入的参数会成为 StartEvent 的属性
    result = await workflow.run(
        text="Event-Driven Architecture",
        source_lang="英文",
        target_lang="中文",
    )

    print("\n" + result)


if __name__ == "__main__":
    asyncio.run(main())

# 预期输出：
# [翻译步骤] 正在将 英文 翻译为 中文...
#   原文: Event-Driven Architecture
#   译文: 事件驱动架构
# [格式化步骤] 生成最终输出
#
# === 翻译结果 ===
# 源语言: 英文
# 目标语言: 中文
# 原文: Event-Driven Architecture
# 译文: 事件驱动架构
# ================

"""
多事件等待示例：等待所有数据源就绪后再合并
"""
import asyncio
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent, Context


class DataSourceReadyEvent(Event):
    """数据源就绪事件"""
    source_name: str
    data: dict

class AllSourcesReadyEvent(Event):
    """所有数据源就绪事件"""
    sources: list

class MergeWorkflow(Workflow):
    @step
    async def load_db(self, ev: StartEvent) -> DataSourceReadyEvent:
        """模拟加载数据库数据"""
        await asyncio.sleep(1)  # 模拟 I/O 延迟
        return DataSourceReadyEvent(
            source_name="database",
            data={"users": 1000, "orders": 5000},
        )

    @step
    async def load_api(self, ev: StartEvent) -> DataSourceReadyEvent:
        """模拟加载 API 数据"""
        await asyncio.sleep(2)  # 模拟更长的 I/O 延迟
        return DataSourceReadyEvent(
            source_name="external_api",
            data={"weather": "sunny", "temperature": 25},
        )

    @step
    async def collect_and_merge(self, ctx: Context, ev: DataSourceReadyEvent) -> StopEvent:
        """
        收集所有数据源的结果并合并。
        使用 Context 累积到达的事件，当所有数据源都就绪后才执行合并。
        """
        # 从 Context 中获取已收集的数据源列表
        collected = await ctx.get("collected_sources", default=[])
        collected.append({
            "source": ev.source_name,
            "data": ev.data,
        })
        await ctx.set("collected_sources", collected)

        # 检查是否已收集到所有数据源（本例中预期 2 个）
        expected_count = 2
        if len(collected) < expected_count:
            # 还没收齐，不返回事件，等待下一个 DataSourceReadyEvent
            return None

        # 所有数据源都已就绪，执行合并
        merged_data = {}
        for source in collected:
            merged_data[source["source"]] = source["data"]

        return StopEvent(result=f"合并完成：{merged_data}")


async def main():
    w = MergeWorkflow(timeout=30, verbose=False)
    result = await w.run()
    print(result)

asyncio.run(main())

from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent, Context


class TaskEvent(Event):
    """任务事件"""
    task_id: int
    task_name: str

class TaskCompleteEvent(Event):
    """任务完成事件"""
    task_id: int
    result: str


class ParallelTaskWorkflow(Workflow):
    @step
    async def dispatch_tasks(self, ctx: Context, ev: StartEvent) -> None:
        """
        分发任务：一次性发出多个 TaskEvent
        使用 ctx.send_event() 手动发送事件
        返回 None 表示不通过返回值发出事件
        """
        tasks = [
            {"task_id": 1, "task_name": "数据清洗"},
            {"task_id": 2, "task_name": "特征提取"},
            {"task_id": 3, "task_name": "模型训练"},
        ]

        for task in tasks:
            # 手动发送事件到工作流
            ctx.send_event(TaskEvent(**task))

        print(f"[分发步骤] 已发送 {len(tasks)} 个任务事件")
        # 返回 None，不产生返回事件

    @step
    async def process_task(self, ev: TaskEvent) -> TaskCompleteEvent:
        """处理单个任务"""
        print(f"[处理步骤] 正在执行任务 {ev.task_id}: {ev.task_name}")
        # 模拟任务处理
        await asyncio.sleep(0.5)
        return TaskCompleteEvent(
            task_id=ev.task_id,
            result=f"任务 '{ev.task_name}' 已完成",
        )

    @step
    async def collect_results(self, ctx: Context, ev: TaskCompleteEvent) -> StopEvent:
        """收集所有任务结果"""
        results = await ctx.get("results", default=[])
        results.append(ev.result)
        await ctx.set("results", results)

        if len(results) < 3:
            return None  # 等待更多结果

        return StopEvent(result=f"所有任务完成：\n" + "\n".join(results))

"""
人机协同示例：文档发布审批流程
流程：
  撰写文档 -> 需要人工审批 -> (暂停等待) -> 收到审批结果 -> 发布或退回修改
"""
import asyncio
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent,
    InputRequiredEvent, HumanResponseEvent, Context
)


class DraftReadyEvent(Event):
    """草稿完成事件"""
    draft_content: str
    author: str

class ApprovalResultEvent(Event):
    """审批结果事件"""
    approved: bool          # 是否批准
    reviewer_comment: str   # 审核意见


class DocumentPublishWorkflow(Workflow):
    """文档发布工作流（含人工审批）"""

    @step
    async def write_draft(self, ev: StartEvent) -> DraftReadyEvent:
        """第一步：撰写文档草稿"""
        topic = ev.get("topic", "LlamaIndex 工作流")
        draft = f"《{topic}》\n\n这是一篇关于 {topic} 的深度技术文章..."
        print(f"[撰写步骤] 草稿已完成")
        return DraftReadyEvent(draft=draft, author="AI 助手")

    @step
    async def request_approval(self, ev: DraftReadyEvent) -> InputRequiredEvent:
        """
        第二步：请求人工审批
        返回 InputRequiredEvent 后，工作流会自动暂停
        等待外部通过 HumanResponseEvent 提供审批结果
        """
        print(f"[审批步骤] 等待人工审核...")
        print(f"  文档作者: {ev.author}")
        print(f"  文档内容: {ev.draft_content[:50]}...")

        return InputRequiredEvent(
            prefix="请审核以下文档是否发布",
            # 你可以携带任意额外信息供审核者参考
        )

    @step
    async def handle_approval(self, ctx: Context, ev: HumanResponseEvent) -> StopEvent:
        """
        第三步：处理审批结果
        接收 HumanResponseEvent 后继续执行
        ev.response 包含人工输入的内容
        """
        response = ev.response.strip().lower()

        if response in ("批准", "通过", "yes", "approve"):
            return StopEvent(result="文档已批准并发布！")
        else:
            return StopEvent(result=f"文档被退回，审核意见：{ev.response}")


async def main():
    workflow = DocumentPublishWorkflow(timeout=300, verbose=False)

    # 启动工作流
    handler = workflow.run(topic="LlamaIndex 最佳实践")

    # 监听事件流，等待 InputRequiredEvent
    async for event in handler.stream_events():
        if isinstance(event, InputRequiredEvent):
            print(f"\n[系统] 收到审批请求: {event.prefix}")
            print("[系统] 等待用户输入...")

            # 模拟人工输入（实际项目中可以是 Web 表单、API 调用等）
            user_input = "批准"  # 实际中通过 input() 或 API 获取

            # 通过 handler 发送 HumanResponseEvent 恢复工作流
            handler.ctx.send_event(
                HumanResponseEvent(response=user_input)
            )

    # 获取最终结果
    result = await handler
    print(f"\n最终结果: {result}")


asyncio.run(main())

"""
逐步执行示例：用于调试和测试
"""
import asyncio
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent


class IntermediateEvent(Event):
    """中间事件"""
    data: str
    step_number: int


class DebugWorkflow(Workflow):
    @step
    async def step_a(self, ev: StartEvent) -> IntermediateEvent:
        print("[Step A] 执行中...")
        return IntermediateEvent(data="Step A 的输出", step_number=1)

    @step
    async def step_b(self, ev: IntermediateEvent) -> IntermediateEvent:
        print(f"[Step B] 收到: {ev.data}")
        return IntermediateEvent(data="Step B 的输出", step_number=2)

    @step
    async def step_c(self, ev: IntermediateEvent) -> StopEvent:
        print(f"[Step C] 收到: {ev.data}")
        return StopEvent(result=f"最终结果: {ev.data}")


async def main():
    workflow = DebugWorkflow(timeout=60, verbose=True)

    # ============================================================
    # 方式一：逐步执行（用于调试）
    # ============================================================
    # run_step() 每次只执行一个 Step，返回执行结果
    # 你可以检查每一步的输入和输出

    handler = workflow.run()  # 获取运行处理器

    # 逐步执行
    print("=== 逐步执行模式 ===")
    step_result_1 = await handler.run_step("step_a")
    print(f"Step A 完成: {step_result_1}")

    step_result_2 = await handler.run_step("step_b")
    print(f"Step B 完成: {step_result_2}")

    step_result_3 = await handler.run_step("step_c")
    print(f"Step C 完成: {step_result_3}")

    # 获取最终结果
    final_result = await handler
    print(f"最终结果: {final_result}")


asyncio.run(main())

"""
检查点示例：带状态持久化的长运行工作流
"""
import asyncio
import json
from pathlib import Path
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)
from llama_index.core.workflow.checkpoint import (
    WorkflowCheckpointer,
)


class ProcessingEvent(Event):
    """处理阶段事件"""
    stage: str       # 当前阶段
    progress: float  # 进度百分比
    data: dict       # 阶段数据


# ============================================================
# 配置检查点器
# ============================================================
# 方式一：使用内存检查点（适合开发/测试）
checkpointer = WorkflowCheckpointer(
    # 内存存储，进程结束后丢失
    checkpoint_store="memory"
)

# 方式二：使用文件检查点（适合单机持久化）
# checkpointer = WorkflowCheckpointer(
#     checkpoint_store="file",
#     checkpoint_dir="./checkpoints"
# )


class LongRunningWorkflow(Workflow):
    """长运行工作流：多阶段数据处理管道"""

    @step
    async def stage_extract(self, ctx: Context, ev: StartEvent) -> ProcessingEvent:
        """阶段一：数据提取"""
        print("[阶段一] 正在提取数据...")
        await asyncio.sleep(1)  # 模拟耗时操作

        # 保存阶段结果到 Context
        await ctx.set("extract_result", {"records": 10000})
        await ctx.set("current_stage", "extract")

        return ProcessingEvent(
            stage="extract",
            progress=33.3,
            data={"records": 10000},
        )

    @step
    async def stage_transform(self, ctx: Context, ev: ProcessingEvent) -> ProcessingEvent:
        """阶段二：数据转换"""
        print(f"[阶段二] 正在转换 {ev.data['records']} 条数据...")
        await asyncio.sleep(2)  # 模拟耗时操作

        extract_result = await ctx.get("extract_result")
        transformed_count = extract_result["records"] * 0.95  # 模拟 5% 数据过滤

        await ctx.set("transform_result", {"transformed": transformed_count})
        await ctx.set("current_stage", "transform")

        return ProcessingEvent(
            stage="transform",
            progress=66.6,
            data={"transformed": transformed_count},
        )

    @step
    async def stage_load(self, ctx: Context, ev: ProcessingEvent) -> StopEvent:
        """阶段三：数据加载"""
        print(f"[阶段三] 正在加载 {ev.data['transformed']:.0f} 条数据...")
        await asyncio.sleep(1)  # 模拟耗时操作

        await ctx.set("current_stage", "load")
        await ctx.set("completed", True)

        return StopEvent(
            result=f"ETL 管道完成！共处理 {ev.data['transformed']:.0f} 条数据"
        )


async def main():
    # 创建带检查点的工作流
    workflow = LongRunningWorkflow(
        timeout=300,
        verbose=True,
    )

    # 运行工作流（检查点会自动保存）
    result = await workflow.run(source="production_db")
    print(f"最终结果: {result}")


asyncio.run(main())

"""
Context 状态管理详解
"""
import asyncio
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent, Context


class AnalyzeEvent(Event):
    """分析事件"""
    text: str

class SummaryEvent(Event):
    """摘要事件"""
    summary: str


class ContextDemoWorkflow(Workflow):
    """Context 使用示例"""

    @step
    async def collect_data(self, ctx: Context, ev: StartEvent) -> AnalyzeEvent:
        """
        Context 写入操作：
        - ctx.set(key, value)：存储键值对
        - 支持任意 Python 对象（列表、字典、自定义对象等）
        - 状态在整个工作流运行期间持久存在
        """
        # 存储简单值
        await ctx.set("start_time", "2026-06-06T10:00:00")
        await ctx.set("request_id", "req-001")

        # 存储复杂数据结构
        await ctx.set("statistics", {
            "total_documents": 0,
            "processed_documents": 0,
            "error_count": 0,
        })

        # 存储列表
        await ctx.set("processing_log", [])

        text = ev.get("text", "需要分析的文本内容")
        print(f"[数据采集] 已初始化上下文状态")

        return AnalyzeEvent(text=text)

    @step
    async def analyze_data(self, ctx: Context, ev: AnalyzeEvent) -> SummaryEvent:
        """
        Context 读写操作：
        - ctx.get(key)：读取值（不存在时抛出异常）
        - ctx.get(key, default)：读取值（不存在时返回默认值）
        """
        # 读取状态
        stats = await ctx.get("statistics")
        log = await ctx.get("processing_log")

        # 更新状态（先读后写模式）
        stats["total_documents"] = 1
        stats["processed_documents"] = 1
        await ctx.set("statistics", stats)

        log.append(f"分析文本: {ev.text[:20]}...")
        await ctx.set("processing_log", log)

        summary = f"分析完成：文本长度 {len(ev.text)} 字符"
        return SummaryEvent(summary=summary)

    @step
    async def generate_report(self, ctx: Context, ev: SummaryEvent) -> StopEvent:
        """
        Context 最终汇总：读取所有状态，生成最终报告
        """
        stats = await ctx.get("statistics")
        log = await ctx.get("processing_log")
        start_time = await ctx.get("start_time")
        request_id = await ctx.get("request_id")

        report = (
            f"=== 执行报告 ===\n"
            f"请求 ID: {request_id}\n"
            f"开始时间: {start_time}\n"
            f"摘要: {ev.summary}\n"
            f"统计: {stats}\n"
            f"日志: {log}\n"
            f"================"
        )

        return StopEvent(result=report)


async def main():
    w = ContextDemoWorkflow(timeout=30, verbose=False)
    result = await w.run(text="这是一段用于演示 Context 状态管理的示例文本")
    print(result)


asyncio.run(main())

# 两次独立运行，各自有独立的 Context
result1 = await workflow.run(text="第一篇文档")  # Context 1
result2 = await workflow.run(text="第二篇文档")  # Context 2（完全独立）

"""
线性工作流：数据清洗 -> 特征提取 -> 模型推理
"""
import asyncio
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent


class CleanedDataEvent(Event):
    """清洗后的数据"""
    clean_text: str
    word_count: int

class FeatureEvent(Event):
    """特征提取结果"""
    features: dict
    clean_text: str


class LinearPipeline(Workflow):
    """线性处理管道"""

    @step
    async def clean_data(self, ev: StartEvent) -> CleanedDataEvent:
        """第一步：数据清洗"""
        raw_text = ev.get("raw_text", "")
        # 模拟数据清洗：去除多余空白、特殊字符
        clean = " ".join(raw_text.split())
        clean = clean.replace("\t", " ").replace("\n", " ")
        print(f"[清洗] 原始长度={len(raw_text)}, 清洗后长度={len(clean)}")
        return CleanedDataEvent(clean_text=clean, word_count=len(clean.split()))

    @step
    async def extract_features(self, ev: CleanedDataEvent) -> FeatureEvent:
        """第二步：特征提取"""
        text = ev.clean_text
        features = {
            "length": len(text),
            "word_count": ev.word_count,
            "has_numbers": any(c.isdigit() for c in text),
            "avg_word_length": sum(len(w) for w in text.split()) / max(ev.word_count, 1),
        }
        print(f"[特征] 提取完成: {features}")
        return FeatureEvent(features=features, clean_text=text)

    @step
    async def predict(self, ev: FeatureEvent) -> StopEvent:
        """第三步：模型推理（模拟）"""
        # 模拟基于特征的简单预测
        score = ev.features["avg_word_length"] * 0.3 + ev.features["word_count"] * 0.01
        label = "正面" if score > 3 else "中性" if score > 1 else "负面"
        result = f"预测结果: {label} (得分: {score:.2f})\n特征: {ev.features}"
        return StopEvent(result=result)


async def main():
    w = LinearPipeline(timeout=30, verbose=False)
    result = await w.run(raw_text="  这是一段   需要清洗的  文本数据！包含123数字。  ")
    print(result)


asyncio.run(main())

"""
分支工作流：文档分类后的差异化处理
根据文档类型（技术文档/商务邮件/其他）走不同的处理分支
"""
import asyncio
from typing import Union
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent


# 定义不同分支的事件
class TechDocEvent(Event):
    """技术文档事件"""
    content: str
    language: str

class BusinessEmailEvent(Event):
    """商务邮件事件"""
    content: str
    sender: str

class GeneralDocEvent(Event):
    """通用文档事件"""
    content: str


class BranchingWorkflow(Workflow):
    """文档分类与分支处理工作流"""

    @step
    async def classify_document(
        self, ev: StartEvent
    ) -> Union[TechDocEvent, BusinessEmailEvent, GeneralDocEvent]:
        """
        文档分类步骤：根据内容特征判断文档类型
        返回不同类型的事件触发不同的下游分支
        """
        content = ev.get("content", "")
        doc_type = ev.get("doc_type", "auto")

        # 简单规则分类（实际项目中可用 LLM 分类）
        if doc_type == "tech" or any(
            kw in content for kw in ["API", "代码", "函数", "class", "def", "import"]
        ):
            print("[分类] 判定为技术文档")
            return TechDocEvent(content=content, language="Python")

        elif doc_type == "email" or any(
            kw in content for kw in ["尊敬的", "此致", "敬礼", "Dear", "Best regards"]
        ):
            print("[分类] 判定为商务邮件")
            return BusinessEmailEvent(content=content, sender="unknown")

        else:
            print("[分类] 判定为通用文档")
            return GeneralDocEvent(content=content)

    # ============ 技术文档分支 ============
    @step
    async def process_tech_doc(self, ev: TechDocEvent) -> StopEvent:
        """技术文档处理：提取代码块、生成文档注释"""
        code_blocks = [line for line in ev.content.split("\n")
                       if line.strip().startswith(("def ", "class ", "import "))]
        summary = (
            f"[技术文档处理报告]\n"
            f"编程语言: {ev.language}\n"
            f"发现代码定义: {len(code_blocks)} 处\n"
            f"代码行: {code_blocks[:3]}"  # 只显示前 3 个
        )
        return StopEvent(result=summary)

    # ============ 商务邮件分支 ============
    @step
    async def process_email(self, ev: BusinessEmailEvent) -> StopEvent:
        """商务邮件处理：提取关键信息、生成回复建议"""
        has_greeting = any(kw in ev.content for kw in ["尊敬的", "Dear"])
        has_closing = any(kw in ev.content for kw in ["此致", "Best regards"])
        summary = (
            f"[商务邮件处理报告]\n"
            f"发件人: {ev.sender}\n"
            f"包含称呼: {has_greeting}\n"
            f"包含落款: {has_closing}\n"
            f"建议: 已分析邮件结构，可自动生成回复"
        )
        return StopEvent(result=summary)

    # ============ 通用文档分支 ============
    @step
    async def process_general(self, ev: GeneralDocEvent) -> StopEvent:
        """通用文档处理：生成摘要、提取关键词"""
        words = ev.content.split()
        summary = (
            f"[通用文档处理报告]\n"
            f"文档长度: {len(words)} 词\n"
            f"前 20 字: {ev.content[:20]}...\n"
            f"建议: 使用 LLM 生成详细摘要"
        )
        return StopEvent(result=summary)


async def main():
    w = BranchingWorkflow(timeout=30, verbose=False)

    # 测试不同分支
    print("=== 测试技术文档分支 ===")
    result1 = await w.run(content="import os\ndef main():\n    print('Hello')\nclass App:\n    pass")
    print(result1)

    print("\n=== 测试商务邮件分支 ===")
    result2 = await w.run(content="尊敬的张总，\n\n感谢您的合作。\n\n此致敬礼")
    print(result2)

    print("\n=== 测试通用文档分支 ===")
    result3 = await w.run(content="这是一篇普通的说明文档，介绍产品的基本功能和使用方法。")
    print(result3)


asyncio.run(main())

"""
循环工作流：LLM 自纠错循环
流程：生成内容 -> 评估质量 -> 不合格则带着反馈重新生成 -> 合格则输出
"""
import asyncio
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)
# from llama_index.llms.openai import OpenAI


class GenerateEvent(Event):
    """生成请求事件"""
    prompt: str           # 生成提示
    feedback: str = ""    # 上轮评估反馈（首次为空）
    attempt: int = 1      # 当前尝试次数

class EvaluateEvent(Event):
    """评估事件"""
    generated_text: str   # 待评估的文本
    attempt: int          # 当前尝试次数

class RetryEvent(Event):
    """重试事件（评估不合格时触发）"""
    original_prompt: str  # 原始提示
    feedback: str         # 评估反馈
    attempt: int          # 已尝试次数


class SelfCorrectionWorkflow(Workflow):
    """LLM 自纠错工作流"""

    # 配置参数
    max_retries: int = 3           # 最大重试次数
    quality_threshold: float = 0.7 # 质量阈值（0-1）

    @step
    async def generate(self, ctx: Context, ev: GenerateEvent) -> EvaluateEvent:
        """
        生成步骤：根据提示和反馈生成内容
        如果有上轮反馈，会在提示中加入改进要求
        """
        prompt = ev.prompt
        if ev.feedback:
            prompt = (
                f"{ev.prompt}\n\n"
                f"[上次生成结果的反馈]：{ev.feedback}\n"
                f"请根据以上反馈改进你的输出。"
            )

        # 实际项目中替换为 LLM 调用：
        # llm = OpenAI(model="gpt-4.1")
        # response = await llm.acomplete(prompt)
        # generated_text = str(response)

        # 模拟生成（演示循环逻辑）
        quality_by_attempt = {1: "较差", 2: "一般", 3: "优秀"}
        quality = quality_by_attempt.get(ev.attempt, "未知")
        generated_text = f"第 {ev.attempt} 次生成的内容（质量：{quality}）"

        print(f"[生成步骤] 第 {ev.attempt} 次尝试")
        print(f"  提示: {prompt[:60]}...")

        # 记录到 Context
        await ctx.set(f"attempt_{ev.attempt}", generated_text)

        return EvaluateEvent(generated_text=generated_text, attempt=ev.attempt)

    @step
    async def evaluate(self, ctx: Context, ev: EvaluateEvent) -> StopEvent | GenerateEvent:
        """
        评估步骤：判断生成质量是否达标
        - 达标 -> 返回 StopEvent（结束工作流）
        - 不达标 -> 返回 GenerateEvent（触发重新生成，形成循环）
        """
        # 模拟质量评分（实际项目中可用另一个 LLM 评估）
        quality_scores = {1: 0.3, 2: 0.6, 3: 0.85}
        score = quality_scores.get(ev.attempt, 0.5)

        print(f"[评估步骤] 质量得分: {score:.2f} (阈值: {self.quality_threshold})")

        if score >= self.quality_threshold:
            # 质量达标，结束工作流
            print(f"[评估步骤] 质量达标！")
            return StopEvent(
                result=f"最终输出（第 {ev.attempt} 次尝试，得分 {score:.2f}）:\n{ev.generated_text}"
            )

        if ev.attempt >= self.max_retries:
            # 超过最大重试次数
            print(f"[评估步骤] 已达最大重试次数，返回最佳结果")
            return StopEvent(
                result=f"达到最大重试次数，返回最后结果:\n{ev.generated_text}"
            )

        # 质量不达标，发出重新生成事件（形成循环！）
        feedback = f"质量得分 {score:.2f}，低于阈值 {self.quality_threshold}。需要更详细、更准确的内容。"
        print(f"[评估步骤] 质量不达标，请求重新生成")
        return GenerateEvent(
            prompt=await ctx.get("original_prompt", "请生成高质量的技术文章"),
            feedback=feedback,
            attempt=ev.attempt + 1,
        )

    @step
    async def init_workflow(self, ctx: Context, ev: StartEvent) -> GenerateEvent:
        """初始化：保存原始提示，发出首次生成事件"""
        prompt = ev.get("prompt", "请生成一篇关于 LlamaIndex 的技术文章")
        await ctx.set("original_prompt", prompt)
        return GenerateEvent(prompt=prompt, feedback="", attempt=1)


async def main():
    w = SelfCorrectionWorkflow(timeout=60, verbose=False)
    result = await w.run(prompt="请写一篇关于 LlamaIndex Workflow 的深度教程")
    print(f"\n{'='*50}")
    print(result)


asyncio.run(main())

# 预期输出：
# [初始化] 保存原始提示
# [生成步骤] 第 1 次尝试
# [评估步骤] 质量得分: 0.30 (阈值: 0.7)
# [评估步骤] 质量不达标，请求重新生成
# [生成步骤] 第 2 次尝试
# [评估步骤] 质量得分: 0.60 (阈值: 0.7)
# [评估步骤] 质量不达标，请求重新生成
# [生成步骤] 第 3 次尝试
# [评估步骤] 质量得分: 0.85 (阈值: 0.7)
# [评估步骤] 质量达标！

"""
并发工作流：多源数据并行采集与汇总
流程：
  分发采集任务 -> 多个采集器并行运行 -> 汇总所有结果
"""
import asyncio
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)


class CollectTaskEvent(Event):
    """采集任务事件"""
    source_name: str
    source_url: str

class CollectResultEvent(Event):
    """采集结果事件"""
    source_name: str
    data: list
    record_count: int


class ParallelCollectWorkflow(Workflow):
    """多源并行采集工作流"""

    # 预期的数据源数量（用于判断是否全部完成）
    expected_sources: int = 3

    @step
    async def dispatch(self, ctx: Context, ev: StartEvent) -> None:
        """
        分发步骤：向多个数据源发出采集任务
        使用 ctx.send_event() 发出多个事件，触发并行执行
        返回 None 表示不通过返回值发出事件
        """
        sources = [
            {"source_name": "新闻网站", "source_url": "https://news.example.com"},
            {"source_name": "技术博客", "source_url": "https://blog.example.com"},
            {"source_name": "学术论文", "source_url": "https://papers.example.com"},
        ]

        await ctx.set("expected_count", len(sources))

        for source in sources:
            # 每个 send_event 都会触发一个 process_source Step 并行执行
            ctx.send_event(CollectTaskEvent(**source))

        print(f"[分发] 已向 {len(sources)} 个数据源发出采集任务")

    @step(num_workers=3)
    async def process_source(self, ev: CollectTaskEvent) -> CollectResultEvent:
        """
        采集步骤：并行采集各数据源
        num_workers=3 表示最多同时运行 3 个此 Step 的实例
        """
        print(f"[采集] 正在从 {ev.source_name} ({ev.source_url}) 采集数据...")

        # 模拟不同数据源的采集延迟
        delays = {"新闻网站": 1, "技术博客": 2, "学术论文": 3}
        await asyncio.sleep(delays.get(ev.source_name, 1))

        # 模拟采集结果
        mock_data = {
            "新闻网站": ["AI 新突破", "量子计算进展", "太空探索"],
            "技术博客": ["LlamaIndex 教程", "RAG 最佳实践", "Agent 设计模式"],
            "学术论文": ["Attention Is All You Need", "BERT 论文", "GPT 技术报告"],
        }
        data = mock_data.get(ev.source_name, [])

        print(f"[采集] {ev.source_name} 采集完成，获取 {len(data)} 条记录")
        return CollectResultEvent(
            source_name=ev.source_name,
            data=data,
            record_count=len(data),
        )

    @step
    async def aggregate(self, ctx: Context, ev: CollectResultEvent) -> StopEvent:
        """
        汇总步骤：等待所有采集结果到达后汇总
        """
        # 累积已收到的结果
        results = await ctx.get("all_results", default=[])
        results.append({
            "source": ev.source_name,
            "data": ev.data,
            "count": ev.record_count,
        })
        await ctx.set("all_results", results)

        expected = await ctx.get("expected_count", 3)

        if len(results) < expected:
            # 还没收齐所有结果，等待下一个 CollectResultEvent
            print(f"[汇总] 已收到 {len(results)}/{expected} 个数据源的结果")
            return None  # 返回 None，不发出事件，继续等待

        # 所有数据源采集完成，汇总
        print(f"[汇总] 所有 {expected} 个数据源采集完成！")
        total_records = sum(r["count"] for r in results)
        all_data = []
        for r in results:
            all_data.extend(r["data"])

        summary = (
            f"=== 采集汇总报告 ===\n"
            f"数据源数量: {len(results)}\n"
            f"总记录数: {total_records}\n"
            f"所有数据:\n"
        )
        for r in results:
            summary += f"  [{r['source']}] {r['data']}\n"

        return StopEvent(result=summary)


async def main():
    w = ParallelCollectWorkflow(timeout=60, verbose=False)
    result = await w.run()
    print(result)


asyncio.run(main())

"""
嵌套工作流：主工作流调用子工作流
场景：
  主工作流负责整体协调
  子工作流负责具体的 RAG 检索和生成任务
"""
import asyncio
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)


# ============================================================
# 子工作流：RAG 检索与生成
# ============================================================
class RetrievalEvent(Event):
    """检索结果事件"""
    documents: list

class RAGSubWorkflow(Workflow):
    """RAG 子工作流：检索 -> 生成"""

    @step
    async def retrieve(self, ev: StartEvent) -> RetrievalEvent:
        """检索步骤"""
        query = ev.get("query", "")
        # 模拟检索（实际项目中使用 VectorStoreIndex 等）
        mock_docs = [
            f"文档1: 关于 {query} 的基础知识",
            f"文档2: {query} 的高级应用",
            f"文档3: {query} 的最佳实践",
        ]
        print(f"  [子工作流-检索] 查询 '{query}'，找到 {len(mock_docs)} 个文档")
        return RetrievalEvent(documents=mock_docs)

    @step
    async def generate(self, ev: RetrievalEvent) -> StopEvent:
        """生成步骤"""
        # 模拟 LLM 生成
        context = "\n".join(ev.documents)
        answer = f"基于检索到的 {len(ev.documents)} 个文档，生成的回答如下..."
        print(f"  [子工作流-生成] 回答已生成")
        return StopEvent(result=answer)


# ============================================================
# 主工作流：协调多个子任务
# ============================================================
class SubTaskEvent(Event):
    """子任务事件"""
    task_type: str
    query: str

class SubTaskResultEvent(Event):
    """子任务结果事件"""
    task_type: str
    result: str


class MainWorkflow(Workflow):
    """主工作流：协调多个子工作流"""

    def __init__(self, **kwargs):
        super().__init__(**kwargs)
        # 在主工作流中实例化子工作流
        self.rag_workflow = RAGSubWorkflow(timeout=30, verbose=False)

    @step
    async def plan_tasks(self, ctx: Context, ev: StartEvent) -> SubTaskEvent:
        """
        规划步骤：分解用户请求为多个子任务
        返回多个 SubTaskEvent 触发并行处理
        """
        user_query = ev.get("query", "")
        print(f"[主工作流-规划] 分解查询: {user_query}")

        # 发送多个子任务事件
        ctx.send_event(SubTaskEvent(task_type="rag", query=f"{user_query} 技术概述"))
        ctx.send_event(SubTaskEvent(task_type="rag", query=f"{user_query} 应用案例"))

        await ctx.set("expected_tasks", 2)
        # 第一个通过返回值发出，第二个通过 send_event 发出
        return SubTaskEvent(task_type="rag", query=f"{user_query} 未来趋势")

    @step
    async def execute_subtask(self, ev: SubTaskEvent) -> SubTaskResultEvent:
        """
        执行子任务：调用子工作流完成具体任务
        """
        print(f"[主工作流-执行] 开始子任务: {ev.task_type} - {ev.query}")

        # 调用子工作流
        sub_result = await self.rag_workflow.run(query=ev.query)

        return SubTaskResultEvent(task_type=ev.task_type, result=str(sub_result))

    @step
    async def synthesize(self, ctx: Context, ev: SubTaskResultEvent) -> StopEvent:
        """汇总所有子任务结果"""
        results = await ctx.get("sub_results", default=[])
        results.append({"type": ev.task_type, "result": ev.result})
        await ctx.set("sub_results", results)

        expected = await ctx.get("expected_tasks", 3)
        if len(results) < expected:
            return None

        # 汇总所有结果
        final = "=== 综合研究报告 ===\n"
        for i, r in enumerate(results, 1):
            final += f"\n[部分 {i}] ({r['type']})\n{r['result']}\n"

        return StopEvent(result=final)


async def main():
    w = MainWorkflow(timeout=60, verbose=False)
    result = await w.run(query="LlamaIndex Workflow")
    print(result)


asyncio.run(main())

"""
条件循环工作流：迭代式文档优化
不断修改文档直到满足所有质量标准
"""
import asyncio
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)


class RevisionEvent(Event):
    """修订事件"""
    document: str
    issues: list       # 当前存在的问题
    revision_number: int


class QualityCheckEvent(Event):
    """质量检查事件"""
    document: str
    revision_number: int
    score: float
    remaining_issues: list


class IterativeRefinementWorkflow(Workflow):
    """迭代优化工作流"""

    max_iterations: int = 5
    target_score: float = 0.9

    @step
    async def init_revision(self, ctx: Context, ev: StartEvent) -> RevisionEvent:
        """初始化修订"""
        document = ev.get("document", "这是一份初稿文档。")
        await ctx.set("original_document", document)
        return RevisionEvent(
            document=document,
            issues=["缺少标题", "内容过短", "缺乏数据支撑"],
            revision_number=0,
        )

    @step
    async def revise_document(self, ev: RevisionEvent) -> QualityCheckEvent:
        """修订文档步骤"""
        doc = ev.document
        issues = ev.issues

        # 模拟根据问题修订文档
        improvements = []
        if "缺少标题" in issues:
            doc = "# 完整标题\n\n" + doc
            improvements.append("已添加标题")
        if "内容过短" in issues:
            doc += "\n\n补充内容：更多详细信息和分析..."
            improvements.append("已扩充内容")
        if "缺乏数据支撑" in issues:
            doc += "\n\n数据：根据统计，效率提升 40%..."
            improvements.append("已添加数据")

        print(f"[修订 #{ev.revision_number + 1}] 改进: {improvements}")

        # 模拟质量评分（每次递增）
        score = min(0.5 + ev.revision_number * 0.2, 1.0)
        remaining = issues[1:] if issues else []  # 每次解决一个问题

        return QualityCheckEvent(
            document=doc,
            revision_number=ev.revision_number + 1,
            score=score,
            remaining_issues=remaining,
        )

    @step
    async def check_quality(
        self, ev: QualityCheckEvent
    ) -> StopEvent | RevisionEvent:
        """
        质量检查步骤：决定是否继续循环
        """
        print(
            f"[质量检查] 第 {ev.revision_number} 版，"
            f"得分: {ev.score:.2f}，"
            f"剩余问题: {len(ev.remaining_issues)}"
        )

        # 终止条件 1：分数达标
        if ev.score >= self.target_score:
            return StopEvent(
                result=f"文档优化完成！最终得分: {ev.score:.2f}\n{ev.document}"
            )

        # 终止条件 2：达到最大迭代次数
        if ev.revision_number >= self.max_iterations:
            return StopEvent(
                result=f"达到最大迭代次数 ({self.max_iterations})，"
                       f"最终得分: {ev.score:.2f}\n{ev.document}"
            )

        # 继续循环：发出新的 RevisionEvent
        return RevisionEvent(
            document=ev.document,
            issues=ev.remaining_issues,
            revision_number=ev.revision_number,
        )


async def main():
    w = IterativeRefinementWorkflow(timeout=60, verbose=False)
    result = await w.run(document="这是一份简单的初稿。")
    print(f"\n{'='*50}")
    print(result)


asyncio.run(main())

"""
复杂状态管理模式示例
"""
from llama_index.core.workflow import Context


async def complex_state_example(ctx: Context):
    # ============================================================
    # 模式 1：列表追加（先读-修改-写回）
    # ============================================================
    logs = await ctx.get("processing_logs", default=[])
    logs.append({"step": "analyze", "status": "success", "timestamp": "2026-06-06T10:00:00"})
    await ctx.set("processing_logs", logs)

    # ============================================================
    # 模式 2：字典更新（部分更新）
    # ============================================================
    config = await ctx.get("config", default={})
    config.update({
        "max_tokens": 4096,
        "temperature": 0.7,
        "model": "gpt-4.1",
    })
    await ctx.set("config", config)

    # ============================================================
    # 模式 3：计数器（原子递增）
    # ============================================================
    count = await ctx.get("retry_count", default=0)
    await ctx.set("retry_count", count + 1)

    # ============================================================
    # 模式 4：嵌套结构
    # ============================================================
    pipeline_state = {
        "stages": {
            "extraction": {"status": "completed", "records": 1000},
            "transformation": {"status": "running", "records": 0},
            "loading": {"status": "pending", "records": 0},
        },
        "errors": [],
        "warnings": [],
    }
    await ctx.set("pipeline_state", pipeline_state)

"""
流式事件处理示例：实时输出 RAG 工作流的中间结果
"""
import asyncio
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)


class StreamTokenEvent(Event):
    """流式 Token 事件（用于实时输出）"""
    token: str
    is_complete: bool = False

class ProgressEvent(Event):
    """进度事件"""
    stage: str
    progress: float  # 0.0 - 1.0
    message: str


class StreamingRAGWorkflow(Workflow):
    """支持流式输出的 RAG 工作流"""

    @step
    async def retrieve(self, ctx: Context, ev: StartEvent) -> ProgressEvent:
        """检索步骤"""
        query = ev.get("query", "")
        print(f"[检索] 查询: {query}")
        await asyncio.sleep(0.5)  # 模拟检索延迟

        # 保存检索结果到 Context
        await ctx.set("retrieved_docs", [
            "文档1: LlamaIndex 是一个数据框架...",
            "文档2: 工作流使用事件驱动模型...",
            "文档3: Context 提供状态管理...",
        ])

        return ProgressEvent(stage="retrieve", progress=0.3, message="检索完成")

    @step
    async def generate_stream(self, ctx: Context, ev: ProgressEvent) -> StopEvent:
        """
        流式生成步骤：模拟逐 Token 生成
        通过 ctx.send_event() 发送中间事件
        """
        docs = await ctx.get("retrieved_docs", [])
        ctx.send_event(ProgressEvent(
            stage="generate", progress=0.5, message="开始生成回答..."
        ))

        # 模拟流式生成（实际项目中使用 LLM streaming API）
        tokens = ["根据", "检索", "到的", "文档，", "LlamaIndex", " 工作流",
                  "的核心", "特点", "包括：", "事件", "驱动、", "状态", "管理",
                  "和", "并发", "执行。"]

        full_response = ""
        for i, token in enumerate(tokens):
            full_response += token
            # 发送流式 Token 事件
            ctx.send_event(StreamTokenEvent(
                token=token,
                is_complete=(i == len(tokens) - 1),
            ))
            await asyncio.sleep(0.05)  # 模拟生成延迟

        return StopEvent(result=full_response)


async def main():
    workflow = StreamingRAGWorkflow(timeout=30, verbose=False)

    # 使用 stream_events() 获取实时事件流
    handler = workflow.run(query="LlamaIndex 工作流有什么特点？")

    print("=== 实时事件流 ===")
    async for event in handler.stream_events():
        if isinstance(event, ProgressEvent):
            progress_bar = "█" * int(event.progress * 20) + "░" * (20 - int(event.progress * 20))
            print(f"[{event.stage}] {progress_bar} {event.progress*100:.0f}% - {event.message}")
        elif isinstance(event, StreamTokenEvent):
            # 实时输出 Token（不换行）
            print(event.token, end="", flush=True)
            if event.is_complete:
                print()  # 生成完成后换行

    # 获取最终结果
    result = await handler
    print(f"\n最终结果: {result}")


asyncio.run(main())

"""
错误处理与重试示例
"""
import asyncio
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)


class ProcessingEvent(Event):
    """处理事件"""
    data: str
    attempt: int = 1

class ErrorEvent(Event):
    """错误事件"""
    error_type: str
    error_message: str
    original_data: str
    attempt: int


class RobustWorkflow(Workflow):
    """带错误处理和重试的工作流"""

    max_retries: int = 3

    @step
    async def process_data(self, ev: StartEvent) -> ProcessingEvent:
        """数据处理步骤"""
        data = ev.get("data", "")
        return ProcessingEvent(data=data, attempt=1)

    @step
    async def risky_operation(
        self, ctx: Context, ev: ProcessingEvent
    ) -> StopEvent | ErrorEvent:
        """
        风险操作步骤：可能失败的操作
        使用 try/except 捕获异常，返回 ErrorEvent
        """
        try:
            # 模拟可能失败的操作
            if ev.attempt <= 2:
                raise ConnectionError(f"网络连接失败（第 {ev.attempt} 次尝试）")

            # 成功路径
            result = f"数据处理成功: {ev.data}"
            print(f"[处理] {result}")
            return StopEvent(result=result)

        except ConnectionError as e:
            print(f"[错误] {e}")
            return ErrorEvent(
                error_type="ConnectionError",
                error_message=str(e),
                original_data=ev.data,
                attempt=ev.attempt,
            )
        except Exception as e:
            print(f"[未知错误] {e}")
            return ErrorEvent(
                error_type=type(e).__name__,
                error_message=str(e),
                original_data=ev.data,
                attempt=ev.attempt,
            )

    @step
    async def handle_error(
        self, ev: ErrorEvent
    ) -> ProcessingEvent | StopEvent:
        """
        错误处理步骤：决定重试还是终止
        """
        if ev.attempt >= self.max_retries:
            return StopEvent(
                result=f"操作失败，已达最大重试次数 ({self.max_retries})。"
                       f"最后错误: {ev.error_message}"
            )

        # 重试：增加等待时间（指数退避）
        wait_time = 2 ** ev.attempt * 0.1
        print(f"[重试] 等待 {wait_time:.1f}s 后重试（第 {ev.attempt + 1} 次）...")
        await asyncio.sleep(wait_time)

        return ProcessingEvent(data=ev.original_data, attempt=ev.attempt + 1)


async def main():
    w = RobustWorkflow(timeout=60, verbose=False)
    result = await w.run(data="重要数据")
    print(f"最终结果: {result}")


asyncio.run(main())

"""
可复用的 RAG 工作流模板
支持自定义 LLM、向量存储、提示模板等
"""
import asyncio
from typing import Optional, Any
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context
)


# ============================================================
# 定义通用事件
# ============================================================
class QueryEvent(Event):
    """查询事件"""
    query: str
    filters: dict = {}

class RetrievedDocsEvent(Event):
    """检索结果事件"""
    documents: list
    query: str

class RerankedDocsEvent(Event):
    """重排后文档事件"""
    documents: list
    query: str

class GeneratedAnswerEvent(Event):
    """生成回答事件"""
    answer: str
    source_documents: list
    query: str


# ============================================================
# 可复用的 RAG 工作流模板
# ============================================================
class BaseRAGWorkflow(Workflow):
    """
    基础 RAG 工作流模板
    支持自定义检索器、重排器、LLM、提示模板
    """

    def __init__(
        self,
        retriever: Any = None,          # 检索器实例
        reranker: Any = None,           # 重排器实例（可选）
        llm: Any = None,                # LLM 实例
        system_prompt: str = "",        # 系统提示
        top_k: int = 5,                 # 检索文档数
        rerank_top_n: int = 3,          # 重排后保留文档数
        enable_reranking: bool = False, # 是否启用重排
        **kwargs,
    ):
        super().__init__(**kwargs)
        self.retriever = retriever
        self.reranker = reranker
        self.llm = llm
        self.system_prompt = system_prompt
        self.top_k = top_k
        self.rerank_top_n = rerank_top_n
        self.enable_reranking = enable_reranking

    @step
    async def retrieve(self, ctx: Context, ev: StartEvent) -> RetrievedDocsEvent:
        """检索步骤"""
        query = ev.get("query", "")
        print(f"[RAG-检索] 查询: {query}, top_k={self.top_k}")

        if self.retriever:
            # 使用真实检索器
            documents = await self.retriever.aretrieve(query)
        else:
            # 模拟检索
            documents = [
                {"content": f"文档 {i}: 关于 {query} 的信息...", "score": 0.9 - i * 0.1}
                for i in range(1, self.top_k + 1)
            ]

        # 保存查询到 Context
        await ctx.set("query", query)

        return RetrievedDocsEvent(documents=documents, query=query)

    @step
    async def maybe_rerank(self, ev: RetrievedDocsEvent) -> RerankedDocsEvent:
        """可选的重排步骤"""
        if self.enable_reranking and self.reranker:
            print(f"[RAG-重排] 重排 {len(ev.documents)} 个文档...")
            # 使用真实重排器
            reranked = await self.reranker.aretrieve(ev.documents, ev.query)
            documents = reranked[:self.rerank_top_n]
        else:
            # 跳过重排，直接使用检索结果
            documents = ev.documents[:self.rerank_top_n] if self.rerank_top_n else ev.documents

        return RerankedDocsEvent(documents=documents, query=ev.query)

    @step
    async def generate(self, ev: RerankedDocsEvent) -> StopEvent:
        """生成步骤"""
        print(f"[RAG-生成] 基于 {len(ev.documents)} 个文档生成回答...")

        # 构建上下文
        context = "\n---\n".join(
            doc.get("content", str(doc)) if isinstance(doc, dict) else str(doc)
            for doc in ev.documents
        )

        if self.llm:
            # 使用真实 LLM
            prompt = f"{self.system_prompt}\n\n上下文:\n{context}\n\n问题: {ev.query}\n\n回答:"
            response = await self.llm.acomplete(prompt)
            answer = str(response)
        else:
            # 模拟生成
            answer = f"基于 {len(ev.documents)} 个检索文档，针对 '{ev.query}' 的回答如下..."

        return StopEvent(result=answer)


# ============================================================
# 使用示例：基于模板创建具体的 RAG 工作流
# ============================================================
class TechDocRAGWorkflow(BaseRAGWorkflow):
    """技术文档 RAG 工作流（继承模板，添加自定义逻辑）"""

    def __init__(self, **kwargs):
        # 设置技术文档专用的默认参数
        kwargs.setdefault("system_prompt", "你是一个技术文档助手，请基于提供的文档准确回答问题。")
        kwargs.setdefault("top_k", 10)
        kwargs.setdefault("enable_reranking", True)
        kwargs.setdefault("rerank_top_n", 5)
        super().__init__(**kwargs)


async def main():
    # 方式一：使用基础模板（模拟模式）
    print("=== 基础 RAG（模拟模式）===")
    basic_rag = BaseRAGWorkflow(timeout=30, verbose=False)
    result = await basic_rag.run(query="什么是事件驱动架构？")
    print(f"结果: {result}\n")

    # 方式二：使用专用子类
    print("=== 技术文档 RAG ===")
    tech_rag = TechDocRAGWorkflow(timeout=30, verbose=False)
    result = await tech_rag.run(query="LlamaIndex Workflow 的核心特性")
    print(f"结果: {result}")


asyncio.run(main())

"""
工作流可视化工具详解
"""
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent
from llama_index.utils.workflow import (
    draw_all_possible_flows,     # 绘制所有可能的流程图
    draw_most_likely_execution,  # 绘制最可能的执行路径
)


# 定义一个有分支和循环的工作流用于演示
class ReviewEvent(Event):
    content: str
    score: float

class ReviseEvent(Event):
    content: str
    feedback: str
    attempt: int


class DemoWorkflow(Workflow):
    @step
    async def generate(self, ev: StartEvent) -> ReviewEvent:
        """生成内容"""
        return ReviewEvent(content="初始内容", score=0.5)

    @step
    async def review(self, ev: ReviewEvent) -> StopEvent | ReviseEvent:
        """评审内容"""
        if ev.score >= 0.8:
            return StopEvent(result=ev.content)
        return ReviseEvent(
            content=ev.content,
            feedback="需要改进",
            attempt=1,
        )

    @step
    async def revise(self, ev: ReviseEvent) -> ReviewEvent:
        """修订内容"""
        return ReviewEvent(content=f"修订版: {ev.content}", score=0.85)


# ============================================================
# 可视化方法 1：绘制所有可能的流程图
# ============================================================
# 这会生成一个包含所有可能执行路径的流程图
# 输出格式支持：PNG、PDF、SVG 等

draw_all_possible_flows(
    DemoWorkflow,
    filename="workflow_all_flows.html",  # 输出文件名
)

# 也可以绘制到 Streamlit 应用
# draw_to_streamlit(DemoWorkflow)

print("流程图已生成: workflow_all_flows.html")


# ============================================================
# 工作流验证：检查连通性和潜在问题
# ============================================================
workflow = DemoWorkflow(timeout=30, verbose=False)

# validate() 检查：
# 1. 所有事件类型是否都有对应的接收 Step
# 2. 是否存在孤立 Step（无法到达的 Step）
# 3. 是否存在无法终止的路径（无限循环风险）
# 4. StartEvent 和 StopEvent 是否正确连接
try:
    # 验证工作流结构
    workflow.validate()
    print("工作流验证通过！")
except Exception as e:
    print(f"工作流验证失败: {e}")

"""
工作流调试最佳实践
"""
import logging
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent, Context

# ============================================================
# 技巧 1：启用详细日志
# ============================================================
logging.basicConfig(level=logging.DEBUG)

# 创建工作流时启用 verbose 模式
workflow = DemoWorkflow(timeout=60, verbose=True)

# ============================================================
# 技巧 2：在 Step 中添加调试输出
# ============================================================
class DebugWorkflow(Workflow):
    @step
    async def debug_step(self, ctx: Context, ev: StartEvent) -> StopEvent:
        # 打印事件内容
        print(f"[DEBUG] 收到事件: {ev.model_dump()}")

        # 打印当前 Context 状态
        # 注意：Context 的所有键可以通过 ctx 的内部属性获取
        print(f"[DEBUG] Context 状态: 已设置")

        # 使用 assert 进行断言检查
        topic = ev.get("topic", "")
        assert len(topic) > 0, "topic 不能为空"

        return StopEvent(result="调试完成")

# ============================================================
# 技巧 3：使用逐步执行模式
# ============================================================
async def debug_with_steps():
    w = DebugWorkflow(timeout=60, verbose=True)
    handler = w.run(topic="调试主题")

    # 一次执行一个 Step，检查每步的输出
    result = await handler.run_step("debug_step")
    print(f"Step 结果: {result}")

    final = await handler
    print(f"最终结果: {final}")

# asyncio.run(debug_with_steps())

# ============================================================
# Dockerfile：LlamaIndex Workflow 应用容器化
# ============================================================
FROM python:3.12-slim

# 设置工作目录
WORKDIR /app

# 安装系统依赖
RUN apt-get update && apt-get install -y \
    build-essential \
    && rm -rf /var/lib/apt/lists/*

# 复制依赖文件
COPY requirements.txt .

# 安装 Python 依赖
RUN pip install --no-cache-dir -r requirements.txt

# 复制应用代码
COPY . .

# 设置环境变量
ENV PYTHONUNBUFFERED=1
ENV OPENAI_API_KEY=""

# 暴露端口
EXPOSE 8000

# 启动命令
CMD ["uvicorn", "app:app", "--host", "0.0.0.0", "--port", "8000"]

# ============================================================
# docker-compose.yml：多服务编排
# ============================================================
version: "3.9"

services:
  # LlamaIndex Workflow API 服务
  llamaindex-api:
    build: .
    ports:
      - "8000:8000"
    environment:
      - OPENAI_API_KEY=${OPENAI_API_KEY}
      - VECTOR_STORE_URL=http://qdrant:6333
    depends_on:
      - qdrant
      - redis
    volumes:
      - ./data:/app/data
    restart: unless-stopped

  # Qdrant 向量数据库
  qdrant:
    image: qdrant/qdrant:latest
    ports:
      - "6333:6333"
    volumes:
      - qdrant_data:/qdrant/storage
    restart: unless-stopped

  # Redis 缓存与会话管理
  redis:
    image: redis:7-alpine
    ports:
      - "6379:6379"
    volumes:
      - redis_data:/data
    restart: unless-stopped

volumes:
  qdrant_data:
  redis_data:

# requirements.txt
llama-index-core>=0.14.0
llama-index-llms-openai>=0.4.0
llama-index-embeddings-openai>=0.3.0
llama-index-vector-stores-qdrant>=0.4.0
fastapi>=0.115.0
uvicorn[standard]>=0.34.0
pydantic>=2.0
python-dotenv>=1.0.0
redis>=5.0

"""
app.py：将 LlamaIndex Workflow 封装为 FastAPI REST API
支持 REST 接口和 WebSocket 流式接口
"""
import asyncio
import uuid
from contextlib import asynccontextmanager
from typing import Optional

from fastapi import FastAPI, HTTPException, WebSocket, WebSocketDisconnect
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field

from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context,
    InputRequiredEvent, HumanResponseEvent,
)


# ============================================================
# 定义 API 请求/响应模型
# ============================================================
class QueryRequest(BaseModel):
    """查询请求"""
    query: str = Field(..., description="用户查询文本", min_length=1, max_length=2000)
    session_id: Optional[str] = Field(None, description="会话 ID")
    filters: Optional[dict] = Field(default_factory=dict, description="过滤条件")

class QueryResponse(BaseModel):
    """查询响应"""
    answer: str
    session_id: str
    sources: list = []
    metadata: dict = {}

class WorkflowStatusResponse(BaseModel):
    """工作流状态响应"""
    status: str       # running, completed, failed
    progress: float   # 0.0 - 1.0
    current_step: str


# ============================================================
# 定义 RAG 工作流
# ============================================================
class RAGQueryEvent(Event):
    query: str
    filters: dict

class RAGRetrieveEvent(Event):
    documents: list
    query: str

class RAGAnswerEvent(Event):
    answer: str
    sources: list


class RAGWorkflow(Workflow):
    """RAG 查询工作流"""

    @step
    async def parse_query(self, ev: StartEvent) -> RAGQueryEvent:
        """解析查询"""
        return RAGQueryEvent(
            query=ev.get("query", ""),
            filters=ev.get("filters", {}),
        )

    @step
    async def retrieve(self, ev: RAGQueryEvent) -> RAGRetrieveEvent:
        """检索相关文档"""
        # 模拟检索（实际项目中使用 VectorStoreIndex）
        docs = [
            {"content": f"关于 {ev.query} 的文档片段 1", "score": 0.95},
            {"content": f"关于 {ev.query} 的文档片段 2", "score": 0.87},
            {"content": f"关于 {ev.query} 的文档片段 3", "score": 0.82},
        ]
        return RAGRetrieveEvent(documents=docs, query=ev.query)

    @step
    async def generate_answer(self, ev: RAGRetrieveEvent) -> StopEvent:
        """生成回答"""
        context = "\n".join(d["content"] for d in ev.documents)
        # 模拟 LLM 生成
        answer = f"基于检索到的 {len(ev.documents)} 个文档，关于 '{ev.query}' 的回答..."
        return StopEvent(result={
            "answer": answer,
            "sources": [d["content"] for d in ev.documents],
        })


# ============================================================
# FastAPI 应用
# ============================================================

# 全局工作流实例
rag_workflow = RAGWorkflow(timeout=60, verbose=False)


@asynccontextmanager
async def lifespan(app: FastAPI):
    """应用生命周期管理"""
    print("应用启动中...")
    # 初始化资源（如加载索引、连接数据库等）
    yield
    print("应用关闭中...")
    # 清理资源


app = FastAPI(
    title="LlamaIndex Workflow API",
    description="基于 LlamaIndex Workflow 的智能问答 API",
    version="1.0.0",
    lifespan=lifespan,
)

# 跨域配置
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


# ============================================================
# REST API 端点
# ============================================================

@app.post("/api/v1/query", response_model=QueryResponse)
async def query_endpoint(request: QueryRequest):
    """
    查询端点：接收用户问题，运行 RAG 工作流，返回回答
    """
    session_id = request.session_id or str(uuid.uuid4())

    try:
        # 运行工作流
        result = await rag_workflow.run(
            query=request.query,
            filters=request.filters,
        )

        # 解析结果
        if isinstance(result, dict):
            return QueryResponse(
                answer=result.get("answer", ""),
                session_id=session_id,
                sources=result.get("sources", []),
            )
        else:
            return QueryResponse(
                answer=str(result),
                session_id=session_id,
            )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"工作流执行失败: {str(e)}")


@app.get("/api/v1/health")
async def health_check():
    """健康检查端点"""
    return {"status": "healthy", "version": "1.0.0"}


# ============================================================
# WebSocket 流式接口
# ============================================================

@app.websocket("/ws/query")
async def websocket_query(websocket: WebSocket):
    """
    WebSocket 流式查询端点
    支持实时推送工作流执行进度和中间结果
    """
    await websocket.accept()

    try:
        while True:
            # 接收客户端消息
            data = await websocket.receive_json()
            query = data.get("query", "")

            if not query:
                await websocket.send_json({"error": "查询内容不能为空"})
                continue

            # 运行工作流并流式输出事件
            handler = rag_workflow.run(query=query)

            # 实时推送每个事件
            async for event in handler.stream_events():
                if isinstance(event, InputRequiredEvent):
                    await websocket.send_json({
                        "type": "input_required",
                        "prefix": event.prefix,
                    })
                else:
                    await websocket.send_json({
                        "type": "progress",
                        "event": event.model_dump(),
                    })

            # 获取最终结果
            result = await handler
            await websocket.send_json({
                "type": "result",
                "data": result if isinstance(result, dict) else {"answer": str(result)},
            })

    except WebSocketDisconnect:
        print("WebSocket 连接已断开")
    except Exception as e:
        await websocket.send_json({"error": str(e)})


# ============================================================
# 启动命令
# ============================================================
# uvicorn app:app --host 0.0.0.0 --port 8000 --reload
#
# API 文档: http://localhost:8000/docs
# ReDoc:    http://localhost:8000/redoc

"""
生产环境性能优化策略
"""
import asyncio
import functools
from typing import Any
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent, Context


# ============================================================
# 优化策略 1：异步并发优化
# ============================================================
# 使用 asyncio.gather() 并行执行多个独立的异步操作

class OptimizedWorkflow(Workflow):
    @step
    async def parallel_retrieve(self, ctx: Context, ev: StartEvent) -> StopEvent:
        """并行检索多个数据源"""
        queries = [
            "查询1: 技术概述",
            "查询2: 应用场景",
            "查询3: 性能对比",
        ]

        # 使用 asyncio.gather 并行执行多个检索
        results = await asyncio.gather(
            *[self._mock_retrieve(q) for q in queries],
            return_exceptions=True,  # 不让单个失败影响整体
        )

        # 处理结果（过滤掉失败的）
        valid_results = [r for r in results if not isinstance(r, Exception)]
        return StopEvent(result=valid_results)

    async def _mock_retrieve(self, query: str) -> dict:
        """模拟检索"""
        await asyncio.sleep(0.5)
        return {"query": query, "docs": ["文档1", "文档2"]}


# ============================================================
# 优化策略 2：缓存策略
# ============================================================
class CachedResult:
    """简单的内存缓存"""
    def __init__(self, ttl: int = 3600):
        self._cache: dict[str, tuple[Any, float]] = {}
        self._ttl = ttl

    async def get_or_compute(self, key: str, compute_fn):
        """获取缓存或计算"""
        import time
        now = time.time()

        if key in self._cache:
            value, expire_at = self._cache[key]
            if now < expire_at:
                return value  # 缓存命中

        # 缓存未命中或已过期
        value = await compute_fn()
        self._cache[key] = (value, now + self._ttl)
        return value


# 全局缓存实例
result_cache = CachedResult(ttl=3600)  # 1 小时 TTL


# ============================================================
# 优化策略 3：连接池管理
# ============================================================
# 对于数据库连接、HTTP 连接等，使用连接池避免频繁创建/销毁

class ConnectionPool:
    """简单的异步连接池"""

    def __init__(self, max_size: int = 10):
        self._pool: asyncio.Queue = asyncio.Queue(maxsize=max_size)
        self._max_size = max_size
        self._current_size = 0

    async def acquire(self):
        """获取连接"""
        if self._current_size < self._max_size:
            self._current_size += 1
            return self._create_connection()

        return await self._pool.get()

    async def release(self, conn):
        """归还连接"""
        await self._pool.put(conn)

    def _create_connection(self):
        """创建新连接"""
        # 实际项目中替换为真实的连接创建逻辑
        return {"id": self._current_size, "status": "connected"}

    async def __aenter__(self):
        self._conn = await self.acquire()
        return self._conn

    async def __aexit__(self, *args):
        await self.release(self._conn)


# 使用连接池
pool = ConnectionPool(max_size=10)

async def use_connection_pool():
    async with pool as conn:
        # 使用连接执行操作
        print(f"使用连接: {conn}")
        pass

"""
可观测性：追踪（Tracing）与指标（Metrics）
"""
import time
import logging
from functools import wraps
from llama_index.core.workflow import Workflow, step, Event, StartEvent, StopEvent, Context

# ============================================================
# 方式 1：LlamaIndex 内置 Observability
# ============================================================
# LlamaIndex 提供了 Observability 模块，可以与多种追踪平台集成

# from llama_index.core import set_global_handler

# 与 Arize Phoenix 集成
# import phoenix as px
# px.launch_app()
# from llama_index.core import set_global_handler
# set_global_handler("arize_phoenix")

# 与 LangSmith 集成
# import os
# os.environ["LANGCHAIN_TRACING_V2"] = "true"
# os.environ["LANGCHAIN_API_KEY"] = "your-api-key"
# os.environ["LANGCHAIN_PROJECT"] = "llamaindex-workflow"


# ============================================================
# 方式 2：自定义追踪装饰器
# ============================================================
logger = logging.getLogger("workflow_metrics")


def trace_step(step_name: str):
    """
    Step 追踪装饰器
    记录每个 Step 的执行时间、输入输出、错误信息
    """
    def decorator(func):
        @wraps(func)
        async def wrapper(*args, **kwargs):
            start_time = time.time()
            span_id = f"{step_name}-{int(start_time * 1000)}"

            logger.info(f"[TRACE] {span_id} 开始执行: {step_name}")

            try:
                result = await func(*args, **kwargs)
                duration = time.time() - start_time

                logger.info(
                    f"[TRACE] {span_id} 执行完成: "
                    f"duration={duration:.3f}s, "
                    f"result_type={type(result).__name__}"
                )

                # 记录到指标系统（实际项目中接入 Prometheus/Grafana 等）
                # metrics.histogram("step_duration", duration, tags={"step": step_name})
                # metrics.increment("step_success", tags={"step": step_name})

                return result

            except Exception as e:
                duration = time.time() - start_time
                logger.error(
                    f"[TRACE] {span_id} 执行失败: "
                    f"duration={duration:.3f}s, "
                    f"error={type(e).__name__}: {e}"
                )
                # metrics.increment("step_error", tags={"step": step_name})
                raise

        return wrapper
    return decorator


# 使用追踪装饰器
class ObservableWorkflow(Workflow):
    @step
    @trace_step("retrieve_documents")
    async def retrieve(self, ev: StartEvent) -> StopEvent:
        """带追踪的检索步骤"""
        await asyncio.sleep(0.5)  # 模拟检索
        return StopEvent(result="检索结果")


# ============================================================
# 方式 3：工作流级别的指标收集
# ============================================================
class WorkflowMetrics:
    """工作流指标收集器"""

    def __init__(self):
        self.step_durations: dict[str, list[float]] = {}
        self.step_errors: dict[str, int] = {}
        self.workflow_runs: int = 0
        self.workflow_errors: int = 0

    def record_step(self, step_name: str, duration: float, error: bool = False):
        """记录 Step 指标"""
        if step_name not in self.step_durations:
            self.step_durations[step_name] = []
        self.step_durations[step_name].append(duration)

        if error:
            self.step_errors[step_name] = self.step_errors.get(step_name, 0) + 1

    def record_workflow(self, error: bool = False):
        """记录工作流指标"""
        self.workflow_runs += 1
        if error:
            self.workflow_errors += 1

    def get_summary(self) -> dict:
        """获取指标汇总"""
        summary = {
            "total_runs": self.workflow_runs,
            "error_rate": self.workflow_errors / max(self.workflow_runs, 1),
            "step_stats": {},
        }
        for step_name, durations in self.step_durations.items():
            summary["step_stats"][step_name] = {
                "count": len(durations),
                "avg_duration": sum(durations) / len(durations),
                "max_duration": max(durations),
                "min_duration": min(durations),
                "error_count": self.step_errors.get(step_name, 0),
            }
        return summary


# 全局指标实例
metrics = WorkflowMetrics()

[用户界面 (Web UI)]
       |
       v
[FastAPI 后端服务]
       |
       v
[LlamaIndex Workflow (RAG 工作流)]
       |
       +---> [文档加载器] ---> [文本分块器] ---> [嵌入模型] ---> [向量数据库]
       |
       +---> [检索器] ---> [重排器] ---> [LLM 生成器]
       |
       +---> [对话记忆管理器]

"""
企业知识库问答系统 - 完整实现
文件: knowledge_base_app.py
"""
import asyncio
import os
import uuid
from typing import Optional, List
from pathlib import Path

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import HTMLResponse
from pydantic import BaseModel, Field

from llama_index.core import (
    VectorStoreIndex,
    SimpleDirectoryReader,
    Document,
    Settings,
    StorageContext,
    load_index_from_storage,
)
from llama_index.core.node_parser import SentenceSplitter
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context,
)

# ============================================================
# 配置
# ============================================================
# 在 .env 文件或环境变量中设置
# OPENAI_API_KEY=your-api-key-here

DATA_DIR = "./knowledge_base_data"
INDEX_DIR = "./index_storage"
UPLOAD_DIR = "./uploads"

# 确保目录存在
for d in [DATA_DIR, INDEX_DIR, UPLOAD_DIR]:
    Path(d).mkdir(parents=True, exist_ok=True)


# ============================================================
# 全局设置
# ============================================================
def init_llama_index_settings():
    """初始化 LlamaIndex 全局设置"""
    # 设置嵌入模型（用于将文本转为向量）
    from llama_index.embeddings.openai import OpenAIEmbedding
    Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")

    # 设置 LLM（用于生成回答）
    from llama_index.llms.openai import OpenAI
    Settings.llm = OpenAI(model="gpt-4.1-mini", temperature=0.1)

    # 设置文本分块器
    Settings.node_parser = SentenceSplitter(
        chunk_size=512,     # 每个块的最大 token 数
        chunk_overlap=50,   # 块之间的重叠 token 数
    )

    # 设置默认检索参数
    Settings.num_output = 512       # LLM 输出最大 token 数
    Settings.context_window = 4096  # 上下文窗口大小


# ============================================================
# 知识库管理器
# ============================================================
class KnowledgeBaseManager:
    """知识库管理器：负责文档加载、索引构建和检索"""

    def __init__(self):
        self.index: Optional[VectorStoreIndex] = None
        self._load_or_create_index()

    def _load_or_create_index(self):
        """加载已有索引或创建新索引"""
        if os.path.exists(INDEX_DIR) and os.listdir(INDEX_DIR):
            try:
                storage_context = StorageContext.from_defaults(
                    persist_dir=INDEX_DIR
                )
                self.index = load_index_from_storage(storage_context)
                print("已加载现有索引")
            except Exception as e:
                print(f"加载索引失败: {e}，将创建新索引")
                self._create_empty_index()
        else:
            self._create_empty_index()

    def _create_empty_index(self):
        """创建空索引"""
        # 用一份占位文档创建初始索引
        placeholder = Document(text="知识库初始化完成。", doc_id="placeholder")
        self.index = VectorStoreIndex.from_documents([placeholder])
        self._persist_index()

    def _persist_index(self):
        """持久化索引到磁盘"""
        self.index.storage_context.persist(persist_dir=INDEX_DIR)

    async def add_documents(self, file_paths: list[str]) -> int:
        """
        添加文档到知识库
        返回添加的文档数量
        """
        documents = []

        for file_path in file_paths:
            path = Path(file_path)
            if not path.exists():
                continue

            # 根据文件类型选择加载器
            if path.suffix.lower() in (".pdf", ".md", ".txt", ".docx"):
                reader = SimpleDirectoryReader(input_files=[str(path)])
                docs = reader.load_data()
                documents.extend(docs)
            else:
                # 尝试作为纯文本读取
                try:
                    content = path.read_text(encoding="utf-8")
                    documents.append(
                        Document(
                            text=content,
                            metadata={"file_name": path.name},
                            doc_id=str(uuid.uuid4()),
                        )
                    )
                except Exception:
                    print(f"跳过不支持的文件: {path}")

        if documents:
            # 将新文档添加到索引
            for doc in documents:
                self.index.insert(doc)
            self._persist_index()

        print(f"已添加 {len(documents)} 个文档到知识库")
        return len(documents)

    def get_retriever(self, similarity_top_k: int = 5):
        """获取检索器"""
        return self.index.as_retriever(similarity_top_k=similarity_top_k)


# ============================================================
# RAG 问答工作流
# ============================================================
class QueryParsedEvent(Event):
    """查询解析事件"""
    query: str
    session_id: str
    chat_history: list

class RetrievedEvent(Event):
    """检索结果事件"""
    nodes: list
    query: str

class AugmentedEvent(Event):
    """增强事件（检索结果 + 对话历史）"""
    context: str
    query: str
    chat_history: list


class KnowledgeBaseRAGWorkflow(Workflow):
    """知识库 RAG 问答工作流"""

    def __init__(self, kb_manager: KnowledgeBaseManager, **kwargs):
        super().__init__(**kwargs)
        self.kb_manager = kb_manager

    @step
    async def parse_query(self, ctx: Context, ev: StartEvent) -> QueryParsedEvent:
        """解析查询参数"""
        query = ev.get("query", "")
        session_id = ev.get("session_id", str(uuid.uuid4()))
        chat_history = ev.get("chat_history", [])

        await ctx.set("session_id", session_id)
        print(f"[RAG] 解析查询: '{query}' (会话: {session_id[:8]}...)")

        return QueryParsedEvent(
            query=query,
            session_id=session_id,
            chat_history=chat_history,
        )

    @step
    async def retrieve_documents(self, ev: QueryParsedEvent) -> RetrievedEvent:
        """检索相关文档"""
        retriever = self.kb_manager.get_retriever(similarity_top_k=5)
        nodes = retriever.retrieve(ev.query)

        print(f"[RAG] 检索到 {len(nodes)} 个相关文档片段")
        return RetrievedEvent(nodes=nodes, query=ev.query)

    @step
    async def augment_context(self, ev: RetrievedEvent) -> AugmentedEvent:
        """构建增强上下文"""
        # 提取检索到的文档内容
        context_parts = []
        for i, node in enumerate(ev.nodes, 1):
            content = node.get_content() if hasattr(node, "get_content") else str(node)
            context_parts.append(f"[文档片段 {i}] {content}")

        context = "\n\n".join(context_parts)

        return AugmentedEvent(
            context=context,
            query=ev.query,
            chat_history=[],  # 简化处理
        )

    @step
    async def generate_answer(self, ev: AugmentedEvent) -> StopEvent:
        """生成回答"""
        # 构建提示
        chat_context = ""
        if ev.chat_history:
            chat_context = "\n对话历史:\n" + "\n".join(
                f"{msg.get('role', 'user')}: {msg.get('content', '')}"
                for msg in ev.chat_history[-4:]  # 只保留最近 4 条
            )

        prompt = f"""你是一个专业的企业知识库助手。请基于以下参考文档回答用户的问题。
如果参考文档中没有相关信息，请如实说明。

参考文档:
{ev.context}
{chat_context}

用户问题: {ev.query}

请用中文给出详细、准确的回答:"""

        # 调用 LLM 生成回答
        from llama_index.core.llms import ChatMessage
        from llama_index.llms.openai import OpenAI

        llm = OpenAI(model="gpt-4.1-mini", temperature=0.1)
        response = await llm.acomplete(prompt)
        answer = str(response)

        # 提取来源信息
        sources = []
        for node in ev.nodes[:3] if hasattr(ev, 'nodes') else []:
            meta = node.metadata if hasattr(node, "metadata") else {}
            sources.append(meta.get("file_name", "未知来源"))

        return StopEvent(result={
            "answer": answer,
            "sources": list(set(sources)),
            "doc_count": len(ev.context.split("[文档片段"))),
        })


# ============================================================
# FastAPI 应用
# ============================================================
app = FastAPI(
    title="企业知识库问答系统",
    description="基于 LlamaIndex 的企业知识库 RAG 系统",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_methods=["*"],
    allow_headers=["*"],
)

# 全局实例
kb_manager = KnowledgeBaseManager()
rag_workflow = KnowledgeBaseRAGWorkflow(
    kb_manager=kb_manager,
    timeout=120,
    verbose=False,
)

# 会话存储（生产环境应使用 Redis）
sessions: dict[str, list] = {}


class ChatRequest(BaseModel):
    query: str = Field(..., min_length=1, max_length=2000)
    session_id: Optional[str] = None


class ChatResponse(BaseModel):
    answer: str
    sources: list[str] = []
    session_id: str


@app.on_event("startup")
async def startup():
    init_llama_index_settings()


@app.post("/api/upload")
async def upload_documents(files: List[UploadFile] = File(...)):
    """上传文档到知识库"""
    saved_paths = []

    for file in files:
        file_path = os.path.join(UPLOAD_DIR, file.filename)
        content = await file.read()
        with open(file_path, "wb") as f:
            f.write(content)
        saved_paths.append(file_path)

    doc_count = await kb_manager.add_documents(saved_paths)

    return {
        "message": f"成功上传并处理 {len(saved_paths)} 个文件",
        "documents_added": doc_count,
        "files": [f.filename for f in files],
    }


@app.post("/api/chat", response_model=ChatResponse)
async def chat(request: ChatRequest):
    """对话接口"""
    session_id = request.session_id or str(uuid.uuid4())

    # 获取对话历史
    chat_history = sessions.get(session_id, [])

    try:
        result = await rag_workflow.run(
            query=request.query,
            session_id=session_id,
            chat_history=chat_history,
        )

        # 更新对话历史
        if session_id not in sessions:
            sessions[session_id] = []
        sessions[session_id].append({"role": "user", "content": request.query})

        if isinstance(result, dict):
            answer = result.get("answer", "")
            sources = result.get("sources", [])
        else:
            answer = str(result)
            sources = []

        sessions[session_id].append({"role": "assistant", "content": answer})

        # 保留最近 20 条对话
        sessions[session_id] = sessions[session_id][-20:]

        return ChatResponse(
            answer=answer,
            sources=sources,
            session_id=session_id,
        )

    except Exception as e:
        raise HTTPException(status_code=500, detail=f"处理失败: {str(e)}")


@app.get("/api/sessions/{session_id}")
async def get_session(session_id: str):
    """获取会话历史"""
    history = sessions.get(session_id, [])
    return {"session_id": session_id, "messages": history}


@app.get("/", response_class=HTMLResponse)
async def index_page():
    """简单的前端页面"""
    return """
    <!DOCTYPE html>
    <html lang="zh-CN">
    <head>
        <meta charset="UTF-8">
        <meta name="viewport" content="width=device-width, initial-scale=1.0">
        <title>企业知识库问答系统</title>
        <style>
            * { margin: 0; padding: 0; box-sizing: border-box; }
            body { font-family: -apple-system, sans-serif; background: #f5f5f5; }
            .container { max-width: 800px; margin: 0 auto; padding: 20px; }
            h1 { text-align: center; margin: 20px 0; color: #333; }
            .chat-box { background: white; border-radius: 12px; padding: 20px;
                        box-shadow: 0 2px 8px rgba(0,0,0,0.1); min-height: 400px;
                        max-height: 600px; overflow-y: auto; margin-bottom: 20px; }
            .message { margin: 10px 0; padding: 10px 15px; border-radius: 8px; }
            .user-msg { background: #e3f2fd; text-align: right; }
            .bot-msg { background: #f1f8e9; }
            .input-area { display: flex; gap: 10px; }
            input[type="text"] { flex: 1; padding: 12px; border: 1px solid #ddd;
                                  border-radius: 8px; font-size: 16px; }
            button { padding: 12px 24px; background: #1976d2; color: white;
                     border: none; border-radius: 8px; cursor: pointer; font-size: 16px; }
            button:hover { background: #1565c0; }
            .sources { font-size: 12px; color: #666; margin-top: 5px; }
        </style>
    </head>
    <body>
        <div class="container">
            <h1>企业知识库问答系统</h1>
            <div class="chat-box" id="chatBox">
                <div class="message bot-msg">你好！我是知识库助手，请问有什么可以帮助你的？</div>
            </div>
            <div class="input-area">
                <input type="text" id="queryInput" placeholder="输入你的问题..."
                       onkeypress="if(event.key==='Enter')sendMessage()">
                <button onclick="sendMessage()">发送</button>
            </div>
        </div>
        <script>
            let sessionId = null;
            async function sendMessage() {
                const input = document.getElementById('queryInput');
                const query = input.value.trim();
                if (!query) return;
                input.value = '';

                const chatBox = document.getElementById('chatBox');
                chatBox.innerHTML += `<div class="message user-msg">${query}</div>`;
                chatBox.innerHTML += `<div class="message bot-msg" id="loading">思考中...</div>`;
                chatBox.scrollTop = chatBox.scrollHeight;

                try {
                    const res = await fetch('/api/chat', {
                        method: 'POST',
                        headers: {'Content-Type': 'application/json'},
                        body: JSON.stringify({query, session_id: sessionId})
                    });
                    const data = await res.json();
                    sessionId = data.session_id;
                    document.getElementById('loading').remove();
                    let msg = `<div class="message bot-msg">${data.answer}`;
                    if (data.sources && data.sources.length > 0) {
                        msg += `<div class="sources">来源: ${data.sources.join(', ')}</div>`;
                    }
                    msg += '</div>';
                    chatBox.innerHTML += msg;
                    chatBox.scrollTop = chatBox.scrollHeight;
                } catch(e) {
                    document.getElementById('loading').innerHTML = '出错了，请重试。';
                }
            }
        </script>
    </body>
    </html>
    """


# 启动命令：
# uvicorn knowledge_base_app:app --host 0.0.0.0 --port 8000 --reload
#
# 访问 http://localhost:8000 使用 Web 界面
# 访问 http://localhost:8000/docs 查看 API 文档

"""
智能研究助手：使用 Workflow 编排多步骤研究任务
支持：搜索 -> 摘要 -> 分析 -> 报告生成
"""
import asyncio
from typing import Union
from llama_index.core.workflow import (
    Workflow, step, Event, StartEvent, StopEvent, Context,
)


# ============================================================
# 事件定义
# ============================================================
class ResearchPlanEvent(Event):
    """研究计划事件"""
    topic: str
    subtopics: list[str]

class SearchResultEvent(Event):
    """搜索结果事件"""
    subtopic: str
    results: list[dict]

class SummaryEvent(Event):
    """摘要事件"""
    subtopic: str
    summary: str
    key_findings: list[str]

class AnalysisEvent(Event):
    """分析事件"""
    report_sections: list[dict]

class ReportReadyEvent(Event):
    """报告就绪事件"""
    full_report: str


# ============================================================
# 研究助手工作流
# ============================================================
class ResearchAssistantWorkflow(Workflow):
    """
    智能研究助手工作流
    流程：
      1. 制定研究计划（分解子主题）
      2. 并行搜索各子主题
      3. 汇总搜索结果
      4. 生成各子主题摘要
      5. 综合分析
      6. 生成完整研究报告
    """

    @step
    async def plan_research(self, ctx: Context, ev: StartEvent) -> SearchResultEvent:
        """
        步骤 1：制定研究计划
        将主题分解为多个子主题，并行搜索
        """
        topic = ev.get("topic", "")
        print(f"[研究助手] 制定研究计划: {topic}")

        # 分解子主题（实际项目中使用 LLM 分解）
        subtopics = [
            f"{topic} - 基本概念与原理",
            f"{topic} - 最新研究进展",
            f"{topic} - 实际应用场景",
            f"{topic} - 未来发展趋势",
        ]

        await ctx.set("topic", topic)
        await ctx.set("subtopics", subtopics)
        await ctx.set("expected_searches", len(subtopics))

        # 并行发出多个搜索事件
        for subtopic in subtopics[1:]:
            ctx.send_event(SearchResultEvent(
                subtopic=subtopic,
                results=[],  # 将在 search_step 中填充
            ))

        # 第一个通过返回值发出
        return SearchResultEvent(subtopic=subtopics[0], results=[])

    @step(num_workers=4)
    async def search_subtopic(self, ev: SearchResultEvent) -> SummaryEvent:
        """
        步骤 2：搜索每个子主题
        并行执行（最多 4 个并发）
        """
        print(f"[搜索] 搜索子主题: {ev.subtopic}")

        # 模拟搜索（实际项目中接入搜索 API 或 RAG）
        await asyncio.sleep(1)
        mock_results = [
            {"title": f"关于 {ev.subtopic} 的论文", "snippet": "重要发现..."},
            {"title": f"{ev.subtopic} 综述", "snippet": "综合分析..."},
            {"title": f"{ev.subtopic} 案例研究", "snippet": "实践案例..."},
        ]

        # 生成摘要（实际项目中使用 LLM）
        summary = f"关于 '{ev.subtopic}' 的研究摘要：基于 {len(mock_results)} 个来源的综合分析。"
        key_findings = [r["snippet"] for r in mock_results]

        return SummaryEvent(
            subtopic=ev.subtopic,
            summary=summary,
            key_findings=key_findings,
        )

    @step
    async def collect_summaries(self, ctx: Context, ev: SummaryEvent) -> AnalysisEvent:
        """
        步骤 3：收集所有子主题的摘要
        """
        summaries = await ctx.get("summaries", default=[])
        summaries.append({
            "subtopic": ev.subtopic,
            "summary": ev.summary,
            "key_findings": ev.key_findings,
        })
        await ctx.set("summaries", summaries)

        expected = await ctx.get("expected_searches", 4)
        if len(summaries) < expected:
            print(f"[收集] 已收到 {len(summaries)}/{expected} 个子主题的摘要")
            return None  # 等待更多摘要

        # 所有摘要已收集完成，进入分析阶段
        print(f"[收集] 所有 {expected} 个子主题摘要已就绪")
        return AnalysisEvent(report_sections=summaries)

    @step
    async def analyze_and_report(self, ctx: Context, ev: AnalysisEvent) -> StopEvent:
        """
        步骤 4：综合分析并生成完整报告
        """
        topic = await ctx.get("topic", "未知主题")
        sections = ev.report_sections

        # 生成报告（实际项目中使用 LLM）
        report = f"""# 研究报告：{topic}

## 摘要
本报告对 {topic} 进行了全面的研究和分析，涵盖基本概念、最新进展、应用场景和未来趋势四个方面。

## 详细分析
"""
        for i, section in enumerate(sections, 1):
            report += f"""
### {i}. {section['subtopic']}

{section['summary']}

**关键发现：**
"""
            for finding in section["key_findings"]:
                report += f"- {finding}\n"

        report += f"""
## 结论
综合以上分析，{topic} 领域正在快速发展，具有广阔的应用前景。
建议持续关注最新研究动态，并结合实际场景进行探索和实践。

---
*本报告由智能研究助手自动生成*
"""

        return StopEvent(result=report)


# ============================================================
# 运行研究助手
# ============================================================
async def main():
    workflow = ResearchAssistantWorkflow(timeout=120, verbose=False)

    print("=== 智能研究助手 ===")
    print("研究主题: LlamaIndex Workflow 事件驱动架构")
    print("=" * 50)

    result = await workflow.run(topic="LlamaIndex Workflow 事件驱动架构")
    print(result)


asyncio.run(main())

# 错误示例：Step 有时不返回任何事件（忘记 return）
@step
async def my_step(self, ev: StartEvent) -> ResultEvent:
    if some_condition:
        return ResultEvent(data="ok")
    # 遗漏了 else 分支的 return！

# 正确做法：确保所有分支都有返回值，或显式返回 None
@step
async def my_step(self, ev: StartEvent) -> ResultEvent:
    if some_condition:
        return ResultEvent(data="ok")
    return None  # 显式表示不发出事件

# 潜在问题：多个并发 Step 同时读写同一个 Context 键
# Step A 和 Step B 同时读取 count=0，都写回 count=1，丢失一次递增

# 解决方案 1：使用收集模式（在单个 Step 中累积）
@step
async def collector(self, ctx: Context, ev: ResultEvent) -> StopEvent:
    results = await ctx.get("results", default=[])
    results.append(ev.data)
    await ctx.set("results", results)
    # 这个 Step 每次只处理一个事件，不会并发

# 解决方案 2：使用 asyncio.Lock
lock = asyncio.Lock()

@step(num_workers=4)
async def safe_step(self, ctx: Context, ev: TaskEvent) -> ResultEvent:
    async with lock:
        count = await ctx.get("count", default=0)
        await ctx.set("count", count + 1)
    return ResultEvent(data="done")

# 问题：工作流 timeout 设置过短，LLM 调用还没完成就超时了

# 解决方案：根据实际耗时合理设置 timeout
workflow = MyWorkflow(
    timeout=300,   # 5 分钟（对于涉及多次 LLM 调用的工作流）
    verbose=True,  # 开发阶段开启，方便排查
)

# 或者在每个 Step 中单独设置超时
@step
async def llm_call(self, ev: InputEvent) -> OutputEvent:
    try:
        result = await asyncio.wait_for(
            self.llm.acomplete(prompt),
            timeout=30,  # 单次 LLM 调用超时 30 秒
        )
    except asyncio.TimeoutError:
        # 处理超时
        result = "生成超时，请稍后重试"
    return OutputEvent(result=result)

# 问题：Step 的返回类型注解与实际返回不匹配
@step
async def wrong_step(self, ev: StartEvent) -> EventA:
    return EventB(data="xxx")  # 返回了 EventB，但注解写的是 EventA！

# 正确做法：使用 Union 标注所有可能的返回类型
from typing import Union

@step
async def correct_step(self, ev: StartEvent) -> Union[EventA, EventB]:
    if condition:
        return EventA(data="路径A")
    return EventB(data="路径B")