feat: 可配置方案引擎通用原型(Mock 检索)
需求未定稿时可先跑通追问/判定/出方案;槽位与判定树外置 YAML,后续可接 Ragflow。 Co-authored-by: haijd <haijd@163.com>
This commit is contained in:
@@ -0,0 +1,94 @@
|
||||
from __future__ import annotations
|
||||
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from .config_loader import business_config
|
||||
from .settings import settings
|
||||
|
||||
|
||||
class RetrievalResult:
|
||||
def __init__(self, chunks: list[dict[str, Any]]):
|
||||
self.chunks = chunks
|
||||
|
||||
|
||||
def retrieve_for_scheme(slots: dict[str, str], bundle: list[str]) -> RetrievalResult:
|
||||
backend = (settings.retrieval_backend or "mock").lower()
|
||||
if backend == "ragflow":
|
||||
return _retrieve_ragflow(slots, bundle)
|
||||
return _retrieve_mock(slots, bundle)
|
||||
|
||||
|
||||
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 c in chunks:
|
||||
if c.get("product") in bundle_set:
|
||||
picked.append(c)
|
||||
# tag overlap with slot values
|
||||
blob = " ".join(slots.values())
|
||||
for c in chunks:
|
||||
if c in picked:
|
||||
continue
|
||||
tags = c.get("tags") or []
|
||||
if any(str(t) in blob for t in tags):
|
||||
picked.append(c)
|
||||
if len(picked) >= 4:
|
||||
break
|
||||
if not picked:
|
||||
picked = chunks[:2]
|
||||
return RetrievalResult(picked[:4])
|
||||
|
||||
|
||||
def _retrieve_ragflow(slots: dict[str, str], bundle: list[str]) -> RetrievalResult:
|
||||
"""Minimal Ragflow retrieval hook. Falls back to mock if not configured/unavailable."""
|
||||
if not settings.ragflow_api_key:
|
||||
return _retrieve_mock(slots, bundle)
|
||||
|
||||
query = " ".join(
|
||||
[
|
||||
slots.get("tech_category", ""),
|
||||
slots.get("scenario", ""),
|
||||
slots.get("key_metric", ""),
|
||||
" ".join(bundle),
|
||||
]
|
||||
).strip() or "产品方案"
|
||||
dataset_ids = [x for x in settings.ragflow_dataset_ids.split(",") if x.strip()]
|
||||
url = settings.ragflow_base_url.rstrip("/") + "/api/v1/retrieval"
|
||||
headers = {"Authorization": f"Bearer {settings.ragflow_api_key}"}
|
||||
payload: dict[str, Any] = {"question": query, "top_k": 4}
|
||||
if dataset_ids:
|
||||
payload["dataset_ids"] = dataset_ids
|
||||
|
||||
try:
|
||||
with httpx.Client(timeout=20.0) as client:
|
||||
resp = client.post(url, headers=headers, json=payload)
|
||||
resp.raise_for_status()
|
||||
data = resp.json()
|
||||
except Exception:
|
||||
return _retrieve_mock(slots, bundle)
|
||||
|
||||
# Ragflow response shapes vary by version; normalize best-effort
|
||||
records = []
|
||||
if isinstance(data, dict):
|
||||
records = data.get("data") or data.get("chunks") or data.get("records") or []
|
||||
if isinstance(records, dict):
|
||||
records = records.get("chunks") or records.get("records") or []
|
||||
normalized = []
|
||||
for i, r in enumerate(records[:4]):
|
||||
if not isinstance(r, dict):
|
||||
continue
|
||||
normalized.append(
|
||||
{
|
||||
"id": r.get("id") or f"rf-{i}",
|
||||
"title": r.get("document_keyword") or r.get("title") or "Ragflow片段",
|
||||
"product": r.get("document_id") or "",
|
||||
"content": r.get("content") or r.get("text") or "",
|
||||
"tags": [],
|
||||
}
|
||||
)
|
||||
if not normalized:
|
||||
return _retrieve_mock(slots, bundle)
|
||||
return RetrievalResult(normalized)
|
||||
Reference in New Issue
Block a user