500 lines
22 KiB
Python
500 lines
22 KiB
Python
from __future__ import annotations
|
|
|
|
import re
|
|
from collections.abc import Callable
|
|
from contextlib import AbstractAsyncContextManager
|
|
from datetime import UTC, datetime, timedelta
|
|
from decimal import Decimal
|
|
from typing import Any
|
|
|
|
from sqlalchemy import text
|
|
from sqlalchemy.ext.asyncio import AsyncSession
|
|
|
|
from app.core.contracts import RequestContext
|
|
from app.core.errors import ForbiddenAgentError, ValidationAgentError
|
|
from app.core.nl2sql_catalog import (
|
|
ALLOWED_TABLES,
|
|
BANNED_SQL,
|
|
CURRENT_ONLY_TABLES,
|
|
DOMAINS,
|
|
JOIN_SQL,
|
|
TABLE_COLUMNS,
|
|
alias,
|
|
)
|
|
from app.core.nl2sql_contracts import (
|
|
FinancialNL2SQLInput,
|
|
FinancialNL2SQLResult,
|
|
FinancialQueryPlan,
|
|
)
|
|
from app.infrastructure.db import SessionFactory
|
|
|
|
SessionFactoryType = Callable[[], AbstractAsyncContextManager[AsyncSession]]
|
|
|
|
|
|
def _range(question: str) -> tuple[str | None, str | None]:
|
|
days = 0
|
|
if "近7天" in question or "最近7天" in question:
|
|
days = 7
|
|
elif "近30天" in question or "最近30天" in question or "近一个月" in question:
|
|
days = 30
|
|
elif "近90天" in question or "最近90天" in question or "近三个月" in question:
|
|
days = 90
|
|
if not days:
|
|
return None, None
|
|
end = datetime.now(UTC).replace(tzinfo=None)
|
|
start = end - timedelta(days=days)
|
|
return start.strftime("%Y-%m-%d %H:%M:%S"), end.strftime("%Y-%m-%d %H:%M:%S")
|
|
|
|
|
|
def _filters(question: str) -> tuple[dict[str, Any], ...]:
|
|
return tuple(
|
|
{"field": "fin_product.product_code", "operator": "=", "value": code}
|
|
for code in sorted(set(re.findall(r"(?<!\d)\d{6}(?!\d)", question)))
|
|
)
|
|
|
|
|
|
class RuleBasedFinancialPlanner:
|
|
def plan(self, query: FinancialNL2SQLInput) -> FinancialQueryPlan:
|
|
question = query.question
|
|
start, end = _range(question)
|
|
time_mode = "range" if start and end else "none"
|
|
if any(word in question for word in ("截至", "历史时点", "当时", "某日")):
|
|
time_mode = "as_of"
|
|
filters = _filters(question)
|
|
for builder in (self._market_plan, self._trade_plan, self._profile_plan):
|
|
plan = builder(query, time_mode, start, end, filters)
|
|
if plan is not None:
|
|
return plan
|
|
return FinancialQueryPlan(
|
|
intent="unknown", domains=("customer_risk",), tables=("fin_customer_profile",),
|
|
limit=query.limit, confidence=0.45, needs_confirmation=True,
|
|
confirmation_question="请明确要查询的客户、产品、时间范围或指标口径。",
|
|
)
|
|
|
|
@staticmethod
|
|
def _market_plan(
|
|
query: FinancialNL2SQLInput,
|
|
time_mode: str,
|
|
start: str | None,
|
|
end: str | None,
|
|
filters: tuple[dict[str, Any], ...],
|
|
) -> FinancialQueryPlan | None:
|
|
question = query.question
|
|
if "净值" in question:
|
|
return FinancialQueryPlan(
|
|
intent="nav_history_query", domains=("market_nav", "product_fee"),
|
|
tables=("fin_nav_history", "fin_product"), metrics=("基金净值",),
|
|
dimensions=("fin_product.product_code", "fin_product.product_name",
|
|
"fin_nav_history.nav_date"),
|
|
filters=filters, time_mode=time_mode, time_column="nav_date",
|
|
start=start, end=end, limit=query.limit, confidence=0.90,
|
|
)
|
|
if any(word in question for word in ("行情", "收盘", "开盘", "最高价", "最低价")):
|
|
return FinancialQueryPlan(
|
|
intent="market_price_query", domains=("market_nav", "product_fee"),
|
|
tables=("fin_market_price", "fin_product"), metrics=("收盘价",),
|
|
dimensions=("fin_product.product_code", "fin_product.product_name",
|
|
"fin_market_price.trade_date"),
|
|
filters=filters, time_mode=time_mode, time_column="trade_date",
|
|
start=start, end=end, limit=query.limit, confidence=0.90,
|
|
)
|
|
return None
|
|
|
|
@staticmethod
|
|
def _trade_plan(
|
|
query: FinancialNL2SQLInput,
|
|
time_mode: str,
|
|
start: str | None,
|
|
end: str | None,
|
|
filters: tuple[dict[str, Any], ...],
|
|
) -> FinancialQueryPlan | None:
|
|
question = query.question
|
|
if "资金" in question or "现金流水" in question:
|
|
cash_metrics = (
|
|
("历史资金变化",)
|
|
if any(w in question for w in ("汇总", "合计", "变化"))
|
|
else ()
|
|
)
|
|
return FinancialQueryPlan(
|
|
intent="cash_ledger_query", domains=("trading_account", "customer_risk"),
|
|
tables=("fin_cash_ledger", "fin_sim_account", "fin_customer_profile"),
|
|
metrics=cash_metrics,
|
|
dimensions=("fin_cash_ledger.occurred_at", "fin_cash_ledger.entry_type"),
|
|
filters=filters, time_mode=time_mode, time_column="occurred_at",
|
|
start=start, end=end, limit=query.limit, confidence=0.88,
|
|
)
|
|
if "持仓" in question:
|
|
return FinancialQueryPlan(
|
|
intent="holding_query", domains=("trading_account", "customer_risk", "product_fee"),
|
|
tables=("fin_holding", "fin_customer_profile", "fin_product"),
|
|
metrics=("持仓市值", "浮动盈亏") if "盈亏" in question else ("持仓市值",),
|
|
dimensions=("fin_customer_profile.customer_id", "fin_product.product_code",
|
|
"fin_product.product_name"),
|
|
filters=filters, time_mode=time_mode, time_column="updated_at",
|
|
start=start, end=end, limit=query.limit, confidence=0.88,
|
|
)
|
|
if any(word in question for word in ("成交", "交易")):
|
|
transaction_metrics = (
|
|
("交易金额", "成交数量")
|
|
if any(w in question for w in ("统计", "汇总", "合计"))
|
|
else ()
|
|
)
|
|
return FinancialQueryPlan(
|
|
intent="transaction_query",
|
|
domains=("trading_account", "customer_risk", "product_fee"),
|
|
tables=("fin_transaction", "fin_customer_profile", "fin_product"),
|
|
metrics=transaction_metrics,
|
|
dimensions=("fin_customer_profile.customer_id", "fin_product.product_code",
|
|
"fin_transaction.order_side"),
|
|
filters=filters, time_mode=time_mode, time_column="executed_at",
|
|
start=start, end=end, limit=query.limit, confidence=0.87,
|
|
)
|
|
if "委托" in question or "订单" in question:
|
|
return FinancialQueryPlan(
|
|
intent="order_query", domains=("trading_account", "customer_risk", "product_fee"),
|
|
tables=("fin_sim_order", "fin_customer_profile", "fin_product"),
|
|
dimensions=("fin_customer_profile.customer_id", "fin_product.product_code",
|
|
"fin_sim_order.status"),
|
|
filters=filters, time_mode=time_mode, time_column="submitted_at",
|
|
start=start, end=end, limit=query.limit, confidence=0.86,
|
|
)
|
|
if "账户" in question or "余额" in question or "现金" in question:
|
|
return FinancialQueryPlan(
|
|
intent="account_query", domains=("trading_account", "customer_risk"),
|
|
tables=("fin_sim_account", "fin_customer_profile"),
|
|
metrics=("现金余额", "可用现金"),
|
|
dimensions=("fin_customer_profile.customer_id",),
|
|
time_mode=time_mode,
|
|
time_column="updated_at",
|
|
start=start,
|
|
end=end,
|
|
limit=query.limit,
|
|
confidence=0.88,
|
|
)
|
|
return None
|
|
|
|
@staticmethod
|
|
def _profile_plan(
|
|
query: FinancialNL2SQLInput,
|
|
time_mode: str,
|
|
start: str | None,
|
|
end: str | None,
|
|
filters: tuple[dict[str, Any], ...],
|
|
) -> FinancialQueryPlan | None:
|
|
question = query.question
|
|
if "怎么样" in question:
|
|
return FinancialQueryPlan(
|
|
intent="unknown", domains=("customer_risk",),
|
|
tables=("fin_customer_profile",), limit=query.limit,
|
|
confidence=0.45, needs_confirmation=True,
|
|
confirmation_question="请明确要查询客户画像、风险测评、交易、持仓还是账户信息。",
|
|
)
|
|
if "费率" in question or "费用" in question:
|
|
return FinancialQueryPlan(
|
|
intent="fee_rule_query", domains=("product_fee",),
|
|
tables=("fin_fee_rule", "fin_product"),
|
|
dimensions=("fin_product.product_code", "fin_fee_rule.order_side",
|
|
"fin_fee_rule.customer_tier"),
|
|
filters=filters, limit=query.limit, confidence=0.87,
|
|
)
|
|
if "内容" in question or "审核" in question or "发布" in question:
|
|
return FinancialQueryPlan(
|
|
intent="client_content_query", domains=("client_content", "customer_risk"),
|
|
tables=("client_facing_content", "fin_customer_profile"),
|
|
dimensions=(
|
|
"fin_customer_profile.customer_id",
|
|
"client_facing_content.content_type",
|
|
"client_facing_content.review_status",
|
|
),
|
|
time_mode=time_mode, time_column="created_at", start=start, end=end,
|
|
limit=query.limit, confidence=0.84, needs_confirmation=True,
|
|
confirmation_question="请确认要查询的是对客内容的审核或发布记录。",
|
|
)
|
|
if "风险" in question or "测评" in question or "画像" in question or "客户" in question:
|
|
return FinancialQueryPlan(
|
|
intent="customer_risk_query", domains=("customer_risk",),
|
|
tables=("fin_risk_assessment", "fin_customer_profile"),
|
|
metrics=("客户数",) if any(w in question for w in ("多少", "数量", "统计")) else (),
|
|
dimensions=(
|
|
"fin_customer_profile.customer_id",
|
|
"fin_customer_profile.investor_type",
|
|
),
|
|
time_mode=time_mode, time_column="assessed_at", start=start, end=end,
|
|
limit=query.limit, confidence=0.86,
|
|
)
|
|
return None
|
|
|
|
|
|
class FinancialNL2SQLService:
|
|
def __init__(
|
|
self,
|
|
planner: RuleBasedFinancialPlanner | None = None,
|
|
session_factory: SessionFactoryType = SessionFactory,
|
|
) -> None:
|
|
self.planner = planner or RuleBasedFinancialPlanner()
|
|
self.session_factory = session_factory
|
|
|
|
async def query(
|
|
self, arguments: FinancialNL2SQLInput, context: RequestContext
|
|
) -> dict[str, Any]:
|
|
self._require_context(context)
|
|
plan = self.planner.plan(arguments)
|
|
if arguments.confirmation and plan.needs_confirmation:
|
|
plan = plan.model_copy(update={
|
|
"needs_confirmation": False, "confidence": max(0.85, plan.confidence),
|
|
})
|
|
valid, message = self._validate_plan(plan, context)
|
|
if not valid:
|
|
return self._result("rejected", message, arguments, context, plan, None, {}, 0)
|
|
if plan.needs_confirmation or plan.confidence < 0.85:
|
|
question = plan.confirmation_question or "请确认查询范围、指标口径和时间条件。"
|
|
return self._result(
|
|
"need_confirmation", question, arguments, context, plan, None, {}, 0
|
|
)
|
|
sql, params = self._compile_sql(plan, context)
|
|
self._safe_sql_check(sql, params, plan)
|
|
if arguments.dry_run:
|
|
return self._result(
|
|
"ready", "SQL 已生成并通过只读校验",
|
|
arguments, context, plan, sql, params, 0,
|
|
)
|
|
rows = await self._execute(sql, params)
|
|
result = self._result(
|
|
"success", "查询成功", arguments, context, plan, sql, params, len(rows)
|
|
)
|
|
result["data"] = {"total": len(rows), "rows": rows}
|
|
return result
|
|
|
|
@staticmethod
|
|
def _require_context(context: RequestContext) -> None:
|
|
if "financial:nl2sql:read" not in context.permissions:
|
|
raise ForbiddenAgentError("缺少金融 NL2SQL 查询权限")
|
|
if not {"advisor", "operator", "admin", "super_admin"}.intersection(context.roles):
|
|
raise ForbiddenAgentError("当前角色不能使用金融 NL2SQL 查询")
|
|
|
|
@staticmethod
|
|
def _validate_plan(plan: FinancialQueryPlan, context: RequestContext) -> tuple[bool, str]:
|
|
if len(plan.domains) > 3:
|
|
return False, "最多支持跨三个业务域查询"
|
|
if not set(plan.tables).issubset(ALLOWED_TABLES):
|
|
return False, "查询包含未纳入 NL2SQL 范围的数据表"
|
|
allowed = set().union(*(DOMAINS[d] for d in plan.domains if d in DOMAINS))
|
|
if not set(plan.tables).issubset(allowed):
|
|
return False, "查询计划的数据表和业务域不匹配"
|
|
if plan.time_mode == "as_of" and CURRENT_ONLY_TABLES.intersection(plan.tables):
|
|
return False, "当前快照表缺少历史版本，暂不支持该历史时点查询"
|
|
if context.data_scope not in {"self", "own_customers", "all"}:
|
|
return False, "未知的数据权限范围"
|
|
return True, "计划校验通过"
|
|
|
|
def _compile_sql(
|
|
self, plan: FinancialQueryPlan, context: RequestContext
|
|
) -> tuple[str, dict[str, Any]]:
|
|
select_parts: list[str] = []
|
|
group_parts: list[str] = []
|
|
for dimension in plan.dimensions:
|
|
table, column = self._split_field(dimension, plan.tables[0])
|
|
self._ensure_column(table, column)
|
|
select_parts.append(f"{alias(table)}.{column} AS {column}")
|
|
group_parts.append(f"{alias(table)}.{column}")
|
|
for metric in plan.metrics:
|
|
select_parts.append(self._metric_expression(metric))
|
|
if not select_parts:
|
|
select_parts = self._default_select(plan.tables[0])
|
|
params: dict[str, Any] = {}
|
|
where = ["1=1"]
|
|
base_alias = alias(plan.tables[0])
|
|
if plan.time_column and plan.start:
|
|
where.append(f"{base_alias}.{plan.time_column} >= :start_time")
|
|
params["start_time"] = plan.start
|
|
if plan.time_column and plan.end:
|
|
where.append(f"{base_alias}.{plan.time_column} <= :end_time")
|
|
params["end_time"] = plan.end
|
|
for index, item in enumerate(plan.filters):
|
|
table, column = self._split_field(str(item.get("field", "")), plan.tables[0])
|
|
self._ensure_column(table, column)
|
|
operator = str(item.get("operator", "=")).upper()
|
|
if operator not in {"=", "!=", ">", ">=", "<", "<=", "LIKE"}:
|
|
raise ValidationAgentError("查询筛选条件不合法")
|
|
key = f"filter_{index}"
|
|
where.append(f"{alias(table)}.{column} {operator} :{key}")
|
|
params[key] = item.get("value")
|
|
self._append_customer_scope(plan, context, where, params)
|
|
sql = (
|
|
f"SELECT {', '.join(select_parts)} FROM {self._join(plan.tables)} "
|
|
f"WHERE {' AND '.join(where)}"
|
|
)
|
|
if plan.metrics and group_parts:
|
|
sql += f" GROUP BY {', '.join(group_parts)}"
|
|
return f"{sql} LIMIT {plan.limit}", params
|
|
|
|
@staticmethod
|
|
def _metric_expression(metric: str) -> str:
|
|
mapping = {
|
|
"交易金额": "SUM(t.gross_amount) AS gross_amount",
|
|
"成交数量": "SUM(t.executed_quantity) AS executed_quantity",
|
|
"持仓市值": "h.market_value AS market_value",
|
|
"浮动盈亏": "h.profit_loss AS profit_loss",
|
|
"现金余额": "a.cash_balance AS cash_balance",
|
|
"可用现金": "a.available_cash AS available_cash",
|
|
"历史资金变化": "SUM(l.amount) AS cash_change",
|
|
"收盘价": "m.close_price AS close_price",
|
|
"基金净值": "n.nav AS nav",
|
|
"客户数": "COUNT(DISTINCT cp.customer_id) AS customer_count",
|
|
}
|
|
if metric not in mapping:
|
|
raise ValidationAgentError(f"暂不支持指标：{metric}")
|
|
return mapping[metric]
|
|
|
|
@staticmethod
|
|
def _default_select(primary: str) -> list[str]:
|
|
defaults = {
|
|
"fin_fee_rule": [
|
|
"p.product_code AS product_code",
|
|
"f.order_side AS order_side",
|
|
"f.fee_rate AS fee_rate",
|
|
],
|
|
"fin_sim_order": [
|
|
"o.order_no AS order_no",
|
|
"o.status AS status",
|
|
"o.submitted_at AS submitted_at",
|
|
],
|
|
"client_facing_content": [
|
|
"c.content_type AS content_type",
|
|
"c.review_status AS review_status",
|
|
],
|
|
"fin_transaction": [
|
|
"t.transaction_no AS transaction_no",
|
|
"t.gross_amount AS gross_amount",
|
|
],
|
|
}
|
|
return defaults.get(primary, [f"{alias(primary)}.id AS id"])
|
|
|
|
@staticmethod
|
|
def _split_field(field: str, default_table: str) -> tuple[str, str]:
|
|
if "." not in field:
|
|
return default_table, field
|
|
table, column = field.split(".", 1)
|
|
return table, column
|
|
|
|
@staticmethod
|
|
def _ensure_column(table: str, column: str) -> None:
|
|
if column not in TABLE_COLUMNS.get(table, set()):
|
|
raise ValidationAgentError("查询字段不在白名单内")
|
|
|
|
@staticmethod
|
|
def _join(tables: tuple[str, ...]) -> str:
|
|
primary = tables[0]
|
|
joined = {primary}
|
|
pending = set(tables) - joined
|
|
sql = f"{primary} {alias(primary)}"
|
|
while pending:
|
|
for table in sorted(pending):
|
|
condition = next(
|
|
(
|
|
JOIN_SQL.get((existing, table)) or JOIN_SQL.get((table, existing))
|
|
for existing in joined
|
|
if JOIN_SQL.get((existing, table)) or JOIN_SQL.get((table, existing))
|
|
),
|
|
None,
|
|
)
|
|
if condition:
|
|
sql += f" JOIN {table} {alias(table)} ON {condition}"
|
|
joined.add(table)
|
|
pending.remove(table)
|
|
break
|
|
else:
|
|
raise ValidationAgentError("查询涉及的表缺少合法关联路径")
|
|
return sql
|
|
|
|
@staticmethod
|
|
def _append_customer_scope(
|
|
plan: FinancialQueryPlan, context: RequestContext, where: list[str], params: dict[str, Any]
|
|
) -> None:
|
|
aliases = [
|
|
alias(table)
|
|
for table in plan.tables
|
|
if "customer_id" in TABLE_COLUMNS.get(table, set())
|
|
]
|
|
if not aliases or context.data_scope == "all":
|
|
return
|
|
customer_ids = tuple(context.customer_ids) or (context.user_id,)
|
|
keys = []
|
|
for index, customer_id in enumerate(customer_ids):
|
|
key = f"scope_customer_{index}"
|
|
keys.append(f":{key}")
|
|
params[key] = int(customer_id)
|
|
where.append(f"{aliases[0]}.customer_id IN ({', '.join(keys)})")
|
|
|
|
@staticmethod
|
|
def _safe_sql_check(sql: str, params: dict[str, Any], plan: FinancialQueryPlan) -> None:
|
|
normalized = re.sub(r"\s+", " ", sql.strip())
|
|
if not normalized.upper().startswith("SELECT "):
|
|
raise ForbiddenAgentError("仅支持 SELECT 查询")
|
|
if ";" in normalized.rstrip(";") or BANNED_SQL.search(normalized):
|
|
raise ForbiddenAgentError("检测到非只读数据库操作")
|
|
if "*" in normalized:
|
|
raise ForbiddenAgentError("禁止使用 SELECT *")
|
|
refs = set(re.findall(r"\b(?:FROM|JOIN)\s+([A-Za-z_][A-Za-z0-9_]*)", normalized, re.I))
|
|
if refs != set(plan.tables):
|
|
raise ForbiddenAgentError("SQL 使用的数据表与查询计划不一致")
|
|
if len(params) > 30:
|
|
raise ForbiddenAgentError("查询参数过多")
|
|
|
|
async def _execute(self, sql: str, params: dict[str, Any]) -> list[dict[str, Any]]:
|
|
async with self.session_factory() as session:
|
|
result = await session.execute(text(sql), params)
|
|
return [self._jsonable(dict(row)) for row in result.mappings().all()]
|
|
|
|
@staticmethod
|
|
def _jsonable(row: dict[str, Any]) -> dict[str, Any]:
|
|
converted: dict[str, Any] = {}
|
|
for key, value in row.items():
|
|
if isinstance(value, Decimal):
|
|
converted[key] = str(value)
|
|
elif isinstance(value, datetime):
|
|
converted[key] = value.isoformat(sep=" ", timespec="seconds")
|
|
else:
|
|
converted[key] = value
|
|
return converted
|
|
|
|
@staticmethod
|
|
def _result(
|
|
status: str,
|
|
message: str,
|
|
arguments: FinancialNL2SQLInput,
|
|
context: RequestContext,
|
|
plan: FinancialQueryPlan,
|
|
sql: str | None,
|
|
params: dict[str, Any],
|
|
row_count: int,
|
|
) -> dict[str, Any]:
|
|
audit = {
|
|
"tool": "query_financial_data",
|
|
"engine_version": "mvp-1",
|
|
"trace_id": context.trace_id,
|
|
"question": arguments.question,
|
|
"query_plan": plan.model_dump(mode="json"),
|
|
"generated_sql": sql,
|
|
"parameters": params,
|
|
"permission_check": {
|
|
"status": "passed" if status in {"ready", "success"} else status,
|
|
"roles": list(context.roles),
|
|
"data_scope": context.data_scope,
|
|
"permission": "financial:nl2sql:read",
|
|
"allowed_tables": list(plan.tables),
|
|
},
|
|
"execution": {"status": status, "row_count": row_count},
|
|
"created_at": datetime.now(UTC).replace(tzinfo=None).isoformat(timespec="seconds"),
|
|
}
|
|
return FinancialNL2SQLResult(
|
|
status=status, message=message, query_plan=plan.model_dump(mode="json"),
|
|
sql=sql, parameters=params, audit=audit,
|
|
).model_dump(mode="json")
|
|
|
|
|
|
async def query_financial_data_tool(
|
|
arguments: FinancialNL2SQLInput, context: RequestContext
|
|
) -> dict[str, Any]:
|
|
return await FinancialNL2SQLService().query(arguments, context)
|