246 lines
9.6 KiB
Python
246 lines
9.6 KiB
Python
"""投顾 Agent 的 NL2SQL 数据查询适配层。"""
|
||
from __future__ import annotations
|
||
|
||
from dataclasses import asdict
|
||
from typing import Any
|
||
from uuid import uuid4
|
||
|
||
from agent.advisor_agent.auth import ensure_customer_access
|
||
from agent.data_query.agent import DataQueryAgent
|
||
from common.common_const import CUSTOMER_REL_STATUS_SIGNED, CUSTOMER_REL_STATUS_UNSIGNED
|
||
from config import database
|
||
from config.settings import settings
|
||
from nl2sql.contracts import DataQueryRequest, DataQueryResult
|
||
from nl2sql.embedding import EmbeddingError
|
||
from nl2sql.retrieval import retrieve_metadata
|
||
from nl2sql.cache import cache_get, cache_set, build_question_cache_key
|
||
from nl2sql.runtime_config import runtime_config
|
||
from nl2sql.schema import load_authoritative_schema
|
||
from repositories.customer_relation import CustomerRelationRepo
|
||
from service.nl2sql.permission_service import load_query_permission
|
||
from service.nl2sql.query_service import QueryServiceError
|
||
from tool.llm import llm as default_llm
|
||
from utils.exceptions import LLMFailError
|
||
from repositories.fin_holdings import FinHoldingsRepo
|
||
from repositories.fin_product import FinProductRepo
|
||
|
||
|
||
def enrich_customer_rows(
|
||
rows: list[dict[str, Any]],
|
||
columns: list[str],
|
||
name_by_id: dict[int, str],
|
||
) -> dict[str, Any]:
|
||
"""为查询结果补充客户姓名，并生成可直接展示的姓名摘要。"""
|
||
id_keys = ("customer_id", "客户ID", "客户编号", "客户_id")
|
||
enriched = [dict(row) for row in rows]
|
||
names: list[str] = []
|
||
for row in enriched:
|
||
customer_id = next((row.get(key) for key in id_keys if row.get(key) is not None), None)
|
||
try:
|
||
customer_id = int(customer_id)
|
||
except (TypeError, ValueError):
|
||
continue
|
||
name = name_by_id.get(customer_id)
|
||
if name:
|
||
row["客户姓名"] = name
|
||
names.append(f"客户{customer_id}（{name}）")
|
||
output_columns = list(columns)
|
||
if any("客户姓名" in row for row in enriched) and "客户姓名" not in output_columns:
|
||
insert_at = next(
|
||
(index + 1 for index, column in enumerate(output_columns) if column in id_keys),
|
||
len(output_columns),
|
||
)
|
||
output_columns.insert(insert_at, "客户姓名")
|
||
return {"rows": enriched, "columns": output_columns, "name_summary": "、".join(dict.fromkeys(names))}
|
||
|
||
|
||
def _is_current_holdings_query(question: str) -> bool:
|
||
text = "".join((question or "").split())
|
||
if any(term in text for term in ("当前持仓", "目前持仓", "现有持仓", "持仓明细", "持仓情况")):
|
||
return True
|
||
return "持仓" in text and not any(term in text for term in ("历史", "曾经", "已卖出", "交易记录"))
|
||
|
||
|
||
async def _query_current_holdings(
|
||
db, *, customer_id: int, trace_id: str, customer_name: str | None = None
|
||
) -> dict[str, Any]:
|
||
holdings = await FinHoldingsRepo(db).list_by_customer(customer_id, status="持有中")
|
||
product_repo = FinProductRepo(db)
|
||
rows: list[dict[str, Any]] = []
|
||
total_value = 0
|
||
for holding in holdings:
|
||
product = await product_repo.get(holding.product_id)
|
||
rows.append(
|
||
{
|
||
"产品代码": product.product_code if product else None,
|
||
"产品名称": product.product_name if product else None,
|
||
"风险等级": product.risk_level if product else None,
|
||
"持有份额": f"{holding.shares:.4f}",
|
||
"成本金额": f"{holding.cost_amount:.2f}",
|
||
"当前市值": f"{holding.current_value:.2f}",
|
||
"盈亏": f"{holding.profit_loss:.2f}",
|
||
"收益率": f"{holding.profit_ratio:.4f}",
|
||
"状态": holding.status,
|
||
}
|
||
)
|
||
total_value += holding.current_value
|
||
names = [row["产品名称"] for row in rows if row["产品名称"]]
|
||
return {
|
||
"query_id": f"holdings-{uuid4().hex}",
|
||
"trace_id": trace_id,
|
||
"columns": list(rows[0].keys()) if rows else ["产品代码", "产品名称", "风险等级", "持有份额", "成本金额", "当前市值", "盈亏", "收益率", "状态"],
|
||
"rows": rows,
|
||
"row_count": len(rows),
|
||
"truncated": False,
|
||
"summary": f"当前持仓共 {len(rows)} 条记录。",
|
||
"answer": (
|
||
(f"客户{customer_id}（{customer_name}）" if customer_name else f"客户{customer_id}")
|
||
+ f"当前持有 {len(rows)} 只基金，总市值约 {total_value:.2f} 元。"
|
||
+ (f"包括：{'、'.join(names[:6])}。" if names else "")
|
||
),
|
||
"sql": None,
|
||
}
|
||
|
||
|
||
async def execute_advisor_data_query(
|
||
db,
|
||
*,
|
||
advisor_id: int,
|
||
customer_id: int | None,
|
||
scope: str = "customer",
|
||
question: str,
|
||
trace_id: str,
|
||
session_id: str | None = None,
|
||
conversation_context: str = "",
|
||
data_scope: dict[str, Any] | None = None,
|
||
max_rows: int | None = None,
|
||
page: int = 1,
|
||
page_size: int = 100,
|
||
sort_by: str | None = None,
|
||
sort_order: str = "asc",
|
||
milvus=None,
|
||
redis=None,
|
||
llm_client=None,
|
||
query_agent=None,
|
||
) -> dict[str, Any]:
|
||
"""在当前投顾或选中客户范围内执行只读自然语言查询。
|
||
|
||
``data_scope`` 即使由调用方传入也不会被信任，服务端始终覆盖为当前
|
||
客户关系范围，避免投顾借助 NL2SQL 查询其他客户数据。
|
||
"""
|
||
if scope == "advisor":
|
||
relations = await CustomerRelationRepo(db).list_by_advisor(advisor_id)
|
||
customer_ids = [
|
||
relation.customer_id
|
||
for relation in relations
|
||
if relation.status in {CUSTOMER_REL_STATUS_UNSIGNED, CUSTOMER_REL_STATUS_SIGNED}
|
||
]
|
||
if not customer_ids:
|
||
raise QueryServiceError("当前投顾名下没有可查询客户")
|
||
else:
|
||
if customer_id is None:
|
||
raise QueryServiceError("单客户查询需要明确客户范围")
|
||
await ensure_customer_access(db, advisor_id=advisor_id, customer_id=customer_id)
|
||
customer_ids = [customer_id]
|
||
|
||
name_by_id: dict[int, str] = {}
|
||
try:
|
||
relation_rows = await CustomerRelationRepo(db).list_customer_rows(
|
||
advisor_id=advisor_id, limit=max(len(customer_ids), 100)
|
||
)
|
||
name_by_id = {
|
||
int(account.id): account.real_name
|
||
for _relation, account, _profile in relation_rows
|
||
if account.real_name
|
||
}
|
||
except Exception: # noqa: BLE001 姓名增强失败不阻断数据查询
|
||
pass
|
||
|
||
if scope == "customer" and _is_current_holdings_query(question):
|
||
redis = redis or database.redis.client()
|
||
return await _query_current_holdings(
|
||
db,
|
||
customer_id=customer_id,
|
||
trace_id=trace_id,
|
||
customer_name=name_by_id.get(customer_id),
|
||
)
|
||
permission = await load_query_permission(db, advisor_id)
|
||
if not permission.get("can_query", False):
|
||
raise QueryServiceError("当前投顾没有 NL2SQL 查询权限")
|
||
|
||
milvus = milvus or database.milvus.client()
|
||
redis = redis or database.redis.client()
|
||
cache_key = build_question_cache_key(
|
||
question,
|
||
permission=permission,
|
||
data_scope={"customer_ids": customer_ids},
|
||
page=page,
|
||
page_size=page_size,
|
||
sort_by=sort_by,
|
||
sort_order=sort_order,
|
||
)
|
||
cached = await cache_get(redis, cache_key)
|
||
if cached is not None:
|
||
cached["trace_id"] = trace_id
|
||
cached["customer_id"] = customer_id
|
||
cached["warnings"] = [*cached.get("warnings", []), "cache_hit"]
|
||
return cached
|
||
llm_client = llm_client or default_llm
|
||
request = DataQueryRequest(
|
||
question=question,
|
||
user_id=advisor_id,
|
||
trace_id=trace_id,
|
||
session_id=session_id,
|
||
caller_agent="advisor_agent",
|
||
data_scope={"customer_ids": customer_ids},
|
||
max_rows=min(max_rows or runtime_config.max_rows, runtime_config.max_rows),
|
||
include_sql=False,
|
||
page=page,
|
||
page_size=page_size,
|
||
sort_by=sort_by,
|
||
sort_order=sort_order,
|
||
)
|
||
|
||
async def permission_loader(_user_id: int):
|
||
return permission
|
||
|
||
async def metadata_retriever(query: str):
|
||
return await retrieve_metadata(
|
||
query,
|
||
milvus,
|
||
top_k=runtime_config.retrieval_top_k,
|
||
)
|
||
|
||
async def schema_loader(table_names: set[str], _permission: dict):
|
||
return await load_authoritative_schema(
|
||
db,
|
||
database=settings.mysql.database,
|
||
candidate_tables=table_names,
|
||
)
|
||
|
||
try:
|
||
result: DataQueryResult = await (query_agent or DataQueryAgent()).query(
|
||
request,
|
||
session=db,
|
||
permission_loader=permission_loader,
|
||
metadata_retriever=metadata_retriever,
|
||
schema_loader=schema_loader,
|
||
llm_client=llm_client,
|
||
summary_llm=llm_client,
|
||
masks=permission.get("masks"),
|
||
conversation_context=conversation_context,
|
||
)
|
||
except (EmbeddingError, LLMFailError) as exc:
|
||
raise QueryServiceError("投顾 Agent 依赖服务不可用，请检查 LLM/Embedding 服务连接") from exc
|
||
payload = asdict(result)
|
||
enriched = enrich_customer_rows(payload.get("rows", []), payload.get("columns", []), name_by_id)
|
||
payload["rows"] = enriched["rows"]
|
||
payload["columns"] = enriched["columns"]
|
||
if enriched["name_summary"]:
|
||
existing_answer = payload.get("answer") or payload.get("summary") or ""
|
||
payload["answer"] = f"{existing_answer.rstrip('。')}。客户姓名：{enriched['name_summary']}。"
|
||
payload["sql"] = None
|
||
payload["customer_id"] = customer_id
|
||
await cache_set(redis, cache_key, payload, ttl=runtime_config.cache_ttl)
|
||
return payload
|