90 lines
3.1 KiB
Python
90 lines
3.1 KiB
Python
"""NL2SQL SQL 生成器：只负责调用 Chat 模型并清理模型输出。"""
|
||
from __future__ import annotations
|
||
|
||
import json
|
||
import re
|
||
from dataclasses import dataclass
|
||
from typing import Any
|
||
|
||
from sqlglot import exp, parse
|
||
from sqlglot.errors import ParseError
|
||
|
||
from tool.llm import llm
|
||
from nl2sql.semantics import build_semantic_context
|
||
|
||
|
||
class SqlGenerationError(ValueError):
|
||
"""模型未返回可用的只读 SQL。"""
|
||
|
||
|
||
@dataclass(frozen=True)
|
||
class GeneratedSql:
|
||
sql: str
|
||
|
||
|
||
def _clean_model_output(output: str) -> str:
|
||
value = (output or "").strip()
|
||
value = re.sub(r"^```(?:sql)?\s*", "", value, flags=re.IGNORECASE)
|
||
value = re.sub(r"\s*```$", "", value).strip()
|
||
return value.removesuffix(";").strip()
|
||
|
||
|
||
def _build_prompt(
|
||
question: str,
|
||
schema: dict[str, Any],
|
||
few_shot: list[dict[str, Any]],
|
||
semantic_context: dict[str, Any] | None = None,
|
||
conversation_context: str = "",
|
||
) -> str:
|
||
return (
|
||
"你是基金业务数据库 SQL 生成器。\n"
|
||
"只根据提供的 Schema 生成一条 MySQL SELECT，禁止写操作、跨库访问和未提供的表字段。\n"
|
||
"只输出 SQL，不要输出解释、Markdown 或代码围栏。\n"
|
||
f"用户问题：{question}\n"
|
||
f"Schema：{json.dumps(schema, ensure_ascii=False, sort_keys=True)}\n"
|
||
f"业务语义：{json.dumps(semantic_context or {}, ensure_ascii=False, sort_keys=True)}\n"
|
||
f"会话上下文：{conversation_context}\n"
|
||
f"Few-shot：{json.dumps(few_shot or [], ensure_ascii=False, sort_keys=True)}"
|
||
)
|
||
|
||
|
||
async def generate_sql(
|
||
question: str,
|
||
schema: dict[str, Any],
|
||
*,
|
||
llm_client=llm,
|
||
few_shot: list[dict[str, Any]] | None = None,
|
||
semantic_context: dict[str, Any] | None = None,
|
||
conversation_context: str = "",
|
||
) -> GeneratedSql:
|
||
"""调用 Chat 模型生成 SQL，并在返回前确认其为单条 SELECT。"""
|
||
if not question or not question.strip():
|
||
raise SqlGenerationError("用户问题不能为空")
|
||
messages = [
|
||
{"role": "system", "content": "你必须严格遵守只输出单条 SELECT SQL。"},
|
||
{
|
||
"role": "user",
|
||
"content": _build_prompt(
|
||
question,
|
||
schema,
|
||
few_shot or [],
|
||
semantic_context or build_semantic_context(question, schema),
|
||
conversation_context,
|
||
),
|
||
},
|
||
]
|
||
try:
|
||
output = await llm_client.chat(messages, temperature=0)
|
||
except Exception as exc: # noqa: BLE001 统一收敛模型异常
|
||
raise SqlGenerationError("SQL 模型调用失败") from exc
|
||
sql = _clean_model_output(output)
|
||
if not sql:
|
||
raise SqlGenerationError("模型未返回 SQL")
|
||
try:
|
||
statements = parse(sql, read="mysql")
|
||
except ParseError as exc:
|
||
raise SqlGenerationError("模型返回的 SQL 无法解析") from exc
|
||
if len(statements) != 1 or not isinstance(statements[0], exp.Select):
|
||
raise SqlGenerationError("模型只允许返回单条 SELECT")
|
||
return GeneratedSql(sql=statements[0].sql(dialect="mysql"))
|