feat:修复客服agent功能

nl2sql/query_rewriter.py

@@ -0,0 +1,95 @@
+"""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。
+只补全指代和省略，不改变用户的查询目标、客户范围、时间范围或排序条件。
+不能扩大权限范围，不能臆造查询结果；上下文不足时保留原问题。
+"""
+
+
+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:
+        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"]