使用 LlamaIndex 构建智能文档查询系统
- 1. 环境准备
- 2. 初始化模型
- 3. 加载文档
- 4. 构建索引和查询引擎
- 5. 生成扩展查询
- 6. 主函数
- 7. 总结
在现代信息检索系统中，如何高效地从大量文档中提取出有用的信息是一个重要的挑战。本文将介绍如何使用 LlamaIndex 构建一个智能文档查询系统，通过多种查询优化技术来提高检索的准确性和效率。
1. 环境准备
首先，我们需要安装必要的 Python 库，并加载环境变量。
import logging
import os
from dotenv import find_dotenv, load_dotenv
from llama_index.core import VectorStoreIndex, SimpleDirectoryReader, PromptTemplate, get_response_synthesizer
from llama_index.core.indices.query.query_transform.base import (
StepDecomposeQueryTransform, HyDEQueryTransform
)
from llama_index.core.node_parser import (
TokenTextSplitter, SentenceWindowNodeParser
)
from llama_index.core.postprocessor import MetadataReplacementPostProcessor
from llama_index.core.query_engine import (
MultiStepQueryEngine, TransformQueryEngine
)
from llama_index.embeddings.openai import OpenAIEmbedding
from llama_index.llms.openai import OpenAI
# 读取本地 .env 文件
load_dotenv(find_dotenv())
# 设置日志级别为 ERROR，避免 WARNING 信息干扰
logging.basicConfig(level=logging.ERROR)
2. 初始化模型
接下来，我们初始化 OpenAI 的 LLM 和 Embedding 模型。
# 初始化 OpenAI 模型
llm_client = OpenAI(
model="gpt-4",
api_base=os.environ["OPENAI_BASE_URL"],
api_key=os.environ["OPENAI_API_KEY"],
is_chat_model=True,
seed=42,
)
# 初始化 OpenAI Embedding 模型
embed_client = OpenAIEmbedding(
model="text-embedding-3-large",
api_base=os.environ["OPENAI_EMBED_BASE_URL"],
api_key=os.environ["OPENAI_API_KEY"],
)
3. 加载文档
我们定义一个函数来加载指定目录下的文档。
# 加载文档
def load_documents(directory):
print('=' * 50)
print('📂 正在加载文档...')
print('=' * 50 + '\n')
documents = SimpleDirectoryReader(directory).load_data()
print(f'✅ 文档加载完成。\n')
return documents
4. 构建索引和查询引擎
我们使用不同的节点解析器来构建索引，并创建查询引擎。
# 构建索引并创建查询引擎
def build_index_and_query_engine(documents, embed_model, llm, node_parser, postprocessors=None):
print(f"\n{'=' * 50}")
print(f"🔍 正在使用 {node_parser.__class__.__name__} 方法进行测试...")
print(f"{'=' * 50}\n")
print("📑 正在处理文档...")
nodes = node_parser.get_nodes_from_documents(documents)
index = VectorStoreIndex(nodes, embed_model=embed_model)
query_engine = index.as_query_engine(
similarity_top_k=5,
streaming=True,
llm=llm,
node_postprocessors=postprocessors if postprocessors else []
)
return query_engine, index
5. 生成扩展查询
为了增强查询的效果，我们使用 LLM 对原始查询进行扩展。
# 生成扩展查询
query_gen_str = """\
系统角色设定:
你是一个专业的问题改写助手。你的任务是将用户的原始问题扩充为一个更完整、更全面的问题。
规则：
1. 将可能的歧义、相关概念和上下文信息整合到一个完整的问题中
2. 使用括号对歧义概念进行补充说明
3. 添加关键的限定词和修饰语
4. 确保改写后的问题清晰且语义完整
5. 对于模糊概念，在括号中列举主要可能性
原始问题:
{query}
请生成一个综合的改写问题，确保：
- 包含原始问题的核心意图
- 涵盖可能的歧义解释
- 使用清晰的逻辑关系词连接不同方面
- 必要时使用括号补充说明
输出格式：
[综合改写] - 改写后的问题
"""
query_gen_prompt = PromptTemplate(query_gen_str)
def generate_queries(query: str):
response = llm_client.predict(query_gen_prompt, query=query)
return response
6. 主函数
在主函数中，我们加载文档并测试不同的查询方法。
# 主函数
def main():
# 加载文档
documents = load_documents('./docs/md')
# 定义问题和节点解析器
question = '张伟是哪个部门的'
node_parsers = [
TokenTextSplitter(chunk_size=1024, chunk_overlap=20),
# SentenceSplitter(chunk_size=512, chunk_overlap=50),
# SentenceWindowNodeParser.from_defaults(
# window_size=3,
# window_metadata_key="window",
# original_text_metadata_key="original_text"
# ),
# SemanticSplitterNodeParser(
# buffer_size=1,
# breakpoint_percentile_threshold=95,
# embed_model=embed_client
# ),
# MarkdownNodeParser()
]
# 遍历不同的节点解析器进行测试
for parser in node_parsers:
if isinstance(parser, SentenceWindowNodeParser):
postprocessors = [MetadataReplacementPostProcessor(target_metadata_key="window")]
else:
postprocessors = None
query_engine, index = build_index_and_query_engine(documents, embed_client, llm_client, parser, postprocessors)
# 方法一：使用大模型扩充用户问题
print("\n🔍 原始问题:")
print(f" {question}")
expanded_query = generate_queries(question)
print("\n📝 扩展查询:")
print(f" {expanded_query}\n")
expanded_response = query_engine.query(expanded_query)
print("💭 AI回答:")
print("-" * 40)
expanded_response.print_response_stream()
print("\n")
# 显示参考文档
print("\n📚 参考依据:")
print("-" * 40)
for i, node in enumerate(expanded_response.source_nodes, 1):
print(f"\n文档片段 {i}:")
print(f"相关度得分: {node.score:.4f}")
print("-" * 30)
print(node.text)
# 方法二：将单一查询改写为多步骤查询
step_decompose_transform = StepDecomposeQueryTransform(verbose=True, llm=llm_client)
response_synthesizer = get_response_synthesizer(llm=llm_client, callback_manager=query_engine.callback_manager,
verbose=True)
multi_step_query_engine = MultiStepQueryEngine(
query_engine=query_engine,
query_transform=step_decompose_transform,
response_synthesizer=response_synthesizer,
index_summary="公司人员信息",
)
print(f"❓ 用户问题: {question}\n")
print("🤖 AI正在进行多步查询...")
multi_step_response = multi_step_query_engine.query(question)
print("\n📚 参考依据:")
print("-" * 40)
for i, node in enumerate(multi_step_response.source_nodes, 1):
print(f"\n文档片段 {i}:")
print("-" * 30)
print(node.text)
# 方法三：用假设文档来增强检索（HyDE）
hyde = HyDEQueryTransform(include_original=True, llm=llm_client)
hyde_query_engine = TransformQueryEngine(
query_engine=query_engine,
query_transform=hyde
)
print(f"❓ 用户问题: {question}\n")
print("🤖 AI正在通过 HyDE 分析...")
hyde_response = hyde_query_engine.query(question)
print("\n💭 AI回答:")
print("-" * 40)
hyde_response.print_response_stream()
# 显示参考文档
print("\n📚 参考依据:")
print("-" * 40)
for i, node in enumerate(hyde_response.source_nodes, 1):
print(f"\n文档片段 {i}:")
print("-" * 30)
print(node.text)
# 显示生成的假设文档
query_bundle = hyde(question)
hyde_doc = query_bundle.embedding_strs[0]
print(f"🤖 AI生成的假想文档:\n{hyde_doc}\n")
if __name__ == "__main__":
main()
7. 总结
通过本文的介绍，我们学习了如何使用 LlamaIndex 构建一个智能文档查询系统。我们探讨了如何加载文档、构建索引、生成扩展查询以及使用多种查询优化技术来提高检索的准确性和效率。希望本文能帮助你更好地理解和应用 LlamaIndex 在实际项目中的使用。
