基于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()
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.

# 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"
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.

# 重构的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()
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.

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

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)
1.
2.
3.
4.
5.
6.
7.
8.

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)  # 向向量数据库中插入数据
1.
2.
3.
4.
5.
6.

重构后的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'.")
1.
2.
3.
4.
5.
6.
7.
8.
9.
10.
11.
12.
13.
14.
15.
16.
17.
18.
19.
20.
21.
22.
23.
24.
25.
26.
27.
28.
29.
30.
31.
32.
33.
34.
35.
36.
37.
38.
39.
40.
41.
42.
43.
44.
45.
46.
47.
48.
49.
50.
51.
52.
53.
54.
55.
56.
57.
58.
59.
60.
61.
62.
63.
64.
65.
66.
67.
68.
69.
70.
71.
72.
73.
74.
75.
76.
77.
78.
79.
80.
81.
82.
83.
84.
85.
86.
87.
88.
89.
90.
91.
92.
93.
94.