Cursor Agentandhaijd
b7bd830da7
咨询台接 Ragflow 检索与示例知识库，支持追问、判定、引用和留资。 Co-authored-by: haijd <haijd@163.com>
81 lines
2.8 KiB
Python
81 lines
2.8 KiB
Python
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)
|