zhanghongyu_0626
0fb7d34d7a
- Enhanced the `AnalystAgent` class to include an `_audit_terminal` method for logging query denials, clarifications, and errors, ensuring compliance and traceability. - Updated error handling paths to call the new audit method, capturing relevant details such as question, user authentication, and SQL context. - Introduced new validation checks in `sql_guard.py` to enforce ownership filters for sensitive queries, improving security measures. - Added unit tests to verify the correct logging behavior and ownership filter enforcement, ensuring robust functionality. This update significantly strengthens the auditing capabilities of the analyst agent, enhancing security and compliance in query handling.
208 lines
7.7 KiB
Python
208 lines
7.7 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 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)
|
|
|
|
|
|
@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)
|