56 lines
1.9 KiB
Python
56 lines
1.9 KiB
Python
from typing import Any, Dict
|
|
|
|
from config import Config
|
|
from services.ragflow_client import RagflowClient, extract_table_name
|
|
|
|
|
|
class TemplateMatcher:
|
|
"""模板匹配器：RAGFlow 检索"""
|
|
|
|
def __init__(self):
|
|
self._ragflow = RagflowClient()
|
|
cfg = Config.get_section("ragflow")
|
|
self._dataset_id = (cfg.get("table_retrieval_dataset_id") or "").strip()
|
|
self._top_k = int(cfg.get("retrieval_top_k", 3))
|
|
|
|
def _validate(self) -> None:
|
|
if not self._dataset_id:
|
|
raise RuntimeError("未配置 ragflow.table_retrieval_dataset_id，无法进行表名检索")
|
|
|
|
def match(self, normalized_text: str) -> Dict[str, Any]:
|
|
"""返回匹配的表名与原始响应"""
|
|
self._validate()
|
|
try:
|
|
response = self._ragflow.retrieve(normalized_text, top_k=self._top_k, dataset_id=self._dataset_id)
|
|
except Exception as e:
|
|
return {"table_name": None, "raw": {"error": str(e)}}
|
|
|
|
candidates = []
|
|
data = response.get("data") if isinstance(response, dict) else None
|
|
records = []
|
|
if isinstance(data, list):
|
|
records = data
|
|
elif isinstance(data, dict):
|
|
chunks = data.get("chunks")
|
|
if isinstance(chunks, list):
|
|
records = chunks
|
|
|
|
for item in records:
|
|
table_name = extract_table_name(item)
|
|
if table_name:
|
|
candidates.append(table_name)
|
|
|
|
matched = candidates[0] if candidates else None
|
|
return {"table_name": matched, "raw": response}
|
|
|
|
|
|
_GLOBAL_TEMPLATE_MATCHER: TemplateMatcher | None = None
|
|
|
|
|
|
def get_template_matcher() -> TemplateMatcher:
|
|
"""获取全局 TemplateMatcher（单例）"""
|
|
global _GLOBAL_TEMPLATE_MATCHER
|
|
if _GLOBAL_TEMPLATE_MATCHER is None:
|
|
_GLOBAL_TEMPLATE_MATCHER = TemplateMatcher()
|
|
return _GLOBAL_TEMPLATE_MATCHER
|