"""投顾 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.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 _is_current_holdings_query(question: str) -> bool: text = "".join((question or "").split()) return any(term in text for term in ("当前持仓", "目前持仓", "现有持仓", "持仓明细", "持仓情况")) async def _query_current_holdings(db, *, customer_id: int, trace_id: str) -> 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"当前持有 {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, 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] if scope == "customer" and _is_current_holdings_query(question): return await _query_current_holdings( db, customer_id=customer_id, trace_id=trace_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() 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"), ) except (EmbeddingError, LLMFailError) as exc: raise QueryServiceError("投顾 Agent 依赖服务不可用,请检查 LLM/Embedding 服务连接") from exc payload = asdict(result) payload["sql"] = None payload["customer_id"] = customer_id return payload