class SimpleRetrieverWrapper():
"""自定義檢索器實(shí)現(xiàn)"""

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')

        chromadb_retriever = self.store.as_retriever()
return chromadb_retriever

docker network create es-net

docker pull elasticsearch:8.6.0

smart-finance-bot \
    |- app \                
    |- docker \
        |- elasticsearch \       # 創(chuàng)建elasticsearch掛載目錄
            |- data \            # 創(chuàng)建數(shù)據(jù)目錄
            |- config \          # 創(chuàng)建配置目錄
            |- plugins  \        # 創(chuàng)建插件目錄

docker run -d \
--restart=always \
--name es \
--network es-net \
-p 9200:9200 \
-p 9300:9300 \
--privileged \
-v /Users/deadwalk/Code/smart-finance-bot/docker/elasticsearch/data:/usr/share/elasticsearch/data \
-v /Users/deadwalk/Code/smart-finance-bot/docker/elasticsearch/plugins:/usr/share/elasticsearch/plugins \
-e "discovery.type=single-node" \
-e "ES_JAVA_OPTS=-Xms512m -Xmx512m" \
elasticsearch:8.6.0

docker exec -it es /bin/bash

bin/elasticsearch-reset-password -i -u elastic

def test_es_connect():
from elasticsearch importElasticsearch

    ELASTIC_PASSWORD ="123abc"
    host ="localhost"
    port =9200
    schema ="https"
    url =f"{schema}://elastic:{ELASTIC_PASSWORD}@{host}:{port}"

    client =Elasticsearch(
        url,
        verify_certs=False,
)

print(client.info())

# 引入
from langchain_core.retrievers importBaseRetriever
from langchain_core.documents importDocument
# ES需要導(dǎo)入的庫(kù)
from typing importList
import re
import jieba
import nltk
from nltk.corpus import stopwords
import time
from elasticsearch importElasticsearch
from elasticsearch.exceptions importConnectionError,AuthenticationException
from elasticsearch import helpers
import settings
from utils.logger_config importLoggerManager
from utils.util_nltk importUtilNltk
import os
import warnings


warnings.simplefilter("ignore")# 屏蔽 ES 的一些Warnings
utilnltk =UtilNltk()



logger =LoggerManager().logger


classTraditionDB:
defadd_documents(self, docs):
"""
        將文檔添加到數(shù)據(jù)庫(kù)
        """
raiseNotImplementedError("Subclasses should implement this method!")

defget_store(self):
"""
        獲得向量數(shù)據(jù)庫(kù)的對(duì)象實(shí)例
        """
raiseNotImplementedError("Subclasses should implement this method!")


classElasticsearchDB(TraditionDB):
def__init__(self,
                 schema=settings.ELASTIC_SCHEMA,
                 host=settings.ELASTIC_HOST,
                 port=settings.ELASTIC_PORT,
                 index_name=settings.ELASTIC_INDEX_NAME,
                 k=3
                 #  docs=docs
                 ):
# 定義索引名稱(chēng)
        self.index_name = index_name
        self.k = k

try:
            url =f"{schema}://elastic:{settings.ELASTIC_PASSWORD}@{host}:{port}"
            logger.info(f'初始化ES服務(wù)連接：{url}')

            self.es =Elasticsearch(
                url,
                verify_certs=False,
# ca_certs="./docker/elasticsearch/certs/ca/ca.crt",
# basic_auth=("elastic", settings.ELASTIC_PASSWORD)
)

            response = self.es.info()# 嘗試獲取信息
            logger.info(f'ES服務(wù)響應(yīng): {response}')
except(ConnectionError,AuthenticationException)as e:
            logger.error(f'連接 Elasticsearch 失敗: {e}')
raise
exceptExceptionas e:
            logger.error(f'發(fā)生其他錯(cuò)誤: {e}')
            logger.error(f'異常類(lèi)型: {type(e).__name__}')# 記錄異常類(lèi)型
raise

defto_keywords(self, input_string):
"""將句子轉(zhuǎn)成檢索關(guān)鍵詞序列"""
# 按搜索引擎模式分詞
        word_tokens = jieba.cut_for_search(input_string)
# 加載停用詞表
        stop_words =set(stopwords.words('chinese'))
# 去除停用詞
        filtered_sentence =[w for w in word_tokens ifnot w in stop_words]
return' '.join(filtered_sentence)

defsent_tokenize(self, input_string):
"""按標(biāo)點(diǎn)斷句,沒(méi)有用到"""
# 按標(biāo)點(diǎn)切分
        sentences = re.split(r'(?<=[。?。?；?!])', input_string)
# 去掉空字符串
return[sentence for sentence in sentences if sentence.strip()]

defcreate_index(self):
"""如果索引不存在，則創(chuàng)建索引"""
ifnot self.es.indices.exists(index=self.index_name):
# 創(chuàng)建索引
            self.es.indices.create(index=self.index_name, ignore=400)

defbluk_data(self, paragraphs):
"""批量進(jìn)行數(shù)據(jù)灌庫(kù)"""
# 灌庫(kù)指令
        actions =[
{
"_index": self.index_name,
"_source":{
"keywords": self.to_keywords(para.page_content),
"text": para.page_content
}
}
for para in paragraphs
]
# 文本灌庫(kù)
        helpers.bulk(self.es, actions)
# # 灌庫(kù)是異步的
# time.sleep(2)

defflush(self):
# 刷新數(shù)據(jù),數(shù)據(jù)入庫(kù)完成以后刷新數(shù)據(jù)
        self.es.indices.flush()

defsearch(self, query_string):
"""關(guān)鍵詞檢索"""
# ES 的查詢(xún)語(yǔ)言
        search_query ={
"match":{
"keywords": self.to_keywords(query_string)
}
}
        res = self.es.search(index=self.index_name, query=search_query, size=self.k)
return[hit["_source"]["text"]for hit in res["hits"]["hits"]]

defdelete(self):
"""如果索引存在，則刪除索引"""
if self.es.indices.exists(index=self.index_name):
# 創(chuàng)建索引
            self.es.indices.delete(index=self.index_name, ignore=400)

defadd_documents(self, docs):
        self.bluk_data(docs)
        self.flush()

"""
ES數(shù)據(jù)庫(kù)相關(guān)的配置
"""
# ES服務(wù)開(kāi)關(guān)：True表示開(kāi)啟ES服務(wù)，F(xiàn)alse表示關(guān)閉ES服務(wù)
ELASTIC_ENABLE_ES = True
ELASTIC_PASSWORD = os.getenv("ELASTIC_PASSWORD", "123abc")
ELASTIC_HOST = os.getenv("ELASTIC_HOST", "localhost")
ELASTIC_PORT = os.getenv("ELASTIC_PORT", 9200)
ELASTIC_SCHEMA = "https"
ELASTIC_INDEX_NAME = "smart_test_index"

def test_import_elasticsearch():
# from rag.elasticsearch_db import TraditionDB
from rag.elasticsearch_db importElasticsearchDB
from rag.pdf_processor importPDFProcessor

    llm, chat, embed = settings.LLM, settings.CHAT, settings.EMBED

# 導(dǎo)入文件的文件目錄
    directory ="./dataset/pdf"

# 創(chuàng)建 Elasticsearch 數(shù)據(jù)庫(kù)實(shí)例
    es_db =ElasticsearchDB()

# 創(chuàng)建 PDFProcessor 實(shí)例
    pdf_processor =PDFProcessor(directory=directory,
                                 db_type="es",
                                 es_client=es_db,
                                 embed=embed)

# 處理 PDF 文件
    pdf_processor.process_pdfs()

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 rag.elasticsearch_db importElasticsearchDB
# ES需要導(dǎo)入的庫(kù)
from typing importList
import logging
import settings


logger =LoggerManager().logger


classSimpleRetrieverWrapper():
"""自定義檢索器實(shí)現(xiàn)"""

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')

# 初始化一個(gè)空的檢索器列表
        retrievers =[]
        weights =[]

# Step1：創(chuàng)建一個(gè) 多路召回檢索器 MultiQueryRetriever
        chromadb_retriever = self.store.as_retriever()
        mq_retriever =MultiQueryRetriever.from_llm(retriever=chromadb_retriever, llm=self.llm)

        retrievers.append(mq_retriever)
        weights.append(0.5)
        logger.info(f'已啟用 MultiQueryRetriever')

# Step2：創(chuàng)建一個(gè) 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

classElasticsearchRetriever(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)

        logger.info(f"ElasticSearch檢索到資料文件個(gè)數(shù)：{len(query_result)}")

if query_result:
return[Document(page_content=doc)for doc in query_result]
return[]

asyncdef_aget_relevant_documents(self, query: str)->List[Document]:
"""(Optional) async native implementation."""
        es_connector =ElasticsearchDB()
        query_result = es_connector.search(query)
if query_result:
return[Document(page_content=doc)for doc in query_result]
return []

def test_financebot_ex():
    from finance_bot_ex import FinanceBotEx
    # 使用Chroma 的向量庫(kù)
    financebot = FinanceBotEx()

    example_query = "根據(jù)武漢興圖新科電子股份有限公司招股意向書(shū)，電子信息行業(yè)的上游涉及哪些企業(yè)？"

    financebot.handle_query(example_query)