起因:我受够了重复造轮子
作为一个每天和YOLO模型、FastAPI打交道的AI架构师,我发现自己陷入了一个诡异的循环:
- 明明三个月前写过一段完美的数据预处理代码,死活想不起来放哪了
- 每次问AI"怎么导出ONNX模型",它给的答案总是差那么点意思,因为它不知道我的项目结构
- 积累了一堆设计文档,结果用的时候还是靠
Ctrl+F大法
于是我决定搞一个本地知识库,把自己的代码片段、设计文档、踩坑记录都扔进去,让AI能够基于我的知识体系来回答问题。
选型:一场短暂的纠结
在正式开始折腾ChromaDB之前,我确实调研了一圈主流方案。简单说说为什么最后选了它:
Pinecone - 企业级向量数据库,性能杠杠的,十亿级向量毫秒响应。但是...我只是想存点自己的代码啊,又不是要给淘宝做推荐系统。更关键的是,数据要上传到云端,我那些写得乱七八糟的代码要是被人看到,社死程度堪比公开我的浏览器历史记录。Pass。
NotebookLM - Google家的产品,上传文档就能问答,还能生成播客风格的摘要,听起来很酷对吧?但它没有API,只能在网页上点点点。我一个写代码的,你让我每次都手动上传文件?而且它对代码的理解约等于我对时尚的理解——基本为零。Pass。
Weaviate / Qdrant - 都是正经的向量数据库,功能强大,部署灵活。但需要跑独立服务,配置文件比我的头发还多。我只是想要一个本地知识库,不是想运维一套分布式系统。先Pass,以后规模大了再说。
FAISS - Facebook出品,性能怪兽,学术圈的最爱。但它只是一个索引库,持久化、元数据管理、API封装全都要自己搞。我只是想用,不是想从零开始写一个向量数据库。敬而远之。
LanceDB - 新秀选手,Rust写的,性能比ChromaDB好。但生态还不够成熟,文档有点少,遇到问题只能靠自己摸索。等它再长大一点我再考虑。观望中。
最后的选择:ChromaDB。理由很简单:
pip install chromadb没错,就这一行,完事。数据存本地,开箱即用,和LangChain配合得天衣无缝。ChromaDB已经跨越了长达数年的alpha阶段,当前可用版本为1.4.x,对于个人知识库这个量级,它完全够用了。
ChromaDB:从入门到真香
第一步:Hello World(入门)
先来个最简单的例子,感受一下它有多丝滑:
import chromadb
# 创建一个持久化的客户端,数据存在本地
client = chromadb.PersistentClient(path="./my_brain")
# 创建一个集合,就像数据库里的表
collection = client.get_or_create_collection(name="code_snippets")
# 往里面塞点东西
collection.add(
documents=[
"FastAPI路由示例:@app.get('/users/{user_id}')",
"YOLO模型加载:model = YOLO('yolov8n.pt')",
"PyTorch保存模型:torch.save(model.state_dict(), 'model.pth')"
],
ids=["doc1", "doc2", "doc3"]
)
# 语义搜索!
results = collection.query(
query_texts=["怎么加载预训练模型"],
n_results=2
)
print(results['documents'])
# 输出会包含YOLO和PyTorch那两条,因为语义相近就这么简单。没有配置文件,没有Docker,没有注册账号,没有填信用卡。它甚至自带了一个embedding模型(all-MiniLM-L6-v2),你什么都不用管就能跑起来。
第二步:正经地组织一下(进阶)
Hello World跑通之后,我们来搞一个正经的项目结构。毕竟以后东西多了,总不能一股脑全扔一个collection里吧。
my_knowledge_base/
├── kb/ # 核心代码
│ ├── __init__.py
│ ├── config.py # 配置
│ ├── processor.py # 文档处理(分块、提取元数据)
│ ├── store.py # ChromaDB封装
│ └── retriever.py # 检索接口
├── data/
│ ├── raw/ # 原始文档,按类型分文件夹
│ │ ├── designs/ # 设计文档
│ │ ├── snippets/ # 代码片段
│ │ └── notes/ # 学习笔记
│ └── chroma_db/ # ChromaDB数据目录
├── ingest.py # 导入脚本
├── query.py # 查询脚本
└── requirements.txtconfig.py - 配置集中管理
from pathlib import Path
class Config:
BASE_DIR = Path(__file__).parent.parent
DATA_DIR = BASE_DIR / "data"
RAW_DIR = DATA_DIR / "raw"
CHROMA_DIR = DATA_DIR / "chroma_db"
# 分块参数,这两个数字我调了好几次才满意
CHUNK_SIZE = 500 # 每块大约500字符
CHUNK_OVERLAP = 50 # 块之间重叠50字符,防止语义被切断
# 集合定义
COLLECTIONS = {
"designs": "系统设计文档",
"snippets": "代码片段",
"notes": "学习笔记和踩坑记录"
}
config = Config()processor.py - 文档处理器
这是最容易踩坑的地方。直接把整个文件扔进去?太长了,检索效果差。切太碎?上下文丢失,AI看了也懵。
import re
from pathlib import Path
from dataclasses import dataclass
from typing import List, Dict, Any
@dataclass
class Document:
content: str
metadata: Dict[str, Any]
doc_id: str
class DocumentProcessor:
"""
文档处理器:把大文件切成小块
踩坑记录:
1. 代码文件别按字数切,要按函数/类切,不然一个函数被切两半,AI看了想打人
2. Markdown文件可以按段落切,但要保留标题作为上下文
3. 元数据很重要!后面过滤全靠它
"""
def __init__(self, chunk_size: int = 500, chunk_overlap: int = 50):
self.chunk_size = chunk_size
self.chunk_overlap = chunk_overlap
def process_file(self, file_path: Path) -> List[Document]:
content = file_path.read_text(encoding='utf-8')
suffix = file_path.suffix.lower()
# 元数据:后面检索过滤全靠这些
metadata = {
"source": str(file_path),
"filename": file_path.name,
"file_type": suffix,
"category": file_path.parent.name,
}
# 根据文件类型选择分块策略
if suffix in ['.py', '.js', '.ts']:
chunks = self._chunk_code(content)
else:
chunks = self._chunk_text(content)
return [
Document(
content=chunk,
metadata={**metadata, "chunk_index": i},
doc_id=f"{file_path.stem}_{i}"
)
for i, chunk in enumerate(chunks)
]
def _chunk_text(self, text: str) -> List[str]:
"""普通文本:按段落切,尽量保持语义完整"""
paragraphs = text.split('\n\n')
chunks = []
current = ""
for para in paragraphs:
if len(current) + len(para) < self.chunk_size:
current += para + "\n\n"
else:
if current.strip():
chunks.append(current.strip())
current = para + "\n\n"
if current.strip():
chunks.append(current.strip())
return chunks if chunks else [text]
def _chunk_code(self, code: str) -> List[str]:
"""代码文件:按函数/类切割"""
# 匹配Python的函数和类定义
pattern = r'((?:def |class |async def )\w+[^\n]*\n(?:[ \t]+[^\n]*\n)*)'
matches = re.findall(pattern, code)
if matches:
return [m.strip() for m in matches if len(m.strip()) > 20]
# 如果没匹配到(可能是脚本文件),回退到普通分块
return self._chunk_text(code)store.py - ChromaDB封装
import chromadb
from chromadb.config import Settings
from pathlib import Path
from typing import List, Dict, Any, Optional
class KnowledgeStore:
"""
ChromaDB的封装层
为什么要封装一层?
1. 统一接口,以后换数据库方便(虽然大概率不会换)
2. 加一些便利方法,比如跨集合搜索
3. 处理一些边界情况,比如空结果
"""
def __init__(self, persist_dir: Path):
self.client = chromadb.PersistentClient(
path=str(persist_dir),
settings=Settings(anonymized_telemetry=False) # 关掉遥测,我不想被统计
)
self._collections = {}
def get_collection(self, name: str, description: str = ""):
if name not in self._collections:
self._collections[name] = self.client.get_or_create_collection(
name=name,
metadata={"description": description}
)
return self._collections[name]
def add_documents(
self,
collection_name: str,
documents: List[str],
metadatas: List[Dict],
ids: List[str]
):
collection = self.get_collection(collection_name)
# ChromaDB会自动生成embedding,省心
collection.add(
documents=documents,
metadatas=metadatas,
ids=ids
)
print(f"✅ [{collection_name}] 新增 {len(documents)} 条文档")
def search(
self,
collection_name: str,
query: str,
n_results: int = 5,
where: Optional[Dict] = None
) -> List[Dict]:
"""单集合搜索"""
collection = self.get_collection(collection_name)
results = collection.query(
query_texts=[query],
n_results=n_results,
where=where
)
# 把结果整理成更友好的格式
return [
{
"content": results['documents'][0][i],
"metadata": results['metadatas'][0][i],
"distance": results['distances'][0][i]
}
for i in range(len(results['documents'][0]))
]
def search_all(self, query: str, n_results: int = 5) -> List[Dict]:
"""跨所有集合搜索,然后按相似度排序"""
all_results = []
for coll in self.client.list_collections():
results = self.search(coll.name, query, n_results)
for r in results:
r['collection'] = coll.name
all_results.append(r)
# 按距离排序(越小越相似)
all_results.sort(key=lambda x: x['distance'])
return all_results[:n_results]
def stats(self) -> Dict[str, int]:
"""看看各个集合有多少文档"""
return {
coll.name: coll.count()
for coll in self.client.list_collections()
}retriever.py - 上下文组装器
这是最终对接AI的地方,把检索结果格式化成AI能理解的上下文。
from typing import List, Optional
from .store import KnowledgeStore
class ContextRetriever:
"""
把检索结果组装成AI友好的上下文
这里有个小技巧:不要把所有结果都塞给AI
1. 太长了AI也消化不了
2. 不相关的内容会干扰AI的判断
3. Token是要钱的(如果你用的是付费API)
"""
def __init__(self, store: KnowledgeStore):
self.store = store
def get_context(
self,
query: str,
max_chars: int = 8000,
collections: Optional[List[str]] = None
) -> str:
"""获取格式化的上下文"""
if collections:
results = []
for coll in collections:
results.extend(self.store.search(coll, query, n_results=5))
results.sort(key=lambda x: x['distance'])
else:
results = self.store.search_all(query, n_results=10)
# 组装上下文,注意控制长度
context_parts = []
total_len = 0
for r in results:
part = self._format_result(r)
if total_len + len(part) > max_chars:
break
context_parts.append(part)
total_len += len(part)
return "\n\n---\n\n".join(context_parts)
def _format_result(self, result: dict) -> str:
meta = result['metadata']
return f"**[{meta.get('category', '')}] {meta.get('filename', '')}**\n\n{result['content']}"
def export_for_ai(self, query: str) -> dict:
"""导出为可以直接喂给AI的格式"""
context = self.get_context(query)
return {
"system_prompt": "你是一个有帮助的AI助手。请基于以下知识库内容回答问题,如果知识库中没有相关信息,请诚实说明。",
"context": f"<knowledge_base>\n{context}\n</knowledge_base>",
"user_query": query
}第三步:跑起来!
ingest.py - 导入文档
#!/usr/bin/env python3
"""
文档导入脚本
用法:python ingest.py
"""
from pathlib import Path
from kb.config import config
from kb.processor import DocumentProcessor
from kb.store import KnowledgeStore
def main():
processor = DocumentProcessor(
chunk_size=config.CHUNK_SIZE,
chunk_overlap=config.CHUNK_OVERLAP
)
store = KnowledgeStore(config.CHROMA_DIR)
print("🚀 开始导入文档...\n")
for category_dir in config.RAW_DIR.iterdir():
if not category_dir.is_dir():
continue
collection_name = category_dir.name
print(f"📁 处理 [{collection_name}]")
all_docs = []
for file_path in category_dir.rglob('*'):
if file_path.suffix in ['.md', '.txt', '.py', '.js', '.ts']:
docs = processor.process_file(file_path)
all_docs.extend(docs)
print(f" - {file_path.name} -> {len(docs)} 块")
if all_docs:
store.add_documents(
collection_name=collection_name,
documents=[d.content for d in all_docs],
metadatas=[d.metadata for d in all_docs],
ids=[d.doc_id for d in all_docs]
)
print("\n📊 导入完成!统计:")
for name, count in store.stats().items():
print(f" {name}: {count} 条")
if __name__ == "__main__":
main()query.py - 查询测试
#!/usr/bin/env python3
"""
查询测试脚本
用法:python query.py "你的问题"
"""
import sys
from kb.config import config
from kb.store import KnowledgeStore
from kb.retriever import ContextRetriever
def main():
if len(sys.argv) < 2:
print("用法:python query.py \"你的问题\"")
sys.exit(1)
query = sys.argv[1]
store = KnowledgeStore(config.CHROMA_DIR)
retriever = ContextRetriever(store)
print(f"🔍 查询:{query}\n")
print("=" * 50)
# 获取上下文
context = retriever.get_context(query)
print(context)
print("\n" + "=" * 50)
print("\n💡 可以把上面的内容作为上下文喂给AI了!")
if __name__ == "__main__":
main()第四步:真香时刻
跑了一段时间之后,我发现这东西是真的好用:
# 导入我的所有文档
$ python ingest.py
🚀 开始导入文档...
📁 处理 [designs]
- payment_system.md -> 12 块
- yolo_pipeline.md -> 8 块
📁 处理 [snippets]
- fastapi_utils.py -> 15 块
- yolo_training.py -> 23 块
📁 处理 [notes]
- autodl_gpu_setup.md -> 5 块
📊 导入完成!统计:
designs: 20 条
snippets: 38 条
notes: 5 条
# 查询一下
$ python query.py "怎么在AutoDL上训练YOLO"
🔍 查询:怎么在AutoDL上训练YOLO
==================================================
**[notes] autodl_gpu_setup.md**
AutoDL RTX 4090 环境配置踩坑记录:
1. 先更新pip:pip install --upgrade pip
2. 安装ultralytics:pip install ultralytics
3. 数据集放到/root/autodl-tmp/,这个目录不限速
...
---
**[snippets] yolo_training.py**
def train_yolo(data_yaml: str, epochs: int = 100):
"""YOLO训练函数,已在4090上验证过"""
from ultralytics import YOLO
model = YOLO('yolov8n.pt')
model.train(data=data_yaml, epochs=epochs, imgsz=640)
...
==================================================
💡 可以把上面的内容作为上下文喂给AI了!再也不用翻历史记录了,问啥都能找到之前的笔记和代码。
一些踩坑经验
- 分块大小很重要:太大了检索不精准,太小了丢失上下文。我试下来500字符比较合适。
- 代码文件别按字数切:一定要按函数/类切,不然一个函数被切成两半,AI看了也没法用。
- 元数据要设计好:文件名、类型、分类这些一定要记录,后面过滤全靠它们。
- 定期清理:过时的文档记得删掉,不然AI会给你返回三年前的错误代码。
- Embedding模型可以升级:ChromaDB默认的模型够用,但如果你的文档主要是中文,可以换成
BAAI/bge-small-zh-v1.5,效果更好。
后续可以折腾的方向
- 加个Web界面:用FastAPI + React搞一个,查询更方便
- 接入IDE:写个VSCode插件,直接在编辑器里查知识库
- 自动同步:用watchdog监听文件变化,自动更新知识库
- 更好的Embedding:接入OpenAI的text-embedding-3-small,效果更好(但要花钱)
最后
折腾了一圈下来,ChromaDB确实是个人知识库的最佳选择。简单、免费、够用。虽然它还是alpha版本,但对于我这个量级的使用场景,完全没问题。
等哪天我的知识库膨胀到百万级别,再考虑升级到Qdrant或者Pinecone吧。
不过那应该是很久以后的事了——毕竟我首先得写出那么多代码来。
写于 2026年1月,MacBook M4 Air 旁边的咖啡已经凉了