前言：想让AI基于你的文档回答问题？想构建一个私有的智能助手？本文将从零开始，用不到200行Python代码，带你实现一个完整的RAG系统！

项目地址：https://github.com/muggle-stack/RAG

本项目仅限于想学习入门大模型应用技术的同学，项目包含RAG的前后处理技术。

什么是RAG？

RAG的全称和核心思想

RAG全称是Retrieval-Augmented Generation（检索增强生成），听起来很高深？其实概念很简单：

• 传统AI问答：AI只能基于训练时学到的知识回答，无法获取最新信息
• RAG问答：AI先从你的文档中检索相关内容，再基于这些内容生成回答

简单来说，RAG = 搜索引擎 + AI对话

举个生活化的例子

假如你是一个学生，考试时遇到不会的题目：

1. 传统方式：只能靠脑子里记住的知识作答（可能答错或答不出）
2. RAG方式：先翻书找到相关章节，再基于书本内容作答（准确率更高）

RAG就是让AI具备了"翻书查资料"的能力！

RAG系统的整体架构

RAG系统需要包含以下核心步骤：

文档 → 文本切块 → 向量化 → 存储 → 检索 → 生成回答

环境准备

1. 安装Ollama（本地AI服务）

# Linux
curl -fsSL https://ollama.com/install.sh | sh

# Windows/Mac：访问 https://ollama.com/download

2. 下载AI模型

# 文本向量化模型
ollama pull nomic-embed-text

# 对话生成模型  
ollama pull qwen3:0.6b

# 启动服务
ollama serve

3. 安装Python依赖

pip install numpy openai pypdf python-docx python-pptx openpyxl beautifulsoup4 lxml tiktoken

代码实现详解

让我们一步步分析完整的RAG实现代码：

第一步：基础配置和导入

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
轻量级 RAG – 支持 TXT / PDF / DOCX / PPTX / XLSX / HTML / MD / PY
"""

import sys, pathlib, re, textwrap
from typing import List
import numpy as np
from openai import OpenAI

# ==== 本地 Ollama OpenAI 兼容接口 ====
client = OpenAI(api_key="ollama", base_url="http://127.0.0.1:11434/v1")

# ==== 配置参数 ====
EMBED_MODEL   = "nomic-embed-text"  # 嵌入模型：将文本转换为向量
LLM_MODEL     = "qwen3:0.6b"        # 语言模型：生成回答
CHUNK_TOKENS  = 350                 # 每个文档块的大小
CHUNK_OVERLAP = 30                  # 文档块之间的重叠部分
EMBED_BATCH   = 64                  # 批量处理向量化的数量
TOP_K         = 4                   # 检索返回的文档片段数量

解释：

• 使用OpenAI客户端连接本地Ollama服务
• 配置了文档处理的各种参数，这些参数直接影响系统性能

第二步：多格式文档解析模块

def file_to_text(path: str) -> str:
    """将各种格式的文件转换为纯文本"""
    p = pathlib.Path(path)
    ext = p.suffix.lower()

    # 纯文本、Markdown、Python
    if ext in {".txt", ".md", ".markdown", ".py"}:
        return p.read_text(encoding="utf-8", errors="ignore")

    # PDF文件处理
    if ext == ".pdf":
        from pypdf import PdfReader
        return "\n".join(page.extract_text() or "" for page in PdfReader(p).pages)

    # Word文档处理
    if ext == ".docx":
        import docx
        return "\n".join(par.text for par in docx.Document(p).paragraphs)

    # PowerPoint处理
    if ext == ".pptx":
        from pptx import Presentation
        return "\n".join(sh.text for s in Presentation(p).slides 
                        for sh in s.shapes if getattr(sh, "text", ""))

    # Excel表格处理
    if ext == ".xlsx":
        import openpyxl
        wb = openpyxl.load_workbook(p, data_only=True)
        return "\n".join(str(c) for ws in wb.worksheets 
                        for row in ws.iter_rows(values_only=True) 
                        for c in row if c)

    # HTML网页处理
    if ext in {".html", ".htm"}:
        from bs4 import BeautifulSoup
        return BeautifulSoup(p.read_text(encoding="utf-8", errors="ignore"), 
                           "lxml").get_text("\n")

    raise ValueError(f"暂不支持格式: {path}")

核心功能：

• 自动识别文件格式：根据文件扩展名选择对应的解析方法
• 统一输出格式：不管输入什么格式，都输出纯文本
• 错误处理：编码错误时使用ignore模式，避免程序崩溃

第三步：智能文本分块模块

def _count_tokens(text: str) -> int:
    """计算文本的token数量"""
    try:
        import tiktoken
        return len(tiktoken.get_encoding("cl100k_base").encode(text))
    except Exception:
        # 备用方案：简单的词数统计
        return len(re.findall(r"\S+", text))

def _split_text(text: str, size=CHUNK_TOKENS, overlap=CHUNK_OVERLAP):
    """智能文本分块"""
    # 预处理：标准化空白字符
    text = re.sub(r"\s+", " ", text).strip()
    if not text:
        return
    
    # 按句子分割（支持中英文标点）
    sentences = re.split(r"(?<=[。.!?])\s+", text)
    
    buf, buf_tokens = [], 0
    for sent in sentences:
        sent_tokens = _count_tokens(sent)
        
        # 如果当前缓冲区加上新句子超过限制，输出当前块
        if buf_tokens + sent_tokens > size:
            yield " ".join(buf)
            
            # 保留重叠部分
            while buf and buf_tokens > overlap:
                buf_tokens -= _count_tokens(buf.pop(0))
        
        buf.append(sent)
        buf_tokens += sent_tokens
    
    # 输出最后一块
    if buf:
        yield " ".join(buf)

为什么要分块？

• 长度限制：AI模型有输入长度限制
• 精确检索：小块更容易匹配用户问题
• 重叠设计：避免重要信息被截断

第四步：向量化模块（RAG的核心！）

def _embed(texts: List[str]) -> np.ndarray:
    """将文本列表转换为向量矩阵"""
    idx_map, non_empty = [], []
    
    # 过滤空文本
    for i, t in enumerate(texts):
        if t.strip():
            idx_map.append(i)
            non_empty.append(t.strip())
    
    # 初始化结果矩阵（768维向量）
    out = np.zeros((len(texts), 768), dtype=np.float32)
    
    # 批量处理向量化
    for i in range(0, len(non_empty), EMBED_BATCH):
        batch = non_empty[i:i+EMBED_BATCH]
        
        # 调用嵌入模型
        vecs = np.array([d.embedding for d in 
                        client.embeddings.create(model=EMBED_MODEL, input=batch).data], 
                       dtype=np.float32)
        
        # 关键步骤：L2标准化（为余弦相似度计算做准备）
        vecs /= np.linalg.norm(vecs, axis=1, keepdims=True)
        
        out[idx_map[i:i+EMBED_BATCH]] = vecs
    
    return out

向量化的作用：

• 数学表示：将文本转换为数字向量
• 相似度计算：向量间的距离代表文本相似度
• 快速检索：向量运算比文本匹配快得多

第五步：构建文档向量库

def build_corpus(files: List[str]):
    """构建完整的文档向量库"""
    chunks, meta = [], []
    
    # 处理每个文件
    for f in files:
        print(f"正在处理: {f}")
        
        # 文件 → 文本 → 分块
        for idx, chunk in enumerate(_split_text(file_to_text(f))):
            chunks.append(chunk)
            # 保存元信息：(文件名, 块索引)
            meta.append((pathlib.Path(f).name, idx))
    
    print(f"共生成 {len(chunks)} 个文本块")
    
    # 批量向量化
    vectors = _embed(chunks)
    
    return chunks, vectors, meta

这一步做了什么？

• 批量处理：一次性处理多个文档
• 元数据管理：记录每个块来自哪个文件
• 内存存储：将向量保存在内存中（轻量级方案）

第六步：智能检索模块

def retrieve(query: str, chunks, vecs, meta, k=TOP_K):
    """基于余弦相似度的检索"""
    if vecs.size == 0:
        return []
    
    # 核心算法：余弦相似度计算
    # 1. 将查询转换为向量
    query_vec = _embed([query])[0]
    
    # 2. 计算相似度（已标准化向量的点积 = 余弦相似度）
    similarities = vecs @ query_vec
    
    # 3. 找到最相似的K个文档块
    top_indices = (-similarities).argsort()[:min(k, len(similarities))]
    
    # 4. 返回结果：(文本内容, 相似度, 文件名, 块索引)
    return [(chunks[i], similarities[i], *meta[i]) for i in top_indices]

检索算法详解：

• 余弦相似度：衡量两个向量的方向相似性
• Top-K策略：返回最相似的K个结果
• 相关性排序：按相似度从高到低排列

第七步：智能回答生成

# 系统提示词
SYS_PROMPT = "你是一位严谨的中文问答助手，只能依据参考资料作答。若资料不足请直接回答"资料不足"。"

def rag_answer_stream(question, chunks, vecs, meta):
    """生成流式回答"""
    # 1. 检索相关文档
    retrieved_docs = retrieve(question, chunks, vecs, meta)
    
    # 2. 构建上下文
    context = "\n".join(
        f"[{i+1}]《{filename}》段落{idx}：{textwrap.shorten(text, width=350, placeholder='…')}"
        for i, (text, similarity, filename, idx) in enumerate(retrieved_docs)
    )
    
    # 3. 构建完整提示词
    prompt = f"{SYS_PROMPT}\n\n参考资料：\n{context}\n\n用户问题：{question}"
    
    # 4. 调用大模型生成回答
    stream = client.chat.completions.create(
        model=LLM_MODEL,
        messages=[{"role": "user", "content": prompt}],
        temperature=0.2,  # 较低的温度，保证回答的准确性
        stream=True       # 流式输出，实时显示
    )
    
    # 5. 逐字输出回答
    for chunk in stream:
        delta = chunk.choices[0].delta
        if delta.content:
            yield delta.content

生成策略：

• 严格约束：只基于检索到的内容回答
• 上下文构建：将相关文档整理成结构化格式
• 流式输出：实时显示生成过程

第八步：用户交互

def main():
    """主程序入口"""
    if len(sys.argv) < 2:
        print("用法: python rag.py <文件1> <文件2> ...")
        sys.exit(0)

    print("正在加载文件并构建向量库…")
    
    # 构建向量库
    chunks, vecs, meta = build_corpus(sys.argv[1:])
    print(f"已生成 {len(chunks)} 段文本向量。Ctrl-C / 回车 退出。")

    # 交互式问答循环
    try:
        while True:
            q = input("\n问题> ").strip()
            if not q:  # 空输入退出
                break
            
            print("回答:", end="", flush=True)
            
            # 流式输出回答
            for s in rag_answer_stream(q, chunks, vecs, meta):
                print(s, end="", flush=True)
            print()  # 换行
            
    except (KeyboardInterrupt, EOFError):
        print("\n再见！")

if __name__ == "__main__":
    main()

实际运行效果

启动系统

python RAG.py document1.pdf document2.txt

运行输出

正在加载文件并构建向量库…
已生成 127 段文本向量。Ctrl-C / 回车 退出。

问题> 总结文档？
回答: 用户让我总结文档...

问题> ^C
再见！

核心技术点解析

1. 余弦相似度计算

# 标准化向量
vecs /= np.linalg.norm(vecs, axis=1, keepdims=True)

# 余弦相似度 = 标准化向量的点积
similarities = vecs @ query_vec

为什么用余弦相似度？

• 方向相似性：关注向量方向，不受长度影响
• 语义匹配：更适合文本语义相似度计算
• 计算高效：向量点积运算很快

2. Top-K检索策略

top_indices = (-similarities).argsort()[:k]

优势：

• 简单有效：总能返回结果
• 速度快：不需要设置阈值
• 相对最优：在现有文档中找最相关的

注意：没有绝对阈值，即使相似度很低也会返回结果

3. 流式输出实现

for chunk in stream:
    if chunk.choices[0].delta.content:
        yield chunk.choices[0].delta.content

用户体验提升：

• 实时反馈：不用等完整回答生成完毕
• 对话感强：模拟真人对话的感觉
• 资源高效：边生成边输出，节省内存

系统特点总结

优势

1. 轻量级：无需复杂的向量数据库
2. 多格式支持：7种常见文档格式
3. 本地化：数据不上传云端，隐私安全
4. 实时性：内存计算，响应速度快
5. 可扩展：代码结构清晰，易于修改

局限性

1. 内存限制：大量文档需要较多内存
2. 无持久化：重启后需要重新构建向量库
3. 无阈值过滤：可能返回不相关的内容
4. 模型依赖：需要本地部署AI模型

通过这篇文章，我们从零开始实现了一个完整的RAG系统：

1. 理解概念：RAG = 检索 + 生成
2. 动手实践：不到200行代码完整实现
3. 掌握原理：向量化、相似度计算、流式生成
4. 实际应用：支持多种文档格式的智能问答

RAG技术让AI具备了"查阅资料"的能力，这为构建专业领域的AI助手打开了大门。无论是企业知识库问答、学术论文分析，还是个人文档管理，RAG都有广阔的应用前景。

下一步，你可以尝试：

• 运行这个代码，用自己的文档测试
• 根据需求调整参数配置
• 扩展功能，比如添加网页爬取
• 针对特定领域优化提示词

希望这篇文章能帮助你理解RAG技术的核心原理，并成功构建属于自己的智能问答系统！

---

技术交流：如果你在实现过程中遇到问题，欢迎在评论区讨论！

觉得有用请点赞收藏，你的支持是我分享技术的动力！