feat:修复投顾agent功能

nl2sql/schema.py

@@ -1,12 +1,23 @@
 """从 MySQL information_schema 加载并校验 NL2SQL 权威 Schema。"""
 from __future__ import annotations
 
+import hashlib
+import json
+import logging
 from typing import Any
 
 from sqlalchemy import bindparam, text
 
 from nl2sql.metadata import normalize_column_row, normalize_table_row
 
+logger = logging.getLogger("nl2sql.schema")
+SCHEMA_CACHE_TTL = 60
+
+
+def _schema_cache_key(database: str, candidates: list[str]) -> str:
+    payload = json.dumps([database, candidates], ensure_ascii=False, separators=(",", ":"))
+    return "nl2sql:schema:" + hashlib.sha256(payload.encode()).hexdigest()
+
 
 TABLES_SQL = text(
     """
@@ -34,13 +45,25 @@ async def load_authoritative_schema(
     *,
     database: str,
     candidate_tables: set[str] | list[str],
+    redis=None,
+    cache_ttl: int = SCHEMA_CACHE_TTL,
 ) -> dict[str, list[dict[str, Any]]]:
     """从权威元数据源加载候选表，并剔除不存在的表和孤立字段。"""
     candidates = {str(name).strip() for name in candidate_tables if str(name).strip()}
     if not candidates:
         return {"tables": [], "columns": []}
 
-    params = {"database": database, "candidate_tables": sorted(candidates)}
+    candidate_list = sorted(candidates)
+    cache_key = _schema_cache_key(database, candidate_list)
+    if redis is not None:
+        try:
+            cached = await redis.get(cache_key)
+            if cached:
+                return json.loads(cached)
+        except Exception:  # noqa: BLE001 缓存异常回退权威查询
+            logger.warning("Schema 缓存读取失败", exc_info=True)
+
+    params = {"database": database, "candidate_tables": candidate_list}
     table_result = await session.execute(TABLES_SQL, params)
     column_result = await session.execute(COLUMNS_SQL, params)
 
@@ -56,4 +79,10 @@ async def load_authoritative_schema(
         normalized = normalize_column_row(dict(row))
         if normalized["table_name"] in valid_table_names and normalized["field_name"]:
             columns.append(normalized)
-    return {"tables": tables, "columns": columns}
+    schema = {"tables": tables, "columns": columns}
+    if redis is not None:
+        try:
+            await redis.set(cache_key, json.dumps(schema, ensure_ascii=False), ex=cache_ttl)
+        except Exception:  # noqa: BLE001 缓存异常不阻断查询
+            logger.warning("Schema 缓存写入失败", exc_info=True)
+    return schema