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)