llama-index做一个简单查询工作流（rag） - qidao123.com技术社区-IT企服评测·应用市场

from llama_index.llms.zhipuai import ZhipuAI
llms = ZhipuAI(model="glm-4")#api要设置

复制代码

from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.llms.huggingface import HuggingFaceLLM
from llama_index.core import Settings
llm=HuggingFaceLLM(
model_name="LLM_model/Qwen2.5-0.5B-Instruct",
tokenizer_name="LLM_model/Qwen2.5-0.5B-Instruct",
context_window=30000,
max_new_tokens=2000,
generate_kwargs={"temperature": 0.7, "top_k": 50, "top_p": 0.95},
device_map="auto")

复制代码

from llama_index.embeddings.huggingface import HuggingFaceEmbedding
from llama_index.core import Settings
Settings.embed_model = HuggingFaceEmbedding(
model_name="LLM_model/embeddding-bge-large-zh-v1.5",trust_remote_code=True)

复制代码

#事件
from llama_index.core.workflow import Event
class QueryPlanItemResult(Event):
"""The result of a query plan item"""
query: str
result: str
class ExecutedPlanEvent(Event):
"""The result of a query plan"""
result: str

复制代码

from pydantic import BaseModel, Field
from llama_index.core.workflow import Event
class QueryPlanItem(Event):
"""A single step in an execution plan for a RAG system."""
name: str = Field(description="The name of the tool to use.")
query: str = Field(description="A natural language search query for a RAG system.")
class QueryPlan(BaseModel):
"""A plan for a RAG system. After running the plan, we should have either enough information to answer the user's original query, or enough information to form a new query plan."""
items: list[QueryPlanItem] = Field(
description="A list of the QueryPlanItem objects in the plan."
)
'''通过astructured传入模型即可得到结构化输出
query_plan = await llm.astructured_predict(
QueryPlan,
self.planning_prompt,
context=context_str,
query=query,
)'''

复制代码

class QueryPlanningWorkflow(Workflow):
#让模型进行规划根据上下文内容和用户输入
planning_prompt = PromptTemplate(
"Think step by step. Given an initial query, as well as information about the indexes you can query, return a plan for a RAG system.\n"
"The plan should be a list of QueryPlanItem objects, where each object contains a query.\n"
"The result of executing an entire plan should provide a result that is a substantial answer to the initial query, "
"or enough information to form a new query plan.\n"
"Sources you can query: {context}\n"
"Initial query: {query}\n"
"Plan:"
)
#根据目前结果，结束还是举行执行任务
decision_prompt = PromptTemplate(
"Given the following information, return a final response that satisfies the original query, or return 'PLAN' if you need to continue planning.\n"
"Original query: {query}\n"
"Current results: {results}\n"
)
@step
async def planning_step(
self, ctx: Context, ev: StartEvent | ExecutedPlanEvent
) -> QueryPlanItem | StopEvent:
if isinstance(ev, StartEvent):
# stratevent的时候才会进行Initially, we need to plan
query = ev.get("query")
tools = ev.get("tools")
await ctx.set("tools", {t.metadata.name: t for t in tools})
await ctx.set("original_query", query)
# 将工具的信息作为上下文导入，使得模型看见工具
context_str = "\n".join(
[
f"{i+1}. {tool.metadata.name}: {tool.metadata.description}"
for i, tool in enumerate(tools)
]
)
await ctx.set("context", context_str)
#planning_prompt导入，要求返回列表，列表里面是使用的工具和查询的自然语言
#这里有需要也可以吧大模型设置成类变量，我这里直接改成上面定义的全局变量了
query_plan = await llm.astructured_predict(
QueryPlan,
self.planning_prompt,
context=context_str,
query=query,
)
ctx.write_event_to_stream(
Event(msg=f"Planning step: {query_plan}")
)
num_items = len(query_plan.items)
await ctx.set("num_items", num_items)
#返回事件
for item in query_plan.items:
ctx.send_event(item)#将事件发送到工作流中的特定步骤。如果 step 为 None，则事件将发送给所有接收者，由它们决定是否丢弃不需要的事件。
else:
#ExecutedPlanEvent时会执行
query = await ctx.get("original_query")
current_results_str = ev.result
##根据目前结果，结束还是举行执行任务
decision = await llm.apredict(
self.decision_prompt,
query=query,
results=current_results_str,
)
# 简单的匹配，如果里面有plan就继续规划，否则就结束，疑惑：如果模型输出不需要plan，怎么结束？
if "PLAN" in decision:
context_str = await ctx.get("context")
query_plan = await llm.astructured_predict(
QueryPlan,
self.planning_prompt,
context=context_str,
query=query,
)
ctx.write_event_to_stream(
Event(msg=f"Re-Planning step: {query_plan}")
)
num_items = len(query_plan.items)
await ctx.set("num_items", num_items)
for item in query_plan.items:
ctx.send_event(item)
else:
#如果没有plan就结束
return StopEvent(result=decision)
@step(num_workers=4)
async def execute_item(
self, ctx: Context, ev: QueryPlanItem
) -> QueryPlanItemResult:
tools = await ctx.get("tools")
tool = tools[ev.name]
ctx.write_event_to_stream(
Event(
msg=f"Querying tool {tool.metadata.name} with query: {ev.query}"
)
)
result = await tool.acall(ev.query)#传入的工具定义的方法，直接把语句放到查询引擎中，然后返回结果
ctx.write_event_to_stream(
Event(msg=f"Tool {tool.metadata.name} returned: {result}")
)
return QueryPlanItemResult(query=ev.query, result=str(result))
@step
async def aggregate_results(
self, ctx: Context, ev: QueryPlanItemResult
) -> ExecutedPlanEvent:
num_items = await ctx.get("num_items")
# 收集全部的事件才会结束，否则会一直等待
results = ctx.collect_events(ev, [QueryPlanItemResult] * num_items)
if results is None:
return
#汇总结果
aggregated_result = "\n------\n".join(
[
f"{i+1}. {result.query}: {result.result}"
for i, result in enumerate(results)
]
)
return ExecutedPlanEvent(result=aggregated_result)

复制代码

# 在这里，我们用于加载和解析文档，调用LlamaParseapi解析文件
from llama_parse import LlamaParse
api_key=''#去llama-cloud注册有免费额度的
parser = LlamaParse(api_key=api_key,fast_mode=True)
from llama_index.core import (
VectorStoreIndex,
StorageContext,
load_index_from_storage,
)
from llama_index.core.tools import QueryEngineTool
folder = "./data_pdf/"
files = os.listdir(folder)
query_engine_tools = []
for file in files:
#取名字第一个字符
number = file.split(" - ")[0]#根据论文名字的分隔符确定，防止文章太长
index_persist_path = f"./storage/paper-{number}/"
#创建储存路径，如果存在就加载，不存在就创建
if os.path.exists(index_persist_path):
storage_context = StorageContext.from_defaults(
persist_dir=index_persist_path
)
index = load_index_from_storage(storage_context)
else:
documents = await parser.aload_data(folder + file)
index = VectorStoreIndex.from_documents(documents)
index.storage_context.persist(index_persist_path)
#这里需要输入llm，不然会自动调用openai的大模型
engine = index.as_query_engine(llm=llm)
query_engine_tools.append(
QueryEngineTool.from_defaults(
engine,
name=f"TimeSeriesAnalysis",
description=f"This is a paper on time series analysis, which includes information on time series forecasting and attention mechanisms."
)
)

复制代码

from llama_index.utils.workflow import draw_all_possible_flows
draw_all_possible_flows(
QueryPlanningWorkflow, filename="twwww.html"
)

复制代码

#测试工作流
workflow = QueryPlanningWorkflow(verbose=False, timeout=120)
# run the workflow
handler = workflow.run(
query="什么是时间序列预测",
tools=query_engine_tools,
)
#流式处理
async for event in handler.stream_events():
if hasattr(event, "msg"):
print(event.msg)
result = await handler

复制代码