import json import re from pathlib import Path from typing import Any, Dict, Optional, Set from config import Config from services.core.sql_prompt_manager import get_sql_prompt_manager from services.integrations.ragflow_client import RagflowClient, extract_table_name IDENTIFIER_RE = re.compile(r"[a-zA-Z_][a-zA-Z0-9_]*") EXPLICIT_FILTER_FIELD_RE = re.compile(r"([a-zA-Z_][a-zA-Z0-9_]*)\s*=") def _normalize_term(term: str) -> str: return str(term or "").strip().lower() class TemplateMatcher: """模板匹配器：RAGFlow 检索""" KEYWORD_MATCH_BONUS = 50 def __init__(self): self._ragflow = RagflowClient() self._sql_prompt_manager = get_sql_prompt_manager() 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)) self._non_empty_tables, self._table_keywords = self._load_table_config() self._table_terms_cache: Dict[str, Set[str]] = {} def _load_table_config(self) -> tuple[Optional[Set[str]], Dict[str, Set[str]]]: """从本地 tables.json 读取非空模板表集合和关键词映射。""" try: tables_path = Path(__file__).resolve().parents[2] / "config" / "table_retrieval_prompts" / "tables.json" with open(tables_path, "r", encoding="utf-8") as f: data = json.load(f) if not isinstance(data, dict): return None, {} non_empty: Set[str] = set() keywords_map: Dict[str, Set[str]] = {} for table_name, templates in data.items(): if not isinstance(table_name, str): continue if isinstance(templates, list) and len(templates) > 0: non_empty.add(table_name) keywords_map[table_name] = {_normalize_term(kw) for kw in templates if isinstance(kw, str)} return non_empty, keywords_map except Exception: return None, {} def _validate(self) -> None: if not self._dataset_id: raise RuntimeError("未配置 ragflow.table_retrieval_dataset_id，无法进行表名检索") @staticmethod def _extract_query_terms(normalized_text: str) -> Set[str]: return { _normalize_term(match.group(0)) for match in IDENTIFIER_RE.finditer(normalized_text or "") } @staticmethod def _extract_explicit_filter_fields(normalized_text: str) -> Set[str]: return { _normalize_term(match.group(1)) for match in EXPLICIT_FILTER_FIELD_RE.finditer(normalized_text or "") } def _load_table_terms(self, table_name: str) -> Set[str]: cached = self._table_terms_cache.get(table_name) if cached is not None: return cached prompt = self._sql_prompt_manager.get_prompt(table_name) or {} terms: Set[str] = set() for field in (prompt.get("data_model_specification") or {}).get("fields_list") or []: if isinstance(field, str): terms.add(_normalize_term(field)) field_ref = prompt.get("field_mapping_reference") or {} self._collect_mapping_terms(field_ref, terms) self._table_terms_cache[table_name] = terms return terms def _collect_mapping_terms(self, node: Any, terms: Set[str]) -> None: if isinstance(node, dict): for key, value in node.items(): if key == "alias" and isinstance(value, list): for alias in value: if isinstance(alias, str): terms.add(_normalize_term(alias)) continue if isinstance(value, dict): if "alias" in value or "type" in value: terms.add(_normalize_term(key)) self._collect_mapping_terms(value, terms) elif isinstance(value, list): # 对字段列表直接入词，增强字段覆盖匹配 if key.endswith("_fields") or key in {"fields_list", "list"}: for item in value: if isinstance(item, str): terms.add(_normalize_term(item)) self._collect_mapping_terms(value, terms) elif isinstance(node, list): for item in node: self._collect_mapping_terms(item, terms) def _rank_candidates(self, normalized_text: str, candidates: list[Dict[str, Any]]) -> list[Dict[str, Any]]: query_terms = self._extract_query_terms(normalized_text) explicit_fields = self._extract_explicit_filter_fields(normalized_text) if not query_terms or not candidates: return candidates ranked: list[Dict[str, Any]] = [] for index, candidate in enumerate(candidates): table_name = candidate.get("table_name") if not table_name: continue table_terms = self._load_table_terms(table_name) overlap = len(query_terms & table_terms) missing_explicit_fields = len([field for field in explicit_fields if field not in table_terms]) keyword_bonus = 0 table_keywords = self._table_keywords.get(table_name, set()) if table_keywords: matched_keywords = query_terms & table_keywords keyword_bonus = len(matched_keywords) * self.KEYWORD_MATCH_BONUS rank_score = overlap + keyword_bonus - (missing_explicit_fields * 5) ranked.append({ **candidate, "rank_score": rank_score, "rank_overlap": overlap, "rank_keyword_bonus": keyword_bonus, "rank_missing_explicit_fields": missing_explicit_fields, "rank_index": index, }) ranked.sort(key=lambda item: (item["rank_score"], item["rank_overlap"], -item["rank_index"]), reverse=True) return ranked 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, "candidates": [], "raw": {"error": str(e)}} candidates = [] seen = set() 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 self._non_empty_tables is not None and table_name not in self._non_empty_tables: continue if table_name and table_name not in seen: seen.add(table_name) candidates.append( { "table_name": table_name, "metadata": item.get("metadata") or {}, "content": item.get("content") or item.get("text") or "", } ) candidates = self._rank_candidates(normalized_text, candidates) matched = candidates[0]["table_name"] if candidates else None return {"table_name": matched, "candidates": candidates, "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