基于Agent的金融问答系统：RAG的检索增强之上下文重排和压缩 原创

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

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 []

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)

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 []

# 定义股票年化收益率计算函数
# 年化收益率定义为：（（有记录的一年的最终收盘价-有记录的一年的年初当天开盘价）/有记录的一年的当天开盘价）* 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_return

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_executor