137 lines
6.3 KiB
Python
137 lines
6.3 KiB
Python
# requires: elasticsearch
|
||
"""Elasticsearch：中文分词、BM25 关键词检索、kNN 向量检索、RRF 混合检索、别名零停机重建索引。
|
||
|
||
运行：ES_URL=http://127.0.0.1:9200 python es_hybrid.py
|
||
本地测试用的单节点 ES 关闭了安全认证；生产环境必须开启认证和 TLS。没有设置环境变量时跳过。
|
||
向量用一个“字符二元组哈希”的玩具嵌入代替真实模型，只为演示流程；索引名带随机后缀，结束后删除。
|
||
"""
|
||
import hashlib
|
||
import math
|
||
import os
|
||
import sys
|
||
import uuid
|
||
|
||
from elasticsearch import Elasticsearch, helpers
|
||
|
||
url = os.environ.get("ES_URL")
|
||
if not url:
|
||
print("未设置 ES_URL，跳过")
|
||
sys.exit(0)
|
||
|
||
es = Elasticsearch(url, request_timeout=30)
|
||
suffix = uuid.uuid4().hex[:6]
|
||
alias = f"kb-{suffix}"
|
||
DIM = 64
|
||
|
||
DOCS = [
|
||
{"id": "d1", "tenant": "t1", "title": "退货政策", "body": "签收后七天内可以申请无理由退货，商品需保持完好。"},
|
||
{"id": "d2", "tenant": "t1", "title": "退款到账时间", "body": "退款审核通过后，原路退回，一般一到三个工作日到账。"},
|
||
{"id": "d3", "tenant": "t1", "title": "发票开具", "body": "订单完成后可以在订单详情页申请电子发票。"},
|
||
{"id": "d4", "tenant": "t1", "title": "会员积分", "body": "每消费一元累计一个积分，积分可以抵扣运费。"},
|
||
{"id": "d5", "tenant": "t2", "title": "退货政策（其他租户）", "body": "本店商品不支持无理由退货。"},
|
||
{"id": "d6", "tenant": "t1", "title": "错误码 E1024", "body": "E1024 表示支付渠道超时，请稍后重试或更换支付方式。"},
|
||
]
|
||
|
||
|
||
def embed(text: str) -> list[float]:
|
||
"""玩具嵌入：字符二元组哈希到 64 维再归一化。真实项目换成嵌入模型。"""
|
||
v = [0.0] * DIM
|
||
for a, b in zip(text, text[1:]):
|
||
h = int(hashlib.md5((a + b).encode()).hexdigest(), 16)
|
||
v[h % DIM] += 1.0
|
||
n = math.sqrt(sum(x * x for x in v)) or 1.0
|
||
return [x / n for x in v]
|
||
|
||
|
||
def create_index(name: str) -> None:
|
||
es.indices.create(index=name, mappings={
|
||
"properties": {
|
||
"tenant": {"type": "keyword"}, # 过滤用，不分词
|
||
"title": {"type": "text", "analyzer": "cjk"}, # 内置 cjk 分析器（中文二元切分）；生产常用 IK 插件，以插件文档为准
|
||
"body": {"type": "text", "analyzer": "cjk"},
|
||
"vec": {"type": "dense_vector", "dims": DIM, "index": True, "similarity": "cosine"},
|
||
}
|
||
})
|
||
|
||
|
||
def load(index: str) -> None:
|
||
# 用业务 id 作为 _id：重复同步是覆盖而不是新增（幂等）
|
||
actions = ({"_index": index, "_id": d["id"], "_source": {**{k: v for k, v in d.items() if k != "id"}, "vec": embed(d["title"] + d["body"])}} for d in DOCS)
|
||
ok, errors = helpers.bulk(es, actions, refresh="wait_for")
|
||
assert ok == len(DOCS) and not errors
|
||
|
||
|
||
def bm25(q: str, tenant: str, k: int = 3) -> list[str]:
|
||
r = es.search(index=alias, size=k, query={"bool": {
|
||
"must": {"multi_match": {"query": q, "fields": ["title^2", "body"]}}, # 标题权重更高
|
||
"filter": {"term": {"tenant": tenant}}, # 租户过滤放 filter：不算分、可缓存
|
||
}})
|
||
return [h["_id"] for h in r["hits"]["hits"]]
|
||
|
||
|
||
def knn(q: str, tenant: str, k: int = 3) -> list[str]:
|
||
r = es.search(index=alias, size=k, knn={
|
||
"field": "vec", "query_vector": embed(q), "k": k, "num_candidates": 50,
|
||
"filter": {"term": {"tenant": tenant}}, # 过滤在向量检索过程中生效，而不是检索后再过滤
|
||
})
|
||
return [h["_id"] for h in r["hits"]["hits"]]
|
||
|
||
|
||
def rrf(*rankings: list[str], k: int = 60) -> list[str]:
|
||
"""倒数排名融合：只看名次，不看分数（BM25 分数和向量相似度不可直接比较）。"""
|
||
score: dict[str, float] = {}
|
||
for ranking in rankings:
|
||
for rank, doc_id in enumerate(ranking, start=1):
|
||
score[doc_id] = score.get(doc_id, 0.0) + 1.0 / (k + rank)
|
||
return sorted(score, key=lambda d: (-score[d], d))
|
||
|
||
|
||
try:
|
||
v1 = f"{alias}-v1"
|
||
create_index(v1)
|
||
load(v1)
|
||
es.indices.put_alias(index=v1, name=alias) # 业务只访问别名
|
||
|
||
exact = bm25("E1024", "t1")
|
||
assert exact[0] == "d6", exact
|
||
print(f"1. 关键词检索擅长精确词（错误码）：E1024 → {exact}")
|
||
|
||
kw = bm25("退货", "t1")
|
||
assert "d5" not in kw and kw[0] == "d1", kw
|
||
print(f"2. 租户过滤生效：t1 搜“退货” → {kw}（没有 t2 的 d5）")
|
||
|
||
sem = knn("商品不想要了怎么退", "t1")
|
||
print(f"3. 向量检索（玩具嵌入）：{sem}")
|
||
assert "d5" not in sem
|
||
|
||
fused = rrf(bm25("退货 退款多久到账", "t1", 5), knn("退货 退款多久到账", "t1", 5))
|
||
assert set(fused[:2]) == {"d1", "d2"}, fused
|
||
print(f"4. RRF 混合检索：{fused[:3]}")
|
||
|
||
# 服务端 RRF（retriever 语法）。许可证和可用版本以官方文档为准，不可用时退回客户端融合
|
||
try:
|
||
r = es.search(index=alias, size=3, retriever={"rrf": {"retrievers": [
|
||
{"standard": {"query": {"bool": {"must": {"match": {"body": "退款"}}, "filter": {"term": {"tenant": "t1"}}}}}},
|
||
{"knn": {"field": "vec", "query_vector": embed("退款到账"), "k": 3, "num_candidates": 50, "filter": {"term": {"tenant": "t1"}}}},
|
||
]}})
|
||
print(f"5. 服务端 RRF：{[h['_id'] for h in r['hits']['hits']]}")
|
||
except Exception as e: # noqa: BLE001 —— 例如 basic 许可下返回 403 授权错误
|
||
print(f"5. 服务端 RRF 不可用（{type(e).__name__}: {str(e)[:80]}），使用客户端 RRF")
|
||
|
||
# 别名切换：新建 v2（例如修改了映射），灌数据，原子切换别名，业务无感知
|
||
v2 = f"{alias}-v2"
|
||
create_index(v2)
|
||
load(v2)
|
||
es.indices.update_aliases(actions=[
|
||
{"remove": {"index": v1, "alias": alias}},
|
||
{"add": {"index": v2, "alias": alias}},
|
||
])
|
||
assert list(es.indices.get_alias(name=alias).keys()) == [v2]
|
||
assert bm25("E1024", "t1")[0] == "d6"
|
||
print("6. 别名原子切换到 v2，查询不中断")
|
||
print("全部检查通过")
|
||
finally:
|
||
# 生产集群通常禁止通配符删除（action.destructive_requires_name），这里逐个写明索引名
|
||
es.indices.delete(index=[f"{alias}-v1", f"{alias}-v2"], ignore_unavailable=True)
|
||
es.close()
|