从脚本到工程项目的转变

前面咱们写了很多的代码，但是代码本质上还是处于零散状态的学生脚本，这种东西和企业输出的工程师项目还是有很大差距的，今天咱们就开始把脚本一步步封装，暴露功能，RAG系统一点点开始有一点工程的味道。下面我们就一步步开始吧，第一步先创建一个RAG系统的类：

class RAGSystem: 
    def __init__(self, chunks):
        self.chunks = chunks
        self.index = None
        self.embeddings = None

有了类以后我们需要在类中封装一些核心的功能，包括存知识数据、相关知识检索以及提供可以完成询问功能的接口：

def build_index(self): # 把FAISS封进去
    embeddings = [get_embedding(c) for c in self.chunks]
    self.embeddings = np.vstack(embeddings)

    dim = self.embeddings.shape[1]
    self.index = faiss.IndexFlatL2(dim)
    self.index.add(self.embeddings)
    
def retrieve(self, query, k=5): #封装retrieve
    query_vec = get_embedding(query).reshape(1, -1)
    distances, indices = self.index.search(query_vec, k)
    return [self.chunks[i] for i in indices[0]]


def ask(self, question):
    retrieved = self.retrieve(question, k=5)
    
    # 可以加 rerank（咱们前面写过）
    context = "\n".join(retrieved)

    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": "Answer based only on context."},
            {"role": "user", "content": context + "\n\n" + question}
        ]
    )

    return response.choices[0].message.content

有了这样一个类，咱们的系统再使用的时候只需要：

rag = RAGSystem(chunks)
rag.build_index()

answer = rag.ask("What is the core contribution?")
print(answer)

这样便实现了一个完整的可调用的功能模块，全部代码我这里给大家贴出来：

import faiss
import numpy as np
from openai import OpenAI

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

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

def get_embedding(text):
    response = client2.embeddings.create(
        model="text-embedding-3-small",  # 改这里
        input=text
    )
    return np.array(response.data[0].embedding, dtype="float32")

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


class RAGSystem:
    def __init__(self, chunks):
        self.chunks = chunks
        self.index = None
        self.embeddings = None

    def build_index(self):
        embeddings = [get_embedding(c) for c in self.chunks]
        self.embeddings = np.vstack(embeddings)

        dim = self.embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dim)
        self.index.add(self.embeddings)

    def retrieve(self, query, k=5):
        query_vec = get_embedding(query).reshape(1, -1)
        distances, indices = self.index.search(query_vec, k)
        return [self.chunks[i] for i in indices[0]]

    def ask(self, question):
        retrieved = self.retrieve(question, k=5)

        # 可以加 rerank（咱们前面写过）
        context = "\n".join(retrieved)

        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": "Answer based only on context."},
                {"role": "user", "content": context + "\n\n" + question}
            ]
        )

        return response.choices[0].message.content


if __name__ == '__main__':
    context = "Voice over IP (VoIP) steganography based on low-bit-rate speech codecs has attracted increasing attention due to its high imperceptibility and large embedding capacity, particularly in the fixed codebook (FCB) parameter domain. However, effective steganalysis remains challenging under extreme embedding conditions. At low embedding rates, steganographic artifacts are weak and sparsely distributed, making them difficult to distinguish from natural speech variations. In contrast, at high embedding rates, the recompression-based calibration process may introduce structural distortions that interfere with reliable feature extraction. To address these challenges, this paper proposes a calibration-aware cross-view steganalysis netword for VoIP steganalysis (CACVAN). An embedding-rate-aware data augmentation (ERADA) strategy is first introduced to construct cross-intensity training samples, which improves the robustness of the model under embedding-rate mismatch scenarios. Furthermore, a cross-view interaction backbone (CVIB) is designed to jointly analyze the original speech stream and its recompressed counterpart, enabling the network to capture subtle inconsistencies introduced by steganographic embedding while suppressing content-related variations. A hybrid attention refinement neck (HARN) is then employed to enhance discriminative feature responses and stabilize the modeling of sparse steganographic artifacts. Extensive experiments on public VoIP steganalysis datasets demonstrate that the proposed method consistently outperforms existing state-of-the-art approaches under various embedding rates and speech durations, especially in challenging scenarios involving low embedding rates and short speech segments. Moreover, the proposed framework achieves high computational efficiency and satisfies the real-time requirements of streaming VoIP steganalysis, indicating its practical applicability."
    chunks = split_text(context, chunk_size=200)

    rag = RAGSystem(chunks)
    rag.build_index()

    answer = rag.ask("What is the core contribution?")
    print(answer)

封装基础上2次优化

增加缓存减少API调用的费用

有了封装后的接口后呢，咱们还可以对封装后的版本进行2次优化，而且这里首先就会遇到一个坑，前面的代码关于嵌入的操作是这样写的：

embeddings = [get_embedding(c) for c in self.chunks]

这样会遇到一个问题，每次调用都会重新调用API（慢+贵），所以咱们需要有点加缓存的意识，加上缓存后，系统会将过去嵌入过的知识存贮在类中，这样下次调用的时候就可以不用再调词嵌入的接口了，可以省下很多小钱钱。

def build_index(self):
    if self.embeddings is None:
        embeddings = [get_embedding(c) for c in self.chunks] 
        self.embeddings = np.vstack(embeddings) #这一步便是缓存，嵌入后的结果被装入类变量

    dim = self.embeddings.shape[1]
    self.index = faiss.IndexFlatL2(dim)
    self.index.add(self.embeddings)

增加Rerank操作

前面讲过为了提高模型的质量咱们可以在模型中增加rerank操作，这样可以让模型结合知识的相关程度回答问题，大大提高模型效果，所以咱们这次将rerank操作封装在这个类内：

    def rerank(self, query, chunks):
        prompt = f"""
                You are a ranking assistant.

                Query:
                {query}

                Rank the following passages from most relevant to least relevant.

                Passages:
                """

        for i, c in enumerate(chunks):
            prompt += f"\n[{i}] {c}\n"

        prompt += "\nReturn ONLY the indices in sorted order, like [2,0,1]."

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

        import ast # 可以将字符串传化为list的工具包
        try:
            return ast.literal_eval(response.choices[0].message.content)
        except:
            return list(range(len(chunks)))

修改ask

前面增加了rerank操作后，咱们的ask操作可以跟着升级，将Rerank后的结果加进去：

def ask(self, question):
    retrieved = self.retrieve(question, k=5)

    # rerank
    sorted_indices = self.rerank(question, retrieved)
    best_chunks = [retrieved[i] for i in sorted_indices[:3]]

    context = "\n".join(best_chunks)

    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=[
            {"role": "system", "content": "Answer based only on context."},
            {"role": "user", "content": context + "\n\n" + question}
        ]
    )

    return response.choices[0].message.content

最后进行一定的参数化

那么这个模块运行起来的就会很灵活：

def __init__(self, chunks, top_k=5, rerank_k=3):
    self.chunks = chunks
    self.top_k = top_k
    self.rerank_k = rerank_k
    self.index = None
    self.embeddings = None

然后修改：

retrieved = self.retrieve(question, k=self.top_k)
best_chunks = [retrieved[i] for i in sorted_indices[:self.rerank_k]]

到这里咱们做的不再是一个可运行的脚本了，是一个可运行可调的系统，完整代码如下：

import faiss
import numpy as np
from openai import OpenAI

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

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

def get_embedding(text):
    response = client2.embeddings.create(
        model="text-embedding-3-small",  # 改这里
        input=text
    )
    return np.array(response.data[0].embedding, dtype="float32")

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


class RAGSystem:
    def __init__(self, chunks, top_k=5, rerank_k=3):
        self.chunks = chunks
        self.top_k = top_k
        self.rerank_k = rerank_k
        self.index = None
        self.embeddings = None

    def build_index(self):
        if self.embeddings is None:   #加缓存避免重复计算消耗API
            embeddings = [get_embedding(c) for c in self.chunks]
            self.embeddings = np.vstack(embeddings)

        dim = self.embeddings.shape[1]
        self.index = faiss.IndexFlatL2(dim)
        self.index.add(self.embeddings)

    def retrieve(self, query, k=5):
        query_vec = get_embedding(query).reshape(1, -1)
        distances, indices = self.index.search(query_vec, k)
        return [self.chunks[i] for i in indices[0]]

    def rerank(self, query, chunks):
        prompt = f"""
                You are a ranking assistant.

                Query:
                {query}

                Rank the following passages from most relevant to least relevant.

                Passages:
                """

        for i, c in enumerate(chunks):
            prompt += f"\n[{i}] {c}\n"

        prompt += "\nReturn ONLY the indices in sorted order, like [2,0,1]."

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

        import ast
        try:
            return ast.literal_eval(response.choices[0].message.content)
        except:
            return list(range(len(chunks)))

    def ask(self, question):
        retrieved = self.retrieve(question, k=self.top_k)

        # 可以加 rerank（你已经有了）
        #context = "\n".join(retrieved)
        sorted_indices = self.rerank(question, retrieved)
        best_chunks = [retrieved[i] for i in sorted_indices[:self.rerank_k]]

        context = "\n".join(best_chunks)

        response = client.chat.completions.create(
            model="deepseek-chat",
            messages=[
                {"role": "system", "content": "Answer based only on context."},
                {"role": "user", "content": context + "\n\n" + question}
            ]
        )

        return response.choices[0].message.content



if __name__ == '__main__':
    context = "Voice over IP (VoIP) steganography based on low-bit-rate speech codecs has attracted increasing attention due to its high imperceptibility and large embedding capacity, particularly in the fixed codebook (FCB) parameter domain. However, effective steganalysis remains challenging under extreme embedding conditions. At low embedding rates, steganographic artifacts are weak and sparsely distributed, making them difficult to distinguish from natural speech variations. In contrast, at high embedding rates, the recompression-based calibration process may introduce structural distortions that interfere with reliable feature extraction. To address these challenges, this paper proposes a calibration-aware cross-view steganalysis netword for VoIP steganalysis (CACVAN). An embedding-rate-aware data augmentation (ERADA) strategy is first introduced to construct cross-intensity training samples, which improves the robustness of the model under embedding-rate mismatch scenarios. Furthermore, a cross-view interaction backbone (CVIB) is designed to jointly analyze the original speech stream and its recompressed counterpart, enabling the network to capture subtle inconsistencies introduced by steganographic embedding while suppressing content-related variations. A hybrid attention refinement neck (HARN) is then employed to enhance discriminative feature responses and stabilize the modeling of sparse steganographic artifacts. Extensive experiments on public VoIP steganalysis datasets demonstrate that the proposed method consistently outperforms existing state-of-the-art approaches under various embedding rates and speech durations, especially in challenging scenarios involving low embedding rates and short speech segments. Moreover, the proposed framework achieves high computational efficiency and satisfies the real-time requirements of streaming VoIP steganalysis, indicating its practical applicability."
    chunks = split_text(context, chunk_size=200)

    rag = RAGSystem(chunks)
    rag.build_index()

    answer = rag.ask("What is the core contribution?")
    print(answer)

如果这篇文章对你有帮助，可以点个赞～
完整代码地址：https://github.com/1186141415/A-Paper-Rag-Agent