more_dots/services/template_matcher.py

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