285 lines
10 KiB
Python
285 lines
10 KiB
Python
import hashlib
|
||
import json
|
||
import re
|
||
import time
|
||
from dataclasses import dataclass
|
||
from typing import Any, Dict, List, Tuple
|
||
|
||
import httpx
|
||
|
||
from config.env import KBConfig
|
||
from utils.log_util import logger
|
||
from utils.match_service import get_match_service
|
||
|
||
|
||
@dataclass
|
||
class KBRetrievalMetrics:
|
||
retrieve_ms: float
|
||
top1_score: float
|
||
score_gap: float
|
||
hit_count: int
|
||
keyword_coverage: float
|
||
cache_hit: bool
|
||
|
||
|
||
class KBESServiceError(Exception):
|
||
pass
|
||
|
||
|
||
class KBESService:
|
||
"""ES/OpenSearch（BM25）检索 + 置信度 + 缓存封装"""
|
||
|
||
CACHE_PREFIX = "kb:es:retrieval"
|
||
|
||
@classmethod
|
||
def _cache_key(cls, question: str) -> str:
|
||
q = (question or "").strip().lower()
|
||
digest = hashlib.sha256(q.encode("utf-8")).hexdigest()
|
||
return f"{cls.CACHE_PREFIX}:{digest}"
|
||
|
||
@classmethod
|
||
async def retrieve(cls, question: str, redis=None) -> Tuple[List[Dict[str, Any]], KBRetrievalMetrics]:
|
||
if not question:
|
||
return [], KBRetrievalMetrics(0.0, 0.0, 0.0, 0, 0.0, False)
|
||
|
||
# 仅对 KB_PROVIDER=es_bm25：禁用 KBES 检索缓存（避免 curl/API 不一致与历史命中）
|
||
enable_retrieval_cache = KBConfig.KB_PROVIDER != 'es_bm25'
|
||
|
||
cache_hit = False
|
||
if enable_retrieval_cache and redis and KBConfig.KB_CACHE_TTL > 0:
|
||
key = cls._cache_key(question)
|
||
cached = await redis.get(key)
|
||
if cached:
|
||
try:
|
||
payload = json.loads(cached)
|
||
chunks = payload.get("chunks", [])
|
||
metrics = payload.get("metrics", {})
|
||
return (
|
||
chunks,
|
||
KBRetrievalMetrics(
|
||
retrieve_ms=metrics.get("retrieve_ms", 0.0),
|
||
top1_score=metrics.get("top1_score", 0.0),
|
||
score_gap=metrics.get("score_gap", 0.0),
|
||
hit_count=metrics.get("hit_count", len(chunks)),
|
||
keyword_coverage=metrics.get("keyword_coverage", 0.0),
|
||
cache_hit=True,
|
||
),
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
if not KBConfig.KB_ES_URL:
|
||
raise KBESServiceError("KB_ES_URL/KB_OS_URL 未配置")
|
||
|
||
logger.info(
|
||
f"[KB_ES] using es_url={KBConfig.KB_ES_URL} | index={KBConfig.KB_ES_INDEX} | top_k={KBConfig.KB_TOP_K}"
|
||
)
|
||
|
||
t0 = time.perf_counter()
|
||
chunks, metrics_raw = await cls._es_search(question)
|
||
retrieve_ms = (time.perf_counter() - t0) * 1000
|
||
|
||
chunks, best_cov = cls._rerank_chunks(question, chunks)
|
||
|
||
top1 = float(chunks[0].get("_score") or 0.0) if chunks else 0.0
|
||
top2 = float(chunks[1].get("_score") or 0.0) if len(chunks) > 1 else 0.0
|
||
gap = top1 - top2
|
||
|
||
metrics = KBRetrievalMetrics(
|
||
retrieve_ms=retrieve_ms,
|
||
top1_score=top1,
|
||
score_gap=gap,
|
||
hit_count=len(chunks),
|
||
keyword_coverage=best_cov,
|
||
cache_hit=cache_hit,
|
||
)
|
||
|
||
if enable_retrieval_cache and redis and KBConfig.KB_CACHE_TTL > 0:
|
||
try:
|
||
key = cls._cache_key(question)
|
||
await redis.set(
|
||
key,
|
||
json.dumps(
|
||
{
|
||
"chunks": chunks,
|
||
"metrics": {
|
||
"retrieve_ms": metrics.retrieve_ms,
|
||
"top1_score": metrics.top1_score,
|
||
"score_gap": metrics.score_gap,
|
||
"hit_count": metrics.hit_count,
|
||
"keyword_coverage": metrics.keyword_coverage,
|
||
},
|
||
},
|
||
ensure_ascii=False,
|
||
),
|
||
ex=KBConfig.KB_CACHE_TTL,
|
||
)
|
||
except Exception as e:
|
||
logger.debug(f"[KB_ES] 写入检索缓存失败: {e}")
|
||
|
||
return chunks, metrics
|
||
|
||
@classmethod
|
||
def _rerank_chunks(cls, question: str, chunks: List[Dict[str, Any]]) -> Tuple[List[Dict[str, Any]], float]:
|
||
"""二次重排（轻量、无LLM），避免 ES top1 误命中导致直接输出低质答案。
|
||
|
||
目标：尽量保持速度（topK 很小），只在内存里做简单特征打分。
|
||
"""
|
||
|
||
if not chunks:
|
||
return chunks, 0.0
|
||
|
||
match_service = get_match_service()
|
||
q = (question or "").strip()
|
||
q_norm = q.lower()
|
||
|
||
# 约束词：如果用户问到公司/年份，命中片段最好也出现对应实体
|
||
must_brand = any(x in q_norm for x in ("康达新材", "康达"))
|
||
years = set(re.findall(r"20\d{2}", q))
|
||
|
||
max_es = max([float(c.get("_score") or 0.0) for c in chunks] or [0.0])
|
||
if max_es <= 0:
|
||
max_es = 1.0
|
||
|
||
scored: List[Tuple[float, float, float, Dict[str, Any]]] = []
|
||
for c in chunks:
|
||
text = "\n".join(
|
||
[
|
||
str(c.get("question") or ""),
|
||
str(c.get("answer") or ""),
|
||
str(c.get("title") or ""),
|
||
str(c.get("content") or ""),
|
||
str(c.get("source") or ""),
|
||
]
|
||
)
|
||
text_norm = text.lower()
|
||
|
||
cov = match_service.calculate_keyword_coverage(q, text)
|
||
es = float(c.get("_score") or 0.0)
|
||
es_norm = es / max_es
|
||
|
||
penalty = 0.0
|
||
if must_brand and ("康达" not in text_norm):
|
||
penalty -= 0.35
|
||
if years:
|
||
hit_years = any(y in text for y in years)
|
||
if not hit_years:
|
||
penalty -= 0.25
|
||
|
||
# 覆盖率更能反映“是否问对文档”，ES 分数反映粗排相关性
|
||
final = 0.55 * cov + 0.45 * es_norm + penalty
|
||
|
||
scored.append((final, cov, es, c))
|
||
|
||
scored.sort(key=lambda x: x[0], reverse=True)
|
||
reranked = [c for _, _, _, c in scored]
|
||
best_cov = float(scored[0][1]) if scored else 0.0
|
||
return reranked, best_cov
|
||
|
||
@classmethod
|
||
def is_confident(cls, metrics: KBRetrievalMetrics) -> bool:
|
||
if metrics.hit_count <= 0:
|
||
return False
|
||
if metrics.top1_score < KBConfig.KB_MIN_SCORE:
|
||
return False
|
||
if metrics.keyword_coverage < KBConfig.KB_MIN_COVERAGE:
|
||
return False
|
||
return True
|
||
|
||
@classmethod
|
||
async def _es_search(cls, question: str) -> Tuple[List[Dict[str, Any]], Dict[str, float]]:
|
||
index = KBConfig.KB_ES_INDEX
|
||
url = f"{KBConfig.KB_ES_URL}/{index}/_search"
|
||
|
||
# 与线上排查/运维 curl 查询保持一致，避免 DSL 差异导致“bash 命中正确、API 命中错误”。
|
||
query = {
|
||
"size": KBConfig.KB_TOP_K,
|
||
"track_total_hits": False,
|
||
"_source": [
|
||
"doc_id",
|
||
"chunk_id",
|
||
"qa_id",
|
||
"question",
|
||
"answer",
|
||
"title",
|
||
"content",
|
||
"source",
|
||
"tags",
|
||
"category",
|
||
"updated_at",
|
||
"version",
|
||
],
|
||
"query": {
|
||
"multi_match": {
|
||
"query": question,
|
||
"fields": [
|
||
"question^3",
|
||
"answer^1.5",
|
||
"title^2",
|
||
"content",
|
||
],
|
||
"type": "best_fields",
|
||
}
|
||
},
|
||
}
|
||
|
||
auth = None
|
||
if KBConfig.KB_ES_USERNAME and KBConfig.KB_ES_PASSWORD:
|
||
auth = (KBConfig.KB_ES_USERNAME, KBConfig.KB_ES_PASSWORD)
|
||
|
||
async with httpx.AsyncClient(timeout=5.0, verify=KBConfig.KB_ES_VERIFY_SSL) as client:
|
||
# 线上排障：打印实际发给 ES 的 request body，确保与手工 curl 完全一致
|
||
logger.info(f"[KB_ES] request_body={json.dumps(query, ensure_ascii=False)}")
|
||
resp = await client.post(url, json=query, auth=auth)
|
||
if resp.status_code >= 300:
|
||
raise KBESServiceError(f"ES检索失败 HTTP {resp.status_code}: {resp.text}")
|
||
data = resp.json()
|
||
|
||
try:
|
||
top3 = ((data.get("hits") or {}).get("hits") or [])[:3]
|
||
logger.info(
|
||
"[KB_ES] es_top3="
|
||
+ json.dumps(
|
||
[
|
||
{
|
||
"_id": h.get("_id"),
|
||
"_score": h.get("_score"),
|
||
"question": (h.get("_source") or {}).get("question"),
|
||
}
|
||
for h in top3
|
||
],
|
||
ensure_ascii=False,
|
||
)
|
||
)
|
||
except Exception:
|
||
pass
|
||
|
||
hits = (data.get("hits") or {}).get("hits") or []
|
||
chunks: List[Dict[str, Any]] = []
|
||
scores: List[float] = []
|
||
for h in hits:
|
||
src = h.get("_source") or {}
|
||
score = float(h.get("_score") or 0.0)
|
||
scores.append(score)
|
||
chunks.append(
|
||
{
|
||
"doc_id": src.get("doc_id"),
|
||
"chunk_id": src.get("chunk_id"),
|
||
"qa_id": src.get("qa_id"),
|
||
"question": src.get("question"),
|
||
"answer": src.get("answer"),
|
||
"title": src.get("title"),
|
||
"content": src.get("content"),
|
||
"source": src.get("source"),
|
||
"tags": src.get("tags"),
|
||
"category": src.get("category"),
|
||
"updated_at": src.get("updated_at"),
|
||
"version": src.get("version"),
|
||
"_score": score,
|
||
}
|
||
)
|
||
|
||
top1 = scores[0] if scores else 0.0
|
||
top2 = scores[1] if len(scores) > 1 else 0.0
|
||
return chunks, {"top1_score": top1, "score_gap": top1 - top2}
|