下面使用前面学过的 API 来封装一个简单的知识库管理器。该管理器将持久化知识库信息到临时目录，提供获取集合、删除集合、重命名集合、集合是否存在等功能。
具体代码如下：
import os
import tempfile
import chromadb
from datetime import datetime
from chromadb.api.models.Collection import Collection
# 让持久化示例写到临时目录，工作目录保持干净
os.chdir(tempfile.mkdtemp(prefix="chroma_demo_"))
class KBManager:
"""Chroma 知识库的集合管理封装"""
def __init__(self, path: str = "./chroma_data"):
self.client = chromadb.PersistentClient(path=path)
def ensure_collection(
self,
name: str,
embedding_function=None,
space: str = "cosine",
description: str = "",
) -> Collection:
"""幂等地拿到一个集合：存在就用，不存在就建"""
return self.client.get_or_create_collection(
name=name,
embedding_function=embedding_function,
configuration={
"hnsw": {
"space": space
}
},
metadata={
"description": description,
"created_at": datetime.now().isoformat(),
},
)
def exists(self, name: str) -> bool:
"""判断集合是否存在"""
return any(c.name == name for c in self.client.list_collections())
def stats(self) -> list[dict]:
"""列出所有集合的概况"""
return [
{
"name": c.name,
"count": c.count(),
"space": (c.configuration.get("hnsw") or {}).get("space"),
"description": (c.metadata or {}).get("description", ""),
}
for c in self.client.list_collections()
]
def rename(self, old: str, new: str):
"""集合重命名"""
# 集合不存在
if not self.exists(old):
return
self.client.get_collection(old).modify(name=new)
def drop(self, name: str, force: bool = False):
"""删除集合"""
# 集合不存在
if not self.exists(name):
return
# 不强制删除，但集合中还存在数据，则抛出异常进行提示
tmp_col = self.client.get_collection(name)
if not force and tmp_col.count() > 0:
raise RuntimeError(f"{name} 还有 {tmp_col.count()} 条数据，force=True 才允许删")
self.client.delete_collection(name)
if __name__ == "__main__":
# 创建实例
kb = KBManager("./_demo_kb")
# 获取 product_docs 集合实例
col = kb.ensure_collection("product_docs", embedding_function=None, description="产品文档")
# 向集合写入数据
col.add(ids=["d1"], embeddings=[[1.0, 0.0]], documents=["快速上手"])
print(kb.stats())
# [{'name': 'product_docs', 'count': 1, 'space': 'cosine', 'description': '产品文档'}]
# 删除集合
kb.drop("product_docs", force=True)封装成工具类后，在删除、重命名等操作时不用重复写集合是否存在判断等逻辑，用起来更方便。
川公网安备51010802032098