Mutual_Fund/nl2sql/query_rewriter.py

"""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"]