RAG 检索增强生成Python 示例

A_miahao

247人浏览 · 2026-08-06 20:57:19

A_miahao · 2026-08-06 20:57:19 发布

RAG 简单理解：读取 docx 文档 → 文本切分 → 向量化 → 存入 Chroma 向量库；用户提问，把问题转向量，从向量库检索相关文档片段，把片段塞给大模型，大模型依据文档片段回答问题，减少幻觉。

依赖安装

bash

pip install chromadb python‑docx openai

整体流程链路： 读取docx文档 → 文本切分split_text → 调用Embedding接口生成向量 → 存入Chroma持久化向量数据库 用户提问 → 问题向量化 → Chroma向量相似度检索 → 把检索出来的文档片段+用户问题组装Prompt → Ollama本地qwen3.5:9b生成答案

1. 导入库

python

运行

import chromadb
from docx import Document
from openai import OpenAI

chromadb：轻量级本地向量数据库，持久化存储文本片段 + 向量，做相似度检索；
docx.Document：读取 Word (.docx) 文档；
OpenAI：OpenAI 兼容 SDK，同时对接阿里云百炼 Embedding 接口和本地 Ollama 大模型。

⚠️注意：这里两个不同 OpenAI 客户端！

Embedding：调用阿里云百炼云端接口 base_url="https://ws‑fyhmdrd51kppowv3.cn‑beijing.maas.aliyuncs.com/compatible‑mode/v1"

生成回答：调用本地 Ollama base_url="http://127.0.0.1:11434/v1"

2.load_file ()：读取 docx 文件

python

运行

def load_file(file_path):
    document = Document(file_path)
    all_text = []
    for para in document.paragraphs:
        all_text.append(para.text)
    return "\n".join(all_text)

Document(file_path) 加载 word 文档；
document.paragraphs 获取文档全部段落对象；
循环取出每个段落的.text文本，存入列表；
使用换行符把所有段落拼接成一整个大字符串返回。

缺陷：当前代码只能读取段落，读不到表格里面的内容，如果 word 有表格，表格内容会丢失。

3.split_text ()：简单文本切分

python

运行

def split_text(all_text):
    chunks = []
    for i in range(0, len(all_text), 200):
        chunks.append(all_text[i:i + 200])
    return chunks

最简单的固定长度切分，每一段 200 字符。

range(0, len(all_text),200)：下标 0、200、400……
all_text[i:i+200] 切片截取 200 个字符，生成多个文本块 (chunk)

⚠️缺点：粗暴按字符切割，会把一句话、一个制度条目直接从中间切断，会破坏语义。生产一般用递归字符分割RecursiveCharacterTextSplitter。

4.text_embedding (chunks)：文本生成向量（Embedding）

python

运行

def text_embedding(chunks):
    client = OpenAI(
        api_key="sk‑xxx",
        base_url="https://ws‑fyhmdrd51kppowv3.cn‑beijing.maas.aliyuncs.com/compatible‑mode/v1",
    )

    if isinstance(chunks, str):
        chunks = [chunks]

    batch_size = 20
    embeddings = []

    for i in range(0, len(chunks), batch_size):
        batch_chunks = chunks[i:i + batch_size]
        completion = client.embeddings.create(
            model="qwen3.7‑text‑embedding",
            input=batch_chunks
        )
        batch_embeddings = [item.embedding for item in completion.data]
        embeddings.extend(batch_embeddings)

    return embeddings

Embedding 作用：把人类文本变成一串浮点数向量，语义越相近，向量距离越小。

OpenAI 客户端指向阿里云百炼 Embedding 接口；
判断入参，如果传入字符串，包装成列表；
batch_size=20：批量请求，一次最多 20 个文本，减少 http 请求次数，提升速度；
循环分批调用client.embeddings.create，模型qwen3.7‑text‑embedding；
返回结果completion.data中每个 item 的.embedding就是向量数组；
将多批向量汇总，返回向量列表。

重要：文档切片、用户查询，全部调用同一个 embedding 模型，向量空间保持一致，检索才有效。

5.save_to_vector_db (chunks)：保存文本 + 向量到 Chroma 向量库

python

运行

def save_to_vector_db(chunks):
    embeddings = text_embedding(chunks)
    client = chromadb.PersistentClient()
    collection = client.get_or_create_collection(
        name="employee_manual_collection",
        embedding_function=None
    )

    ids = [f"chunk_{i}" for i in range(len(chunks))]

    collection.add(
        ids=ids,
        embeddings=embeddings,
        documents=chunks
    )

chromadb.PersistentClient()：持久化客户端，会在本地生成chroma/文件夹，向量数据落磁盘，程序重启不会丢失；
get_or_create_collection：获取或者新建一个向量集合（相当于一张数据表），集合名字employee_manual_collection；
embedding_function=None：关闭 chroma 自带 embedding，向量由我们自己外部调用阿里云接口生成，手动传入；
ids：给每个文本块生成唯一 ID chunk_0、chunk_1；
collection.add()写入数据：
- ids：每条数据唯一标识
- embeddings：我们算好的向量
- documents：原始文本片段

执行完后本地目录会多出 chroma 文件夹，后续查询直接读取这个文件夹，不需要重复解析文档、重复算 embedding。

6.query_by_chroma (user_query_text)：用户问题向量检索

python

运行

def query_by_chroma(user_query_text):
    user_query_embeddings = text_embedding(user_query_text)
    client = chromadb.PersistentClient()
    collection = client.get_or_create_collection(
        name="employee_manual_collection",
        embedding_function=None
    )

    results = collection.query(
        query_embeddings=user_query_embeddings,
        n_results=5,
        include=[
            "documents",
            "distances",
            "metadatas",
        ]
    )

    return results

RAG 检索阶段核心函数：

用户的问题字符串调用text_embedding转为向量；
连接本地 Chroma，打开同一个集合；
collection.query()向量相似度查询：
- query_embeddings 用户问题的向量；
- n_results=5 返回相似度最高 5 条文档片段；
- include 指定返回内容：原文 documents、距离 distances、元数据 metadatas；
返回 results 字典，结构：

python

运行

{
 "ids":[["chunk_xx",...]],
 "documents":[["片段1","片段2"...]],
 "distances":[[0.12,0.21...]]
}

distances 距离数值越小，代表和用户提问语义越接近。

7.generate_answer ()：组装 Prompt，交给大模型生成答案

python

运行

def generate_answer(query_results, user_query_text):
    results_chunks = query_results["documents"][0]
    content = "\n".join([f"--: {chunk}" for chunk in results_chunks])

    prompt = f"""
    ## 角色设定
    你是一个专业的人力资源专家,
    ## 任务描述
    根据用户的问题和公司的制度内容,回答问题
    ## 输入数据
    用户的问题:{user_query_text}
    公司制度内容:{content}
    ## 约束
    1:严格基于公司的制度内容回答问题,不要胡编乱造
    2:如果公司的制度内容没有相关信息,请明确说明“暂无此知识”
    3:回答简洁明了,条理清晰
    """
    client = OpenAI(
        api_key="xxx",
        base_url="http://127.0.0.1:11434/v1",
    )

    completion = client.chat.completions.create(
        model="qwen3.5:9b",
        messages=[{"role": "user", "content": prompt}],
        temperature=0.75,
    )
    print(completion.model_dump_json())
    return completion.choices[0].message.content

query_results["documents"][0]：取检索回来的 5 条相关文档片段；
拼接参考上下文content；
构造 Prompt 提示词，包含角色、任务、用户问题、从向量库检索到的文档、强制约束；
- 强制大模型只允许参考给的公司制度内容；
- 找不到资料输出 “暂无此知识”，用来抑制大模型幻觉；
OpenAI 客户端指向本地 Ollama qwen3.5:9b，传入完整 prompt，获取大模型输出返回给调用方。

8. 程序入口 main

python

运行

if __name__ == '__main__':
    # all_text = load_file("D:\积云教育\C2502A\demo00730\llm\employee_manual.docx")
    # chunks = split_text(all_text)
    # embeddings = text_embedding(chunks)
    # save_to_vector_db(chunks)
    user_query_text = "员工上班时间超过多久算迟到"
    results = query_by_chroma(user_query_text)
    # print(results)

    final_answer = generate_answer(results, user_query_text)
    print(f"大模型生成的最终答案:{final_answer}")