zyi
b19a2415f9
- 新增 app/api、app/service 数据分析 Agent 全套服务与接口 - schemas.py 重构为 schemas 包（analyst schema） - 新增 SQL 防注入、guardrail、缓存、字典、LLM 等服务 - 新增 tests 测试套件与 scripts/dev、scripts/setup 脚本 - 补充需求规格、架构说明书、开发清单、表设计等文档
107 lines
3.6 KiB
Python
107 lines
3.6 KiB
Python
"""analyst_agent 编排测试：确定性假件 + 一条真实端到端。"""
|
|
import unittest
|
|
|
|
from app.service.analyst_agent import AnalystAgent
|
|
from app.utils.auth import AuthContext
|
|
|
|
|
|
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 ctx(roles, subject="STAFF-A"):
|
|
return AuthContext(subject_id=subject, token_type="staff", roles=roles, staff_type=roles[0] if roles else "")
|
|
|
|
|
|
class TestAgentOrchestration(unittest.TestCase):
|
|
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(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(len(repo.logged), 1)
|
|
|
|
def test_deny_bad_sql(self):
|
|
agent = AnalystAgent(llm=FakeLLM("INSERT INTO core_customer VALUES (1)", []), repo=FakeRepo())
|
|
resp = agent.run("删库", ctx(["analyst"]))
|
|
self.assertEqual(resp.status, "deny")
|
|
self.assertEqual(resp.error_code, "SQL_NOT_SELECT")
|
|
|
|
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"]))
|
|
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")
|
|
|
|
|
|
class TestAgentReal(unittest.TestCase):
|
|
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)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|