xianjing-ai/engine/app/retrieval.py

from __future__ import annotations

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_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)