chromadb向量数据库使用 (1)
目录
- 完整代码
- 代码解释
完整代码
import chromadb
chroma_client = chromadb.Client()
collection = chroma_client.create_collection(name="my_collection")
collection.add(
documents=[
"This is a document about pineapple",
"This is a document about oranges"
],
ids=["id1", "id2"]
)
results = collection.query(
query_texts=["This is a query document about hawaii"],
n_results=2
)
print(results)
{'ids': [['id1', 'id2']], 'embeddings': None, 'documents': [['This is a document about pineapple', 'This is a document about oranges']], 'uris': None, 'data': None, 'metadatas': [[None, None]], 'distances': [[1.0404009819030762, 1.2430799007415771]], 'included': [<IncludeEnum.distances: 'distances'>, <IncludeEnum.documents: 'documents'>, <IncludeEnum.metadatas: 'metadatas'>]}
import chromadb
chroma_client = chromadb.Client()
collection = chroma_client.get_or_create_collection(name="my_collection")
collection.upsert(
documents=[
"This is a document about pineapple",
"This is a document about oranges"
],
ids=["id1", "id2"]
)
results = collection.query(
query_texts=["..."],
n_results=2
)
print(results)
{'ids': [['id2', 'id1']], 'embeddings': None, 'documents': [['This is a document about oranges', 'This is a document about pineapple']], 'uris': None, 'data': None, 'metadatas': [[None, None]], 'distances': [[1.8110723495483398, 1.846815824508667]], 'included': [<IncludeEnum.distances: 'distances'>, <IncludeEnum.documents: 'documents'>, <IncludeEnum.metadatas: 'metadatas'>]}
代码解释
以下是逐行代码的中文解释:
# 导入ChromaDB客户端库
import chromadb
# 创建ChromaDB客户端实例(默认使用内存存储)
chroma_client = chromadb.Client()
# 创建名为"my_collection"的集合(类似数据库表)
collection = chroma_client.create_collection(name="my_collection")
# 向集合添加文档数据
collection.add(
documents=[
"This is a document about pineapple", # 文档1:关于菠萝
"This is a document about oranges" # 文档2:关于橙子
],
ids=["id1", "id2"] # 为每个文档指定唯一ID
)
# 执行相似性查询
results = collection.query(
query_texts=["This is a query document about hawaii"], # 查询文本(夏威夷相关)
n_results=2 # 返回最相似的2个结果
)
# 打印查询结果(包含相似文档及其距离分数)
print(results)
# 获取或创建集合(如果已存在则直接获取)
collection = chroma_client.get_or_create_collection(name="my_collection")
# 使用upsert方法添加/更新文档(存在则更新,不存在则插入)
collection.upsert(
documents=[
"This is a document about pineapple", # 文档内容与之前相同
"This is a document about oranges"
],
ids=["id1", "id2"] # 使用相同ID
)
# 执行空查询(使用"..."作为占位符)
results = collection.query(
query_texts=["..."], # 无效查询文本示例
n_results=2
)
# 打印不同查询条件的结果对比
print(results)
关键点解析:
- 存储方式:默认使用内存存储,重启后数据会丢失
- 集合操作:
create_collection()
严格创建新集合get_or_create_collection()
更安全的获取方式
- 文档操作:
add()
单纯添加新文档upsert()
支持更新已有文档(基于ID)
- 查询结果:
- distances越小表示相似度越高
- 无效查询可能返回随机/全部结果
- 结果排序基于相似度得分
典型使用场景:构建简单的文本相似性搜索系统,适用于知识库检索、FAQ问答等场景。建议后续添加文本向量化模型(如Sentence-BERT)来提升搜索质量。
参考链接:https://docs.trychroma.com/docs/overview/getting-started