101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
"""NL2SQL 查询问题改写与会话上下文补全。"""
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import re
|
||
|
||
from tool.llm import llm
|
||
|
||
|
||
logger = logging.getLogger("nl2sql.query_rewriter")
|
||
|
||
_SYSTEM_PROMPT = """你是基金平台 NL2SQL 的问题改写器。
|
||
将当前问题改写为不依赖上下文、可直接交给数据库查询系统理解的完整中文问题。
|
||
只输出改写后的问题，不要输出解释、JSON、SQL 或 Markdown。
|
||
只补全指代和省略，不改变用户的查询目标、客户范围、时间范围或排序条件。
|
||
不能扩大权限范围，不能臆造查询结果；上下文不足时保留原问题。
|
||
"""
|
||
|
||
_REFERENCE_TERMS = (
|
||
"他", "她", "它", "他们", "她们", "它们", "这个客户", "该客户", "那个客户",
|
||
"这个基金", "该基金", "那只基金", "这只基金", "这个结果", "上一轮", "刚才", "前面",
|
||
)
|
||
|
||
|
||
def _rewrite_explicit_customer_identity(question: str) -> str | None:
|
||
normalized = re.sub(r"\s+", "", question)
|
||
match = re.fullmatch(
|
||
r"(?:查询|请问|告诉我)?客户(?:编号|号|#)?(\d+)(是谁|姓名是什么|姓名|基本信息|联系方式|手机号是什么|手机号)",
|
||
normalized,
|
||
)
|
||
if not match:
|
||
return None
|
||
return f"查询客户编号{match.group(1)}的姓名"
|
||
|
||
|
||
def _rewrite_customer_name_follow_up(question: str, context: str) -> str | None:
|
||
"""从最近的客户身份问答中恢复客户编号，保证追问可独立查询。"""
|
||
normalized_question = re.sub(r"\s+", "", question)
|
||
if not any(term in normalized_question for term in ("持仓", "资产", "基金", "产品", "收益", "盈亏", "余额", "交易")):
|
||
return None
|
||
|
||
lines = context.splitlines()
|
||
for index, line in enumerate(lines):
|
||
customer_match = re.search(r"客户(?:编号|号|#)?(\d+)", line)
|
||
if not customer_match or not any(term in line for term in ("是谁", "姓名")):
|
||
continue
|
||
for answer_line in lines[index + 1 : index + 4]:
|
||
if not answer_line.startswith("assistant:"):
|
||
continue
|
||
answer = answer_line.split(":", 1)[1].strip()
|
||
name_match = re.search(r"(?:姓名|是)\s*[:：]?\s*([\u4e00-\u9fff]{2,6})", answer)
|
||
candidates = [name_match.group(1)] if name_match else re.findall(r"[\u4e00-\u9fff]{2,6}", answer)
|
||
for name in reversed(candidates):
|
||
if name in normalized_question:
|
||
return f"查询客户编号{customer_match.group(1)}的持仓和资产信息"
|
||
return None
|
||
|
||
|
||
async def rewrite_query(
|
||
question: str,
|
||
conversation_context: str = "",
|
||
*,
|
||
llm_client=llm,
|
||
) -> str:
|
||
"""使用有限会话上下文补全问题；改写失败时安全返回原问题。"""
|
||
original = (question or "").strip()
|
||
context = (conversation_context or "").strip()
|
||
if not original:
|
||
return original
|
||
explicit_identity = _rewrite_explicit_customer_identity(original)
|
||
if explicit_identity:
|
||
return explicit_identity
|
||
contextual_identity = _rewrite_customer_name_follow_up(original, context)
|
||
if contextual_identity:
|
||
return contextual_identity
|
||
if not context or llm_client is None or not any(term in original for term in _REFERENCE_TERMS):
|
||
return original
|
||
|
||
prompt = (
|
||
f"会话上下文：\n{context[:4000]}\n\n"
|
||
f"当前问题：\n{original[:2000]}\n\n"
|
||
"改写后的独立问题："
|
||
)
|
||
try:
|
||
rewritten = await llm_client.chat(
|
||
[
|
||
{"role": "system", "content": _SYSTEM_PROMPT},
|
||
{"role": "user", "content": prompt},
|
||
],
|
||
temperature=0,
|
||
max_tokens=512,
|
||
)
|
||
except Exception: # noqa: BLE001 改写失败不阻断原查询
|
||
logger.warning("NL2SQL 查询问题改写失败，继续使用原问题", exc_info=True)
|
||
return original
|
||
value = str(rewritten or "").strip()
|
||
return value or original
|
||
|
||
|
||
__all__ = ["rewrite_query"]
|