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：創(chuàng)建一個 多路召回檢索器 MultiQueryRetriever
        chromadb_retriever = self.store.as_retriever()
        mq_retriever =MultiQueryRetrieverWrapper.from_llm(retriever=chromadb_retriever, llm=self.llm)

# Step2：創(chuàng)建一個 上下文壓縮檢索器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：創(chuàng)建一個 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)

# 創(chuàng)建系統(tǒng)Prompt提示語
        system_prompt = self.create_sys_prompt()

# 創(chuàng)建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