程序员转行学习 AI 大模型： RAG（检索增强生成）｜ 附清晰业务流程示例

本文是程序员转行学习AI大模型的第14个核心知识点笔记，附清晰业务流程示例。
当前阶段：还在学习知识点，由点及面，从 0 到 1 搭建 AI  大模型知识体系中。
系列更新，关注我，后续会持续记录分享转行经历～

概念

RAG：Retrieval-Augmented Generation（检索-增强 生成），通过检索（外部知识库），来增强（大模型）生成（能力），是一种 LLM+检索技术组合。

RAG 优势：

1. 减少幻觉，提高生成准确性；
2. 可解释性强，可以追溯到具体文档；
3. 易于更新，更新知识库即可；
4. 成本可控，可以控制检索的文档数量。

核心流程

用户问题
    ↓
【步骤1】问题向量化
    ↓
    问题向量 [0.1, 0.2, 0.3, ..., 0.1536]
    ↓
【步骤2】向量检索
    ↓
    在向量数据库中搜索相似文档
    ↓
    找到Top-K个相关文档
    ↓
    【步骤3】文档排序
    ↓
    按相似度排序
    ↓
    【步骤4】上下文构建
    ↓
    将相关文档组合成上下文
    ↓
    【步骤5】Prompt构建
    ↓
    构建包含上下文的Prompt
    ↓
    【步骤6】LLM生成
    ↓
    LLM基于上下文生成回答
    ↓
    最终回答

步骤 1：问题向量化

"""
步骤1：问题向量化
"""

from openai import OpenAI
import numpy as np

class Step1_QueryEmbedding:
    def __init__(self, api_key=None):
        self.client = OpenAI(api_key=api_key)
        self.model = "text-embedding-ada-002"

    def embed_query(self, query: str) -> np.ndarray:
        """
        将用户问题转换为向量
        
        输入：用户问题（文本）
        输出：问题向量（1536维）
        """
        print(f"\n=== 步骤1：问题向量化 ===")
        print(f"原始问题: {query}")

        # 调用OpenAI的embedding API
        response = self.client.embeddings.create(
            model=self.model,
            input=query
        )

        # 提取向量
        embedding = np.array(response.data[0].embedding)

        print(f"向量维度: {embedding.shape}")
        print(f"向量前10维: {embedding[:10]}")
        print(f"向量范数: {np.linalg.norm(embedding):.4f}")

        return embedding

# 测试
def test_step1():
    embedder = Step1_QueryEmbedding()
    query = "什么是量子计算？"
    embedding = embedder.embed_query(query)
    print(f"\n问题向量已生成，可以用于检索")

# test_step1()

对应输出

=== 步骤1：问题向量化 ===
原始问题: 什么是量子计算？
向量维度: (1536,)
向量前10维: [ 0.0123, -0.0456,  0.0789, ...,  0.0234]
向量范数: 0.9876

问题向量已生成，可以用于检索

步骤 2：向量检索

"""
步骤2：向量检索
"""

class Step2_VectorRetrieval:
    def __init__(self, vector_db):
        self.vector_db = vector_db

    def retrieve(self, query_embedding: np.ndarray, top_k: int = 3) -> list:
        """
        在向量数据库中检索相似文档
        
        输入：问题向量
        输出：Top-K个最相似的文档
        """
        print(f"\n=== 步骤2：向量检索 ===")
        print(f"检索Top-{top_k}个最相似的文档")

        # 计算与所有文档的相似度
        similarities = []
        for i, doc_embedding in enumerate(self.vector_db.embeddings):
            # 计算余弦相似度
            similarity = self._cosine_similarity(
                query_embedding, 
                doc_embedding
            )
            similarities.append((i, similarity))

        # 按相似度排序
        similarities.sort(key=lambda x: x[1], reverse=True)

        # 取Top-K
        top_k_results = similarities[:top_k]

        print(f"\n检索结果（共{len(top_k_results)}个）:")
        for rank, (idx, similarity) in enumerate(top_k_results, 1):
            doc = self.vector_db.get_document(idx)
            print(f"\n排名 {rank}:")
            print(f"  相似度: {similarity:.4f}")
            print(f"  文档: {doc['text'][:100]}...")
            print(f"  元数据: {doc['metadata']}")

        return top_k_results

    def _cosine_similarity(self, vec1: np.ndarray, vec2: np.ndarray) -> float:
        """
        计算余弦相似度
        """
        dot_product = np.dot(vec1, vec2)
        norm1 = np.linalg.norm(vec1)
        norm2 = np.linalg.norm(vec2)
        return dot_product / (norm1 * norm2 + 1e-8)

# 测试
def test_step2():
    # 假设已经有向量数据库
    vector_db = SimpleVectorDB()

    # 添加一些测试文档
    test_docs = [
        "量子计算是一种基于量子力学原理的计算方式。",
        "量子比特可以同时处于0和1的状态。",
        "量子计算机在特定问题上比传统计算机快得多。"
    ]

    for doc in test_docs:
        # 注意：实际使用时需要调用embedding API
        vector_db.add_document(doc, {"topic": "量子计算"})

    retriever = Step2_VectorRetrieval(vector_db)

    # 模拟问题向量
    query_embedding = np.random.randn(1536)
    results = retriever.retrieve(query_embedding, top_k=3)

# test_step2()

对应输出

=== 步骤2：向量检索 ===
检索Top-3个最相似的文档

检索结果（共3个）:

排名 1:
  相似度: 0.8234
  文档: 量子计算是一种基于量子力学原理的计算方式。
  元数据: {'topic': '量子计算'}

排名 2:
  相似度: 0.7567
  文档: 量子比特可以同时处于0和1的状态。
  元数据: {'topic': '量子计算'}

排名 3:
  相似度: 0.6891
  文档: 量子计算机在特定问题上比传统计算机快得多。
  元数据: {'topic': '量子计算'}

步骤 3：文档排序

"""
步骤3：文档排序
"""

class Step3_DocumentRanking:
    def __init__(self):
        pass
    
    def rank_documents(self, retrieved_docs: list) -> list:
        """
        对检索到的文档进行排序
        
        输入：检索到的文档列表
        输出：排序后的文档列表
        """
        print(f"\n=== 步骤3：文档排序 ===")
        
        # 方法1：按相似度排序（已经完成）
        print("排序方法1: 按相似度排序")
        sorted_by_similarity = sorted(
            retrieved_docs, 
            key=lambda x: x[1], 
            reverse=True
        )
        
        # 方法2：按多样性排序（避免重复）
        print("排序方法2: 按多样性排序")
        sorted_by_diversity = self._rank_by_diversity(retrieved_docs)
        
        # 方法3：综合排序（相似度+多样性）
        print("排序方法3: 综合排序")
        sorted_final = self._rank_combined(
            sorted_by_similarity, 
            sorted_by_diversity
        )
        
        print(f"\n最终排序结果:")
        for rank, (idx, similarity) in enumerate(sorted_final, 1):
            print(f"  {rank}. 相似度: {similarity:.4f}")
        
        return sorted_final
    
    def _rank_by_diversity(self, docs: list) -> list:
        """
        按多样性排序
        """
        # 简单实现：避免选择过于相似的文档
        ranked = []
        used_indices = set()
        
        for idx, similarity in docs:
            if idx not in used_indices:
                ranked.append((idx, similarity))
                used_indices.add(idx)
        
        return ranked
    
    def _rank_combined(self, similarity_ranked: list, diversity_ranked: list) -> list:
        """
        综合排序
        """
        # 简单实现：结合相似度和多样性
        combined = []
        for i, (sim_idx, sim_score) in enumerate(similarity_ranked):
            # 综合得分 = 相似度 * 0.7 + 多样性 * 0.3
            combined_score = sim_score * 0.7 + (1 - i * 0.1) * 0.3
            combined.append((sim_idx, combined_score))
        
        return sorted(combined, key=lambda x: x[1], reverse=True)

# 测试
def test_step3():
    ranker = Step3_DocumentRanking()
    
    # 模拟检索结果
    retrieved_docs = [
        (0, 0.8234),
        (1, 0.7567),
        (2, 0.6891)
    ]
    
    ranked_docs = ranker.rank_documents(retrieved_docs)

# test_step3()

对应输出

=== 步骤3：文档排序 ===
排序方法1: 按相似度排序
排序方法2: 按多样性排序
排序方法3: 综合排序

最终排序结果:
  1. 相似度: 0.8234
  2. 相似度: 0.7567
  3. 相似度: 0.6891

步骤 4：上下文构建

"""
步骤4：上下文构建
"""

class Step4_ContextBuilding:
    def __init__(self):
        pass

    def build_context(self, retrieved_docs: list, vector_db) -> str:
        """
        构建上下文
        
        输入：检索到的文档
        输出：上下文字符串
        """
        print(f"\n=== 步骤4：上下文构建 ===")

        # 方法1：简单拼接
        print("方法1: 简单拼接")
        context_simple = self._build_simple_context(retrieved_docs, vector_db)
        print(f"上下文长度: {len(context_simple)} 字符")

        # 方法2：带元数据的拼接
        print("方法2: 带元数据拼接")
        context_with_metadata = self._build_metadata_context(retrieved_docs, vector_db)
        print(f"上下文长度: {len(context_with_metadata)} 字符")

        # 方法3：摘要式上下文
        print("方法3: 摘要式上下文")
        context_summary = self._build_summary_context(retrieved_docs, vector_db)
        print(f"上下文长度: {len(context_summary)} 字符")

        # 选择最佳方法
        context = context_with_metadata

        print(f"\n最终上下文:")
        print(context[:500] + "...")

        return context

    def _build_simple_context(self, docs: list, vector_db) -> str:
        """
        简单拼接
        """
        context_parts = []
        for rank, (idx, similarity) in enumerate(docs, 1):
            doc = vector_db.get_document(idx)
            context_parts.append(f"文档{rank}: {doc['text']}")

        return "\n\n".join(context_parts)

    def _build_metadata_context(self, docs: list, vector_db) -> str:
        """
        带元数据拼接
        """
        context_parts = []
        for rank, (idx, similarity) in enumerate(docs, 1):
            doc = vector_db.get_document(idx)
            metadata_str = ", ".join([f"{k}={v}" for k, v in doc['metadata'].items()])
            context_parts.append(
                f"文档{rank} (相似度:{similarity:.4f}, {metadata_str}):\n{doc['text']}"
            )

        return "\n\n".join(context_parts)

    def _build_summary_context(self, docs: list, vector_db) -> str:
        """
        摘要式上下文
        """
        context_parts = []
        for rank, (idx, similarity) in enumerate(docs, 1):
            doc = vector_db.get_document(idx)
            # 提取前100字作为摘要
            summary = doc['text'][:100]
            context_parts.append(
                f"文档{rank}摘要 (相似度:{similarity:.4f}):\n{summary}..."
            )

        return "\n\n".join(context_parts)

# 测试
def test_step4():
    builder = Step4_ContextBuilding()

    # 模拟检索结果
    retrieved_docs = [(0, 0.8234), (1, 0.7567)]

    # 模拟向量数据库
    vector_db = SimpleVectorDB()
    vector_db.add_document("量子计算是一种基于量子力学原理的计算方式。", {"topic": "量子计算"})
    vector_db.add_document("量子比特可以同时处于0和1的状态。", {"topic": "量子计算"})

    context = builder.build_context(retrieved_docs, vector_db)

# test_step4()

对应输出

=== 步骤4：上下文构建 ===
方法1: 简单拼接
上下文长度: 156 字符
方法2: 带元数据拼接
上下文长度: 189 字符
方法3: 摘要式上下文
上下文长度: 134 字符

最终上下文:
文档1 (相似度:0.8234, topic=量子计算):
量子计算是一种基于量子力学原理的计算方式。

文档2 (相似度:0.7567, topic=量子计算):
量子比特可以同时处于0和1的状态。

步骤 5：Prompt 构建

"""
步骤5：Prompt构建
"""

class Step5_PromptBuilding:
    def __init__(self):
        pass

    def build_prompt(self, query: str, context: str) -> str:
        """
        构建包含上下文的Prompt
        
        输入：用户问题、上下文
        输出：完整Prompt
        """
        print(f"\n=== 步骤5：Prompt构建 ===")

        # 方法1：基础Prompt
        print("方法1: 基础Prompt")
        prompt_basic = self._build_basic_prompt(query, context)

        # 方法2：结构化Prompt
        print("方法2: 结构化Prompt")
        prompt_structured = self._build_structured_prompt(query, context)

        # 方法3：角色Prompt
        print("方法3: 角色Prompt")
        prompt_role = self._build_role_prompt(query, context)

        # 选择最佳方法
        prompt = prompt_structured

        print(f"\n最终Prompt:")
        print(prompt[:500] + "...")

        return prompt

    def _build_basic_prompt(self, query: str, context: str) -> str:
        """
        基础Prompt
        """
        prompt = f"""
基于以下文档内容回答问题：

{context}

问题: {query}

请基于文档内容回答。
"""
        return prompt

    def _build_structured_prompt(self, query: str, context: str) -> str:
        """
        结构化Prompt
        """
        prompt = f"""
你是一个专业的问答助手，擅长基于文档内容回答问题。

文档内容：
{context}

用户问题：
{query}

请按照以下要求回答：
1. 基于文档内容回答，不要编造
2. 如果文档中没有相关信息，请明确说明
3. 回答要准确、清晰、有条理
4. 可以引用文档中的具体内容

回答：
"""
        return prompt

    def _build_role_prompt(self, query: str, context: str) -> str:
        """
        角色Prompt
        """
        prompt = f"""
你是一个量子计算领域的专家，拥有深厚的专业知识。

参考文档：
{context}

用户问题：
{query}

请以专家的身份，基于参考文档，给出专业、准确的回答。如果文档中没有足够的信息，可以适当补充专业知识，但要明确区分文档内容和补充内容。

专家回答：
"""
        return prompt

# 测试
def test_step5():
    builder = Step5_PromptBuilding()

    query = "什么是量子计算？"
    context = """
文档1 (相似度:0.8234, topic=量子计算):
量子计算是一种基于量子力学原理的计算方式。

文档2 (相似度:0.7567, topic=量子计算):
量子比特可以同时处于0和1的状态。
"""

    prompt = builder.build_prompt(query, context)

# test_step5()

对应输出

=== 步骤5：Prompt构建 ===
方法1: 基础Prompt
方法2: 结构化Prompt
方法3: 角色Prompt

最终Prompt:
你是一个专业的问答助手，擅长基于文档内容回答问题。

文档内容：

文档1 (相似度:0.8234, topic=量子计算):
量子计算是一种基于量子力学原理的计算方式。

文档2 (相似度:0.7567, topic=量子计算):
量子比特可以同时处于0和1的状态。

用户问题：
什么是量子计算？

请按照以下要求回答：
1. 基于文档内容回答，不要编造
2. 如果文档中没有相关信息，请明确说明
3. 回答要准确、清晰、有条理
4. 可以引用文档中的具体内容

回答：

步骤 6：LLM 生成

"""
步骤6：LLM生成
"""

class Step6_LLMGeneration:
    def __init__(self, api_key=None):
        self.client = OpenAI(api_key=api_key)
        self.model = "gpt-3.5-turbo"

    def generate_answer(self, prompt: str) -> str:
        """
        基于Prompt生成回答
        
        输入：完整Prompt
        输出：LLM生成的回答
        """
        print(f"\n=== 步骤6：LLM生成 ===")
        print(f"使用模型: {self.model}")
        print(f"Prompt长度: {len(prompt)} 字符")

        # 调用OpenAI API
        response = self.client.chat.completions.create(
            model=self.model,
            messages=[{"role": "user", "content": prompt}],
            temperature=0.3,  # 降低温度，提高准确性
            max_tokens=500
        )

        # 提取回答
        answer = response.choices[0].message.content

        print(f"\n生成的回答:")
        print(answer)

        # 显示使用情况
        usage = response.usage
        print(f"\nToken使用:")
        print(f"  输入: {usage.prompt_tokens}")
        print(f"  输出: {usage.completion_tokens}")
        print(f"  总计: {usage.total_tokens}")

        return answer

# 测试
def test_step6():
    generator = Step6_LLMGeneration()

    prompt = """
你是一个专业的问答助手，擅长基于文档内容回答问题。

文档内容：

文档1 (相似度:0.8234, topic=量子计算):
量子计算是一种基于量子力学原理的计算方式。

文档2 (相似度:0.7567, topic=量子计算):
量子比特可以同时处于0和1的状态。

用户问题：
什么是量子计算？

请按照以下要求回答：
1. 基于文档内容回答，不要编造
2. 如果文档中没有相关信息，请明确说明
3. 回答要准确、清晰、有条理
4. 可以引用文档中的具体内容

回答：
"""

    answer = generator.generate_answer(prompt)

# test_step6()

对应输出

=== 步骤6：LLM生成 ===
使用模型: gpt-3.5-turbo
Prompt长度: 456 字符

生成的回答:
根据提供的文档内容，量子计算是一种基于量子力学原理的计算方式。

与经典计算不同，量子计算利用量子比特作为基本信息单位。文档2指出，量子比特具有独特的特性，可以同时处于0和1的状态，这被称为量子叠加态。这种特性使得量子计算机在处理某些特定问题时，能够比传统计算机展现出更强大的计算能力。

总结来说，量子计算是基于量子力学原理的新型计算方式，其核心是利用量子比特的叠加特性来实现计算。

Token使用:
  输入: 256
  输出: 187
  总计: 443

完整 RAG 系统

"""
完整的RAG系统
"""

class CompleteRAGSystem:
    def __init__(self, api_key=None):
        # 初始化各个组件
        self.embedder = Step1_QueryEmbedding(api_key)
        self.vector_db = SimpleVectorDB()
        self.retriever = Step2_VectorRetrieval(self.vector_db)
        self.ranker = Step3_DocumentRanking()
        self.context_builder = Step4_ContextBuilding()
        self.prompt_builder = Step5_PromptBuilding()
        self.generator = Step6_LLMGeneration(api_key)

    def add_knowledge(self, texts: list, metadatas: list = None):
        """
        添加知识到向量数据库
        """
        print("\n=== 添加知识 ===")
        print(f"共 {len(texts)} 条文档")

        # 向量化所有文档
        embeddings = []
        for text in texts:
            embedding = self.embedder.embed_query(text)
            embeddings.append(embedding)

        # 添加到向量数据库
        for i, (text, embedding) in enumerate(zip(texts, embeddings)):
            metadata = metadatas[i] if metadatas else None
            self.vector_db.documents.append({
                "text": text,
                "metadata": metadata or {}
            })
            self.vector_db.embeddings.append(embedding)

        print(f"知识库构建完成，共 {len(self.vector_db)} 条文档")

    def query(self, question: str, top_k: int = 3) -> dict:
        """
        完整的RAG查询流程
        """
        print(f"\n{'='*60}")
        print(f"RAG查询开始")
        print(f"问题: {question}")
        print(f"{'='*60}")

        # 步骤1：问题向量化
        query_embedding = self.embedder.embed_query(question)

        # 步骤2：向量检索
        retrieved_docs = self.retriever.retrieve(query_embedding, top_k)

        # 步骤3：文档排序
        ranked_docs = self.ranker.rank_documents(retrieved_docs)

        # 步骤4：上下文构建
        context = self.context_builder.build_context(ranked_docs, self.vector_db)

        # 步骤5：Prompt构建
        prompt = self.prompt_builder.build_prompt(question, context)

        # 步骤6：LLM生成
        answer = self.generator.generate_answer(prompt)

        # 返回结果
        result = {
            "question": question,
            "retrieved_docs": ranked_docs,
            "context": context,
            "answer": answer
        }

        print(f"\n{'='*60}")
        print(f"RAG查询完成")
        print(f"{'='*60}")

        return result

# 测试
def test_complete_rag():
    rag = CompleteRAGSystem()

    # 添加知识
    knowledge = [
        "量子计算是一种基于量子力学原理的计算方式。",
        "量子比特可以同时处于0和1的状态。",
        "量子计算机在特定问题上比传统计算机快得多。",
        "量子纠缠是量子力学中的奇特现象。",
        "量子纠错是量子计算的关键技术。"
    ]

    metadatas = [
        {"topic": "量子计算", "category": "基础"},
        {"topic": "量子计算", "category": "原理"},
        {"topic": "量子计算", "category": "优势"},
        {"topic": "量子计算", "category": "现象"},
        {"topic": "量子计算", "category": "技术"}
    ]

    rag.add_knowledge(knowledge, metadatas)

    # 查询
    questions = [
        "什么是量子计算？",
        "量子比特有什么特点？",
        "量子计算的优势是什么？"
    ]
    
    for question in questions:
        result = rag.query(question)
        print(f"\n最终回答: {result['answer']}")
        print("\n" + "="*80 + "\n")

# test_complete_rag()