8.构建可维护的RAG系统:代码分层与模块化设计

张开发
2026/4/8 3:11:08 15 分钟阅读

分享文章

8.构建可维护的RAG系统:代码分层与模块化设计
写在前面咱们后面的功能将会迈进智能体在这之前我感觉有必要对以前写的代码进行分层重构这样当后面项目越来越大的时候咱们维护代码的时候会很清晰。第一类纯工具函数这里函数有一个特点不依赖FAISS和Memory机制能力是通用的也不用self包括split_text、clean_text、load_pdfs、load_documents、process_documents这类函数我们可以不放进类可以创建一个data_loader.py然后统统将他们搬进去这样有一个好处那就是解耦这样别的项目也能用使得这些函数不属于RAG系统本身。# data_loader.pyimportosfrompypdfimportPdfReaderimportredefsplit_text(text,chunk_size200,overlap50):chunks[]foriinrange(0,len(text),chunk_size-overlap):chunks.append(text[i:ichunk_size])returnchunksdefclean_text(text):textre.sub(r\n,\n,text)# 多行空白换成一行textre.sub(r\s, ,text)# 将所有连续空白字符空格、制表符、换行等替换成单个空格实现“规范化空白”。returntextdefload_pdfs(folder_path):documents[]forfilenameinos.listdir(folder_path):iffilename.endswith(.pdf):pathos.path.join(folder_path,filename)readerPdfReader(path)textforpageinreader.pages:textpage.extract_text()ortextclean_text(text)documents.append({text:text,source:filename})returndocumentsdefload_documents(folder_path):documents[]forfilenameinos.listdir(folder_path):iffilename.endswith(.txt):withopen(os.path.join(folder_path,filename),r,encodingutf-8)asf:textf.read()documents.append({text:text,source:filename})returndocumentsdefprocess_documents(documents):all_chunks[]fordocindocuments:chunkssplit_text(doc[text],chunk_size200,overlap50)forcinchunks:all_chunks.append({text:c,source:doc[source]})returnall_chunks第二类模型相关函数这类函数的特点是需要用到API但是不依赖于RAG内部的结构这话有点绕嘴换一种更好理解的方式来说这类函数不需要用到类变量但是用到了大模型API包括get_embedding、decide_tool这些函数我们可以另建一个llm_utils.py函数将他们都存下# llm_utils.pyimportnumpyasnpimporttimefromopenaiimportOpenAIimportosfromdotenvimportload_dotenv load_dotenv()clientOpenAI(api_keyos.getenv(DEEPSEEK_API_KEY),base_urlhttps://api.deepseek.com)client2OpenAI(api_keyos.getenv(SHUBIAOBIAO_API_KEY),base_urlhttps://api.shubiaobiao.com/v1)defget_embedding(text,max_retries3):forattemptinrange(max_retries):try:responseclient2.embeddings.create(modeltext-embedding-3-small,inputtext)returnnp.array(response.data[0].embedding,dtypefloat32)exceptExceptionase:print(fEmbedding失败第{attempt1}次重试...)time.sleep(2)print(Embedding最终失败返回零向量)returnnp.zeros(1536,dtypefloat32)# embedding维度defdecide_tool(query):promptf You are an AI assistant. Decide whether the following question needs document retrieval. Question:{query}Answer ONLY: - RAG if it needs document-based answer - LLM if it can be answered directly responseclient.chat.completions.create(modeldeepseek-chat,messages[{role:user,content:prompt}])returnresponse.choices[0].message.content.strip()不知道你们注意到没有我这回在调用API_KEY的时候和以往的方式不太一样我是用了这种方式os.getenv(DEEPSEEK_API_KEY)为什么能这样用呢其实目的是为了更好的管理项目的API我们用到了一个管理工具dotenv使用这个工具前需要安装依赖pip install dotenv安装完成后在根目录下创建一个.env的文件然后将你的API_KEY像这样都写进去DEEPSEEK_API_KEYxxxxxxxxxxxxxxxxxxxxxxxx SHUBIAOBIAO_API_KEYxxxxxxxxxxxxxxxxxxxxx这样在需要使用的地方用fromdotenvimportload_dotenv load_dotenv()就能加载我们的API_KEY了这样还有一个好处还可以避免我们的项目中有大量用到API的地方暴露我们的API_KEY。所以最终我们整理好的项目就会像这样,项目分层后咱们就可以正式迈入后面的智能体开发了RagAgent/ │ ├── app.py # FastAPI入口 ├── rag_system.py # RAG核心类 ├── llm_utils.py # embedding / decide_tool ├── data_loader.py # PDF / 文本处理 └── data/ # 你的论文如果这篇文章对你有帮助可以点个赞完整代码地址https://github.com/1186141415/A-Paper-Rag-Agent

更多文章