zhanghongyu_0626
8556da489a
- Introduced new endpoints `/api/analyst/query/{trace_id}/sample` and `/api/analyst/escalate` for sampling query results and escalating issues to human analysts, respectively.
- Enhanced `AnalystAgent` to support sampling of SQL results based on trace ID and to handle escalation requests, improving user experience in error scenarios.
- Updated `analyst_schemas.py` to include `EscalateRequest` for structured escalation requests.
- Added corresponding frontend API calls and UI components to facilitate user interactions with the new features.
- Implemented unit tests to ensure the reliability of the new functionalities.
This update significantly enhances the analytical capabilities of the application, allowing users to retrieve detailed query samples and escalate issues effectively.
256 lines
9.3 KiB
Python
256 lines
9.3 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
|
|
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 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_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)
|