基于Agent的金融問答系統：RAG的檢索增強之上下文重排和壓縮 原創
前言
在上一章??【項目實戰】基于Agent的金融問答系統：RAG的檢索增強之ElasticSearch??中，我們在已經構建的Agent框架上，通過集成檢索器將ES檢索器和多路召回檢索器集成進來，提升了檢索的召回率。本章，我們將基于項目問題，進一步優化Agent的檢索能力。
問題
隨著該項目檢索能力的增強，我們計劃對天池大賽提供的1000個問題進行執行，依次評估我們的系統能力如何。 在這個測試過程中，我們發現了多個待優化的問題，其中有2個問題的解決值得分享，在此作為記錄以供讀者參考。
問題1：執行過程中，偶爾會出現RAG檢索結果內容過長，超出大模型能夠接收的范圍(如下圖中顯示的status_code=400)，導致執行中斷
問題2：Agent遇到計算漲幅、收益率等問題時，會反復構造SQL語句，試圖從數據庫中直接查詢出對應的數據，從而導致思考迭代此處超過限制，程序異常
優化方案
分析上述問題后，我們計劃對整體程序進行優化，優化方案如下：
說明：
- 針對問題1，我們計劃給集成檢索器增加上下文重排和壓縮，以解決長文本問題。
- 針對問題2，我們計劃為Agent掛載更多的工具，進而讓Agent使用工具計算漲幅、收益率等。
優化步驟
1、檢索器增加上下文壓縮
優化代碼文件：??app/rag/retrievers.py??
from langchain_core.callbacks importCallbackManagerForRetrieverRun
from utils.logger_config importLoggerManager
from langchain_core.retrievers importBaseRetriever
from langchain_core.documents importDocument
from langchain.retrievers importEnsembleRetriever
from langchain.retrievers.multi_query importMultiQueryRetriever
from langchain.retrievers.contextual_compression importContextualCompressionRetriever
from langchain.retrievers.document_compressors importLLMChainExtractor
from rag.elasticsearch_db importElasticsearchDB
# ES需要導入的庫
from typing importList
import logging
import settings
from langchain_community.document_transformers import(
LongContextReorder,
)
from utils.util import get_rerank_model
logger =LoggerManager().logger
classSimpleRetrieverWrapper():
"""自定義檢索器實現"""
def__init__(self, store, llm, **kwargs):
self.store = store
self.llm = llm
logger.info(f'檢索器所使用的Chat模型：{self.llm}')
defcreate_retriever(self):
logger.info(f'初始化自定義的Retriever')
# 初始化一個空的檢索器列表
retrievers =[]
weights =[]
# Step1：創建一個 多路召回檢索器 MultiQueryRetriever
chromadb_retriever = self.store.as_retriever()
mq_retriever =MultiQueryRetrieverWrapper.from_llm(retriever=chromadb_retriever, llm=self.llm)
# Step2：創建一個 上下文壓縮檢索器ContextualCompressionRetriever
if settings.COMPRESSOR_ENABLE isTrue:
compressor =LLMChainExtractor.from_llm(llm=self.llm)
compression_retriever =ContextualCompressionRetrieverWrapper(
base_compressor=compressor, base_retriever=mq_retriever
)
# 開啟開關就使用壓縮檢索器
retrievers.append(compression_retriever)
weights.append(0.5)
logger.info(f'已啟用 ContextualCompressionRetriever')
else:
# 關閉開關就使用多路召回檢索器
retrievers.append(mq_retriever)
weights.append(0.5)
logger.info(f'已啟用 MultiQueryRetriever')
# Step3：創建一個 ES 檢索器
if settings.ELASTIC_ENABLE_ES isTrue:
es_retriever =ElasticsearchRetriever()
retrievers.append(es_retriever)
weights.append(0.5)
logger.info(f'已啟用 ElasticsearchRetriever')
# 使用集成檢索器，將所有啟用的檢索器集合在一起
ensemble_retriever =EnsembleRetriever(retrievers=retrievers, weights=weights)
return ensemble_retriever說明：
- 首先，我們創建一個空的檢索器列表，用于存儲啟用的檢索器。
- 其次，創建一個 ?
?MultiQueryRetrieverWrapper?? ，用于將ES檢索器與多路召回檢索器集成。 - 然后，通過 ?
?ContextualCompressionRetrieverWrapper?? ，為多路檢索器添加上下文壓縮功能。 - 最后，將檢索器列表與權重列表傳入集成檢索器，完成集成。
此處，MultiQueryRetrieverWrapper 和 ContextualCompressionRetrieverWrapper 分別是基于 ?
?MultiQueryRetriever??? 和 ??ContextualCompressionRetriever?? 進一步封裝實現的，在3、中會詳細介紹。
2、檢索器增加上下文重排
對之前實現的ElasticsearchRetriever增加上下文重排功能，具體代碼如下：
優化代碼文件：??app/rag/retrievers.py??
class ElasticsearchRetriever(BaseRetriever):
def_get_relevant_documents(self, query: str, )->List[Document]:
"""Return the first k documents from the list of documents"""
es_connector =ElasticsearchDB()
query_result = es_connector.search(query)
# 增加長上下文重排序
reordering =LongContextReorder()
reordered_docs = reordering.transform_documents(query_result)
# logger.info(f"ElasticSearch檢索到的原始文檔：")
# for poriginal in query_result:
# logger.info(f"{poriginal}")
logger.info(f"ElasticSearch檢索重排后的文檔：")
for preordered in reordered_docs:
logger.info(f"{preordered}")
logger.info(f"ElasticSearch檢索到資料文件個數：{len(query_result)}")
if reordered_docs:
return[Document(page_cnotallow=doc)for doc in reordered_docs]
return []3、優化日志輸出
由于 ??MultiQueryRetriever?? 是langchain已經封裝好的檢索器，如果我們需要在其基礎上增加一些功能，比如：增加日志，我們需要對其進行重寫，具體方法：
重寫MultiQueryRetriever
創建一個新的Class ??MultiQueryRetrieverWrapper???，繼承 ??MultiQueryRetriever??? ，重寫 ??_get_relevant_documents?? 方法，具體代碼如下：
class MultiQueryRetrieverWrapper(MultiQueryRetriever):
def_get_relevant_documents(
self,
query: str,
*,
run_manager: CallbackManagerForRetrieverRun,
)->List[Document]:
"""
對MultiQueryRetriever進行重寫，增加日志打印
"""
queries = self.generate_queries(query, run_manager)
if self.include_original:
queries.append(query)
documents = self.retrieve_documents(queries, run_manager)
# 增加長上下文重排序
reordering =LongContextReorder()
reordered_docs = reordering.transform_documents(documents)
logger.info(f'MultiQuery生成的檢索語句：')
for q in queries:
logger.info(f"{q}")
logger.info(f'MultiQuery檢索到的資料文件：')
for doc in documents:
logger.info(f"{doc}")
logger.info(f"MultiQuery檢索到資料文件個數：{len(documents)}")
return self.unique_union(reordered_docs)重寫ContextualCompressionRetriever
創建一個新的Class ??ContextualCompressionRetrieverWrapper???，繼承 ??ContextualCompressionRetriever??? ，重寫 ??_get_relevant_documents?? 方法，具體代碼如下：
class ContextualCompressionRetrieverWrapper(ContextualCompressionRetriever):
from typing importAny,List
def_get_relevant_documents(
self,
query: str,
*,
run_manager: CallbackManagerForRetrieverRun,
**kwargs: Any,
)->List[Document]:
"""
對ContextualCompressionRetriever進行重寫，增加日志打印
"""
docs = self.base_retriever.invoke(
query, cnotallow={"callbacks": run_manager.get_child()},**kwargs
)
if docs:
compressed_docs = self.base_compressor.compress_documents(
docs, query, callbacks=run_manager.get_child()
)
logger.info(f'壓縮后的文檔長度：{len(compressed_docs)}')
logger.info(f'壓縮后的文檔：{compressed_docs}')
returnlist(compressed_docs)
else:
return []4、給Agent增加更多工具
實現工具函數：計算股票年化收益率
代碼文件：??app/finance_bot_ex.py??
# 定義股票年化收益率計算函數
# 年化收益率定義為：（（有記錄的一年的最終收盤價-有記錄的一年的年初當天開盤價）/有記錄的一年的當天開盤價）* 100%。
def calculate_stock_annualized_return(final_closing_price: float, initial_opening_price: float) -> float:
"""
計算股票年化收益率
"""
annualized_return = ((final_closing_price - initial_opening_price) / initial_opening_price) * 100
return annualized_returnAgent增加工具
在??init_agent??函數中，增加工具函數，具體代碼如下：
代碼優化文件：??app/finance_bot_ex.py??
definit_agent(self):
# 初始化 RAG 工具
retriever_tool = self.init_rag_tools()
# 初始化 SQL 工具
sql_tools = self.init_sql_tool(settings.SQLDATABASE_URI)
# 創建系統Prompt提示語
system_prompt = self.create_sys_prompt()
# 創建Agent
agent_executor = create_react_agent(
self.chat,
tools=[
get_datetime,
calculate_stock_annualized_return,# 增加自定義的計算年化收益率工具
retriever_tool]+ sql_tools,
state_modifier=system_prompt,
checkpointer=MemorySaver()
# state_modifier=modify_state_messages,
)
return agent_executor4、驗證測試
完成上述檢索器的優化之后，我們使用test_framework.py進行驗證，驗證結果如下：
問題1解決效果
1. 用戶輸入問題，觸發RAG檢索的 MultiQueryRetriever 檢索器
2. 通過上下文壓縮之后，原本 MultiQueryRetriever 檢索到的資料文件數量由12個減少到1個。
3. 通過上下文壓縮之后，原本 ElasticsearchRetriever 檢索到的資料文件數量由3個減少到2個。
4. 最后，將兩個檢索器的結果進行整合后，大模型給出最終答案。
問題2解決效果
1. 我們輸入問題：計算代碼000798股票在2020年的年化收益率，保留兩位小數。
2. Agent分析問題后，形成SQL查詢語句：獲取2020年第一天開盤價以及2020年最后一天的收盤價
3. Agent獲得開盤價和收盤價之后，調用我們提供的 calculate_stock_annualized_return 函數計算年化收益率，并返回結果。
結束語
基于Agent的金融問答系統系列文章在此告一段落了。 在這個項目中，我們不只利用AI技術完成了項目課題，其中也不乏應用了軟件工程的一些方法論，而最為重要的是：我們解決了一個又一個的問題。 最后，附帶一張黑神話悟空的圖片，希望看到此篇文章的你我一同共勉。
?
本文轉載自公眾號一起AI技術 作者：Dongming
