# 原始的rag.py
class RagManager:
    def __init__(self,
                 chroma_server_type="http",
                 host="localhost", port=8000,
                 persist_path="chroma_db",
                 llm=None, embed=None):
        self.llm = llm
        self.embed = embed

        chrom_db = ChromaDB(chroma_server_type=chroma_server_type,
                            host=host, port=port,
                            persist_path=persist_path,
                            embed=embed)
        self.store = chrom_db.get_store()

# setttings.py

"""
Chroma向量數(shù)據(jù)庫使用時(shí)的相關(guān)的配置
"""
# 默認(rèn)的ChromaDB的服務(wù)器類別
CHROMA_SERVER_TYPE ="http"
# 默認(rèn)本地?cái)?shù)據(jù)庫的持久化目錄
CHROMA_PERSIST_DB_PATH ="chroma_db"

CHROMA_HOST = os.getenv("CHROMA_HOST","localhost")
CHROMA_PORT =int(os.getenv("CHROMA_PORT",8000))
CHROMA_COLLECTION_NAME ="langchain"

# 重構(gòu)的rag.py
import settings


classRagManager:
def__init__(self,
                 vector_db_class=ChromaDB,  # 默認(rèn)使用 ChromaDB
                 db_config=None,  # 數(shù)據(jù)庫配置參數(shù)
                 llm=None, embed=None,
                 retriever_cls=SimpleRetrieverWrapper, **retriever_kwargs):
        self.llm = llm
        self.embed = embed
        logger.info(f'初始化llm大模型：{self.llm}')
        logger.info(f'初始化embed模型：{self.embed}')

# 如果沒有提供 db_config，使用默認(rèn)配置
if db_config isNone:
            db_config ={
"chroma_server_type": settings.CHROMA_SERVER_TYPE,
"host": settings.CHROMA_HOST,
"port": settings.CHROMA_PORT,
"persist_path": settings.CHROMA_PERSIST_DB_PATH,
"collection_name": settings.CHROMA_COLLECTION_NAME
}
            logger.info(f'初始化向量數(shù)據(jù)庫配置：{db_config}')

# 創(chuàng)建向量數(shù)據(jù)庫實(shí)例
        self.vector_db = vector_db_class(**db_config, embed=self.embed)
        self.store = self.vector_db.get_store()

RagManager(chroma_server_type="http", host="localhost", port=8000, persist_path="chroma_db", collection_name="langchain",llm , embed)

db_config = {
"chroma_server_type": settings.CHROMA_SERVER_TYPE,
"host": settings.CHROMA_HOST,
"port": settings.CHROMA_PORT,
"persist_path": settings.CHROMA_PERSIST_DB_PATH,
"collection_name": settings.CHROMA_COLLECTION_NAME,
}
RagManager(vector_db_class=ChromaDB, db_config=db_config, llm=self.llm, embed=self.embed)

def process_pdfs(self)# 處理pdf文件
defprocess_pdfs_group(self, pdf_files_group)# 分組處理pdf文件
defload_pdf_files(self)# 加載pdf文件
defload_pdf_content(self, pdf_path)# 讀取pdf文件內(nèi)容
defsplit_text(self, documents)# 分割讀取到的文本
definsert_docs_chromadb(self, docs, batch_size)  # 向向量數(shù)據(jù)庫中插入數(shù)據(jù)

重構(gòu)后的pdf_processor.py
import os
import logging
import time
from tqdm import tqdm
from langchain_community.document_loaders importPyMuPDFLoader
from langchain_text_splitters importRecursiveCharacterTextSplitter
from rag.vector_db importVectorDB
from rag.elasticsearch_db importTraditionDB
from utils.logger_config importLoggerManager
logger =LoggerManager().logger
classPDFProcessor:
def__init__(self, directory, db_type='vector', **kwargs):
"""
初始化 PDF 處理器
:param directory: PDF 文件所在目錄
:param db_type: 數(shù)據(jù)庫類型 ('vector' 或 'es')
:param kwargs: 其他參數(shù)
"""
self.directory = directory  # PDF 文件所在目錄
self.db_type = db_type  # 數(shù)據(jù)庫類型
self.file_group_num = kwargs.get('file_group_num',20)# 每組處理的文件數(shù)
self.batch_num = kwargs.get('batch_num',6)# 每次插入的批次數(shù)量
self.chunksize = kwargs.get('chunksize',500)# 切分文本的大小
self.overlap = kwargs.get('overlap',100)# 切分文本的重疊大小
logger.info(f"""
初始化PDF文件導(dǎo)入器:
配置參數(shù)：
- 導(dǎo)入的文件路徑：{self.directory}
- 每次處理文件數(shù)：{self.file_group_num}
- 每批次處理樣本數(shù)：{self.batch_num}
- 切分文本的大?。簕self.chunksize}
- 切分文本重疊大?。簕self.overlap}
""")
根據(jù)數(shù)據(jù)庫類型初始化相應(yīng)的客戶端
if db_type =='vector':
self.vector_db = kwargs.get('vector_db')# 向量數(shù)據(jù)庫實(shí)例
self.es_client =None
logger.info(f'導(dǎo)入的目標(biāo)數(shù)據(jù)庫為：向量數(shù)據(jù)庫')elif db_type =='es':
self.vector_db =None
self.es_client = kwargs.get('es_client')# Elasticsearch 客戶端
logger.info(f'導(dǎo)入的目標(biāo)數(shù)據(jù)庫為：ES數(shù)據(jù)庫')else:
raiseValueError("db_type must be either 'vector' or 'es'.")
defload_pdf_files(self):
這部分代碼未做修改，具體內(nèi)容省略
defload_pdf_content(self, pdf_path):
這部分代碼未做修改，具體內(nèi)容省略
defsplit_text(self, documents):
這部分代碼未做修改，具體內(nèi)容省略
defprocess_pdfs(self):
這部分代碼未做修改，具體內(nèi)容省略
definsert_docs(self, docs, insert_function, batch_size=None):
"""
將文檔插入到指定的數(shù)據(jù)庫，并顯示進(jìn)度
:param docs: 要插入的文檔列表
:param insert_function: 插入函數(shù)
:param batch_size: 批次大小
"""
if batch_size isNone:
batch_size = self.batch_num
logging.info(f"Inserting {len(docs)} documents.")
    start_time = time.time()
    total_docs_inserted =0

    total_batches =(len(docs)+ batch_size -1)// batch_sizewith tqdm(total=total_batches, desc="Inserting batches", unit="batch")as pbar:
for i inrange(0,len(docs), batch_size):
batch = docs[i:i + batch_size]
insert_function(batch)# 調(diào)用傳入的插入函數(shù)
total_docs_inserted +=len(batch)計(jì)算并顯示當(dāng)前的TPM
elapsed_time = time.time()- start_timeif elapsed_time >0:
tpm =(total_docs_inserted / elapsed_time)*60
pbar.set_postfix({"TPM":f"{tpm:.2f}"})
pbar.update(1)definsert_to_vector_db(self, docs):
"""
將文檔插入到 VectorDB
"""
self.vector_db.add_with_langchain(docs)
definsert_to_elasticsearch(self, docs):
"""
將文檔插入到 Elasticsearch
"""
self.es_client.add_documents(docs)
defprocess_pdfs_group(self, pdf_files_group):
讀取PDF文件內(nèi)容
pdf_contents =[]for pdf_path in pdf_files_group:
讀取PDF文件內(nèi)容
documents = self.load_pdf_content(pdf_path)將documents 逐一添加到pdf_contents
pdf_contents.extend(documents)將文本切分成小段
docs = self.split_text(pdf_contents)if self.db_type =='vector':
將文檔插入到 VectorDB
self.insert_docs(docs, self.insert_to_vector_db)elif self.db_type =='es':
將文檔插入到 Elasticsearch
self.insert_docs(docs, self.insert_to_elasticsearch)else:
raiseValueError("db_type must be either 'vector' or 'es'.")