12. FastAPI + Chroma 最小原型 - Enkidu 的学习笔记

目标：把双存储架构从”脚本”变成”API 接口”，让前端或其他服务能通过网络调用 RAG 能力。

本章地图#

1
上一章：双存储架构（SQLite + ChromaDB 协作）
2
        ↓
3
本章：  把双存储封装成 FastAPI 接口
4
        ├── 为什么需要 API 化？
5
        ├── 接口设计（RESTful 风格）
6
        ├── 代码结构（router 分层）
7
        ├── 持久化 ChromaDB（数据不丢失）
8
        └── 测试接口（curl / 浏览器）
9
        ↓
10
下一章：手搓最小 RAG 闭环（检索 → Prompt → LLM → 回答）

1. 为什么需要 API 化？#

上一章的问题#

dual_storage_demo.py 是脚本：

✅ 能跑通双存储流程
❌ 只能命令行执行
❌ 前端无法调用
❌ 每次重启 ChromaDB 数据丢失（内存模式）

API 化的好处#

对比	脚本	API
调用方式	`python demo.py`	`POST /rag/search`
谁可以用	你自己	前端、手机 App、其他服务
数据持久	重启清空	存磁盘，重启还在
复用性	低	高

类比：

脚本 = 家里自己做菜，只有自己能吃
API = 开餐厅，任何人都能点菜

2. 接口设计#

只暴露两个接口#

1
POST /rag/documents    → 存入文档（双存储）
2
POST /rag/search       → 语义检索（双存储协作）

为什么只设计两个？

最小原型的核心思想：先跑通主流程，再丰富边缘功能。删除、修改、列表查询可以后面加。

请求/响应格式#

存入文档#

1
// 请求
2
POST /rag/documents
3
{
4
  "title": "Python 入门指南",
5
  "content": "Python 是一种高级编程语言...",
6
  "source": "官方文档"
7
}
8

9
// 响应
10
{
11
  "message": "文档存入成功",
12
  "document_id": 1,
13
  "title": "Python 入门指南",
14
  "chunks_count": 3
15
}

语义检索#

1
// 请求
2
POST /rag/search
3
{
4
  "query": "我想学编程",
5
  "n_results": 3
6
}
7

8
// 响应
9
[
10
  {
11
    "title": "Python 入门指南",
12
    "chunk_content": "Python 是一种高级编程语言...",
13
    "similarity": 0.8234,
14
    "document_id": 1,
15
    "chunk_index": 0
16
  }
17
]

3. 代码结构#

新增文件#

1
routers/
2
├── app/routers/rag.py          ← 新增：RAG 路由（双存储 API）
3
└── ...
4

5
main.py                     ← 修改：注册 rag_router
6
chroma_db/                  ← 新增：ChromaDB 持久化数据目录

router 分层（复习）#

1
# main.py 只负责"注册路由"，不写业务逻辑
2
from routers import rag_router
3
app.include_router(rag_router.router, prefix="/rag")
4

5
# app/routers/rag.py 负责"RAG 业务"
6
router = APIRouter(prefix="/rag", tags=["RAG 向量检索"])
7

8
@router.post("/documents")
9
def add_document(doc: DocumentIn, db: Session = Depends(get_db)):
10
    # 存入双存储...
11

12
@router.post("/search")
13
def search(query: SearchIn, db: Session = Depends(get_db)):
14
    # 语义检索...

分层的好处：main.py 干净，app/routers/rag.py 专注 RAG，以后加新功能互不干扰。

4. 关键变化：ChromaDB 持久化#

上一章（内存模式）#

1
client_db = chromadb.Client()  # 数据存在内存，重启清空

本章（持久化模式）#

1
chroma_client = chromadb.PersistentClient(path="./chroma_db")
2
# 数据存在 ./chroma_db 目录，重启后还在

模式	代码	特点
内存模式	`chromadb.Client()`	快，但重启数据丢
持久化模式	`chromadb.PersistentClient(path="...")`	稍慢，但数据持久

类比：

内存模式 = 草稿纸，写完就扔
持久化模式 = 笔记本，写完了还能翻

5. 依赖注入（复习 + 应用）#

数据库会话注入#

1
def get_db():
2
    db = SessionLocal()
3
    try:
4
        yield db          # ← 把 db 交给路由函数用
5
    finally:
6
        db.close()        # ← 用完自动关闭
7

8
@router.post("/documents")
9
def add_document(doc: DocumentIn, db: Session = Depends(get_db)):
10
    # db 自动传进来，不用手动创建
11
    document = Document(title=doc.title, ...)
12
    db.add(document)
13
    db.commit()

为什么用 yield？

保证数据库连接用完就关，不泄漏。即使代码报错，finally 也会执行关闭。

6. 完整代码#

app/routers/rag.py#

1
from fastapi import APIRouter, Depends
2
from pydantic import BaseModel
3
from sqlalchemy.orm import Session
4
from typing import List
5

6
from database import SessionLocal
7
from models import Document, DocumentChunk
8

9
import chromadb
10
from openai import OpenAI
11
from dotenv import load_dotenv
12
import os
13

14
# ========== 初始化 ==========
15
load_dotenv()
16

17
client = OpenAI(
18
    base_url="https://api-inference.modelscope.cn/v1",
19
    api_key=os.getenv("MODELSCOPE_API_KEY")
20
)
21

22
# 持久化 ChromaDB
23
chroma_client = chromadb.PersistentClient(path="./chroma_db")
24
collection = chroma_client.get_or_create_collection(
25
    name="documents",
26
    metadata={"hnsw:space": "cosine"}
27
)
28

29
# ========== 依赖 ==========
30
def get_db():
31
    db = SessionLocal()
32
    try:
33
        yield db
34
    finally:
35
        db.close()
36

37
# ========== Pydantic 模型 ==========
38
class DocumentIn(BaseModel):
39
    title: str
40
    content: str
41
    source: str = ""
42

43
class SearchIn(BaseModel):
44
    query: str
45
    n_results: int = 3
46

47
class SearchResult(BaseModel):
48
    title: str
49
    chunk_content: str
50
    similarity: float
51
    document_id: int
52
    chunk_index: int
53

54
# ========== 工具函数 ==========
55
def get_embedding(text: str) -> list[float]:
56
    response = client.embeddings.create(
57
        model="Qwen/Qwen3-Embedding-8B",
58
        input=text,
59
        encoding_format="float"
60
    )
61
    return response.data[0].embedding
62

63
def split_text(text: str, chunk_size: int = 80, overlap: int = 10) -> list[str]:
64
    chunks = []
65
    start = 0
66
    while start < len(text):
67
        end = start + chunk_size
68
        chunks.append(text[start:end])
69
        start += (chunk_size - overlap)
70
    return chunks
71

72
# ========== 路由 ==========
73
router = APIRouter(prefix="/rag", tags=["RAG 向量检索"])
74

75
@router.post("/documents")
76
def add_document(doc: DocumentIn, db: Session = Depends(get_db)):
77
    # 1. SQLite 存完整文档
78
    document = Document(title=doc.title, content=doc.content, source=doc.source)
79
    db.add(document)
80
    db.commit()
81
    db.refresh(document)
82

83
    # 2. 切片
84
    chunks = split_text(doc.content)
85

86
    # 3. 准备 ChromaDB 数据
87
    chroma_ids, chroma_embeddings, chroma_documents, chroma_metadatas = [], [], [], []
88

89
    for idx, chunk_text in enumerate(chunks):
90
        # SQLite 切片
91
        chunk = DocumentChunk(
92
            document_id=document.id,
93
            chunk_index=idx,
94
            content=chunk_text,
95
            embedding_id=f"doc{document.id}_chunk{idx}"
96
        )
97
        db.add(chunk)
98

99
        # 向量 + ChromaDB 数据
100
        vec = get_embedding(chunk_text)
101
        chroma_ids.append(f"doc{document.id}_chunk{idx}")
102
        chroma_embeddings.append(vec)
103
        chroma_documents.append(chunk_text)
104
        chroma_metadatas.append({
105
            "document_id": document.id,
106
            "chunk_index": idx,
107
            "title": doc.title
108
        })
109

110
    db.commit()
111

112
    # 4. ChromaDB 存向量
113
    collection.add(
114
        ids=chroma_ids,
115
        embeddings=chroma_embeddings,
116
        documents=chroma_documents,
117
        metadatas=chroma_metadatas
118
    )
119

120
    return {
121
        "message": "文档存入成功",
122
        "document_id": document.id,
123
        "chunks_count": len(chunks)
124
    }
125

126
@router.post("/search", response_model=List[SearchResult])
127
def search(query: SearchIn, db: Session = Depends(get_db)):
128
    # 1. 查询转向量
129
    query_vec = get_embedding(query.query)
130

131
    # 2. ChromaDB 检索
132
    chroma_results = collection.query(
133
        query_embeddings=[query_vec],
134
        n_results=query.n_results,
135
        include=["metadatas", "distances"]
136
    )
137

138
    # 3. 回查 SQLite
139
    results = []
140
    for i in range(query.n_results):
141
        metadata = chroma_results["metadatas"][0][i]
142
        distance = chroma_results["distances"][0][i]
143
        similarity = round(1 - distance, 4)
144

145
        doc_id = metadata["document_id"]
146
        chunk_idx = metadata["chunk_index"]
147

148
        chunk = db.query(DocumentChunk).filter(
149
            DocumentChunk.document_id == doc_id,
150
            DocumentChunk.chunk_index == chunk_idx
151
        ).first()
152

153
        if chunk:
154
            document = db.query(Document).filter(Document.id == doc_id).first()
155
            results.append(SearchResult(
156
                title=document.title if document else "未知",
157
                chunk_content=chunk.content,
158
                similarity=similarity,
159
                document_id=doc_id,
160
                chunk_index=chunk_idx
161
            ))
162

163
    return results

main.py（修改）#

1
from fastapi import FastAPI
2
from routers import ai_router, todos_routers, chat_memory_router, rag_router
3

4
app = FastAPI()
5

6
app.include_router(ai_router)
7
app.include_router(todos_routers)
8
app.include_router(chat_memory_router)
9
app.include_router(rag_router.router)   # ← 新增

7. 测试接口#

启动服务#

1
cd /home/enkidu/study_python
2
python main.py

用 curl 测试#

1
# 存入文档
2
curl -X POST "http://127.0.0.1:8000/rag/documents" \
3
  -H "Content-Type: application/json" \
4
  -d '{
5
    "title": "Python 入门",
6
    "content": "Python 是一种高级编程语言，语法简洁...",
7
    "source": "教程"
8
  }'
9

10
# 语义检索
11
curl -X POST "http://127.0.0.1:8000/rag/search" \
12
  -H "Content-Type: application/json" \
13
  -d '{
14
    "query": "编程语言",
15
    "n_results": 2
16
  }'

用浏览器测试#

访问 http://127.0.0.1:8000/docs → FastAPI 自动生成的 Swagger UI，可以直接点”Try it out”测试。

8. 本章小结#

知识点	一句话总结
API 化	把脚本变成接口，让任何人都能调用
持久化 ChromaDB	`PersistentClient` 存磁盘，重启不丢数据
router 分层	`main.py` 注册，`app/routers/rag.py` 写业务
依赖注入	`Depends(get_db)` 自动管理数据库连接
最小原型	先跑通两个核心接口，再扩展功能

🔧 准确术语速查#

术语	准确含义	本章对应
API 化	把脚本能力包装成 HTTP 接口	`/rag/documents`、`/rag/search`
PersistentClient	ChromaDB 持久化客户端	数据写入 `./chroma_db`
Router	FastAPI 路由模块	`app/routers/rag.py`
Dependency injection	依赖注入	`Depends(get_db)`
Prototype	最小原型，先验证主流程	只保留存入和检索两个核心接口
Swagger UI	FastAPI 自动交互文档	`http://127.0.0.1:8000/docs`

⚠️ 常见坑#

坑	现象	正确做法
原型阶段做太多功能	主流程还没跑通就复杂化	先保留存入和检索两个接口
Chroma 用内存模式	重启后数据消失	用 `PersistentClient(path="./chroma_db")`
router 写了但没注册	Swagger 没有 RAG 接口	`main.py` 里 `app.include_router(rag_router.router)`
curl JSON 引号错误	请求直接 422 或 shell 报错	用单引号包住 JSON，字段用双引号

✅ 四条理解标准#

思想是什么：先把向量检索脚本变成能被外部调用的 API。
干什么：提供文档存入和语义检索两个最小接口。
为什么这么干：API 化后前端、测试脚本、其他服务才能复用 RAG 能力。
怎么干：能启动服务，用 curl 调 /rag/documents 和 /rag/search，并知道数据会持久化到哪里。

下一步#

手搓最小 RAG 闭环：

1
用户提问 → Embedding → ChromaDB 检索 → SQLite 回查 → 拼成 Prompt → LLM 生成回答

把”检索到的碎片”喂给 AI，让 AI 基于文档内容回答，而不是瞎编。