Mutual_Fund/agent/advisor_agent/data_query.py

"""投顾 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 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())
    return 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):
        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()
    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
    return payload