more_dots/services/sql_prompt_manager.py

import json
import os
from typing import Any, Dict, Optional


class SqlPromptManager:
    """按表名读取 SQL 提示词"""

    def __init__(self, base_dir: Optional[str] = None):
        root_dir = os.path.dirname(os.path.dirname(__file__))
        self._base_dir = base_dir or os.path.join(root_dir, "config", "sql_gen_prompts")

    @staticmethod
    def _safe_filename(name: str) -> str:
        return name.replace("..", "").replace("/", "_").replace("\\", "_")

    def get_prompt(self, table_name: str) -> Optional[Dict[str, Any]]:
        """读取指定表的提示词 JSON"""
        if not table_name:
            return None
        safe_name = self._safe_filename(table_name)
        filename = safe_name + ".json"
        path = os.path.join(self._base_dir, filename)
        if not os.path.exists(path):
            return None

        with open(path, "r", encoding="utf-8") as f:
            prompt = json.load(f)

        return prompt


_GLOBAL_SQL_PROMPT_MANAGER: Optional[SqlPromptManager] = None


def get_sql_prompt_manager(base_dir: Optional[str] = None) -> SqlPromptManager:
    """获取全局 SqlPromptManager（单例）"""
    global _GLOBAL_SQL_PROMPT_MANAGER
    if _GLOBAL_SQL_PROMPT_MANAGER is None:
        _GLOBAL_SQL_PROMPT_MANAGER = SqlPromptManager(base_dir=base_dir)
    return _GLOBAL_SQL_PROMPT_MANAGER