通过前面的学习，我们已经学会了如何从 Chroma 查询数据、分页查询、以及进行相识度查询。下面将结合前面学习到的知识，封装一个 Searcher 类，实现对指定集合进行分页查询和相识度查询，以及返回集合所有符合条件数据的工具方法。
注意，嵌入方法还是复用前面自定义的 AliQwenEmbeddingFunction。
代码如下：
import chromadb
from dataclasses import dataclass
from typing import Optional, Sequence
from common.custom_embedding import AliQwenEmbeddingFunction
# 使用 @dataclass 定义数据类
# 用来封装向量库检索单条结果（SearchHit，检索命中项）
@dataclass
class SearchHit:
id: str
document: Optional[str]
metadata: Optional[dict]
distance: Optional[float]
similarity: Optional[float] = None
# 自己封装的检索类
class Searcher:
def __init__(self, cli, collection_name: str, embedding_function=None):
"""构造方法
:param cli Chroma数据库客户端实例
:param collection_name 集合名称
:param embedding_function 嵌入函数
"""
self.col = cli.get_collection(
collection_name, embedding_function=embedding_function
)
# 获取集合的 hnsw 配置，用来判断如何进行距离转换为相似度
self._space = (self.col.configuration.get("hnsw") or {}).get("space", "l2")
def _to_similarity(self, distance: float) -> Optional[float]:
"""将距离转换为相似度"""
if distance is None:
return None
if self._space == "cosine":
# 余弦距离
return 1.0 - distance
elif self._space == "l2":
# 平方欧式距离
return 1.0 / (1.0 + distance)
return None
def search(
self,
query: str | Sequence[float],
top_k: int = 10,
where: Optional[dict] = None,
where_document: Optional[dict] = None,
ids: Optional[list[str]] = None,
) -> list[SearchHit]:
"""检索方法，先过滤，在进行相似度检索
:param query 语义检索，可以是文本，也可以是向量数组
:param top_k 返回前N个记录
:param where 过滤条件，根据 metadatas 过滤
:param where_document 过滤条件，根据文档内容过滤
:param ids 过滤条件，根据ID进行过滤
"""
if isinstance(query, str):
# 如果是字符串，则由 embedding_function 帮我们计算向量
kwargs = {"query_texts": [query]}
else:
# 向量检索，不需要计算向量
kwargs = {"query_embeddings": [list(query)]}
res = self.col.query(
**kwargs,
n_results=top_k,
where=where,
where_document=where_document,
ids=ids,
)
# 使用 zip 将列返回进行打包成行数据，每行包含对应的 ids、documents、metadatas、distances
hits = []
for id_, doc, meta, dist in zip(
res["ids"][0],
res["documents"][0],
res["metadatas"][0],
res["distances"][0],
):
hits.append(SearchHit(
id=id_,
document=doc,
metadata=meta,
distance=dist,
similarity=self._to_similarity(dist),
))
return hits
def page(
self,
where: Optional[dict] = None,
limit: int = 50,
offset: int = 0,
) -> list[SearchHit]:
"""分页查询"""
res = self.col.get(
where=where,
limit=limit,
offset=offset,
include=["documents", "metadatas"],
)
return [
SearchHit(id=i, document=d, metadata=m, distance=None)
for i, d, m in zip(res["ids"], res["documents"], res["metadatas"])
]
def iter_all(self, batch_size: int = 500, where=None):
"""分页扫全量，写导出/迁移脚本时很好用"""
offset = 0
while True:
batch = self.page(where=where, limit=batch_size, offset=offset)
if not batch:
break
yield from batch
if len(batch) < batch_size:
break
offset += batch_size
if __name__ == "__main__":
# 自定义嵌入函数实现
ef = AliQwenEmbeddingFunction()
client = chromadb.EphemeralClient()
col = client.create_collection(
"faq", embedding_function=ef, configuration={"hnsw": {"space": "cosine"}}
)
# 灌注数据
col.add(
ids=["r1", "r2", "r3", "r4", "r5"],
documents=[
"退款流程：在订单页点申请退款，3 个工作日内到账",
"发票开具：在订单详情页申请电子发票",
"物流查询：发货后可在物流详情里看到进度",
"退货政策：7 天无理由退货，运费买家承担",
"账号安全：如何修改密码和绑定手机",
],
metadatas=[
{"cat": "售后", "year": 2024, "score": 0.9, "tags": ["退款"]},
{"cat": "财务", "year": 2023, "score": 0.5, "tags": ["发票"]},
{"cat": "物流", "year": 2024, "score": 0.7, "tags": ["物流"]},
{"cat": "售后", "year": 2022, "score": 0.3, "tags": ["退货"]},
{"cat": "账号", "year": 2024, "score": 0.6, "tags": ["安全"]},
],
)
# 相识度查询验证
print("--- 相识度查询 ---")
s = Searcher(client, collection_name="faq", embedding_function=ef)
for hit in s.search("不要了", top_k=2, where={"cat": "售后"}):
print(f" id={hit.id}, 相似度={hit.similarity}, 距离={hit.distance}, document={hit.document}")
# 演示 page 和 iter_all 函数的用法
# page：按页取，用 limit / offset 控制，适合做前端列表翻页
print("\n--- page：第 1 页（limit=2, offset=0）---")
for hit in s.page(limit=2, offset=0):
print(f" id={hit.id}, cat={hit.metadata.get('cat')}, document={hit.document}")
print("\n--- page：第 2 页（limit=2, offset=2）---")
for hit in s.page(limit=2, offset=2):
print(f" id={hit.id}, cat={hit.metadata.get('cat')}, document={hit.document}")
# page 也支持 where 过滤：只翻“售后”这一类
print("\n--- page：只翻 cat=售后（limit=10）---")
for hit in s.page(where={"cat": "售后"}, limit=10):
print(f" id={hit.id}, cat={hit.metadata.get('cat')}, document={hit.document}")
# iter_all：不用管页码，一次把所有命中项吐出来，写导出 / 迁移脚本很方便
print("\n--- iter_all：遍历全部（batch_size=2，方便看分批拉取效果）---")
all_hits = list(s.iter_all(batch_size=2))
print(f" 共 {len(all_hits)} 条：")
for hit in all_hits:
print(f" id={hit.id}, cat={hit.metadata.get('cat')}, document={hit.document}")运行示例，输出如下：
--- 相识度查询 ---
id=r4, 相似度=0.3523508906364441, 距离=0.6476491093635559, document=退货政策：7 天无理由退货，运费买家承担
id=r1, 相似度=0.28433358669281006, 距离=0.7156664133071899, document=退款流程：在订单页点申请退款，3 个工作日内到账
--- page：第 1 页（limit=2, offset=0）---
id=r1, cat=售后, document=退款流程：在订单页点申请退款，3 个工作日内到账
id=r2, cat=财务, document=发票开具：在订单详情页申请电子发票
--- page：第 2 页（limit=2, offset=2）---
id=r3, cat=物流, document=物流查询：发货后可在物流详情里看到进度
id=r4, cat=售后, document=退货政策：7 天无理由退货，运费买家承担
--- page：只翻 cat=售后（limit=10）---
id=r1, cat=售后, document=退款流程：在订单页点申请退款，3 个工作日内到账
id=r4, cat=售后, document=退货政策：7 天无理由退货，运费买家承担
--- iter_all：遍历全部（batch_size=2，方便看分批拉取效果）---
共 5 条：
id=r1, cat=售后, document=退款流程：在订单页点申请退款，3 个工作日内到账
id=r2, cat=财务, document=发票开具：在订单详情页申请电子发票
id=r3, cat=物流, document=物流查询：发货后可在物流详情里看到进度
id=r4, cat=售后, document=退货政策：7 天无理由退货，运费买家承担
id=r5, cat=账号, document=账号安全：如何修改密码和绑定手机注意，在其他文件中，可以通过 from...import 导入上面的类，然后使用，避免每次都编写重复的模板代码。
川公网安备51010802032098