feat:投顾agent接入nl2sql能力

agent/advisor_agent/data_query.py

@@ -0,0 +1,98 @@
+"""投顾 Agent 的 NL2SQL 数据查询适配层。"""
+from __future__ import annotations
+
+from dataclasses import asdict
+from typing import Any
+
+from agent.advisor_agent.auth import ensure_customer_access
+from agent.data_query.agent import DataQueryAgent
+from config import database
+from config.settings import settings
+from nl2sql.contracts import DataQueryRequest, DataQueryResult
+from nl2sql.retrieval import retrieve_metadata
+from nl2sql.runtime_config import runtime_config
+from nl2sql.schema import load_authoritative_schema
+from service.nl2sql.permission_service import load_query_permission
+from service.nl2sql.query_service import QueryServiceError
+from tool.llm import llm as default_llm
+
+
+async def execute_advisor_data_query(
+    db,
+    *,
+    advisor_id: int,
+    customer_id: int,
+    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`` 即使由调用方传入也不会被信任，服务端始终覆盖为当前
+    ``customer_id``，避免投顾借助 NL2SQL 查询其他客户数据。
+    """
+    await ensure_customer_access(db, advisor_id=advisor_id, customer_id=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_id]},
+        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,
+        )
+
+    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"),
+        redis=redis,
+    )
+    payload = asdict(result)
+    payload["sql"] = None
+    payload["customer_id"] = customer_id
+    return payload