aps-agent/server/knowledge/retrieval.py

# ============================================================
# 轻量混合检索（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