从零开始实现一个 RAG API 服务:ChromaDB + FastAPI 实战
写在前面
上一篇内容聊过 LangChain,它是一个很棒的编排框架。但我真正手写一遍 RAG 后,我对 RAG 底层原理的理解又加深了。
这段时间我用 FastAPI + ChromaDB 从头做了一个 RAG API 服务,没有框架封装,每一步都是自己实现的。写完之后,对 RAG 的理解又上了一个台阶。
这篇文章就把我的实现思路和代码拆开来讲,希望能帮到正在学 RAG 的朋友。
整体架构
先看项目结构,很简洁:
my_RAG_API/
├── main.py # FastAPI 入口,定义路由
├── config.py # 全局配置(API Key、模型、分块参数)
├── ChromaDB.py # 向量数据库操作(嵌入生成、存储、检索)
├── RAG_engine.py # RAG 核心引擎(分块 + 检索 + 生成)
├── llm_client.py # LLM 客户端(封装 DeepSeek API)
├── tools.py # 工具函数(计算器 Function Calling)
├── data/ # 知识库文档
└── chroma_db/ # 向量数据持久化目录
整个流程长这样:
用户提问
│
▼
将问题转为 Embedding ──► 去 ChromaDB 语义检索
│
▼
找到 Top-K 最相关文本块
│
▼
拼装 Prompt(上下文 + 问题)
│
▼
发给 DeepSeek LLM 生成回答
│
▼
返回结果
第一步:配置文件,隐藏敏感信息
所有可变的参数统一放在 config.py,方便管理:
# config.py
SILICONFLOW_API_KEY = os.getenv("SILICONFLOW_API_KEY") # 向量模型 API Key
DEEPSEEK_API_KEY = os.getenv("DEEPSEEK_API_KEY") # LLM API Key
EMBEDDING_MODEL = "BAAI/bge-m3" # 开源中文嵌入模型
LLM_MODEL = "deepseek-chat" # DeepSeek 对话模型
CHROMA_PERSIST_DIR = "./chroma_db" # 向量库持久化目录
CHUNK_SIZE = 500 # 文本分块大小
CHUNK_OVERLAP = 50 # 分块重叠大小
DOCUMENT_PATH = "./data/tech_article.txt" # 知识库文件路径
API Key 存在 .env 里,load_dotenv() 自动加载,不上传到 Git,安全省心。
第二步:文本分块,RAG 的第一步门槛
原始的文档不能直接丢给向量数据库,需要先切分成小块。这里我实现了一个智能分块函数:
def chunk_text(text, chunk_size=CHUNK_SIZE, overlap=CHUNK_OVERLAP):
chunks = []
start = 0
text_len = len(text)
while start < text_len:
end = min(start + chunk_size, text_len)
# 尽量在句号/问号/感叹号处断开,保持语义完整
if end < text_len:
for sep in ["。", "?", "!", "\n", ".", "?", "!"]:
last_sep = text.rfind(sep, start, end)
if last_sep > start + chunk_size // 2:
end = last_sep + 1
break
chunk = text[start:end].strip()
if chunk:
chunks.append(chunk)
start = end - overlap if end < text_len else end
return chunks
核心思路是:
- 从文本开头开始,每 500 字符切一块
- 但不会生硬地切在中间——它会往回找最近的句号,保证每个 chunk 是完整的句子
- 相邻块之间重叠 50 字符,避免信息刚好被切散
这个细节很重要,我一开始没做重叠,结果有些跨 chunk 的关键信息检索不到,改完立马好多了。
第三步:Embedding + ChromaDB,搭建向量知识库
嵌入模型我选的是硅基流动(SiliconFlow)托管的 BAAI/bge-m3,一款国产开源的中文嵌入模型,在中文场景的使用效果还算不错,而且调用有免费额度。
import chromadb
import requests
from config import SILICONFLOW_API_KEY, EMBEDDING_API_URL, EMBEDDING_MODEL
# 调用嵌入 API
def get_embeddings(texts):
headers = {
"Authorization": f"Bearer {SILICONFLOW_API_KEY}",
"Content-Type": "application/json",
}
payload = {
"model": EMBEDDING_MODEL,
"input": texts if isinstance(texts, list) else [texts],
"encoding_format": "float",
}
resp = requests.post(EMBEDDING_API_URL, json=payload, headers=headers, timeout=30)
resp.raise_for_status()
data = resp.json()
return [item["embedding"] for item in data["data"]]
# 初始化持久化集合
def init_collection():
chromadb_client = chromadb.PersistentClient(path=CHROMA_PERSIST_DIR)
collection = chromadb_client.get_or_create_collection(name="rag_docs")
return collection
# 添加文档(自动生成嵌入向量)
def add_documents(collection, documents, ids, metadatas):
embeddings = get_embeddings(documents)
collection.add(documents=documents, embeddings=embeddings, ids=ids, metadatas=metadatas)
ChromaDB 的 PersistentClient 会把向量数据存到本地磁盘,下次启动不用重新生成,省时又省力。
第四步:RAG 引擎,把一切串起来
RAGEngine 是我整个服务的核心类,负责初始化知识库、检索上下文、生成回答三件事:
class RAGEngine:
def __init__(self, collection=None):
self.collection = collection if collection else init_collection()
self.is_initialized = False
# 从文本文件初始化向量库(自动去重)
def init_from_text_file(self, file_path: str):
if self.collection.count() > 0:
print(f"向量库已存在,共 {self.collection.count()} 个向量")
return
with open(file_path, "r", encoding="UTF-8") as f:
raw_text = f.read()
chunks = chunk_text(raw_text)
ids = [f"chunk_{i}" for i in range(len(chunks))]
metadatas = [{"source": file_path, "chunk_index": i} for i in range(len(chunks))]
add_documents(self.collection, chunks, ids, metadatas)
# 语义检索 Top-3
def retrieve(self, query: str, n_results: int = 3):
results = self.collection.query(
query_embeddings=[get_embeddings(query)[0]],
n_results=n_results,
include=["documents", "distances", "metadatas"],
)
return results
# 根据检索结果生成回答
def generate_answer(self, query: str, context_chunks: list) -> str:
context = "\n\n---\n\n".join(context_chunks)
prompt = f"""
你是一个专业的技术助手。请基于以下【参考资料】回答用户的问题。
如果参考资料中没有相关信息,请如实说"资料中没有提及",不要编造。
【参考资料】
{context}
【用户问题】
{query}
【回答】"""
return get_LLM_response(prompt, temperature=0.3)
def query(self, query: str) -> str:
results = self.collection.query(
query_embeddings=[get_embeddings(query)[0]],
n_results=3,
include=["documents"],
)
chunks = results["documents"][0]
if not chunks:
return "没有找到相关信息"
return self.generate_answer(query, chunks)
有一个很实用的细节:每次服务重启时,init_from_text_file 会先检查向量库是否已存在,避免重复生成嵌入向量。提高了效率。
第五步:FastAPI 包装成 API 接口
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
RAG = RAGEngine()
RAG.init_from_text_file(DOCUMENT_PATH)
app = FastAPI(title="RAG API", description="基于 ChromaDB + DeepSeek 的问答服务")
class Question(BaseModel):
question: str
@app.post("/chat")
def chat(q: Question):
try:
answer = RAG.query(q.question)
return {"question": q.question, "answer": answer}
except Exception as e:
raise HTTPException(status_code=500, detail=str(e))
@app.get("/health")
def health():
return {"status": "ok"}
两个接口,清晰明了:
POST /chat— 传入问题,返回基于知识库的答案GET /health— 健康检查,便于部署时排查错误
进阶:Function Calling,让 LLM 学会用工具
除了问答外,我还加了一个 计算器工具 的端点。当用户问"123 × 456 等于多少"时,LLM 会主动调用工具来算,而不是自己瞎编:
# tools.py — 定义工具
calculate_tool = {
"type": "function",
"function": {
"name": "calculator",
"description": "执行数学计算,支持加减乘除、幂运算等",
"parameters": {
"type": "object",
"properties": {
"expression": {
"type": "string",
"description": "数学表达式,例如 '123 * 456'",
}
},
"required": ["expression"],
},
},
}
def run_calculator(expression: str) -> str:
allowed_chars = set("0123456789+-*/().%**")
if not all(c in allowed_chars for c in expression):
return "错误:表达式包含非法字符"
try:
result = eval(expression)
return str(result)
except Exception as e:
return f"计算错误: {e}"
在 /agent 端点中,我实现了标准的 工具调用循环:
@app.post("/agent")
def agent_endpoint(req: Question):
messages = [{"role": "user", "content": req.question}]
# 第一轮:LLM 决定是否调用工具
first_response = LLM_client.chat.completions.create(
model="deepseek-chat",
messages=messages,
tools=[calculate_tool],
tool_choice="auto",
)
message = first_response.choices[0].message
# 如果没调用工具,直接返回
if not message.tool_calls:
return {"question": req.question, "answer": message.content}
# 如果有工具调用,执行工具并把结果发回给 LLM
for tool_call in message.tool_calls:
if tool_call.function.name == "calculator":
args = json.loads(tool_call.function.arguments)
result = run_calculator(args.get("expression", ""))
messages.append(message)
messages.append({
"role": "tool",
"tool_call_id": tool_call.id,
"content": result,
})
# 第二轮:LLM 结合工具结果给出最终回答
final_response = LLM_client.chat.completions.create(
model="deepseek-chat",
messages=messages,
)
return {"question": req.question, "answer": final_response.choices[0].message.content}
流程就是:用户提问 → LLM 判断需要计算 → 提取表达式 → 执行计算 → 结果注入上下文 → LLM 组织最终答案。这一套下来,感觉离真正的 Agent 又近了一步。
启动服务
uvicorn main:app --reload
启动后访问 http://127.0.0.1:8000/docs,Swagger 文档就出来了,可以直接在页面上测试接口。因为只是个人练手的项目,所以没有部署到云端,等今后业务精进了,把它变成能投入生产的服务后再部署吧
总结与心得
在写这个 RAG 服务的过程中,有几个体会很深:
- 分块策略直接影响检索质量——无重叠的硬切会让信息刚好断在两块之间,加上重叠和语义断句后效果提升明显
- Embedding 模型选对语言很重要——中文场景用 BAAI/bge-m3 比通用模型准确不少
- 向量库去重是个容易被忽略的细节——服务重启时重复写入会导致检索结果被稀释,提前判断
count()能避免 - Function Calling 让 RAG 不止于问答——加上工具调用后,LLM 的边界被大大扩展了
当然这个版本也比较简陋:没有流式输出、没有多轮对话记忆、没有异步处理。不过作为理解 RAG 原理的起点够用了。
下一步我打算加上对话历史和流式响应,并用LangChain重写,让它更像一个真正的 AI 助手。欢迎提出意见