feat:修改nl2sql功能

agent/advisor_agent/data_query.py

@@ -7,15 +7,19 @@ 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
 
@@ -67,7 +71,8 @@ async def execute_advisor_data_query(
     db,
     *,
     advisor_id: int,
-    customer_id: int,
+    customer_id: int | None,
+    scope: str = "customer",
     question: str,
     trace_id: str,
     session_id: str | None = None,
@@ -82,13 +87,27 @@ async def execute_advisor_data_query(
     llm_client=None,
     query_agent=None,
 ) -> dict[str, Any]:
-    """在当前投顾和选中客户范围内执行只读自然语言查询。
+    """在当前投顾或选中客户范围内执行只读自然语言查询。
 
     ``data_scope`` 即使由调用方传入也不会被信任，服务端始终覆盖为当前
-    ``customer_id``，避免投顾借助 NL2SQL 查询其他客户数据。
+    客户关系范围，避免投顾借助 NL2SQL 查询其他客户数据。
     """
-    await ensure_customer_access(db, advisor_id=advisor_id, customer_id=customer_id)
-    if _is_current_holdings_query(question):
+    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,
@@ -107,7 +126,7 @@ async def execute_advisor_data_query(
         trace_id=trace_id,
         session_id=session_id,
         caller_agent="advisor_agent",
-        data_scope={"customer_ids": [customer_id]},
+        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,
@@ -133,17 +152,19 @@ async def execute_advisor_data_query(
             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,
-    )
+    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