99 lines
3.2 KiB
Python
99 lines
3.2 KiB
Python
"""投顾 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
|