LlamaIndex 实现 RAG （一）

理解过 LlamaIndex 的功能之后，本文通过 LlamaIndex 快速实现一个简单的 RAG 应用，主要包括以下几个部分：

1. 创建知识库，并进行 Embedding
2. 集成本地 Ollama 模型或者 Qwen 模型
3. 通过 Streamlit 可视化 RAG

文末提供了源代码地址

创建 LlamaIndex 环境

安装 LlamaIndex 依赖，LlamaIndex是模块化设计，每个集成点在 LlamaIndex 中都是

pip install llama-index

集成 LLM 模型

本地模型

1. 集成 Ollama 本地模型

集成 Ollama 本地大语言模型，需要安装 Ollama 依赖

pip install llama-index-llms-ollama

2. 集成 Ollama Qwen2:7b

如需要安装 Ollama 环境，请移步至安装文档。首先下载 Qwen2:7b

ollama pull qwen2:7b

下载完成之后，运行Qwen，确保安装成功。

ollama run qwen2:7b

3. 通过 LlamaIndex 集成 Ollama 并调用模型，初始化 Ollama 时，如果使用的不是本地的 Ollama，需要指定 base_url，如果使用的是本地的 Ollama 可以忽略该参数，使用 LlamaIndex 阅读源码能快速理解 API，文档写的很清晰。调用模型时可以使用 Complete 或者 Chat，区别单轮和多轮对话，API 带有前缀 “stream” 时，代表流式输出。

from llama_index.llms.ollama import Ollama
from llama_index.core.llms import ChatMessage

# 创建 Ollama
llm = Ollama(base_url="http://10.91.3.116:11434", model="qwen2:7b", request_timeout=30.0)

# # 调用模型

response = llm.complete("介绍一下故宫")

print (response)
## 流式
response = llm.stream_complete("介绍一下故宫")
for r in response:
    print(r.delta, end="")


# # 多轮对话
messages = [
    ChatMessage(
        role="system", content="你是一个幽默风趣的导游"
    ),
    ChatMessage(role="user", content="介绍一下故宫"),
]
response = llm.chat(messages)
print (response)

## 流式
response = llm.stream_chat(messages)
for r in response:
    print(r.delta, end="")

通义千问模型

本例中 Ollama 使用的是 7B 模型，有时本地模型精确度达不到应用的要求，需要用更高参数的模型，可以是使用通义千问的商用模型，具体费用可以参考官网。要使用通义千问，需要安装 Dashscope 依赖，并申请 API KEY，在项目目录下创建 .env, 并配置 DASHSCOPE_API_KEY。

pip install load_dotenv
pip install llama-index-llms-dashscope

初始化通义千问模型

from llama_index.llms.dashscope import DashScope, DashScopeGenerationModels
from dotenv import load_dotenv
def get_tongyi():
    load_dotenv()
    return DashScope(model_name=DashScopeGenerationModels.QWEN_TURBO)

模型用法同本地模型。

文档索引

实现 RAG 之前，我们需要创建本地知识库，这里我们拿单个文档进行举例，一个功能完善的知识库需要包括很多功能，但召回的底层逻辑是相似的。

首先导入文档，并对文档进行切块，随后进行 Embedding 并保存到文件中。生产环境中我们需要将文档保存到向量数据中，以保证数据可靠性，后面会单独去讲向量数据库存储的方式。知识库的创建完成之后，就可以开始进行 Retrieve 阶段的代码开发：

from llama_index.core import SimpleDirectoryReader
from llama_index.core import VectorStoreIndex
from llama_index.embeddings.ollama import OllamaEmbedding

# 创建 OllamaEmbedding 实例，用于指定嵌入模型和服务的基本 URL
ollama_embedding = OllamaEmbedding(
    model_name="nomic-embed-text",
    base_url="http://10.91.3.116:11434"
)

# 读取 "./data" 目录中的数据并加载为文档对象
documents = SimpleDirectoryReader("./data").load_data()

# 从文档中创建 VectorStoreIndex，并使用 OllamaEmbedding 作为嵌入模型
vector_index = VectorStoreIndex.from_documents(documents, embed_model=ollama_embedding)
vector_index.set_index_id("vector_index")  # 设置索引 ID
vector_index.storage_context.persist("./storage")  # 将索引持久化到 "./storage"

# 创建 retriever
query_engine = vector_index.as_retriever()

response = query_engine.retrieve("当前社保情况?")

LLM 增强

数据从本地知识库查询结束之后，需要构建查询上下文，并将 Query 一同传给大模型进行回答。LlamaIndex 通过 Setting 进行模型大模型设置，LlamaIndex 运行时会通过 Settings.llm 获取当前使用的 LLM，并进一步对模型进行的操作。index.as_query_engine() 这行始终点，LlamaIndex 把 RAG 的主流程都封装在这个方法中了。

from llm import get_llm
from docs import get_index
from llama_index.core import Settings

Settings.llm = get_llm()

index = get_index()

query  = index.as_query_engine()

print (query.query("当前社保情况?"))

可视化 RAG

通过 StreamLit 可视化 RAG，StreamLit 是一个 Python 可是化框架，和 Gradio 类似，通过可视化可以直接让我们上面完成代码转为可视化 Demo。首先安装 StreamLit 依赖包

pip install streamlit

streamlit 组件都已经内置了，自己需要我们对其进行设计就好。

from llm import get_llm
from docs import get_index
from llama_index.core import Settings
import streamlit as st
# 设置 LLM
Settings.llm = get_llm()



@st.cache_resource(show_spinner=False)
def load_data():
    return get_index()

index = load_data()

# 初始化消息历史
if "messages" not in st.session_state.keys():  
    st.session_state.messages = [
        {
            "role": "assistant",
            "content": "问我一些关于2023 年度人力资源和社会保障事业发展统计公报的问题!",
        }
    ]

if "chat_engine" not in st.session_state.keys():  # 初始化引擎
    st.session_state.chat_engine = index.as_chat_engine(
        chat_mode="condense_question", verbose=True, streaming=True
    )

if prompt := st.chat_input(
    "请提问"
):  # 用户输入
    st.session_state.messages.append({"role": "user", "content": prompt})
    
for message in st.session_state.messages:  # 消息历史
    with st.chat_message(message["role"]):
        st.write(message["content"])

# 发送请求到 Engine 进行 Prompt 回答
if st.session_state.messages[-1]["role"] != "assistant":
    with st.chat_message("assistant"):
        response_stream = st.session_state.chat_engine.stream_chat(prompt)
        st.write_stream(response_stream.response_gen)
        message = {"role": "assistant", "content": response_stream.response}
        # 添加历史消息
        st.session_state.messages.append(message)

启动 Stream

streamlit run simple_rag.py    

总结

本文简单实现 RAG 的流程，并实现了可视化页面，关于文档的切分以及 RAG 不同组件的详细讲解，会在后续文章继续分享。

代码地址：https://gitcode.com/hawk2014bj/llamaindex/overview

源代码也可以从资源中下载。