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)