96 lines
3.6 KiB
Python
96 lines
3.6 KiB
Python
|
"""20 题批量问答测试：调 /api/analyst/chat，记录 status/SQL/rows/answer。
|
||
|
|
|
||
|
|
需本地 uvicorn + MySQL + DeepSeek Key（真实 LLM）。
|
||
|
|
用法：python scripts/dev/run_query_battery.py
|
||
|
|
"""
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import json
|
||
|
|
import os
|
||
|
|
import sys
|
||
|
|
import urllib.request
|
||
|
|
|
||
|
|
ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
|
|
sys.path.insert(0, ROOT)
|
||
|
|
|
||
|
|
from app.service.auth_service import issue_dev_token
|
||
|
|
|
||
|
|
URL = "http://127.0.0.1:8000/api/analyst/chat"
|
||
|
|
|
||
|
|
QUERIES = [
|
||
|
|
("Q01", "各风险等级（C1~C5）客户各有多少？占比多少？"),
|
||
|
|
("Q02", "C3 客户大概占多少？"),
|
||
|
|
("Q03", "高净值 vs 非高净值客户各多少、占比多少？"),
|
||
|
|
("Q04", "客户总体规模（总客户数）怎样？"),
|
||
|
|
("Q05", "近 30 天业务活跃度和再往前 30 天比怎样？按交易笔数"),
|
||
|
|
("Q06", "各服务等级（普通/VIP/钻石）客户分布？"),
|
||
|
|
("Q07", "各城市/地区客户大概怎么分布？汇总"),
|
||
|
|
("Q08", "各风险等级产品（R1~R5）各有多少只？"),
|
||
|
|
("Q09", "钱主要集中在哪类产品？按持仓市值"),
|
||
|
|
("Q10", "最近 30 天申购金额和笔数，和再往前 30 天比怎样？"),
|
||
|
|
("Q11", "近 30 天赎回大概多少？"),
|
||
|
|
("Q12", "哪几类产品申购最猛？按产品类型汇总，不要列单个客户"),
|
||
|
|
("Q13", "还有多少条预警没处理（待审核）？"),
|
||
|
|
("Q14", "预警按状态怎么分布？"),
|
||
|
|
("Q15", "高等级预警（反洗钱/适当性）大概多少？"),
|
||
|
|
("Q16", "大额交易类预警大概多少？"),
|
||
|
|
("Q17", "近一段时间适当性不匹配/被拦大概多少笔？汇总"),
|
||
|
|
("Q18", "理财顾问名下客户规模对比"),
|
||
|
|
("Q19", "近 30 天新开户大概多少？"),
|
||
|
|
("Q20", "AML 风险等级（高/中/低）客户占比？"),
|
||
|
|
]
|
||
|
|
|
||
|
|
|
||
|
|
def call(question: str, idx: int, token: str) -> dict:
|
||
|
|
payload = json.dumps({
|
||
|
|
"question": question,
|
||
|
|
"session_id": f"battery-{idx}",
|
||
|
|
"trace_id": f"trace-{idx}",
|
||
|
|
}).encode("utf-8")
|
||
|
|
req = urllib.request.Request(
|
||
|
|
URL,
|
||
|
|
data=payload,
|
||
|
|
headers={
|
||
|
|
"Authorization": f"Bearer {token}",
|
||
|
|
"Content-Type": "application/json; charset=utf-8",
|
||
|
|
},
|
||
|
|
method="POST",
|
||
|
|
)
|
||
|
|
with urllib.request.urlopen(req, timeout=180) as resp:
|
||
|
|
return json.loads(resp.read().decode("utf-8"))
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
token = issue_dev_token(sub="STAFF-20001", roles=["analyst"])
|
||
|
|
out = []
|
||
|
|
for idx, (tag, q) in enumerate(QUERIES, 1):
|
||
|
|
rec: dict = {"tag": tag, "question": q}
|
||
|
|
try:
|
||
|
|
r = call(q, idx, token)
|
||
|
|
except Exception as exc:
|
||
|
|
rec.update({"status": "EXC", "error": str(exc)})
|
||
|
|
out.append(rec)
|
||
|
|
print(json.dumps(rec, ensure_ascii=False))
|
||
|
|
continue
|
||
|
|
t = r.get("table") or {}
|
||
|
|
rec["status"] = r.get("status")
|
||
|
|
rec["answer"] = r.get("answer")
|
||
|
|
rec["sql"] = (r.get("sql") or "").replace("\n", " ")
|
||
|
|
rec["columns"] = t.get("columns")
|
||
|
|
rec["rows"] = (t.get("rows") or [])[:6]
|
||
|
|
rec["row_count"] = len(t.get("rows") or [])
|
||
|
|
rec["error_code"] = r.get("error_code")
|
||
|
|
out.append(rec)
|
||
|
|
print(json.dumps(rec, ensure_ascii=False))
|
||
|
|
sys.stdout.flush()
|
||
|
|
|
||
|
|
report_path = os.path.join(ROOT, "scripts", "dev", "battery_report.json")
|
||
|
|
with open(report_path, "w", encoding="utf-8") as f:
|
||
|
|
json.dump(out, f, ensure_ascii=False, indent=2)
|
||
|
|
print(f"\nWrote {report_path}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
sys.stdout.reconfigure(encoding="utf-8", errors="replace")
|
||
|
|
main()
|