使用 LlamaIndex 构建多轮文档问答系统

模块与库

导入所需模块

1import streamlit as st
2from llama_index.core import VectorStoreIndex, SimpleDirectoryReader
3from llama_index.embeddings.ollama import OllamaEmbedding
4from llama_index.llms.groq import Groq
5import os
6import tempfile
7import asyncio

初始化嵌入模型和大语言模型

1# 初始化 Groq 大语言模型
2llm = Groq(
3    model="llama3-70b-8192",
4    api_key="YOUR_GROQ_API_KEY"  # 请替换为实际的 API 密钥
5)
6
7# 初始化嵌入模型
8embedding_model = OllamaEmbedding(model_name="nomic-embed-text")

设置 Streamlit 界面

1st.set_page_config(page_title="多轮文档问答")
2st.title("基于PDF文档的多轮问答系统")
3st.markdown("上传一个或多个PDF文档，然后提出问题或请求生成摘要。")
4
5uploaded_files = st.file_uploader("上传PDF文件", type=["pdf"], accept_multiple_files=True)
6query = st.text_input("请输入您的问题：")
7summarize_clicked = st.button("生成文档摘要")

每会话加载并索引一次上传的文档

1if uploaded_files and "documents" not in st.session_state:
2    with st.spinner("正在加载并索引文档..."):
3        with tempfile.TemporaryDirectory() as temp_dir:
4            for uploaded_file in uploaded_files:
5                file_path = os.path.join(temp_dir, uploaded_file.name)
6                with open(file_path, "wb") as f:
7                    f.write(uploaded_file.getbuffer())
8            docs = SimpleDirectoryReader(input_dir=temp_dir).load_data()
9            idx = VectorStoreIndex.from_documents(docs, embed_model=embedding_model)
10
11            st.session_state.documents = docs
12            st.session_state.index = idx
13
14documents = st.session_state.get("documents")
15index = st.session_state.get("index")

导入记忆模块

1from llama_index.core.memory import ChatMemoryBuffer, VectorMemory, SimpleComposableMemory
2from llama_index.core.llms import ChatMessage

初始化主记忆和辅助记忆

1# 短期对话记忆
2chat_memory = ChatMemoryBuffer.from_defaults()
3
4# 基于向量的语义记忆
5vector_memory = VectorMemory.from_defaults(
6    embed_model=embedding_model,
7    retriever_kwargs={"similarity_top_k": 3}
8)
9
10# 组合两种记忆机制
11composable_memory = SimpleComposableMemory.from_defaults(
12    primary_memory=chat_memory,
13    secondary_memory_sources=[vector_memory]
14)

存储与检索对话记录

1# 存储用户消息
2user_msg = ChatMessage(role="user", content=query)
3composable_memory.put(user_msg)
4
5# 运行智能体处理请求
6agent_response = await agent.run(query)
7
8# 存储助手回复
9assistant_msg = ChatMessage(role="assistant", content=str(agent_response))
10composable_memory.put(assistant_msg)
11
12# 显示结果
13st.subheader("回答")
14st.write(str(agent_response))

将完整文档内容发送至大语言模型，并附上摘要提示词，展示生成结果

1if summarize_clicked and documents:
2    with st.spinner("正在生成摘要..."):
3        full_text = "\n".join(doc.text for doc in documents)
4        prompt = f"请对以下文档内容进行总结：\n\n{full_text}"
5        response = llm.complete(prompt)
6        st.subheader("文档摘要")
7        st.write(response.text)

结合记忆与RAG处理问题：将召回的记忆信息与检索到的文档内容整合进提示词，生成回复，并保存助手的回答

1if query and index:
2    async def process_query():
3        # 存储用户消息
4        user_msg = ChatMessage(role="user", content=query)
5        composable_memory.put(user_msg)
6
7        # 回忆相关记忆
8        memory_items = composable_memory.get(query)
9        recalled_memory = "\n".join(f"{msg.role.capitalize()}: {msg.content}" for msg in memory_items)
10
11        # 检索相关文档片段
12        retriever = index.as_retriever(similarity_top_k=3)
13        retrieved_nodes = retriever.retrieve(query)
14        document_context = "\n".join(node.get_content() for node in retrieved_nodes)
15
16        # 构建完整提示词
17        prompt = f"""您需要根据上传的文档内容和记忆信息回答问题。
18
19记忆内容：
20{recalled_memory}
21
22相关文档片段：
23{document_context}
24
25用户问题：
26{query}
27"""
28
29        # 生成回复
30        response = llm.complete(prompt)
31
32        # 存储助手回复
33        assistant_msg = ChatMessage(role="assistant", content=response.text)
34        composable_memory.put(assistant_msg)
35
36        # 显示结果
37        st.subheader("回答")
38        st.write(response.text)
39
40    asyncio.run(process_query())

注册简单追踪处理器

1import llama_index.core
2
3llama_index.core.set_global_handler("simple")