feat:修改nl2sql功能
This commit is contained in:
@@ -25,6 +25,8 @@ from tool.llm import llm
|
||||
async def run_real_query(user_id: int, question: str, *, query_id: str | None = None) -> dict:
|
||||
"""使用已有员工权限执行一条真实查询,只输出脱敏统计。"""
|
||||
from config.database.mysql import get_session_factory
|
||||
from common.common_const import CUSTOMER_REL_STATUS_SIGNED, CUSTOMER_REL_STATUS_UNSIGNED
|
||||
from repositories.customer_relation import CustomerRelationRepo
|
||||
|
||||
query_id = query_id or uuid4().hex
|
||||
try:
|
||||
@@ -33,6 +35,22 @@ async def run_real_query(user_id: int, question: str, *, query_id: str | None =
|
||||
if not permission.get("can_query"):
|
||||
return {"query_status": "permission_denied", "user_id": user_id}
|
||||
|
||||
data_scope = None
|
||||
if any(
|
||||
scope.get("type") == "customer_ids"
|
||||
for scope in (permission.get("row_scopes") or {}).values()
|
||||
):
|
||||
relations = await CustomerRelationRepo(db).list_by_advisor(user_id)
|
||||
customer_ids = [
|
||||
relation.customer_id
|
||||
for relation in relations
|
||||
if relation.status in {
|
||||
CUSTOMER_REL_STATUS_UNSIGNED,
|
||||
CUSTOMER_REL_STATUS_SIGNED,
|
||||
}
|
||||
]
|
||||
data_scope = {"customer_ids": customer_ids}
|
||||
|
||||
async def permission_loader(_user_id: int):
|
||||
return permission
|
||||
|
||||
@@ -51,6 +69,7 @@ async def run_real_query(user_id: int, question: str, *, query_id: str | None =
|
||||
question=question,
|
||||
user_id=user_id,
|
||||
trace_id=f"nl2sql-e2e-{query_id}",
|
||||
data_scope=data_scope,
|
||||
include_sql=False,
|
||||
),
|
||||
session=db,
|
||||
|
||||
Reference in New Issue
Block a user