当前位置:首页>python>手把手搭建一个最小 RAG Demo:Python + Chroma,50 行代码跑通

手把手搭建一个最小 RAG Demo:Python + Chroma,50 行代码跑通

  • 2026-10-11 07:49:46
手把手搭建一个最小 RAG Demo:Python + Chroma,50 行代码跑通

一个Java 程序员的 RAG 学习笔记。不是教程,是"我踩过的坑你不用再踩"。

上篇讲了 RAG 的三件事:建索引、检索、生成。评论区有朋友说"看懂了,但不知道怎么动手"。

我也是。看完概念觉得简单,打开编辑器愣了半天——从哪开始?

这篇解决这个问题。我们用 Python + ChromaDB,50 行代码搭一个能跑的 RAG。跑通了你就有体感了,后面换什么组件都好说。

为什么选ChromaDB

选型的时候我对比了几种方案:

方案
优点
缺点
ChromaDB
pip 装一下就能用,内置 Embedding
不适合生产环境
Milvus
性能好,支持分布式
部署复杂,学习成本高
Pinecone
全托管,不用运维
要注册,要联网,有费用
FAISS
Facebook 出品,性能强
只做检索,不管存储和 Embedding

我的想法是:先跑起来再说。ChromaDB 装个包就能用,不用注册任何东西,不用启动额外服务。等跑通了,再换 Milvus 也不迟。

环境准备

pip install chromadb openai

两个包。chromadb 是向量数据库,openai 用来调大模型。

如果你跟我一样用 DeepSeek(便宜),设两个环境变量:

export OPENAI_API_KEY="你的key"

export OPENAI_BASE_URL="https://api.deepseek.com"

DeepSeek 的接口和 OpenAI 兼容,代码不用改,换个地址就行。

完整代码

一共 50 行出头,我直接贴了:

import chromadb

from openai import OpenAI

---- 第一步:建索引 ----

client = chromadb.Client() # 内存模式,重启就没了

collection = client.create_collection("my_docs")

docs = [

"RAG 是检索增强生成的缩写,核心思路是先从知识库检索相关文档,再交给大模型生成回答。",

"ChromaDB 是一个轻量级向量数据库,pip install 就能用,适合本地开发和原型验证。",

"Embedding 是把文本转换成向量的过程,语义相近的文本在向量空间里距离更近。",

"Java 程序员转 AI Agent 开发,建议先过一遍 Python 基础,再看 LangChain 或 Spring AI。",

"向量检索的原理是算余弦相似度,找到和问题最相关的文档片段喂给大模型。",

]

collection.add(

documents=docs,

ids=[f"doc_{i}" for i in range(len(docs))]

)

---- 第二步:检索 ----

question = "什么是 RAG?"

results = collection.query(query_texts=[question], n_results=3)

---- 第三步:生成 ----

context = "\n".join(results["documents"][0])

prompt = f"""基于以下参考资料回答问题。如果资料里没有相关内容,就说"我不确定"。

参考资料:

{context}

问题:{question}

"""

llm = OpenAI()

resp = llm.chat.completions.create(

model="deepseek-chat",

messages=[{"role": "user", "content": prompt}]

)

print("Q:", question)

print("A:", resp.choices[0].message.content)

跑一下,输出大概是这样:

Q: 什么是 RAG?

A: RAG 是"检索增强生成"的缩写。它的核心思路是先从知识库中检索与问题相关的文档片段,

再把这些片段交给大模型,让大模型基于这些事实来生成回答。这样可以减少大模型"编答案"的问题。

50 行,三步,跑通了。

拆解一下每步在干什么

建索引:把文档喂进去

collection.add(documents=docs, ids=[...])

就这一行。ChromaDB 在背后自动帮你做了 Embedding——把每段文本转成一个几百维的向量,存到内存里。

你不需要自己调 Embedding 模型,ChromaDB 默认用 all-MiniLM-L6-v2。先跑通,后面再换。

检索:找到最相关的几段

results = collection.query(query_texts=[question], n_results=3)

ChromaDB 把你的问题也转成向量,然后和所有文档算相似度,返回最像的 3 段。

这就是上篇说的"检索"。本质就是数学——两个向量越接近,语义越相似。

生成:交给大模型

prompt = f"基于以下参考资料回答问题..."

resp = llm.chat.completions.create(...)

把检索到的 3 段文本拼进 Prompt,丢给 DeepSeek,完事。

注意那句"如果资料里没有相关内容,就说'我不确定'"。这句很关键,不加的话 AI 会编。

跑完你可能会有的几个想法

"Embedding 居然不用我管?"

对,ChromaDB 帮你做了。但代价是默认模型是英文的,中文效果一般。后面你可能会想换中文 Embedding,比如 text2vec-base-chinese。

"检索也太简单了吧?"

确实,一行代码。但真实项目里你不会把整篇文章丢进去——你得先切分,200-500 字一段,切太大检索不精准,切太小上下文丢失。

"Prompt 写法影响很大"

同一份资料,Prompt 写得好,AI 回答得准;写得差,AI 能把子虚乌有的东西说得头头是道。RAG 的质量,一大半在 Prompt 上。

我踩过的几个坑

文档切分

Demo 里的 docs 是我手动写的 5 句话,每句刚好一段。但真实场景你要处理长文档——PDF、Markdown、网页。切分策略直接决定检索质量。

我的经验:200-500 字一段,相邻段之间重叠 50-100 字。别贪多,宁可多切几段也别把不相关的内容塞到同一段里。

中文 Embedding

ChromaDB 默认的 all-MiniLM-L6-v2是英文模型。我拿中文文档试过,检索结果经常对不上。

后来换了text2vec-base-chinese,效果好很多。换法也很简单,初始化的时候指定就行:

collection = client.create_collection(

"my_docs",

metadata={"hf:all-MiniLM-L6-v2": None}   不用默认的

)

自己先 Embedding 好再 add

n_results 调参

n_results=3是随便写的。实际你得试——给多了,AI 被无关信息干扰;给少了,关键信息可能漏掉。

我一般从 5 开始试,看回答质量再调。

"我不知道"

Prompt 里不写这句,AI 面对资料里没有的问题,会非常自信地编一个答案。RAG 最怕这个——你做了检索增强,结果 AI 还是在编,那检索了个寂寞。

下一步

这个 Demo 是最小可运行版本。如果你想继续深入,我觉得按这个顺序比较合理:

1.接真实数据源——读 PDF 或 Markdown,自动切分存入

2.换中文 Embedding——text2vec-base-chinese或bge-large-zh

3.持久化——ChromaDB 可以存到本地目录,重启不丢数据

4.评估——准备一批问答对,测一下检索准不准、回答对不对

最新文章

随机文章