116 lines
4.2 KiB
Python
116 lines
4.2 KiB
Python
# ============================================================
|
||
# 轻量混合检索（moduleId: knowledge-retrieval, 可重生 ✅）
|
||
# plan.md §8.2：向量（可选）+ 字符 bigram + 标签/标题加权；强制带出处
|
||
# ============================================================
|
||
from __future__ import annotations
|
||
|
||
import math
|
||
import re
|
||
from typing import Any
|
||
|
||
_CLEAN = re.compile(r"[^\w\u4e00-\u9fff]+")
|
||
|
||
|
||
def _bigrams(text: str) -> set[str]:
|
||
t = _CLEAN.sub("", text.lower())
|
||
if len(t) < 2:
|
||
return set(t)
|
||
return {t[i:i + 2] for i in range(len(t) - 1)}
|
||
|
||
|
||
def _bigram_sim(q: set[str], d: set[str]) -> float:
|
||
if not q or not d:
|
||
return 0.0
|
||
return 2 * len(q & d) / (len(q) + len(d))
|
||
|
||
|
||
def search(assets: list[dict[str, Any]], query: str, top_k: int = 3,
|
||
min_score: float = 0.05) -> list[dict[str, Any]]:
|
||
"""兼容旧接口：对整篇资产 bigram 检索。"""
|
||
units = []
|
||
for a in assets:
|
||
if not a.get("approved", True):
|
||
continue
|
||
units.append({
|
||
"assetId": a["assetId"], "title": a["title"], "kind": a["kind"],
|
||
"version": a["version"], "tags": a.get("tags") or [],
|
||
"content": a.get("content") or "",
|
||
"chunkId": None, "heading": None, "page": None,
|
||
})
|
||
return hybrid_search(units, query, top_k=top_k, min_score=min_score)
|
||
|
||
|
||
def hybrid_search(units: list[dict[str, Any]], query: str, top_k: int = 3,
|
||
min_score: float = 0.05,
|
||
alpha_vector: float = 0.55, beta_bigram: float = 0.45) -> list[dict[str, Any]]:
|
||
"""混合检索：向量余弦（若可用）+ bigram + 标签/标题加成。"""
|
||
if not query.strip() or not units:
|
||
return []
|
||
|
||
q_bi = _bigrams(query)
|
||
docs_bi = [_bigrams((u.get("title") or "") + (u.get("content") or "")
|
||
+ " ".join(u.get("tags") or [])) for u in units]
|
||
n_docs = max(len(units), 1)
|
||
df: dict[str, int] = {}
|
||
for bi in docs_bi:
|
||
for g in bi:
|
||
df[g] = df.get(g, 0) + 1
|
||
|
||
# 向量分
|
||
vec_scores: dict[str, float] = {}
|
||
try:
|
||
from server.knowledge.embedding import get_embedding_provider, get_embedding_store
|
||
prov = get_embedding_provider()
|
||
if prov.enabled:
|
||
qv = prov.embed_one(query)
|
||
if qv:
|
||
store = get_embedding_store()
|
||
for uid, sc in store.search(qv, top_k=max(top_k * 5, 20)):
|
||
vec_scores[uid] = sc
|
||
except Exception:
|
||
pass
|
||
|
||
query_has_cjk = any("\u4e00" <= c <= "\u9fff" for c in query)
|
||
scored: list[tuple[float, dict[str, Any]]] = []
|
||
for u, d_bi in zip(units, docs_bi):
|
||
shared = q_bi & d_bi
|
||
base = _bigram_sim(q_bi, d_bi) if shared else 0.0
|
||
if not query_has_cjk and base < 0.12 and not vec_scores:
|
||
continue
|
||
idf_bonus = (sum(math.log(1 + n_docs / df.get(g, 1)) for g in shared) / (len(q_bi) + 1)
|
||
if shared else 0.0)
|
||
tag_bonus = 0.3 if any(t in query for t in (u.get("tags") or [])) else 0.0
|
||
title = u.get("title") or ""
|
||
title_bonus = 0.4 if title and (title in query or query in title) else 0.0
|
||
bigram_score = base + 0.1 * idf_bonus + tag_bonus + title_bonus
|
||
|
||
key = f"{u['assetId']}:{u.get('chunkId') or u['assetId']}"
|
||
v_score = vec_scores.get(key, 0.0)
|
||
if v_score > 0 and bigram_score > 0:
|
||
score = alpha_vector * v_score + beta_bigram * bigram_score
|
||
elif v_score > 0:
|
||
score = v_score
|
||
else:
|
||
score = bigram_score
|
||
|
||
if score >= min_score:
|
||
scored.append((score, u))
|
||
|
||
scored.sort(key=lambda x: -x[0])
|
||
hits: list[dict[str, Any]] = []
|
||
for score, u in scored[:top_k]:
|
||
content = u.get("content") or ""
|
||
hits.append({
|
||
"assetId": u["assetId"],
|
||
"title": u.get("title"),
|
||
"kind": u.get("kind"),
|
||
"version": u.get("version"),
|
||
"score": round(score, 3),
|
||
"snippet": content[:120] + ("…" if len(content) > 120 else ""),
|
||
"content": content,
|
||
"chunkId": u.get("chunkId"),
|
||
"heading": u.get("heading"),
|
||
"page": u.get("page"),
|
||
})
|
||
return hits
|