zhanghongyu_0626
188bf6a438
- Introduced `_recent_days_series_hint` to handle queries related to "最近N天" for time series data aggregation. - Added `_cn_num_to_int` function to convert Chinese numerals to integers for better query parsing. - Updated `_nl_sql_hints` to incorporate the new hint generation logic, ensuring accurate SQL output for recent days queries. - Enhanced documentation to reflect these changes and improve clarity on the new functionalities.
430 lines
17 KiB
Python
430 lines
17 KiB
Python
"""analyst_agent 编排测试（Wave 6）。"""
|
|
import unittest
|
|
|
|
import pytest
|
|
|
|
from app.api.analyst_auth_adapter import AnalystAuthContext
|
|
from app.service.analyst_agent import AnalystAgent, _nl_sql_hints
|
|
from app.service.template_service import QueryTemplate, TemplateService
|
|
|
|
|
|
class FakeLLM:
|
|
def __init__(self, sql, answers):
|
|
self.sql = sql
|
|
self.answers = list(answers)
|
|
self.calls = 0
|
|
|
|
def complete(self, messages, temperature=0, max_tokens=2048):
|
|
self.calls += 1
|
|
usage = {"prompt_tokens": 10, "completion_tokens": 10}
|
|
if self.calls == 1:
|
|
return self.sql, usage
|
|
ans = self.answers.pop(0) if self.answers else "无解读"
|
|
return ans, usage
|
|
|
|
|
|
class FakeRepo:
|
|
def __init__(self, rows=(), columns=(), scope=None):
|
|
self.rows = list(rows)
|
|
self.columns = list(columns)
|
|
self.scope = scope or []
|
|
self.logged = []
|
|
|
|
def resolve_advisor_scope(self, sid):
|
|
return self.scope
|
|
|
|
def execute_readonly(self, sql):
|
|
return {"columns": self.columns, "rows": self.rows}
|
|
|
|
def get_data_as_of(self):
|
|
return "2026-09-04"
|
|
|
|
def log_query(self, **kw):
|
|
self.logged.append(kw)
|
|
|
|
def log_audit(self, **kw):
|
|
pass
|
|
|
|
def list_published_templates(self):
|
|
return []
|
|
|
|
def list_published_metrics(self):
|
|
return []
|
|
|
|
def list_published_few_shots(self, limit=5):
|
|
return []
|
|
|
|
def get_query_log_by_trace(self, trace_id):
|
|
return {
|
|
"staff_id": "STAFF-A",
|
|
"generated_sql": "SELECT COUNT(*) AS c FROM core_customer",
|
|
"exec_status": "success",
|
|
}
|
|
|
|
|
|
def ctx(roles, subject="STAFF-A", *, token_type="staff", customer_id=None):
|
|
return AnalystAuthContext(
|
|
subject_id=subject,
|
|
token_type=token_type,
|
|
roles=roles,
|
|
customer_id=customer_id,
|
|
)
|
|
|
|
|
|
class TestAgentOrchestration(unittest.TestCase):
|
|
def test_clarify_trade_flow_ambiguity(self):
|
|
agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=FakeRepo())
|
|
resp = agent.run("近30天交易流水是多少", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "clarify")
|
|
self.assertIn("流水", resp.answer)
|
|
|
|
def test_clarify_ambiguity(self):
|
|
agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=FakeRepo())
|
|
resp = agent.run("我名下的规模是多少", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "clarify")
|
|
|
|
def test_success_query_only(self):
|
|
repo = FakeRepo(rows=[[33]], columns=["c"])
|
|
agent = AnalystAgent(
|
|
llm=FakeLLM("SELECT COUNT(*) AS c FROM core_customer", ["共 33 个客户"]),
|
|
repo=repo,
|
|
)
|
|
resp = agent.run("客户总数是多少", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertEqual(resp.table.rows, [[33]])
|
|
self.assertEqual(resp.answer, "")
|
|
self.assertEqual(len(repo.logged), 1)
|
|
self.assertEqual(agent.llm.calls, 1)
|
|
|
|
def test_success_with_interpret_flag(self):
|
|
repo = FakeRepo(rows=[[33]], columns=["c"])
|
|
llm = FakeLLM("SELECT COUNT(*) AS c FROM core_customer", ["共 33 个客户"])
|
|
agent = AnalystAgent(llm=llm, repo=repo)
|
|
resp = agent.run("客户总数是多少", ctx(["analyst"]), interpret=True)
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertIn("33", resp.answer)
|
|
self.assertEqual(llm.calls, 2)
|
|
|
|
def test_interpret_from_snapshot(self):
|
|
from app.model.analyst_schemas import InterpretRequest, Meta, TableData
|
|
|
|
class AnswerOnlyLLM:
|
|
def complete(self, messages, temperature=0, max_tokens=2048):
|
|
return "共 33 个客户", {"prompt_tokens": 1, "completion_tokens": 1}
|
|
|
|
repo = FakeRepo(rows=[[33]], columns=["c"])
|
|
agent = AnalystAgent(llm=AnswerOnlyLLM(), repo=repo)
|
|
req = InterpretRequest(
|
|
question="客户总数是多少",
|
|
status="success",
|
|
table=TableData(columns=["c"], rows=[[33]]),
|
|
sql="SELECT COUNT(*) AS c FROM core_customer",
|
|
meta=Meta(row_count=1),
|
|
)
|
|
resp = agent.interpret(req, ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertIn("33", resp.answer)
|
|
|
|
def test_interpret_echo_deny_without_llm(self):
|
|
from app.model.analyst_schemas import InterpretRequest
|
|
|
|
llm = FakeLLM("SELECT 1", [])
|
|
agent = AnalystAgent(llm=llm, repo=FakeRepo())
|
|
req = InterpretRequest(
|
|
question="查别人",
|
|
status="deny",
|
|
answer="无法执行：权限不足",
|
|
)
|
|
resp = agent.interpret(req, ctx(["customer"], "CUST-1", customer_id="CUST-1"))
|
|
self.assertEqual(resp.status, "deny")
|
|
self.assertEqual(resp.answer, "无法执行：权限不足")
|
|
self.assertEqual(llm.calls, 0)
|
|
|
|
def test_deny_bad_sql(self):
|
|
repo = FakeRepo()
|
|
agent = AnalystAgent(llm=FakeLLM("INSERT INTO core_customer VALUES (1)", []), repo=repo)
|
|
resp = agent.run("删库", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "deny")
|
|
self.assertEqual(resp.error_code, "SQL_NOT_SELECT")
|
|
self.assertEqual(len(repo.logged), 1)
|
|
self.assertEqual(repo.logged[0].get("exec_status"), "blocked")
|
|
|
|
def test_degrade_wrong_number(self):
|
|
repo = FakeRepo(rows=[[33]], columns=["c"])
|
|
agent = AnalystAgent(
|
|
llm=FakeLLM("SELECT COUNT(*) FROM core_customer", ["共 999 个客户", "共 999 个客户"]),
|
|
repo=repo,
|
|
)
|
|
resp = agent.run("客户总数", ctx(["analyst"]), interpret=True)
|
|
self.assertEqual(resp.status, "degrade")
|
|
|
|
def test_advisor_out_of_scope_deny(self):
|
|
repo = FakeRepo(scope=["CUST-1001"])
|
|
agent = AnalystAgent(
|
|
llm=FakeLLM("SELECT * FROM core_holding WHERE customer_id='CUST-1004'", []),
|
|
repo=repo,
|
|
)
|
|
resp = agent.run("查 CUST-1004 持仓", ctx(["advisor"], "STAFF-B"))
|
|
self.assertEqual(resp.status, "deny")
|
|
self.assertEqual(resp.error_code, "AUTH_403_NOT_ASSIGNED")
|
|
|
|
def test_customer_self_success(self):
|
|
repo = FakeRepo(rows=[[2]], columns=["cnt"])
|
|
agent = AnalystAgent(
|
|
llm=FakeLLM(
|
|
"SELECT COUNT(*) AS cnt FROM core_trade WHERE customer_id='CUST-9527'",
|
|
["近阶段共有 2 笔交易"],
|
|
),
|
|
repo=repo,
|
|
)
|
|
resp = agent.run(
|
|
"我有多少笔交易",
|
|
ctx(["customer"], "CUST-9527", token_type="customer", customer_id="CUST-9527"),
|
|
interpret=True,
|
|
)
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertIn("AI 分析有风险", resp.answer)
|
|
|
|
def test_template_hit_skips_llm_sql(self):
|
|
tpl = QueryTemplate(
|
|
template_key="customer_total_count",
|
|
template_sql="SELECT COUNT(*) AS c FROM core_customer",
|
|
params_schema={"match_all": ["客户", "总数"]},
|
|
)
|
|
repo = FakeRepo(rows=[[33]], columns=["c"])
|
|
llm = FakeLLM("SELECT 1", ["共 33 个客户"])
|
|
agent = AnalystAgent(
|
|
llm=llm,
|
|
repo=repo,
|
|
templates=TemplateService(templates=[tpl]),
|
|
)
|
|
resp = agent.run("客户总数是多少", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertTrue(resp.meta.template_hit)
|
|
self.assertEqual(resp.meta.template_key, "customer_total_count")
|
|
self.assertIn("COUNT(*)", resp.sql)
|
|
self.assertEqual(llm.calls, 0)
|
|
|
|
def test_dual_extreme_question_enriches_llm_prompt(self):
|
|
q = "最大和最低涨幅的产品分别是哪两个"
|
|
enriched = _nl_sql_hints(q)
|
|
self.assertIn("UNION ALL", enriched)
|
|
self.assertIn("禁止仅 LIMIT 1", enriched)
|
|
|
|
def test_recent_days_amount_question_enriches_llm_prompt(self):
|
|
"""时间序列问法（最近N天+金额/流水）→ 提示按日聚合，且天数与问题一致。"""
|
|
enriched = _nl_sql_hints("最近十天的交易流水金额？")
|
|
self.assertIn("DATE(traded_at)", enriched)
|
|
self.assertIn("GROUP BY", enriched)
|
|
self.assertIn("INTERVAL 10 DAY", enriched)
|
|
self.assertIn("禁止", enriched)
|
|
|
|
def test_recent_days_amount_question_arabic_and_cn_numerals(self):
|
|
self.assertIn("INTERVAL 7 DAY", _nl_sql_hints("近7天申购金额"))
|
|
self.assertIn("INTERVAL 30 DAY", _nl_sql_hints("最近三十天的赎回金额"))
|
|
|
|
def test_recent_days_detail_question_not_enriched(self):
|
|
"""要明细/清单的问法不按日聚合（保留原始行）。"""
|
|
q = "最近十天的交易明细"
|
|
self.assertEqual(_nl_sql_hints(q), q)
|
|
|
|
def test_recent_days_count_question_not_enriched(self):
|
|
"""「有多少笔」是计数问法（单行合计即正确答案），不按日聚合。"""
|
|
q = "我近30天有多少笔交易"
|
|
self.assertEqual(_nl_sql_hints(q), q)
|
|
|
|
def test_generate_sql_uses_db_few_shots_only(self):
|
|
class CaptureLLM:
|
|
last_user = ""
|
|
|
|
def complete(self, messages, temperature=0.0, max_tokens=2048):
|
|
self.last_user = messages[-1]["content"]
|
|
return "SELECT 1 AS x", {"prompt_tokens": 1, "completion_tokens": 1}
|
|
|
|
class FewShotRepo(FakeRepo):
|
|
def list_published_few_shots(self, limit=5):
|
|
return [
|
|
{
|
|
"question": "平台持仓市值最大的客户是谁，其持有产品的净值走势如何",
|
|
"sql_text": "SELECT 1 FROM core_product_nav",
|
|
}
|
|
]
|
|
|
|
llm = CaptureLLM()
|
|
agent = AnalystAgent(llm=llm, repo=FewShotRepo())
|
|
agent._generate_sql("x", "full", [])
|
|
self.assertIn("平台持仓市值最大的客户", llm.last_user)
|
|
self.assertIn("core_product_nav", llm.last_user)
|
|
|
|
def test_generate_sql_empty_when_no_db_few_shots(self):
|
|
class CaptureLLM:
|
|
last_user = ""
|
|
|
|
def complete(self, messages, temperature=0.0, max_tokens=2048):
|
|
self.last_user = messages[-1]["content"]
|
|
return "SELECT 1 AS x", {"prompt_tokens": 1, "completion_tokens": 1}
|
|
|
|
llm = CaptureLLM()
|
|
agent = AnalystAgent(llm=llm, repo=FakeRepo())
|
|
agent._generate_sql("客户总数", "full", [])
|
|
self.assertNotIn("参考 few-shot", llm.last_user)
|
|
|
|
def test_clarify_includes_structured_payload(self):
|
|
agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=FakeRepo())
|
|
resp = agent.run("我名下的规模是多少", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "clarify")
|
|
self.assertIsNotNone(resp.clarify)
|
|
assert resp.clarify is not None
|
|
self.assertTrue(resp.clarify.options)
|
|
self.assertIn("规模", resp.clarify.term)
|
|
|
|
def test_generate_sql_prefers_db_few_shots(self):
|
|
class CaptureLLM:
|
|
last_user = ""
|
|
|
|
def complete(self, messages, temperature=0.0, max_tokens=2048):
|
|
self.last_user = messages[-1]["content"]
|
|
return "SELECT 1 AS x", {"prompt_tokens": 1, "completion_tokens": 1}
|
|
|
|
class FewShotRepo(FakeRepo):
|
|
def list_published_few_shots(self, limit=5):
|
|
return [{"question": "库内示例问法", "sql_text": "SELECT 2 AS y FROM core_product"}]
|
|
|
|
llm = CaptureLLM()
|
|
agent = AnalystAgent(llm=llm, repo=FewShotRepo())
|
|
agent._generate_sql("x", "full", [])
|
|
self.assertIn("库内示例问法", llm.last_user)
|
|
self.assertIn("core_product", llm.last_user)
|
|
|
|
def test_template_nav_max_min_skips_llm(self):
|
|
tpl = QueryTemplate(
|
|
template_key="product_nav_max_min_latest",
|
|
template_sql=(
|
|
"(SELECT 'max' AS extremum, n.product_id, n.daily_chg_pct FROM core_product_nav n "
|
|
"ORDER BY n.daily_chg_pct DESC LIMIT 1) "
|
|
"UNION ALL "
|
|
"(SELECT 'min' AS extremum, n.product_id, n.daily_chg_pct FROM core_product_nav n "
|
|
"ORDER BY n.daily_chg_pct ASC LIMIT 1)"
|
|
),
|
|
params_schema={"match_phrases": ["目前净值最高和最低的两个产品"]},
|
|
)
|
|
repo = FakeRepo(
|
|
rows=[["max", "P1", 1.2], ["min", "P2", -0.5]],
|
|
columns=["extremum", "product_id", "daily_chg_pct"],
|
|
)
|
|
llm = FakeLLM("SELECT 1 LIMIT 1", [])
|
|
agent = AnalystAgent(
|
|
llm=llm,
|
|
repo=repo,
|
|
templates=TemplateService(templates=[tpl]),
|
|
)
|
|
resp = agent.run("目前净值最高和最低的两个产品", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertTrue(resp.meta.template_hit)
|
|
self.assertEqual(resp.meta.row_count, 2)
|
|
self.assertIn("UNION ALL", resp.sql.upper())
|
|
self.assertEqual(llm.calls, 0)
|
|
|
|
def test_analyze_text_only_json(self):
|
|
json_out = '{"answer":"产品净值整体平稳，未见异常波动。","chart_type":"none","title":"","x_field":null,"y_fields":[]}'
|
|
repo = FakeRepo(rows=[["P1", 1.1], ["P2", 1.2]], columns=["product_id", "nav"])
|
|
agent = AnalystAgent(llm=FakeLLM(json_out, []), repo=repo)
|
|
from app.model.analyst_schemas import AnalyzeRequest, Meta, TableData
|
|
|
|
req = AnalyzeRequest(
|
|
question="产品净值",
|
|
analysis_prompt="三句话总结",
|
|
status="success",
|
|
table=TableData(columns=["product_id", "nav"], rows=[["P1", 1.1], ["P2", 1.2]]),
|
|
sql="SELECT 1",
|
|
meta=Meta(row_count=2),
|
|
)
|
|
resp = agent.analyze(req, ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertEqual(resp.analysis_kind, "text")
|
|
self.assertIsNone(resp.chart)
|
|
|
|
def test_analyze_chart_only_json(self):
|
|
json_out = (
|
|
'{"answer":"","chart_type":"line","title":"趋势","reason":"日期+数值",'
|
|
'"x_field":"d","y_fields":["v"],"series_field":"s"}'
|
|
)
|
|
repo = FakeRepo(
|
|
rows=[["2026-01-01", "A", 1.0], ["2026-01-02", "A", 1.1]],
|
|
columns=["d", "s", "v"],
|
|
)
|
|
agent = AnalystAgent(llm=FakeLLM(json_out, []), repo=repo)
|
|
from app.model.analyst_schemas import AnalyzeRequest, Meta, TableData
|
|
|
|
req = AnalyzeRequest(
|
|
question="趋势",
|
|
analysis_prompt="画折线图",
|
|
status="success",
|
|
table=TableData(
|
|
columns=["d", "s", "v"],
|
|
rows=[["2026-01-01", "A", 1.0], ["2026-01-02", "A", 1.1]],
|
|
),
|
|
sql="SELECT 1",
|
|
meta=Meta(row_count=2),
|
|
)
|
|
resp = agent.analyze(req, ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertEqual(resp.analysis_kind, "chart")
|
|
self.assertIsNotNone(resp.chart)
|
|
self.assertEqual(resp.chart.chart_type, "line")
|
|
|
|
def test_sql_exec_failure_escalates(self):
|
|
class FailRepo(FakeRepo):
|
|
def execute_readonly(self, sql):
|
|
raise TimeoutError("read timed out")
|
|
|
|
class NoCache:
|
|
def permission_fingerprint(self, *a, **k):
|
|
return "fp"
|
|
|
|
def sql_hash(self, sql):
|
|
return "hash"
|
|
|
|
def get_result(self, *a, **k):
|
|
return None
|
|
|
|
def set_result(self, *a, **k):
|
|
pass
|
|
|
|
agent = AnalystAgent(
|
|
llm=FakeLLM("SELECT COUNT(*) AS c FROM core_customer", []),
|
|
repo=FailRepo(),
|
|
templates=TemplateService(templates=[]),
|
|
cache=NoCache(),
|
|
)
|
|
resp = agent.run("客户总数", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "escalate")
|
|
self.assertIn(resp.error_code, ("EXEC_TIMEOUT", "EXEC_ERROR"))
|
|
|
|
def test_sample_by_trace(self):
|
|
repo = FakeRepo(rows=[[1]], columns=["c"])
|
|
agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=repo)
|
|
resp = agent.sample_by_trace("trace-1", ctx(["analyst"], subject="STAFF-A"))
|
|
self.assertEqual(resp.status, "success")
|
|
self.assertIn("LIMIT", resp.sql)
|
|
|
|
def test_escalate_audit(self):
|
|
repo = FakeRepo()
|
|
agent = AnalystAgent(llm=FakeLLM("SELECT 1", []), repo=repo)
|
|
out = agent.escalate("trace-x", "q", "timeout", ctx(["analyst"]))
|
|
self.assertTrue(out["ok"])
|
|
|
|
|
|
@pytest.mark.integration
|
|
class TestAgentReal(unittest.TestCase):
|
|
@pytest.mark.skip(reason="需要真实 MySQL + DeepSeek Key")
|
|
def test_real_end_to_end(self):
|
|
from app.service.analytics_repo import AnalyticsRepo
|
|
from app.service.llm import DeepSeekLLM
|
|
|
|
agent = AnalystAgent(llm=DeepSeekLLM(), repo=AnalyticsRepo())
|
|
resp = agent.run("客户总数是多少", ctx(["analyst"]))
|
|
self.assertIn(resp.status, ("success", "degrade"))
|
|
self.assertTrue(resp.sql)
|
|
self.assertGreater(resp.meta.row_count, 0)
|