xianjing-ai/engine/app/retrieval.py

from __future__ import annotations

import re
from typing import Any

from .config_loader import business_config
from .ragflow_client import RagflowClient, RagflowError
from .settings import settings


class RetrievalResult:
    def __init__(self, chunks: list[dict[str, Any]], backend: str, warning: str = ""):
        self.chunks = chunks
        self.backend = backend
        self.warning = warning


def retrieve_for_scheme(slots: dict[str, str], bundle: list[str]) -> RetrievalResult:
    backend = (settings.retrieval_backend or "ragflow").lower()
    if backend == "mock":
        return _retrieve_mock(slots, bundle)
    try:
        return _retrieve_ragflow(slots, bundle)
    except RagflowError as exc:
        return RetrievalResult([], "ragflow", warning=str(exc))


def retrieve_question(question: str, *, page_size: int = 3, min_score: float = 0.22) -> RetrievalResult:
    """自由问答检索；分数过低视为知识库未命中。"""
    backend = (settings.retrieval_backend or "ragflow").lower()
    q = (question or "").strip()
    if not q:
        return RetrievalResult([], backend or "none", warning="empty_query")
    if backend == "mock":
        return _retrieve_mock_question(q, min_score=min_score)
    try:
        client = RagflowClient()
        dataset_ids = [item.strip() for item in settings.ragflow_dataset_ids.split(",") if item.strip()]
        raw = client.retrieve(q, dataset_ids, page_size=page_size)
        chunks = []
        for index, item in enumerate(raw[:page_size]):
            if not isinstance(item, dict):
                continue
            score = item.get("similarity")
            try:
                score_f = float(score) if score is not None else 0.0
            except (TypeError, ValueError):
                score_f = 0.0
            if score_f < min_score:
                continue
            chunks.append(
                {
                    "id": item.get("id") or f"rf-q-{index}",
                    "title": item.get("document_keyword")
                    or item.get("document_name")
                    or item.get("docnm_kwd")
                    or "知识片段",
                    "content": item.get("content") or item.get("content_with_weight") or "",
                    "score": score_f,
                }
            )
        warning = "" if chunks else "知识库未召回到足够相关的内容"
        return RetrievalResult(chunks, "ragflow", warning=warning)
    except RagflowError as exc:
        return RetrievalResult([], "ragflow", warning=str(exc))


def _retrieve_mock_question(question: str, *, min_score: float = 0.22) -> RetrievalResult:
    picked: list[dict[str, Any]] = []
    for chunk in business_config.chunks:
        blob = f"{chunk.get('title', '')} {chunk.get('content', '')} {' '.join(chunk.get('tags') or [])}"
        hits = sum(1 for token in re_split_tokens(question) if token and token in blob)
        if hits:
            score = min(0.9, 0.2 + 0.15 * hits)
            if score >= min_score:
                item = dict(chunk)
                item["score"] = score
                picked.append(item)
        if len(picked) >= 3:
            break
    warning = "" if picked else "知识库未召回到足够相关的内容"
    return RetrievalResult(picked, "mock", warning=warning)


def re_split_tokens(text: str) -> list[str]:
    return [t for t in re.split(r"[\s,，。；;？?！!/|]+", text) if len(t) >= 2]


def _retrieve_mock(slots: dict[str, str], bundle: list[str]) -> RetrievalResult:
    chunks = business_config.chunks
    picked: list[dict[str, Any]] = []
    bundle_set = set(bundle)
    for chunk in chunks:
        if chunk.get("product") in bundle_set:
            picked.append(chunk)
    blob = " ".join(slots.values())
    for chunk in chunks:
        if chunk in picked:
            continue
        tags = chunk.get("tags") or []
        if any(str(tag) in blob for tag in tags):
            picked.append(chunk)
        if len(picked) >= 4:
            break
    if not picked:
        picked = chunks[:2]
    return RetrievalResult(picked[:4], "mock")


def _retrieve_ragflow(slots: dict[str, str], bundle: list[str]) -> RetrievalResult:
    client = RagflowClient()
    query = " ".join(
        [
            slots.get("tech_category", ""),
            slots.get("scenario", ""),
            slots.get("key_metric", ""),
            slots.get("constraints", ""),
            " ".join(bundle),
            "产品组合 适用场景 注意事项",
        ]
    ).strip()
    dataset_ids = [item.strip() for item in settings.ragflow_dataset_ids.split(",") if item.strip()]
    raw = client.retrieve(query, dataset_ids, page_size=4)
    normalized = []
    for index, item in enumerate(raw[:4]):
        if not isinstance(item, dict):
            continue
        normalized.append(
            {
                "id": item.get("id") or f"rf-{index}",
                "title": item.get("document_keyword")
                or item.get("document_name")
                or item.get("docnm_kwd")
                or "知识片段",
                "product": item.get("document_id") or "",
                "content": item.get("content") or item.get("content_with_weight") or "",
                "tags": [],
                "score": item.get("similarity"),
            }
        )
    warning = "" if normalized else "知识库没有召回到可用片段，方案仅供草稿，建议转人工。"
    return RetrievalResult(normalized, "ragflow", warning=warning)