LLM应用开发——LlamaIndex框架

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

# Document: 原始文档
doc = Document(
    text="这是一篇关于人工智能的文章...",
    metadata={
        "title": "AI 入门",
        "author": "张三",
        "date": "2024-01-01"
    }
)

# Node: 文档切分后的块
parser = SentenceSplitter(chunk_size=512, chunk_overlap=50)
nodes = parser.get_nodes_from_documents([doc])

# 查看节点
for node in nodes:
    print(f"Node ID: {node.node_id}")
    print(f"Text: {node.text[:100]}...")
    print(f"Metadata: {node.metadata}")

from llama_index.core import SimpleDirectoryReader

# 加载目录下所有支持的文件
documents = SimpleDirectoryReader(
    input_dir="./data",
    recursive=True,  # 递归子目录
    required_exts=[".pdf", ".docx", ".md", ".txt"]
).load_data()

print(f"加载了 {len(documents)} 个文档")

# PDF 加载
from llama_index.readers.file import PDFReader
pdf_reader = PDFReader()
docs = pdf_reader.load_data(file="document.pdf")

# Web 页面加载
from llama_index.readers.web import SimpleWebPageReader
web_reader = SimpleWebPageReader()
docs = web_reader.load_data(urls=["https://example.com"])

# 数据库加载
from llama_index.readers.database import DatabaseReader
db_reader = DatabaseReader(
    sql_database=database,  # SQLAlchemy engine
)
docs = db_reader.load_data(query="SELECT * FROM articles")

# Notion 加载
from llama_index.readers.notion import NotionPageReader
notion_reader = NotionPageReader(integration_token="your-token")
docs = notion_reader.load_data(page_ids=["page-id"])

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

class CustomAPIReader(BaseReader):
    """从自定义 API 加载数据"""
    
    def __init__(self, api_url: str, api_key: str):
        self.api_url = api_url
        self.api_key = api_key
    
    def load_data(self, query: str = None) -> List[Document]:
        import requests
        
        response = requests.get(
            self.api_url,
            headers={"Authorization": f"Bearer {self.api_key}"},
            params={"q": query} if query else {}
        )
        
        data = response.json()
        
        documents = []
        for item in data["results"]:
            doc = Document(
                text=item["content"],
                metadata={
                    "id": item["id"],
                    "title": item["title"],
                    "source": self.api_url
                }
            )
            documents.append(doc)
        
        return documents

# 使用自定义加载器
reader = CustomAPIReader(
    api_url="https://api.example.com/articles",
    api_key="your-key"
)
docs = reader.load_data(query="AI")

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

# 配置模型
embed_model = OpenAIEmbedding(model="text-embedding-3-small")
llm = OpenAI(model="gpt-4")

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

# 创建向量索引
index = VectorStoreIndex.from_documents(
    documents,
    embed_model=embed_model
)

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

# 查询
response = query_engine.query("什么是机器学习？")
print(response)

from llama_index.core import SummaryIndex

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

# 查询（会遍历所有节点）
query_engine = index.as_query_engine(
    response_mode="tree_summarize"  # 树形摘要
)
response = query_engine.query("总结这些文档的主要内容")

from llama_index.core import TreeIndex

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

# 查询
query_engine = index.as_query_engine(
    child_branch_factor=2  # 每次选择的分支数
)

from llama_index.core import KeywordTableIndex

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

# 查询
query_engine = index.as_query_engine()
response = query_engine.query("人工智能相关内容")

from llama_index.core import VectorStoreIndex, SummaryIndex
from llama_index.core.query_engine import RouterQueryEngine
from llama_index.core.selectors import LLMSingleSelector

# 创建多个索引
vector_index = VectorStoreIndex.from_documents(documents)
summary_index = SummaryIndex.from_documents(documents)

# 创建查询引擎
vector_query_engine = vector_index.as_query_engine()
summary_query_engine = summary_index.as_query_engine()

# 路由查询引擎
query_engine = RouterQueryEngine(
    selector=LLMSingleSelector.from_defaults(),
    query_engine_tools=[
        {
            "query_engine": vector_query_engine,
            "description": "用于具体问题的精确检索"
        },
        {
            "query_engine": summary_query_engine,
            "description": "用于总结和概述类问题"
        }
    ]
)

response = query_engine.query("总结文档的主要观点")

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")
collection = chroma_client.get_or_create_collection("my_collection")

# 创建向量存储
vector_store = ChromaVectorStore(chroma_collection=collection)

# 创建存储上下文
storage_context = StorageContext.from_defaults(vector_store=vector_store)

# 创建索引（数据会自动存入 Chroma）
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context
)

# 后续可以直接从 Chroma 加载
index = VectorStoreIndex.from_vector_store(vector_store)

from pinecone import Pinecone
from llama_index.vector_stores.pinecone import PineconeVectorStore

# 初始化 Pinecone
pc = Pinecone(api_key="your-api-key")
pinecone_index = pc.Index("my-index")

# 创建向量存储
vector_store = PineconeVectorStore(pinecone_index=pinecone_index)

# 创建索引
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context
)

import faiss
from llama_index.vector_stores.faiss import FaissVectorStore

# 创建 FAISS 索引
d = 1536  # OpenAI embedding 维度
faiss_index = faiss.IndexFlatL2(d)

# 创建向量存储
vector_store = FaissVectorStore(faiss_index=faiss_index)

# 创建 LlamaIndex 索引
storage_context = StorageContext.from_defaults(vector_store=vector_store)
index = VectorStoreIndex.from_documents(
    documents,
    storage_context=storage_context
)

# 保存到磁盘
index.storage_context.persist(persist_dir="./faiss_index")

# 从磁盘加载
from llama_index.core import load_index_from_storage
storage_context = StorageContext.from_defaults(persist_dir="./faiss_index")
index = load_index_from_storage(storage_context)

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)

# 创建查询引擎
query_engine = index.as_query_engine(
    similarity_top_k=5,          # 返回 top-5 相似结果
    response_mode="compact",      # 响应模式
    streaming=False               # 是否流式
)

# 执行查询
response = query_engine.query("解释机器学习的基本原理")

# 查看源文档
for node in response.source_nodes:
    print(f"Score: {node.score}")
    print(f"Text: {node.text[:200]}...")
    print("---")

# 1. refine: 逐个节点精炼答案
query_engine = index.as_query_engine(response_mode="refine")

# 2. compact: 压缩后一次性处理
query_engine = index.as_query_engine(response_mode="compact")

# 3. tree_summarize: 树形摘要
query_engine = index.as_query_engine(response_mode="tree_summarize")

# 4. simple_summarize: 简单摘要
query_engine = index.as_query_engine(response_mode="simple_summarize")

# 5. no_text: 只返回检索结果，不生成答案
query_engine = index.as_query_engine(response_mode="no_text")

from llama_index.core.query_engine import CustomQueryEngine
from llama_index.core.retrievers import BaseRetriever
from llama_index.llms.openai import OpenAI

class RAGQueryEngine(CustomQueryEngine):
    """自定义 RAG 查询引擎"""
    
    retriever: BaseRetriever
    llm: OpenAI
    
    def custom_query(self, query_str: str) -> str:
        # 检索相关节点
        nodes = self.retriever.retrieve(query_str)
        
        # 构建上下文
        context = "\n\n".join([n.node.text for n in nodes])
        
        # 构建提示
        prompt = f"""基于以下上下文回答问题。如果上下文中没有相关信息，请说明。

上下文：
{context}

问题：{query_str}

回答："""
        
        # 调用 LLM
        response = self.llm.complete(prompt)
        
        return str(response)

# 使用自定义引擎
retriever = index.as_retriever(similarity_top_k=5)
query_engine = RAGQueryEngine(
    retriever=retriever,
    llm=OpenAI(model="gpt-4")
)

response = query_engine.query("什么是深度学习？")

from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)

# 创建对话引擎
chat_engine = index.as_chat_engine(
    chat_mode="condense_question",  # 对话模式
    verbose=True
)

# 多轮对话
response1 = chat_engine.chat("什么是机器学习？")
print(response1)

response2 = chat_engine.chat("它有哪些应用场景？")
print(response2)

response3 = chat_engine.chat("举一个具体的例子")
print(response3)

# 重置对话
chat_engine.reset()

# 1. simple: 简单模式，直接将历史拼接
chat_engine = index.as_chat_engine(chat_mode="simple")

# 2. condense_question: 将问题压缩为独立问题
chat_engine = index.as_chat_engine(chat_mode="condense_question")

# 3. context: 保持上下文的对话
chat_engine = index.as_chat_engine(chat_mode="context")

# 4. react: ReAct 模式，支持工具调用
chat_engine = index.as_chat_engine(chat_mode="react")

# 创建流式对话引擎
chat_engine = index.as_chat_engine(streaming=True)

# 流式响应
streaming_response = chat_engine.stream_chat("解释深度学习")

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

from llama_index.core.retrievers import QueryFusionRetriever
from llama_index.retrievers.bm25 import BM25Retriever

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

# BM25 检索器
bm25_retriever = BM25Retriever.from_defaults(
    nodes=nodes,
    similarity_top_k=5
)

# 混合检索器
hybrid_retriever = QueryFusionRetriever(
    retrievers=[vector_retriever, bm25_retriever],
    similarity_top_k=10,
    num_queries=4,  # 生成多个查询变体
    mode="reciprocal_rerank"  # 倒数排名融合
)

# 使用混合检索
nodes = hybrid_retriever.retrieve("机器学习的应用")

from llama_index.core.postprocessor import SentenceTransformerRerank

# 创建重排序器
reranker = SentenceTransformerRerank(
    model="cross-encoder/ms-marco-MiniLM-L-6-v2",
    top_n=3
)

# 带重排序的查询引擎
query_engine = index.as_query_engine(
    similarity_top_k=10,
    node_postprocessors=[reranker]
)

response = query_engine.query("什么是 Transformer？")

from llama_index.core.vector_stores import MetadataFilters, FilterCondition

# 创建过滤条件
filters = MetadataFilters(
    filters=[
        {"key": "author", "value": "张三", "operator": "=="},
        {"key": "year", "value": 2023, "operator": ">="}
    ],
    condition=FilterCondition.AND
)

# 带过滤的检索
retriever = index.as_retriever(
    similarity_top_k=5,
    filters=filters
)

# 或在查询引擎中使用
query_engine = index.as_query_engine(
    similarity_top_k=5,
    filters=filters
)

from llama_index.core.node_parser import (
    SentenceSplitter,
    TokenTextSplitter,
    SemanticSplitterNodeParser
)
from llama_index.embeddings.openai import OpenAIEmbedding

# 1. 句子分割器（推荐）
splitter = SentenceSplitter(
    chunk_size=512,
    chunk_overlap=50,
    paragraph_separator="\n\n"
)

# 2. Token 分割器
splitter = TokenTextSplitter(
    chunk_size=512,
    chunk_overlap=50,
    separator=" "
)

# 3. 语义分割器（基于嵌入的相似度）
splitter = SemanticSplitterNodeParser(
    embed_model=OpenAIEmbedding(),
    breakpoint_percentile_threshold=95
)

# 使用分割器
nodes = splitter.get_nodes_from_documents(documents)

from llama_index.core.node_parser import HierarchicalNodeParser

# 创建分层解析器
parser = HierarchicalNodeParser.from_defaults(
    chunk_sizes=[2048, 512, 128]  # 从大到小的块大小
)

# 解析文档
nodes = parser.get_nodes_from_documents(documents)

# 查看层级关系
for node in nodes:
    print(f"Node: {node.node_id}")
    print(f"Parent: {node.relationships.get('parent')}")
    print(f"Children: {node.relationships.get('children')}")

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.extractors import (
    TitleExtractor,
    KeywordsExtractor,
    SummaryExtractor
)

# 创建摄取管道
pipeline = IngestionPipeline(
    transformations=[
        # 文本分割
        SentenceSplitter(chunk_size=512, chunk_overlap=50),
        # 标题提取
        TitleExtractor(nodes=5),
        # 关键词提取
        KeywordsExtractor(keywords=10),
        # 摘要提取
        SummaryExtractor(summaries=["self"]),
        # 嵌入生成
        OpenAIEmbedding()
    ]
)

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

# 查看提取的元数据
for node in nodes[:3]:
    print(f"Title: {node.metadata.get('document_title')}")
    print(f"Keywords: {node.metadata.get('excerpt_keywords')}")
    print(f"Summary: {node.metadata.get('section_summary')}")
    print("---")

from llama_index.core.ingestion import IngestionPipeline, IngestionCache
from llama_index.core.storage.docstore import SimpleDocumentStore

# 创建缓存
cache = IngestionCache(
    collection="my_cache"
)

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

# 创建管道
pipeline = IngestionPipeline(
    transformations=[
        SentenceSplitter(chunk_size=512),
        OpenAIEmbedding()
    ],
    cache=cache,
    docstore=docstore
)

# 首次运行
nodes = pipeline.run(documents=documents)

# 增量更新（只处理新文档）
new_documents = SimpleDirectoryReader("./new_data").load_data()
new_nodes = pipeline.run(documents=new_documents)

from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, Settings
from llama_index.llms.openai import OpenAI
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.core.node_parser import SentenceSplitter
import chromadb
from llama_index.vector_stores.chroma import ChromaVectorStore
from llama_index.core import StorageContext

class KnowledgeBase:
    def __init__(self, persist_dir: str = "./kb_data"):
        self.persist_dir = persist_dir
        
        # 配置模型
        Settings.llm = OpenAI(model="gpt-4")
        Settings.embed_model = OpenAIEmbedding(model="text-embedding-3-small")
        Settings.node_parser = SentenceSplitter(chunk_size=512, chunk_overlap=50)
        
        # 初始化 Chroma
        self.chroma_client = chromadb.PersistentClient(path=persist_dir)
        self.collection = self.chroma_client.get_or_create_collection("knowledge_base")
        self.vector_store = ChromaVectorStore(chroma_collection=self.collection)
        
        # 加载或创建索引
        self.index = self._load_or_create_index()
    
    def _load_or_create_index(self) -> VectorStoreIndex:
        """加载或创建索引"""
        try:
            return VectorStoreIndex.from_vector_store(self.vector_store)
        except:
            return VectorStoreIndex(
                nodes=[],
                storage_context=StorageContext.from_defaults(
                    vector_store=self.vector_store
                )
            )
    
    def add_documents(self, directory: str):
        """添加文档"""
        documents = SimpleDirectoryReader(directory).load_data()
        
        for doc in documents:
            self.index.insert(doc)
        
        print(f"已添加 {len(documents)} 个文档")
    
    def query(self, question: str, top_k: int = 5) -> dict:
        """查询知识库"""
        query_engine = self.index.as_query_engine(
            similarity_top_k=top_k,
            response_mode="compact"
        )
        
        response = query_engine.query(question)
        
        return {
            "answer": str(response),
            "sources": [
                {
                    "text": node.text[:200],
                    "score": node.score,
                    "metadata": node.metadata
                }
                for node in response.source_nodes
            ]
        }
    
    def chat(self, message: str):
        """对话接口"""
        if not hasattr(self, 'chat_engine'):
            self.chat_engine = self.index.as_chat_engine(
                chat_mode="condense_question",
                verbose=True
            )
        
        response = self.chat_engine.chat(message)
        return str(response)

# 使用示例
kb = KnowledgeBase()
kb.add_documents("./docs")

# 查询
result = kb.query("什么是深度学习？")
print(result["answer"])

# 对话
print(kb.chat("它和机器学习有什么区别？"))

# 1. 使用异步
from llama_index.core import VectorStoreIndex

index = VectorStoreIndex.from_documents(documents)
query_engine = index.as_query_engine()

# 异步查询
response = await query_engine.aquery("问题")

# 2. 批量处理
from llama_index.core.ingestion import IngestionPipeline

pipeline = IngestionPipeline(
    transformations=[...],
    num_workers=4  # 并行处理
)

# 3. 缓存嵌入
from llama_index.core import Settings
from llama_index.core.callbacks import CallbackManager, TokenCountingHandler

# 启用 Token 计数
token_counter = TokenCountingHandler()
Settings.callback_manager = CallbackManager([token_counter])