group_fqcd_jr/nl2sql_yc.py

"""金融 NL2SQL MVP。

查询数据只允许使用参数化 SELECT；审计记录可通过回调写入
conversation_message.tool_calls。该文件可被运营和投顾 Agent 直接导入调用。
"""
from __future__ import annotations

import json
import logging
import os
import re
import uuid
from dataclasses import asdict, dataclass, field
from datetime import datetime, time, timedelta
from pathlib import Path
from typing import Any, Callable

logger = logging.getLogger(__name__)

ALLOWED_TABLES = {
    "sys_customer_assignment", "fin_customer_profile", "fin_risk_assessment",
    "fin_product", "fin_fee_rule", "fin_market_price", "fin_nav_history",
    "fin_holding", "fin_transaction", "fin_sim_order", "fin_sim_account",
    "fin_cash_ledger", "client_facing_content",
}

DOMAINS = {
    "market_nav": {"fin_market_price", "fin_nav_history", "fin_product"},
    "customer_risk": {
        "sys_customer_assignment", "fin_customer_profile", "fin_risk_assessment",
    },
    "product_fee": {"fin_product", "fin_fee_rule"},
    "trading_account": {
        "fin_sim_order", "fin_transaction", "fin_sim_account",
        "fin_cash_ledger", "fin_holding", "fin_product",
    },
    "client_content": {"client_facing_content"},
}

TABLE_COLUMNS = {
    "sys_customer_assignment": {
        "id", "customer_id", "employee_id", "employee_role", "assigned_at", "unassigned_at",
    },
    "fin_customer_profile": {
        "customer_id", "trade_account", "real_name", "birth_date", "occupation",
        "mobile_masked", "investor_type", "investment_horizon", "preferred_asset_class",
        "trading_frequency", "last_active_at", "total_asset", "behavior_score",
        "risk_tags", "opened_at", "updated_at",
    },
    "fin_risk_assessment": {
        "id", "customer_id", "questionnaire_version", "answers", "total_score",
        "investor_type", "assessed_at", "valid_until", "created_at",
    },
    "fin_product": {
        "id", "product_code", "product_name", "exchange_code", "product_category",
        "risk_level", "fund_manager", "currency", "lot_size", "price_tick",
        "current_nav", "current_nav_at", "min_amount", "open_start_at", "open_end_at",
        "open_period_start", "open_period_end", "transaction_fee_rate",
        "single_investor_max_holding_ratio", "management_fee_rate", "custodian_fee_rate",
        "risk_disclosure_required", "second_confirmation_required", "recording_required",
        "status", "created_at", "updated_at",
    },
    "fin_fee_rule": {
        "id", "rule_code", "product_id", "exchange_code", "order_side", "customer_tier",
        "min_trade_amount", "max_trade_amount", "fee_rate", "minimum_fee", "fixed_fee",
        "priority", "effective_from", "effective_until", "status", "created_at", "updated_at",
    },
    "fin_market_price": {
        "id", "product_id", "trade_date", "open_price", "high_price", "low_price",
        "close_price", "volume", "turnover_amount", "total_fund_shares", "source",
        "source_updated_at", "created_at",
    },
    "fin_nav_history": {"id", "product_id", "nav_date", "nav", "created_at"},
    "fin_holding": {
        "id", "customer_id", "trade_account", "product_id", "total_quantity",
        "shares", "available_quantity", "frozen_quantity", "average_cost", "cost_amount",
        "market_value", "current_value", "profit_loss", "profit_loss_ratio", "status",
        "first_acquired_at", "version", "updated_at",
    },
    "fin_transaction": {
        "id", "transaction_no", "order_id", "work_order_id", "customer_id", "account_id",
        "product_id", "order_side", "transaction_type", "executed_price", "nav", "executed_quantity",
        "shares", "gross_amount", "amount", "fee_rule_id", "fee_rate_snapshot", "fee_amount",
        "fee", "net_amount", "quote_at", "quote_source", "executed_at", "confirmed_at",
        "confirmed_by", "auto_confirmed", "created_at",
    },
    "fin_sim_order": {
        "id", "order_no", "customer_id", "account_id", "product_id", "order_side",
        "price_type", "quantity", "limit_price", "quote_price", "quote_at", "quote_source",
        "channel", "advisor_id", "filled_quantity", "average_executed_price", "status",
        "risk_rule_hits", "risk_disclosure_ack_at", "second_confirmation_at",
        "recording_reference", "ops_handler_id", "ops_handled_at", "compliance_handler_id",
        "compliance_handled_at", "reject_reason", "submitted_at", "cancelled_at",
        "created_at", "updated_at",
    },
    "fin_sim_account": {
        "id", "account_no", "customer_id", "currency", "cash_balance", "available_cash",
        "frozen_cash", "initial_balance", "status", "version", "created_at", "updated_at",
    },
    "fin_cash_ledger": {
        "id", "ledger_no", "account_id", "transaction_id", "entry_type", "amount",
        "balance_after", "available_cash_after", "frozen_cash_after", "idempotency_key",
        "occurred_at", "created_at",
    },
    "client_facing_content": {
        "id", "customer_id", "content_type", "draft_content", "generated_by_portal",
        "review_status", "reviewer_user_id", "reviewed_at", "published_at",
        "created_at", "updated_at",
    },
}

JOIN_SQL = {
    ("fin_transaction", "fin_product"): "t.product_id = p.id",
    ("fin_holding", "fin_product"): "h.product_id = p.id",
    ("fin_market_price", "fin_product"): "m.product_id = p.id",
    ("fin_nav_history", "fin_product"): "n.product_id = p.id",
    ("fin_sim_order", "fin_product"): "o.product_id = p.id",
    ("fin_fee_rule", "fin_product"): "f.product_id = p.id",
    ("fin_transaction", "fin_customer_profile"): "t.customer_id = cp.customer_id",
    ("fin_holding", "fin_customer_profile"): "h.customer_id = cp.customer_id",
    ("fin_sim_account", "fin_customer_profile"): "a.customer_id = cp.customer_id",
    ("fin_cash_ledger", "fin_sim_account"): "l.account_id = a.id",
    ("fin_sim_account", "fin_customer_profile"): "a.customer_id = cp.customer_id",
    ("fin_cash_ledger", "fin_sim_account"): "l.account_id = a.id",
    ("fin_transaction", "fin_sim_account"): "t.account_id = a.id",
    ("fin_transaction", "sys_customer_assignment"): "t.customer_id = ca.customer_id",
    ("fin_holding", "sys_customer_assignment"): "h.customer_id = ca.customer_id",
    ("fin_customer_profile", "sys_customer_assignment"): "cp.customer_id = ca.customer_id",
}


def _load_local_config() -> dict[str, str]:
    """加载标准环境变量，并兼容 0901/.evn 的中文标签密钥格式。"""
    values = dict(os.environ)
    path = Path(__file__).with_name(".evn")
    if not path.exists():
        return values
    raw_lines = [line.strip() for line in path.read_text(encoding="utf-8").splitlines() if line.strip()]
    for line in raw_lines:
        if "=" in line:
            key, value = line.split("=", 1)
            values.setdefault(key.strip(), value.strip().strip('"').strip("'"))
    if raw_lines:
        first_value = raw_lines[0].split("：", 1)[-1].strip() if "：" in raw_lines[0] else raw_lines[0]
        values.setdefault("DEEPSEEK_API_KEY_NL2SQL", first_value)
    if len(raw_lines) > 1:
        second_value = raw_lines[1].split("：", 1)[-1].strip() if "：" in raw_lines[1] else raw_lines[1]
        values.setdefault("ALIYUN_API_KEY", second_value)
    return values


CONFIG = _load_local_config()
DATABASE_URL = CONFIG.get("DATABASE_URL", "")
LLM_BASE_URL = CONFIG.get("DEEPSEEK_BASE_URL") or "https://api.deepseek.com"
LLM_MODEL = CONFIG.get("DEEPSEEK_NL2SQL_MODEL") or CONFIG.get("DEEPSEEK_MODEL") or "deepseek-chat"
LLM_KEY = CONFIG.get("DEEPSEEK_API_KEY_NL2SQL") or CONFIG.get("DEEPSEEK_API_KEY", "")


@dataclass
class AuthContext:
    user_id: int | None = None
    roles: list[str] = field(default_factory=list)
    allowed_domains: list[str] = field(default_factory=lambda: list(DOMAINS))
    customer_scope: str = "all"
    allowed_fields: list[str] | None = None
    masked_fields: list[str] = field(default_factory=list)
    max_rows: int = 500
    max_query_seconds: int = 10


@dataclass
class QueryRequest:
    question: str
    auth_context: AuthContext
    conversation_id: str | None = None
    request_id: str | None = None
    timezone: str = "Asia/Shanghai"
    confirmation: str | None = None


@dataclass
class QueryPlan:
    intent: str
    domains: list[str]
    tables: list[str]
    time_mode: str = "none"
    time_column: str | None = None
    start: str | None = None
    end: str | None = None
    metrics: list[str] = field(default_factory=list)
    dimensions: list[str] = field(default_factory=list)
    filters: list[dict[str, Any]] = field(default_factory=list)
    sort: list[dict[str, str]] = field(default_factory=list)
    limit: int = 50
    confidence: float = 0.0
    needs_confirmation: bool = False
    confirmation_question: str | None = None
    unsupported_reason: str | None = None


def _json_object(text_value: str) -> dict[str, Any]:
    cleaned = re.sub(r"```(?:json)?|```", "", text_value or "").strip()
    match = re.search(r"\{.*\}", cleaned, flags=re.S)
    if not match:
        raise ValueError("模型未返回结构化查询计划")
    value = json.loads(match.group(0))
    if not isinstance(value, dict):
        raise ValueError("查询计划必须是 JSON 对象")
    return value


def _days_range(days: int) -> tuple[str, str]:
    end = datetime.now()
    start = end - timedelta(days=days)
    return start.strftime("%Y-%m-%d %H:%M:%S"), end.strftime("%Y-%m-%d %H:%M:%S")


def _catalog_text(tables: set[str]) -> str:
    lines = ["只允许使用以下表和字段："]
    for table in sorted(tables):
        lines.append(f"{table}: {', '.join(sorted(TABLE_COLUMNS[table]))}")
    lines.append("指标口径：交易金额=gross_amount；净交易金额=net_amount；收盘价=close_price；基金净值=nav；")
    lines.append("成交价格=executed_price；当前持仓市值=market_value；历史资金变化=fin_cash_ledger.amount。")
    return "\n".join(lines)


def _route_keywords(question: str) -> list[str]:
    routes = []
    if any(word in question for word in ("行情", "收盘", "开盘", "最高价", "最低价")):
        routes.append("market_nav")
    if "净值" in question:
        routes.append("market_nav")
    if any(word in question for word in ("客户", "风险", "测评", "画像", "投顾", "运营")):
        routes.append("customer_risk")
    if any(word in question for word in ("产品", "基金", "费率", "适配", "风险等级")):
        routes.append("product_fee")
    if any(word in question for word in ("委托", "成交", "交易", "持仓", "账户", "资金", "现金", "盈亏")):
        routes.append("trading_account")
    if any(word in question for word in ("内容", "审核", "发布")):
        routes.append("client_content")
    return list(dict.fromkeys(routes)) or ["trading_account"]


def _llm_plan(request: QueryRequest, domains: list[str]) -> dict[str, Any]:
    try:
        from openai import OpenAI
    except ImportError as exc:
        raise RuntimeError("缺少 openai 依赖，无法调用模型") from exc
    tables = set().union(*(DOMAINS[d] for d in domains))
    system = f"""你是金融查询计划生成器。只输出 JSON，不输出解释。
用户身份已由后端确认，customer_scope={request.auth_context.customer_scope}。
{_catalog_text(tables)}
返回字段：intent, domains, tables, time_mode, time_column, start, end, metrics,
dimensions, filters, sort, limit, confidence, needs_confirmation, confirmation_question。
只从给定表和字段中选择；无法准确判断时降低 confidence。历史时点查询不能使用当前
fin_holding、fin_sim_account 或当前归属表冒充历史数据。"""
    client = OpenAI(api_key=LLM_KEY, base_url=LLM_BASE_URL)
    response = client.chat.completions.create(
        model=LLM_MODEL,
        messages=[{"role": "system", "content": system}, {"role": "user", "content": request.question}],
        temperature=0,
        response_format={"type": "json_object"},
        timeout=request.auth_context.max_query_seconds,
    )
    return _json_object(response.choices[0].message.content)


def _normalize_plan(raw: dict[str, Any], request: QueryRequest, domains: list[str]) -> QueryPlan:
    tables = [name for name in raw.get("tables", []) if name in ALLOWED_TABLES]
    if not tables:
        tables = sorted(set().union(*(DOMAINS[d] for d in domains)))
    confidence = float(raw.get("confidence", 0.0) or 0.0)
    confidence = max(0.0, min(confidence, 1.0))
    plan = QueryPlan(
        intent=str(raw.get("intent") or "unknown"),
        domains=[d for d in raw.get("domains", domains) if d in DOMAINS] or domains,
        tables=tables,
        time_mode=str(raw.get("time_mode") or raw.get("temporal", {}).get("mode") or "none"),
        time_column=raw.get("time_column") or raw.get("temporal", {}).get("time_column"),
        start=raw.get("start") or raw.get("temporal", {}).get("start"),
        end=raw.get("end") or raw.get("temporal", {}).get("end"),
        metrics=list(raw.get("metrics") or []),
        dimensions=list(raw.get("dimensions") or []),
        filters=list(raw.get("filters") or []),
        sort=list(raw.get("sort") or []),
        limit=min(max(int(raw.get("limit", 50) or 50), 1), request.auth_context.max_rows),
        confidence=confidence,
        needs_confirmation=bool(raw.get("needs_confirmation", False)),
        confirmation_question=raw.get("confirmation_question"),
        unsupported_reason=raw.get("unsupported_reason"),
    )
    if plan.confidence < 0.60:
        plan.needs_confirmation = True
        plan.confirmation_question = plan.confirmation_question or "请补充客户、产品、时间范围或指标口径。"
    elif plan.confidence < 0.85:
        plan.needs_confirmation = True
        plan.confirmation_question = plan.confirmation_question or "请确认我对查询范围和指标口径的理解是否正确。"
    return plan


def _validate_plan(plan: QueryPlan, auth: AuthContext) -> tuple[bool, str]:
    allowed_tables = set().union(*(DOMAINS[d] for d in auth.allowed_domains if d in DOMAINS))
    if not set(plan.tables).issubset(allowed_tables):
        return False, "查询包含当前角色未授权的数据表"
    if len(plan.tables) > 8:
        return False, "查询涉及的数据表过多"
    if plan.time_mode in {"range", "as_of"} and not plan.time_column:
        return False, "历史查询缺少时间字段"
    temporal_current = {"fin_holding", "fin_sim_account", "sys_customer_assignment"}
    if plan.time_mode == "as_of" and temporal_current.intersection(plan.tables):
        return False, "当前表缺少历史时点来源，暂不支持该历史时点查询"
    if plan.start and plan.end and plan.start > plan.end:
        return False, "查询开始时间不能晚于结束时间"
    return True, "计划校验通过"


def _field_allowed(table: str, field_name: str, auth: AuthContext) -> bool:
    if field_name not in TABLE_COLUMNS.get(table, set()):
        return False
    if auth.allowed_fields is None:
        return True
    return f"{table}.{field_name}" in auth.allowed_fields or field_name in auth.allowed_fields


def _safe_sql_check(sql: str, params: dict[str, Any], plan: QueryPlan, auth: AuthContext) -> tuple[bool, str]:
    normalized = re.sub(r"\s+", " ", sql.strip())
    upper = normalized.upper()
    if not upper.startswith("SELECT "):
        return False, "仅支持 SELECT 查询"
    if ";" in normalized.rstrip(";"):
        return False, "禁止执行多语句"
    if re.search(r"\b(INSERT|UPDATE|DELETE|DROP|ALTER|TRUNCATE|CREATE|GRANT|REVOKE|CALL|INTO\s+OUTFILE)\b", upper):
        return False, "检测到禁止的数据库操作"
    if "*" in normalized:
        return False, "禁止使用 SELECT *"
    table_refs = set(re.findall(r"\b(?:FROM|JOIN)\s+([A-Za-z_][A-Za-z0-9_]*)", normalized, flags=re.I))
    if not table_refs.issubset(set(plan.tables)):
        return False, "SQL 使用了计划外的数据表"
    for table in table_refs:
        aliases = re.findall(rf"\b{re.escape(table)}\s+(?:AS\s+)?([A-Za-z_][A-Za-z0-9_]*)", normalized, flags=re.I)
        del aliases
    if plan.time_mode in {"range", "as_of"} and plan.time_column and plan.time_column not in normalized:
        return False, "历史查询缺少计划中的时间条件"
    if len(params) > 30:
        return False, "查询参数过多"
    return True, "SQL 安全校验通过"


def _alias(table: str) -> str:
    return {
        "fin_transaction": "t", "fin_product": "p", "fin_holding": "h",
        "fin_market_price": "m", "fin_nav_history": "n", "fin_sim_order": "o",
        "fin_sim_account": "a", "fin_cash_ledger": "l",
        "fin_customer_profile": "cp", "sys_customer_assignment": "ca",
        "fin_risk_assessment": "ra", "fin_fee_rule": "f",
        "client_facing_content": "c",
    }.get(table, table[:1])


def _compile_sql(plan: QueryPlan, auth: AuthContext) -> tuple[str, dict[str, Any]]:
    if plan.unsupported_reason:
        raise ValueError(plan.unsupported_reason)
    primary = plan.tables[0]
    alias = _alias(primary)
    select_parts = []
    group_parts = []
    for dimension in plan.dimensions:
        if "." not in dimension:
            dimension = f"{primary}.{dimension}"
        table, column = dimension.split(".", 1)
        if not _field_allowed(table, column, auth):
            raise PermissionError(f"字段未授权：{dimension}")
        select_parts.append(f"{_alias(table)}.{column} AS {column}")
        group_parts.append(f"{_alias(table)}.{column}")
    metric_map = {
        "交易金额": "SUM(t.gross_amount) AS gross_amount",
        "gross_amount": "SUM(t.gross_amount) AS gross_amount",
        "净交易金额": "SUM(t.net_amount) AS net_amount",
        "net_amount": "SUM(t.net_amount) AS net_amount",
        "成交数量": "SUM(t.executed_quantity) AS executed_quantity",
        "持仓数量": "h.total_quantity AS total_quantity",
        "持仓市值": "h.market_value AS market_value",
        "浮动盈亏": "h.profit_loss AS profit_loss",
        "可用份额": "h.available_quantity AS available_quantity",
        "可用现金": "a.available_cash AS available_cash",
        "现金余额": "a.cash_balance AS cash_balance",
        "历史资金变化": "SUM(l.amount) AS cash_change",
        "收盘价": "m.close_price AS close_price",
        "基金总份额": "m.total_fund_shares AS total_fund_shares",
        "基金净值": "n.nav AS nav",
        "成交价格": "t.executed_price AS executed_price",
        "客户数": "COUNT(DISTINCT cp.customer_id) AS customer_count",
    }
    for metric in plan.metrics or ["客户数"]:
        expression = metric_map.get(str(metric))
        if not expression:
            raise ValueError(f"暂不支持指标：{metric}")
        select_parts.append(expression)
    if not select_parts:
        raise ValueError("查询计划没有可返回的字段")
    params: dict[str, Any] = {}
    where = ["1=1"]
    if primary == "fin_transaction":
        where.append("t.order_side IN ('买入', '卖出')")
    if plan.time_column and plan.start:
        params["start_time"] = plan.start
        where.append(f"{alias}.{plan.time_column} >= :start_time")
    if plan.time_column and plan.end:
        params["end_time"] = plan.end
        where.append(f"{alias}.{plan.time_column} <= :end_time")
    for index, item in enumerate(plan.filters):
        field_name = item.get("field")
        operator = str(item.get("operator", "=")).upper()
        if not isinstance(field_name, str) or "." not in field_name or operator not in {"=", "!=", ">", ">=", "<", "<=", "LIKE"}:
            raise ValueError("查询筛选条件不合法")
        table, column = field_name.split(".", 1)
        if not _field_allowed(table, column, auth):
            raise PermissionError(f"字段未授权：{field_name}")
        key = f"filter_{index}"
        params[key] = item.get("value")
        where.append(f"{_alias(table)}.{column} {operator} :{key}")
    from_sql = f"{primary} {alias}"
    pending = set(plan.tables) - {primary}
    joined = {primary}
    while pending:
        progress = False
        for table in sorted(pending):
            join = None
            for existing in joined:
                join = JOIN_SQL.get((existing, table)) or JOIN_SQL.get((table, existing))
                if join:
                    break
            if not join:
                continue
            from_sql += f" JOIN {table} {_alias(table)} ON {join}"
            joined.add(table)
            pending.remove(table)
            progress = True
        if not progress:
            break
    if pending:
        raise ValueError(f"缺少合法 Join 路径：{', '.join(sorted(pending))}")
    sql = f"SELECT {', '.join(select_parts)} FROM {from_sql} WHERE {' AND '.join(where)}"
    if group_parts:
        sql += f" GROUP BY {', '.join(group_parts)}"
    if plan.sort:
        sort_items = []
        for item in plan.sort:
            field = str(item.get("field", "")).split(".")[-1]
            direction = "DESC" if str(item.get("direction", "desc")).lower() == "desc" else "ASC"
            if field not in {part.split(" AS ")[-1] for part in select_parts}:
                continue
            sort_items.append(f"{field} {direction}")
        if sort_items:
            sql += " ORDER BY " + ", ".join(sort_items)
    sql += f" LIMIT {min(plan.limit, auth.max_rows)}"
    return sql, params


def _mock_plan(request: QueryRequest, domains: list[str]) -> dict[str, Any]:
    """无模型或测试时的保守规则，便于单元测试和离线联调。"""
    question = request.question
    if "近30天" in question or "最近30天" in question:
        start, end = _days_range(30)
    else:
        start = end = None
    if "收盘" in question or "行情" in question:
        return {"intent": "market_price", "domains": ["market_nav"], "tables": ["fin_market_price", "fin_product"],
                "time_mode": "range", "time_column": "trade_date", "start": start, "end": end,
                "metrics": ["收盘价"], "dimensions": ["fin_product.product_name", "fin_market_price.trade_date"],
                "confidence": 0.88}
    if "净值" in question:
        return {"intent": "nav_history", "domains": ["market_nav"], "tables": ["fin_nav_history", "fin_product"],
                "time_mode": "range", "time_column": "nav_date", "start": start, "end": end,
                "metrics": ["基金净值"], "dimensions": ["fin_product.product_name", "fin_nav_history.nav_date"],
                "confidence": 0.88}
    if "当前持仓" in question or "持仓市值" in question:
        return {"intent": "current_holding", "domains": ["customer_risk", "trading_account"],
                "tables": ["fin_holding", "fin_product", "fin_customer_profile"],
                "metrics": ["持仓市值"], "dimensions": ["fin_customer_profile.real_name", "fin_product.product_name"],
                "confidence": 0.87}
    if "账户余额" in question or "可用现金" in question:
        return {"intent": "current_account", "domains": ["trading_account"],
                "tables": ["fin_sim_account", "fin_customer_profile"], "metrics": ["可用现金"],
                "dimensions": ["fin_customer_profile.real_name"], "confidence": 0.87}
    if "资金变动" in question or "资金流水" in question:
        return {"intent": "cash_ledger", "domains": ["trading_account"],
                "tables": ["fin_cash_ledger", "fin_sim_account"], "time_mode": "range",
                "time_column": "occurred_at", "start": start, "end": end,
                "metrics": ["历史资金变化"], "dimensions": ["fin_cash_ledger.occurred_at"], "confidence": 0.87}
    return {"intent": "unknown", "domains": domains, "tables": sorted(set().union(*(DOMAINS[d] for d in domains))),
            "confidence": 0.45, "needs_confirmation": True,
            "confirmation_question": "请明确要查询的业务对象、指标和时间范围。"}


def _offsite_mock_plan(request: QueryRequest, domains: list[str]) -> dict[str, Any] | None:
    """为场外核对补充确定性查询计划，仍复用统一 SQL 安全校验。"""
    question = request.question
    if "基金代码" not in question:
        return None
    fund_match = re.search(r"基金代码(?:为|是)?\s*([A-Za-z0-9_-]+)", question)
    if fund_match is None:
        return None
    filters: list[dict[str, Any]] = [{
        "field": "fin_product.product_code",
        "operator": "=",
        "value": fund_match.group(1),
    }]
    account_match = re.search(r"账户标识(?:为|是)?\s*([^，,。；;\s]+)", question)
    account_filter = {
            "field": "fin_holding.trade_account",
            "operator": "=",
            "value": account_match.group(1),
        } if account_match is not None else None
    common = {
        "domains": ["market_nav", "trading_account"],
        "filters": filters,
        "confidence": 0.99,
        "needs_confirmation": False,
        "limit": 1,
    }
    if "最新总份额" in question and "申请前持有份额" in question:
        holding_filters = filters + ([account_filter] if account_filter else [])
        return {
            **common,
            "intent": "offsite_subscription_holding_check",
            "tables": ["fin_market_price", "fin_nav_history", "fin_holding", "fin_product"],
            "metrics": ["基金总份额", "基金净值", "持仓数量"],
            "filters": holding_filters,
        }
    if "最新总份额" in question and "最新净值" in question:
        return {
            **common,
            "intent": "offsite_fund_limit_check",
            "tables": ["fin_market_price", "fin_nav_history", "fin_product"],
            "metrics": ["基金总份额", "基金净值"],
        }
    if "当前最新可用份额" in question or "可用份额" in question:
        holding_filters = filters + ([account_filter] if account_filter else [])
        return {
            **common,
            "intent": "offsite_redemption_available_check",
            "tables": ["fin_holding", "fin_product"],
            "metrics": ["可用份额"],
            "filters": holding_filters,
        }
    if "最新净值" in question:
        return {
            **common,
            "intent": "offsite_latest_nav",
            "tables": ["fin_nav_history", "fin_product"],
            "metrics": ["基金净值"],
        }
    if "最新总份额" in question:
        return {
            **common,
            "intent": "offsite_latest_total_shares",
            "tables": ["fin_market_price", "fin_product"],
            "metrics": ["基金总份额"],
        }
    return None


def _audit_payload(request: QueryRequest, plan: QueryPlan, sql: str | None, params: dict[str, Any],
                   validation: dict[str, Any], execution: dict[str, Any]) -> dict[str, Any]:
    return {
        "tool": "nl2sql", "engine_version": "mvp-1", "request_id": request.request_id,
        "conversation_id": request.conversation_id, "intent": plan.intent,
        "confidence": plan.confidence, "domains": plan.domains,
        "query_plan": asdict(plan), "generated_sql": sql,
        "parameters": {key: "<redacted>" if "password" in key.lower() else value for key, value in params.items()},
        "authorized_context": {"user_id": request.auth_context.user_id, "roles": request.auth_context.roles,
                               "customer_scope": request.auth_context.customer_scope},
        "validation": validation, "execution": execution,
        "created_at": datetime.now().isoformat(timespec="seconds"),
    }


def write_conversation_tool_calls(connection: Any, request: QueryRequest, audit: dict[str, Any]) -> None:
    """将审计对象写入已有 conversation_message；宿主也可使用 audit_writer 自行落库。"""
    from sqlalchemy import text
    connection.execute(
        text("""INSERT INTO conversation_message
        (session_id, customer_id, portal, role, content, tool_calls, intent, confidence, created_at)
        VALUES (:session_id, :customer_id, :portal, 'assistant', :content, :tool_calls,
                :intent, :confidence, :created_at)"""),
        {
            "session_id": request.conversation_id or request.request_id,
            "customer_id": request.auth_context.user_id,
            "portal": "nl2sql",
            "content": audit.get("execution", {}).get("status", "nl2sql"),
            "tool_calls": json.dumps(audit, ensure_ascii=False, default=str),
            "intent": audit.get("intent"),
            "confidence": audit.get("confidence"),
            "created_at": datetime.now(),
        },
    )


def query(request: QueryRequest, *, db_engine: Any = None,
          audit_writer: Callable[[dict[str, Any]], None] | None = None,
          use_llm: bool = True, persist_audit: bool = False) -> dict[str, Any]:
    """统一入口：返回结构化结果，供运营和投顾 Agent 直接调用。"""
    if not request.question or not request.question.strip():
        return {"status": "rejected", "message": "查询问题不能为空"}
    request.request_id = request.request_id or str(uuid.uuid4())
    domains = _route_keywords(request.question)
    try:
        raw = (
            _llm_plan(request, domains)
            if use_llm and LLM_KEY
            else _offsite_mock_plan(request, domains) or _mock_plan(request, domains)
        )
        plan = _normalize_plan(raw, request, domains)
        if request.confirmation and plan.needs_confirmation:
            plan.needs_confirmation = False
            plan.confidence = max(plan.confidence, 0.85)
        valid, message = _validate_plan(plan, request.auth_context)
        if not valid:
            audit = _audit_payload(request, plan, None, {}, {"valid": False, "message": message}, {"status": "rejected"})
            if audit_writer:
                audit_writer(audit)
            if persist_audit and db_engine is not None:
                with db_engine.begin() as connection:
                    write_conversation_tool_calls(connection, request, audit)
            return {"status": "rejected", "message": message, "audit": audit}
        if plan.needs_confirmation:
            audit = _audit_payload(request, plan, None, {}, {"valid": True, "message": "等待确认"}, {"status": "waiting"})
            if audit_writer:
                audit_writer(audit)
            if persist_audit and db_engine is not None:
                with db_engine.begin() as connection:
                    write_conversation_tool_calls(connection, request, audit)
            return {"status": "need_confirmation", "message": plan.confirmation_question, "query_plan": asdict(plan), "audit": audit}
        sql, params = _compile_sql(plan, request.auth_context)
        safe, safety_message = _safe_sql_check(sql, params, plan, request.auth_context)
        if not safe:
            raise PermissionError(safety_message)
        if db_engine is None:
            execution = {"status": "not_executed", "reason": "未提供数据库连接，仅返回已校验 SQL"}
            audit = _audit_payload(request, plan, sql, params, {"valid": True, "message": safety_message}, execution)
            if audit_writer:
                audit_writer(audit)
            if persist_audit and db_engine is not None:
                with db_engine.begin() as connection:
                    write_conversation_tool_calls(connection, request, audit)
            return {"status": "ready", "sql": sql, "parameters": params, "query_plan": asdict(plan), "audit": audit}
        from sqlalchemy import text
        with db_engine.connect() as connection:
            result = connection.execute(text(sql), params)
            columns = list(result.keys())
            rows = [dict(zip(columns, row)) for row in result.fetchmany(request.auth_context.max_rows)]
        execution = {"status": "success", "row_count": len(rows), "truncated": len(rows) >= request.auth_context.max_rows}
        audit = _audit_payload(request, plan, sql, params, {"valid": True, "message": safety_message}, execution)
        if audit_writer:
            audit_writer(audit)
        if persist_audit:
            with db_engine.begin() as connection:
                write_conversation_tool_calls(connection, request, audit)
        return {"status": "success", "data": {"total": len(rows), "rows": rows},
                "query_plan": asdict(plan), "sql": sql, "audit": audit}
    except (ValueError, PermissionError, RuntimeError) as exc:
        logger.warning("NL2SQL业务失败：%s", exc)
        return {"status": "error", "message": str(exc), "request_id": request.request_id}
    except Exception:
        logger.exception("NL2SQL执行失败")
        return {"status": "error", "message": "查询执行失败，请稍后重试", "request_id": request.request_id}


def build_request(question: str, auth_context: dict[str, Any], **kwargs: Any) -> QueryRequest:
    """将 Agent 的字典请求转换为统一请求对象。"""
    return QueryRequest(question=question, auth_context=AuthContext(**auth_context), **kwargs)


def query_dict(question: str, auth_context: dict[str, Any], *,
               db_engine: Any = None,
               audit_writer: Callable[[dict[str, Any]], None] | None = None,
               use_llm: bool = True, persist_audit: bool = False,
               **request_kwargs: Any) -> dict[str, Any]:
    """给 Agent 使用的字典式快捷入口。"""
    return query(
        build_request(question, auth_context, **request_kwargs),
        db_engine=db_engine,
        audit_writer=audit_writer,
        use_llm=use_llm,
        persist_audit=persist_audit,
    )


def get_tool_definition() -> dict[str, Any]:
    """返回可注册到运营或投顾 Agent 的统一工具定义。"""
    return {
        "name": "financial_nl2sql",
        "description": "对金融业务数据执行只读自然语言查询；低置信度时先确认。",
        "input_schema": {
            "type": "object",
            "required": ["question", "auth_context"],
            "properties": {
                "question": {"type": "string"},
                "auth_context": {
                    "type": "object",
                    "required": ["roles", "customer_scope"],
                    "properties": {
                        "user_id": {"type": ["integer", "null"]},
                        "roles": {"type": "array", "items": {"type": "string"}},
                        "customer_scope": {"type": "string", "enum": ["self", "own_customers", "all"]},
                        "allowed_domains": {"type": "array", "items": {"type": "string"}},
                    },
                },
                "conversation_id": {"type": ["string", "null"]},
                "request_id": {"type": ["string", "null"]},
                "timezone": {"type": "string"},
                "confirmation": {"type": ["string", "null"]},
            },
        },
    }


if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO)
    demo = build_request("查询最近30天的行情收盘价", {"roles": ["advisor"], "customer_scope": "all"})
    print(json.dumps(query(demo, use_llm=False), ensure_ascii=False, indent=2, default=str))