基于Agent的金融问答系统：代码重构 原创

# 原始的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向量数据库使用时的相关的配置
"""
# 默认的ChromaDB的服务器类别
CHROMA_SERVER_TYPE ="http"
# 默认本地数据库的持久化目录
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"

# 重构的rag.py
import settings


classRagManager:
def__init__(self,
                 vector_db_class=ChromaDB,  # 默认使用 ChromaDB
                 db_config=None,  # 数据库配置参数
                 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，使用默认配置
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'初始化向量数据库配置：{db_config}')

# 创建向量数据库实例
        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文件内容
defsplit_text(self, documents)# 分割读取到的文本
definsert_docs_chromadb(self, docs, batch_size)  # 向向量数据库中插入数据

重构后的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: 数据库类型 ('vector' 或 'es')
:param kwargs: 其他参数
"""
self.directory = directory  # PDF 文件所在目录
self.db_type = db_type  # 数据库类型
self.file_group_num = kwargs.get('file_group_num',20)# 每组处理的文件数
self.batch_num = kwargs.get('batch_num',6)# 每次插入的批次数量
self.chunksize = kwargs.get('chunksize',500)# 切分文本的大小
self.overlap = kwargs.get('overlap',100)# 切分文本的重叠大小
logger.info(f"""
初始化PDF文件导入器:
配置参数：
- 导入的文件路径：{self.directory}
- 每次处理文件数：{self.file_group_num}
- 每批次处理样本数：{self.batch_num}
- 切分文本的大小：{self.chunksize}
- 切分文本重叠大小：{self.overlap}
""")
根据数据库类型初始化相应的客户端
if db_type =='vector':
self.vector_db = kwargs.get('vector_db')# 向量数据库实例
self.es_client =None
logger.info(f'导入的目标数据库为：向量数据库')elif db_type =='es':
self.vector_db =None
self.es_client = kwargs.get('es_client')# Elasticsearch 客户端
logger.info(f'导入的目标数据库为：ES数据库')else:
raiseValueError("db_type must be either 'vector' or 'es'.")
defload_pdf_files(self):
这部分代码未做修改，具体内容省略
defload_pdf_content(self, pdf_path):
这部分代码未做修改，具体内容省略
defsplit_text(self, documents):
这部分代码未做修改，具体内容省略
defprocess_pdfs(self):
这部分代码未做修改，具体内容省略
definsert_docs(self, docs, insert_function, batch_size=None):
"""
将文档插入到指定的数据库，并显示进度
:param docs: 要插入的文档列表
:param insert_function: 插入函数
: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)# 调用传入的插入函数
total_docs_inserted +=len(batch)计算并显示当前的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文件内容
pdf_contents =[]for pdf_path in pdf_files_group:
读取PDF文件内容
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'.")