import json import os from typing import Any, Dict, Optional from config import Config from services.cache import NoopCache, RedisCache 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") self._cache = self._init_cache() self._cache_ttl = self._get_cache_ttl() @staticmethod def _get_cache_ttl() -> int: redis_cfg = Config.get_section("redis") try: return int(redis_cfg.get("sql_prompt_ttl", 600)) except Exception: return 600 @staticmethod def _init_cache(): redis_cfg = Config.get_section("redis") enabled = str(redis_cfg.get("enabled", "false")).lower() in ("1", "true", "yes") if not enabled: return NoopCache() # 优先使用完整 URL;否则使用 host/port/password/database 拼接 url = redis_cfg.get("url") db = int(redis_cfg.get("db", redis_cfg.get("database", 0))) if not url: host = redis_cfg.get("host") port = redis_cfg.get("port", "6379") password = redis_cfg.get("password", "") database = redis_cfg.get("database", str(db)) if host: auth = f":{password}@" if password else "" url = f"redis://{auth}{host}:{port}/{database}" if not url: return NoopCache() try: return RedisCache(url=url, db=db) except Exception: return NoopCache() @staticmethod def _safe_filename(name: str) -> str: return name.replace("..", "").replace("/", "_").replace("\\", "_") @staticmethod def _cache_key(table_name: str, mtime: float) -> str: return f"sql_prompt:{table_name}:{int(mtime)}" 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 mtime = os.path.getmtime(path) key = self._cache_key(safe_name, mtime) cached = self._cache.get(key) if cached: try: return json.loads(cached) except Exception: pass with open(path, "r", encoding="utf-8") as f: prompt = json.load(f) self._cache.set(key, json.dumps(prompt, ensure_ascii=False), self._cache_ttl) 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