自己動(dòng)手實(shí)現(xiàn)一個(gè)RAG應(yīng)用
我們知道 RAG 有兩個(gè)核心的過(guò)程，一個(gè)是把信息存放起來(lái)的索引過(guò)程，一個(gè)是利用找到相關(guān)信息生成內(nèi)容的檢索生成過(guò)程。所以，我們這個(gè) RAG 應(yīng)用也要分成兩個(gè)部分：索引和檢索生成。
RAG 是為了讓大模型知道更多的東西，所以，接下來(lái)要實(shí)現(xiàn)的 RAG 應(yīng)用，用來(lái)增強(qiáng)的信息就是我們這門課程的內(nèi)容，我會(huì)把開(kāi)篇詞做成一個(gè)文件，這樣，我們就可以和大模型討論我們的課程了。LangChain 已經(jīng)提供了一些基礎(chǔ)設(shè)施，我們可以利用這些基礎(chǔ)設(shè)施構(gòu)建我們的應(yīng)用。
我們先從索引的過(guò)程開(kāi)始！
圖片
下面是實(shí)現(xiàn)這個(gè)索引過(guò)程的代碼：
from langchain_community.document_loaders import TextLoader
loader = TextLoader("introduction.txt")
docs = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)
vectorstore = Chroma(
collection_name="ai_learning",
embedding_functinotallow=OpenAIEmbeddings(),
persist_directory="vectordb"
)
vectorstore.add_documents(splits)出于簡(jiǎn)化的目的，我這里直接從文本內(nèi)容中加載信息源，而且選擇了 Chroma 作為向量數(shù)據(jù)庫(kù)，它對(duì)開(kāi)發(fā)很友好，可以把向量數(shù)據(jù)存儲(chǔ)在本地的指定目錄下。
我們結(jié)合代碼來(lái)看一下。首先是 TextLoader，它負(fù)責(zé)加載文本信息。
loader = TextLoader("introduction.txt")
docs = loader.load()這里的 TextLoader 屬于 DocumentLoader。在 LangChain 中，有一個(gè)很重要的概念叫文檔（Document），它包括文檔的內(nèi)容（page_content）以及相關(guān)的元數(shù)據(jù)（metadata）。所有原始信息都是文檔，索引信息的第一步就是把這些文檔加載進(jìn)來(lái)，這就是 DocumentLoader 的作用。
除了這里用到的 TextLoader，LangChain 社區(qū)里已經(jīng)實(shí)現(xiàn)了大量的 DocumentLoader，比如，從數(shù)據(jù)庫(kù)里加載數(shù)據(jù)的 SQLDatabaseLoader，從亞馬遜 S3 加載文件的 S3FileLoader?；旧?，大部分我們需要的文檔加載器都可以找到直接的實(shí)現(xiàn)。
拆分加載進(jìn)來(lái)的文檔是 TextSplitter 的主要職責(zé)。
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
splits = text_splitter.split_documents(docs)雖然都是文本，但怎樣拆分還是有講究的，拆分源代碼和拆分普通文本，處理方法就是不一樣的。LangChain 社區(qū)里同樣實(shí)現(xiàn)了大量的 TextSplitter，我們可以根據(jù)自己的業(yè)務(wù)特點(diǎn)進(jìn)行選擇。我們這里使用了 RecursiveCharacterTextSplitter，它會(huì)根據(jù)常見(jiàn)的分隔符（比如換行符）遞歸地分割文檔，直到把每個(gè)塊拆分成適當(dāng)?shù)拇笮 ?/span>
做好基礎(chǔ)的準(zhǔn)備之后，就要把拆分的文檔存放到向量數(shù)據(jù)庫(kù)里了：
vectorstore = Chroma(
collection_name="ai_learning",
embedding_functinotallow=OpenAIEmbeddings(),
persist_directory="vectordb"
)
vectorstore.add_documents(splits)LangChain 支持了很多的向量數(shù)據(jù)庫(kù)，它們都有一個(gè)統(tǒng)一的接口：VectorStore，在這個(gè)接口中包含了向量數(shù)據(jù)庫(kù)的統(tǒng)一操作，比如添加、查詢之類的。這個(gè)接口屏蔽了向量數(shù)據(jù)庫(kù)的差異，在向量數(shù)據(jù)庫(kù)并不為所有程序員熟知的情況下，給嘗試不同的向量數(shù)據(jù)庫(kù)留下了空間。各個(gè)具體實(shí)現(xiàn)負(fù)責(zé)實(shí)現(xiàn)這些接口，我們這里采用的實(shí)現(xiàn)是 Chroma。
在 Chroma 初始化的過(guò)程中，我們指定了 Embedding 函數(shù)，它負(fù)責(zé)把文本變成向量。這里我們采用了 OpenAI 的 Embeddings 實(shí)現(xiàn)，你完全可以根據(jù)自己的需要選擇相應(yīng)的實(shí)現(xiàn)，LangChain 社區(qū)同樣提供了大量的實(shí)現(xiàn)，比如，你可以指定 Hugging Face 這個(gè)模型社區(qū)中的特定模型來(lái)做 Embedding。
到這里，我們就完成了索引的過(guò)程，看上去還是比較簡(jiǎn)單的。為了驗(yàn)證我們索引的結(jié)果，我們可以調(diào)用 similarity_search 檢索向量數(shù)據(jù)庫(kù)的數(shù)據(jù)：
vectorstore = Chroma(
collection_name="ai_learning",
embedding_functinotallow=OpenAIEmbeddings(),
persist_directory="vectordb"
)
documents = vectorstore.similarity_search("專欄的作者是誰(shuí)？")
print(documents)我們這里用的 similarity_search 表示的是根據(jù)相似度進(jìn)行搜索，還可以使用 max_marginal_relevance_search，它會(huì)采用 MMR（Maximal Marginal Relevance，最大邊際相關(guān)性）算法。這個(gè)算法可以在保持結(jié)果相關(guān)性的同時(shí)，盡量選擇與已選結(jié)果不相似的內(nèi)容，以增加結(jié)果的多樣性。
檢索生成
現(xiàn)在，我們已經(jīng)為我們 RAG 應(yīng)用準(zhǔn)備好了數(shù)據(jù)。接下來(lái)，就該正式地構(gòu)建我們的 RAG 應(yīng)用了。我在之前的聊天機(jī)器上做了一些修改，讓它能夠支持 RAG，代碼如下：
from operator import itemgetter
from typing import List
import tiktoken
from langchain_core.messages import BaseMessage, HumanMessage, AIMessage, ToolMessage, SystemMessage, trim_messages
from langchain_core.chat_history import BaseChatMessageHistory, InMemoryChatMessageHistory
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.runnables import RunnablePassthrough
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_openai import OpenAIEmbeddings
from langchain_openai.chat_models import ChatOpenAI
from langchain_chroma import Chroma
vectorstore = Chroma(
collection_name="ai_learning",
embedding_functinotallow=OpenAIEmbeddings(),
persist_directory="vectordb"
)
retriever = vectorstore.as_retriever(search_type="similarity")
def str_token_counter(text: str) -> int:
enc = tiktoken.get_encoding("o200k_base")
return len(enc.encode(text))
def tiktoken_counter(messages: List[BaseMessage]) -> int:
num_tokens = 3
tokens_per_message = 3
tokens_per_name = 1
for msg in messages:
if isinstance(msg, HumanMessage):
role = "user"
elif isinstance(msg, AIMessage):
role = "assistant"
elif isinstance(msg, ToolMessage):
role = "tool"
elif isinstance(msg, SystemMessage):
role = "system"
else:
raise ValueError(f"Unsupported messages type {msg.__class__}")
num_tokens += (
tokens_per_message
+ str_token_counter(role)
+ str_token_counter(msg.content)
)
if msg.name:
num_tokens += tokens_per_name + str_token_counter(msg.name)
return num_tokens
trimmer = trim_messages(
max_tokens=4096,
strategy="last",
token_counter=tiktoken_counter,
include_system=True,
)
store = {}
def get_session_history(session_id: str) -> BaseChatMessageHistory:
if session_id not in store:
store[session_id] = InMemoryChatMessageHistory()
return store[session_id]
model = ChatOpenAI()
prompt = ChatPromptTemplate.from_messages(
[
(
"system",
"""You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
Context: {context}""",
),
MessagesPlaceholder(variable_name="history"),
("human", "{question}"),
]
)
def format_docs(docs):
return "\n\n".join(doc.page_content for doc in docs)
context = itemgetter("question") | retriever | format_docs
first_step = RunnablePassthrough.assign(cnotallow=context)
chain = first_step | prompt | trimmer | model
with_message_history = RunnableWithMessageHistory(
chain,
get_session_history=get_session_history,
input_messages_key="question",
history_messages_key="history",
)
config = {"configurable": {"session_id": "dreamhead"}}
while True:
user_input = input("You:> ")
if user_input.lower() == 'exit':
break
if user_input.strip() == "":
continue
stream = with_message_history.stream(
{"question": user_input},
cnotallow=config
)
for chunk in stream:
print(chunk.content, end='', flush=True)
print()為了進(jìn)行檢索，我們需要指定數(shù)據(jù)源，這里就是我們的向量數(shù)據(jù)庫(kù)，其中存放著我們前面已經(jīng)索引過(guò)的數(shù)據(jù)：
vectorstore = Chroma(
collection_name="ai_learning",
embedding_functinotallow=OpenAIEmbeddings(),
persist_directory="vectordb"
)
retriever = vectorstore.as_retriever(search_type="similarity")這段代碼引入了一個(gè)新的概念：Retriever。從名字不難看出，它就是充當(dāng) RAG 中的 R。Retriever 的核心能力就是根據(jù)文本查詢出對(duì)應(yīng)的文檔（Document）。
為什么不直接使用向量數(shù)據(jù)庫(kù)呢？因?yàn)?Retriever 并不只有向量數(shù)據(jù)庫(kù)一種實(shí)現(xiàn)，比如，WikipediaRetriever 可以從 Wikipedia 上進(jìn)行搜索。所以，一個(gè) Retriever 接口就把具體的實(shí)現(xiàn)隔離開(kāi)來(lái)。
回到向量數(shù)據(jù)庫(kù)上，當(dāng)我們調(diào)用 as_retriever 創(chuàng)建 Retriever 時(shí)，還傳入了搜索類型（search_type），這里的搜索類型和前面講到向量數(shù)據(jù)庫(kù)的檢索方式是一致的，這里我們傳入的是 similarity，當(dāng)然也可以傳入 mmr
You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
Context: {context}在這段提示詞里，我們告訴大模型，根據(jù)提供的上下文回答問(wèn)題，不知道就說(shuō)不知道。這是一個(gè)提示詞模板，在提示詞的最后是我們給出的上下文（Context）。這里上下文是根據(jù)問(wèn)題檢索出來(lái)的內(nèi)容。
有了這個(gè)提示詞，再加上聊天歷史和我們的問(wèn)題，就構(gòu)成了一個(gè)完整的提示詞模板：
prompt = ChatPromptTemplate.from_messages(
[
(
"system",
"""You are an assistant for question-answering tasks. Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. Use three sentences maximum and keep the answer concise.
Context: {context}""",
),
MessagesPlaceholder(variable_name="history"),
("human", "{question}"),
]
)好，我們已經(jīng)理解了這一講的新內(nèi)容，接下來(lái)，就是把各個(gè)組件組裝到一起，構(gòu)成一條完整的鏈：
context = itemgetter("question") | retriever | format_docs
first_step = RunnablePassthrough.assign(cnotallow=context)
chain = first_step | prompt | trimmer | model
with_message_history = RunnableWithMessageHistory(
chain,
get_session_history=get_session_history,
input_messages_key="question",
history_messages_key="history",
)在這段代碼里，我們首先構(gòu)建了一個(gè) context 變量，它也一條鏈。第一步是從傳入?yún)?shù)中獲取到 question 屬性，也就是我們的問(wèn)題，然后把它傳給 retriever。retriever 會(huì)根據(jù)問(wèn)題去做檢索，對(duì)應(yīng)到我們這里的實(shí)現(xiàn)，就是到向量數(shù)據(jù)庫(kù)中檢索，檢索的結(jié)果是一個(gè)文檔列表。
文檔是 LangChain 應(yīng)用內(nèi)部的表示，要傳給大模型，我們需要把它轉(zhuǎn)成文本，這就是 format_docs 做的事情，它主要是把文檔內(nèi)容取出來(lái)拼接到一起：
def format_docs(docs):
return "\n\n".join(doc.page_content for doc in docs)這里補(bǔ)充幾句實(shí)現(xiàn)細(xì)節(jié)。在 LangChain 代碼里， | 運(yùn)算符被用作不同組件之間的連接，其實(shí)現(xiàn)的關(guān)鍵就是大部分組件都實(shí)現(xiàn)了 Runnable 接口，在這個(gè)接口里實(shí)現(xiàn)了 __or__ 和 __ror__。__or__ 表示這個(gè)對(duì)象出現(xiàn)在| 左邊時(shí)的處理，相應(yīng)的 __ror__ 表示這個(gè)對(duì)象出現(xiàn)在右邊時(shí)的處理。
Python 在處理 a | b 這個(gè)表達(dá)式時(shí)，它會(huì)先嘗試找 a 的 __or__，如果找不到，它會(huì)嘗試找 b 的 __ror__。所以，在 context 的處理中， 來(lái)自標(biāo)準(zhǔn)庫(kù)的 itemgetter 雖然沒(méi)有實(shí)現(xiàn)
__or__，但 retriever 因?yàn)閷?shí)現(xiàn)了 Runnable 接口，所以，它也實(shí)現(xiàn)了 __ror__。所以，這段代碼才能組裝出我們所需的鏈。
有了 context 變量，我們可以用它構(gòu)建了另一個(gè)變量 first_step：
first_step = RunnablePassthrough.assign(cnotallow=context)還記得我們的提示詞模板里有一個(gè) context 變量嗎？它就是從這里來(lái)的。
RunnablePassthrough.assign 這個(gè)函數(shù)就是在不改變鏈當(dāng)前狀態(tài)值的前提下，添加新的狀態(tài)值。前面我們說(shuō)了，這里賦給 context 變量的值是一個(gè)鏈，我們可以把它理解成一個(gè)函數(shù)，它會(huì)在運(yùn)行期執(zhí)行，其參數(shù)就是我們當(dāng)前的狀態(tài)值?，F(xiàn)在你可以理解 itemgetter(“question”) 的參數(shù)是從哪來(lái)的了。這個(gè)函數(shù)的返回值會(huì)用來(lái)在當(dāng)前的狀態(tài)里添加一個(gè)叫 context 的變量，以便在后續(xù)使用。
其余的代碼我們之前已經(jīng)講解過(guò)了，這里就不再贅述了。至此，我們擁有了一個(gè)可以運(yùn)行的 RAG 應(yīng)用，我們可以運(yùn)行一下看看效果：
You:> 專欄的作者是誰(shuí)？
專欄的作者是鄭曄。
You:> 作者還寫過(guò)哪些專欄？
作者鄭曄還寫過(guò)《10x程序員工作法》、《軟件設(shè)計(jì)之美》、《代碼之丑》和《程序員的測(cè)試課》這四個(gè)專欄。
