group_fqcd_jr/tools/risk_agent_business_e2e.py

"""奶龙风控智能助手业务对话端到端验收。"""

from __future__ import annotations

import asyncio
import datetime as dt
import json
import uuid
from pathlib import Path
from typing import Any

import httpx
import jwt
from sqlalchemy import text

from app.core.config import get_settings
from app.infrastructure.db import SessionFactory
from app.main import create_app
from app.worker.runtime import WorkerRuntime

PRIVATE_KEY = Path("config/jwt/jwt-private.pem").read_text(encoding="utf-8")
RISK_USER = "9002"
AGENT_TYPE = "risk"

CASES = (
    {
        "name": "风险概览",
        "message": "查看当前风险概览",
        "expected_tools": {"get_risk_overview"},
    },
    {
        "name": "完整筛选回答",
        "message": "当前低风险预警都是哪些客户的？他们买的都是什么产品？",
        "expected_tools": {"search_risk_alerts"},
    },
    {
        "name": "误报研判",
        "message": "哪些预警可以按误报复核？",
        "expected_tools": {"search_risk_alerts", "get_alert_evidence"},
    },
)


def token(subject: str) -> str:
    now = dt.datetime.now(dt.UTC)
    return jwt.encode(
        {
            "sub": subject,
            "iss": get_settings().jwt_issuer,
            "aud": get_settings().jwt_audience,
            "exp": now + dt.timedelta(minutes=30),
            "nbf": now - dt.timedelta(seconds=5),
            "jti": str(uuid.uuid4()),
        },
        PRIVATE_KEY,
        algorithm="RS256",
    )


async def audit_rows(trace_id: str, run_id: str) -> list[dict[str, object]]:
    async with SessionFactory() as session:
        rows = (
            await session.execute(
                text(
                    "SELECT action_type, detail FROM interaction_audit "
                    "WHERE JSON_UNQUOTE(JSON_EXTRACT(detail,'$.trace_id'))=:trace "
                    "OR JSON_UNQUOTE(JSON_EXTRACT(detail,'$.run_id'))=:run "
                    "ORDER BY id"
                ),
                {"trace": trace_id, "run": run_id},
            )
        ).mappings().all()
    result = []
    for row in rows:
        detail = row["detail"]
        if isinstance(detail, str):
            detail = json.loads(detail)
        result.append({"action_type": row["action_type"], "detail": detail})
    return result


async def run_case(
    client: httpx.AsyncClient,
    auth: dict[str, str],
    case: dict[str, Any],
) -> dict[str, Any]:
    accepted = await client.post(
        "/api/v1/agent-runs",
        headers=auth,
        json={
            "agent_type": AGENT_TYPE,
            "message": case["message"],
            "session_id": f"risk-business-e2e-{uuid.uuid4()}",
            "idempotency_key": uuid.uuid4().hex,
        },
    )
    if accepted.status_code != 202:
        raise SystemExit(
            f"{case['name']} 受理失败：{accepted.status_code} {accepted.text}"
        )
    accepted_data = accepted.json()["data"]
    run_id = str(accepted_data["run_id"])
    trace_id = str(accepted_data["trace_id"])

    await WorkerRuntime().execute(run_id)
    detail_response = await client.get(f"/api/v1/agent-runs/{run_id}", headers=auth)
    if detail_response.status_code != 200:
        raise SystemExit(
            f"{case['name']} 结果查询失败："
            f"{detail_response.status_code} {detail_response.text}"
        )
    detail = detail_response.json()["data"]
    events = await client.get(
        f"/api/v1/agent-runs/{run_id}/events",
        headers={**auth, "Accept": "text/event-stream"},
    )
    audits = await audit_rows(trace_id, run_id)
    calls = ((detail.get("result") or {}).get("tool_calls") or {}).get("calls") or []
    return {
        "name": case["name"],
        "run_id": run_id,
        "trace_id": trace_id,
        "status": detail.get("status"),
        "error_code": detail.get("error_code"),
        "content": (detail.get("result") or {}).get("content") or "",
        "tool_calls": calls,
        "events": events,
        "audits": audits,
        "expected_tools": case["expected_tools"],
    }


def validate_result(result: dict[str, Any]) -> None:
    name = str(result["name"])
    if result["status"] != "succeeded":
        raise SystemExit(
            f"{name} 执行失败：status={result['status']} "
            f"error_code={result['error_code']}"
        )
    if not str(result["content"]).strip():
        raise SystemExit(f"{name} 返回内容为空")
    tool_names = {
        str(call.get("tool_name"))
        for call in result["tool_calls"]
        if isinstance(call, dict)
    }
    if not tool_names.intersection(result["expected_tools"]):
        raise SystemExit(
            f"{name} 未调用预期工具：expected={result['expected_tools']} "
            f"actual={tool_names}"
        )
    events = result["events"]
    if not str(events.headers.get("content-type") or "").startswith("text/event-stream"):
        raise SystemExit(f"{name} SSE Content-Type 异常：{events.headers.get('content-type')}")
    if "event: done" not in events.text:
        raise SystemExit(f"{name} SSE 未返回 done")
    actions = {str(row["action_type"]) for row in result["audits"]}
    if "agent.tool_executed" not in actions:
        raise SystemExit(f"{name} 未记录工具审计")
    if "agent.run_completed" not in actions:
        raise SystemExit(f"{name} 未记录运行完成审计")


async def main() -> None:
    app = create_app()
    auth = {"Authorization": f"Bearer {token(RISK_USER)}"}
    async with httpx.AsyncClient(
        transport=httpx.ASGITransport(app=app),
        base_url="http://test",
        timeout=120,
    ) as client:
        for case in CASES:
            result = await run_case(client, auth, case)
            validate_result(result)
            print(
                f"{result['name']} PASSED "
                f"run_id={result['run_id']} "
                f"tools={[call.get('tool_name') for call in result['tool_calls']]} "
                f"content={result['content'][:120]!r}"
            )


if __name__ == "__main__":
    asyncio.run(main())