写在前面

咱们后面的功能将会迈进智能体，在这之前我感觉有必要对以前写的代码进行分层重构，这样当后面项目越来越大的时候，咱们维护代码的时候会很清晰。

第一类：纯工具函数

这里函数有一个特点，不依赖FAISS和Memory机制，能力是通用的也不用self，包括split_text、clean_text、load_pdfs、load_documents、process_documents这类函数我们可以不放进类，可以创建一个data_loader.py，然后统统将他们搬进去，这样有一个好处，那就是解耦，这样别的项目也能用，使得这些函数不属于RAG系统本身。

# data_loader.py

import os
from pypdf import PdfReader
import re

def split_text(text, chunk_size=200, overlap=50):
    chunks = []
    for i in range(0, len(text), chunk_size - overlap):
        chunks.append(text[i:i + chunk_size])
    return chunks

def clean_text(text):
    text = re.sub(r'\n+', '\n', text)  # 多行空白换成一行
    text = re.sub(r'\s+', ' ', text)  # 将所有连续空白字符（空格、制表符、换行等）替换成单个空格，实现“规范化空白”。
    return text

def load_pdfs(folder_path):
    documents = []

    for filename in os.listdir(folder_path):
        if filename.endswith(".pdf"):
            path = os.path.join(folder_path, filename)
            reader = PdfReader(path)

            text = ""
            for page in reader.pages:
                text += page.extract_text() or ""

            text = clean_text(text)

            documents.append({
                "text": text,
                "source": filename
            })

    return documents

def load_documents(folder_path):
    documents = []

    for filename in os.listdir(folder_path):
        if filename.endswith(".txt"):
            with open(os.path.join(folder_path, filename), "r", encoding="utf-8") as f:
                text = f.read()
                documents.append({
                    "text": text,
                    "source": filename
                })

    return documents

def process_documents(documents):
    all_chunks = []

    for doc in documents:
        chunks = split_text(doc["text"], chunk_size=200, overlap=50)

        for c in chunks:
            all_chunks.append({
                "text": c,
                "source": doc["source"]
            })

    return all_chunks

第二类：模型相关函数

这类函数的特点是需要用到API，但是不依赖于RAG内部的结构，这话有点绕嘴，换一种更好理解的方式来说，这类函数不需要用到类变量但是用到了大模型API，包括：get_embedding、decide_tool，这些函数我们可以另建一个llm_utils.py函数将他们都存下：

# llm_utils.py
import numpy as np
import time
from openai import OpenAI
import os
from dotenv import load_dotenv
load_dotenv()

client = OpenAI(
    api_key=os.getenv("DEEPSEEK_API_KEY"),
    base_url="https://api.deepseek.com"
)

client2 = OpenAI(
    api_key=os.getenv("SHUBIAOBIAO_API_KEY"),
    base_url="https://api.shubiaobiao.com/v1"
)


def get_embedding(text, max_retries=3):
    for attempt in range(max_retries):
        try:
            response = client2.embeddings.create(
                model="text-embedding-3-small",
                input=text
            )
            return np.array(response.data[0].embedding, dtype="float32")

        except Exception as e:
            print(f"Embedding失败，第{attempt + 1}次重试...")
            time.sleep(2)

    print("Embedding最终失败，返回零向量")
    return np.zeros(1536, dtype="float32")  # embedding维度



def decide_tool(query):
    prompt = f"""  
    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   
              """

    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[{"role": "user", "content": prompt}]
    )

    return response.choices[0].message.content.strip()

不知道你们注意到没有，我这回在调用API_KEY的时候和以往的方式不太一样，我是用了这种方式os.getenv("DEEPSEEK_API_KEY") ，为什么能这样用呢？其实目的是为了更好的管理项目的API，我们用到了一个管理工具dotenv，使用这个工具前需要安装依赖

pip install dotenv

安装完成后，在根目录下创建一个.env的文件，然后将你的API_KEY像这样都写进去：

DEEPSEEK_API_KEY=xxxxxxxxxxxxxxxxxxxxxxxx
SHUBIAOBIAO_API_KEY=xxxxxxxxxxxxxxxxxxxxx

这样在需要使用的地方用：

from dotenv import load_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