kangda-robot-backend/ruoyi-fastapi-backend/module_admin/service/kb_es_service.py

import hashlib
import json
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)

        cache_hit = False
        if 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 未配置")

        t0 = time.perf_counter()
        chunks, metrics_raw = await cls._es_search(question)
        retrieve_ms = (time.perf_counter() - t0) * 1000

        match_service = get_match_service()
        combined = "\n".join([(c.get("title") or "") + "\n" + (c.get("content") or "") for c in chunks])
        coverage = match_service.calculate_keyword_coverage(question, combined)

        top1 = metrics_raw.get("top1_score", 0.0)
        gap = metrics_raw.get("score_gap", 0.0)

        metrics = KBRetrievalMetrics(
            retrieve_ms=retrieve_ms,
            top1_score=top1,
            score_gap=gap,
            hit_count=len(chunks),
            keyword_coverage=coverage,
            cache_hit=cache_hit,
        )

        if 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 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"

        query = {
            "size": KBConfig.KB_TOP_K,
            "track_total_hits": False,
            "_source": [
                "doc_id",
                "chunk_id",
                "title",
                "content",
                "source",
                "tags",
                "category",
                "updated_at",
                "version",
            ],
            "query": {
                "bool": {
                    "should": [
                        {
                            "multi_match": {
                                "query": question,
                                "fields": ["title^2", "content", "tags^1.2", "category^1.1"],
                                "type": "best_fields",
                                "minimum_should_match": "60%",
                            }
                        },
                        {"match_phrase": {"title": {"query": question, "boost": 3}}},
                        {"match_phrase": {"content": {"query": question, "boost": 2}}},
                    ],
                    "minimum_should_match": 1,
                }
            },
        }

        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:
            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()

        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"),
                    "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}