aiagent/test_agent_execution.py

#!/usr/bin/env python3
"""
Agent工作流执行测试脚本
用于测试Agent工作流的正常执行

与《工作流调用测试总结》一致：input_data 仅包含 query、USER_INPUT，便于 LLM 正确提取 user_query。

用法示例:
  python test_agent_execution.py
  python test_agent_execution.py <agent_id>
  python test_agent_execution.py <agent_id> "你好"
  python test_agent_execution.py --homework
  python test_agent_execution.py --homework --base-url http://127.0.0.1:8037
"""
from __future__ import annotations

import argparse
import json
import os
import sys
import time
from typing import Any, Dict, List, Optional

import requests

DEFAULT_BASE_URL = os.environ.get("API_BASE_URL", "http://localhost:8037")


def _ensure_utf8_stdio() -> None:
    """Windows 默认 GBK 控制台打印含 emoji 的模型回复会报错，尽量切到 UTF-8。"""
    if sys.platform != "win32":
        return
    for name in ("stdout", "stderr"):
        stream = getattr(sys, name, None)
        if stream is not None and hasattr(stream, "reconfigure"):
            try:
                stream.reconfigure(encoding="utf-8", errors="replace")
            except Exception:
                pass


_ensure_utf8_stdio()

DEFAULT_ANDROID_PROMPT = "生成一个导出androidlog的脚本"
HOMEWORK_AGENT_NAME = "学生作业管理助手"
HOMEWORK_DEFAULT_MESSAGE = "你好"


def print_section(title: str) -> None:
    print("\n" + "=" * 80)
    print(f"  {title}")
    print("=" * 80)


def _login(
    base_url: str,
    username: str,
    password: str,
    timeout: int,
) -> Optional[Dict[str, str]]:
    login_data = {"username": username, "password": password}
    try:
        response = requests.post(
            f"{base_url}/api/v1/auth/login",
            data=login_data,
            headers={"Content-Type": "application/x-www-form-urlencoded"},
            timeout=timeout,
        )
        if response.status_code != 200:
            print(f"[FAIL] 登录失败: {response.status_code}")
            print(f"响应: {response.text[:800]}")
            return None
        token = response.json().get("access_token")
        if not token:
            print("[FAIL] 登录失败: 未获取到 token")
            return None
        print("[OK] 登录成功")
        return {"Authorization": f"Bearer {token}"}
    except Exception as e:
        print(f"[FAIL] 登录异常: {e}")
        return None


def _find_agent_by_name(
    base_url: str,
    headers: Dict[str, str],
    name: str,
    timeout: int,
) -> Optional[str]:
    try:
        response = requests.get(
            f"{base_url}/api/v1/agents",
            headers=headers,
            params={"search": name, "limit": 100},
            timeout=timeout,
        )
        if response.status_code != 200:
            print(f"[FAIL] 按名称查找 Agent 失败: {response.status_code}")
            print(response.text[:800])
            return None
        agents: List[Dict[str, Any]] = response.json() or []
        exact = [a for a in agents if (a.get("name") or "").strip() == name]
        pick = exact[0] if exact else (agents[0] if agents else None)
        if not pick:
            print(f"[FAIL] 未找到名为「{name}」的 Agent（search 无结果）")
            return None
        print(
            f"[OK] 使用 Agent: {pick.get('name')} (ID: {pick['id']}) "
            f"状态: {pick.get('status')}"
        )
        return str(pick["id"])
    except Exception as e:
        print(f"[FAIL] 查找 Agent 异常: {e}")
        return None


def _find_first_published_agent(
    base_url: str,
    headers: Dict[str, str],
    timeout: int,
) -> Optional[str]:
    try:
        response = requests.get(
            f"{base_url}/api/v1/agents",
            headers=headers,
            params={"status": "published", "limit": 10},
            timeout=timeout,
        )
        if response.status_code != 200:
            print(f"[FAIL] 获取 Agent 列表失败: {response.status_code}")
            print(response.text[:800])
            return None
        agents: List[Dict[str, Any]] = response.json() or []
        if not agents:
            print("[FAIL] 未找到可用的 Agent")
            print("请先创建并发布 Agent，或指定 agent_id / --agent-name / --homework")
            return None
        published_agents = [a for a in agents if a.get("status") == "published"]
        if published_agents:
            a = published_agents[0]
            print(f"[OK] 找到已发布的 Agent: {a['name']} (ID: {a['id']})")
            return str(a["id"])
        a = agents[0]
        print(
            f"[WARN] 使用 Agent: {a['name']} (ID: {a['id']}) - 状态: {a.get('status')}"
        )
        return str(a["id"])
    except Exception as e:
        print(f"[FAIL] 获取 Agent 列表异常: {e}")
        return None


def test_agent_execution(
    agent_id: Optional[str] = None,
    user_input: str = DEFAULT_ANDROID_PROMPT,
    *,
    base_url: str = DEFAULT_BASE_URL,
    username: str = "admin",
    password: str = "123456",
    agent_name: Optional[str] = None,
    request_timeout: int = 120,
    max_wait_time: int = 300,
    poll_interval: float = 2.0,
) -> None:
    """
    测试 Agent 执行

    Args:
        agent_id: Agent ID；为 None 时按 agent_name 或已发布列表解析
        user_input: 用户输入（写入 query / USER_INPUT）
        base_url: API 根地址
        agent_name: 按名称精确匹配查找（配合 search 参数）
        request_timeout: 单次 HTTP 超时（秒）
        max_wait_time: 轮询最长等待（秒）
        poll_interval: 轮询间隔（秒）
    """
    base_url = base_url.rstrip("/")
    print_section("Agent工作流执行测试")
    print(f"API: {base_url}")

    print_section("1. 用户登录")
    headers = _login(base_url, username, password, request_timeout)
    if not headers:
        return

    if not agent_id:
        print_section("2. 查找可用的 Agent")
        if agent_name:
            agent_id = _find_agent_by_name(
                base_url, headers, agent_name, request_timeout
            )
        else:
            agent_id = _find_first_published_agent(
                base_url, headers, request_timeout
            )
        if not agent_id:
            return
    else:
        print_section("2. 使用指定的 Agent")
        print(f"Agent ID: {agent_id}")

    print_section("3. 执行 Agent 工作流")
    print(f"用户输入: {user_input}")

    input_data = {"query": user_input, "USER_INPUT": user_input}
    execution_data = {"agent_id": agent_id, "input_data": input_data}

    try:
        response = requests.post(
            f"{base_url}/api/v1/executions",
            headers=headers,
            json=execution_data,
            timeout=request_timeout,
        )
        if response.status_code != 201:
            print(f"[FAIL] 创建执行任务失败: {response.status_code}")
            print(f"响应: {response.text[:2000]}")
            return
        execution = response.json()
        execution_id = execution["id"]
        print(f"[OK] 执行任务已创建: {execution_id}")
        print(f"状态: {execution.get('status')}")
    except Exception as e:
        print(f"[FAIL] 创建执行任务异常: {e}")
        return

    print_section("4. 等待执行完成")
    start_time = time.time()
    final_loop_status: Optional[str] = None

    while True:
        elapsed_time = time.time() - start_time
        if elapsed_time > max_wait_time:
            print(f"[FAIL] 执行超时（超过 {max_wait_time} 秒）")
            break
        try:
            status_response = requests.get(
                f"{base_url}/api/v1/executions/{execution_id}/status",
                headers=headers,
                timeout=request_timeout,
            )
            if status_response.status_code == 200:
                status = status_response.json()
                current_status = status.get("status")
                final_loop_status = current_status
                progress = status.get("progress", 0)
                print(
                    f"[...] 执行中 状态={current_status} 进度={progress}%",
                    end="\r",
                )
                if current_status == "completed":
                    print("\n[OK] 执行完成")
                    break
                if current_status == "failed":
                    print("\n[FAIL] 执行失败")
                    err = status.get("error") or status.get("error_message")
                    if not err and status.get("failed_nodes"):
                        fn = status["failed_nodes"][0]
                        err = fn.get("error_message") or fn.get("error_type")
                    print(f"错误信息: {err or '未知错误'}")
                    break
                if current_status in ("cancelled", "awaiting_approval"):
                    print(f"\n[WARN] 结束轮询: 状态={current_status}")
                    break
                current_node = status.get("current_node")
                if current_node:
                    nid = current_node.get("node_id")
                    ntype = current_node.get("node_type")
                    print(f"\n   当前节点: {nid} ({ntype})")
            time.sleep(poll_interval)
        except Exception as e:
            print(f"\n[FAIL] 获取执行状态异常: {e}")
            time.sleep(poll_interval)

    print_section("5. 获取执行结果")
    try:
        response = requests.get(
            f"{base_url}/api/v1/executions/{execution_id}",
            headers=headers,
            timeout=request_timeout,
        )
        if response.status_code == 200:
            execution = response.json()
            status = execution.get("status")
            output_data = execution.get("output_data")
            execution_time = execution.get("execution_time")
            err_msg = execution.get("error_message")

            print(f"执行状态: {status}")
            if err_msg and status != "completed":
                print(f"服务端错误信息: {err_msg[:2000]}")
            if execution_time is not None:
                print(f"执行时间: {execution_time}ms")

            print("\n输出结果:")
            print("-" * 80)
            if output_data:
                if isinstance(output_data, dict):
                    text_output = (
                        output_data.get("result")
                        or output_data.get("output")
                        or output_data.get("text")
                        or output_data.get("content")
                        or json.dumps(output_data, ensure_ascii=False, indent=2)
                    )
                    print(text_output)
                else:
                    print(output_data)
            else:
                print("(无输出数据)")
            print("-" * 80)

            if execution.get("logs"):
                print("\n执行日志:")
                for log in execution.get("logs", []):
                    print(f"  [{log.get('timestamp')}] {log.get('message')}")
        else:
            print(f"[FAIL] 获取执行结果失败: {response.status_code}")
            print(f"响应: {response.text[:2000]}")
    except Exception as e:
        print(f"[FAIL] 获取执行结果异常: {e}")

    print_section("测试完成")
    if final_loop_status and final_loop_status != "completed":
        sys.exit(1)


def _parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(
        description="Agent 工作流执行测试（input_data 与总结文档一致：query + USER_INPUT）"
    )
    p.add_argument("agent_id", nargs="?", default=None, help="Agent UUID（可选）")
    p.add_argument("user_input", nargs="?", default=None, help="用户消息（可选）")
    p.add_argument(
        "--homework",
        action="store_true",
        help=f"测试「{HOMEWORK_AGENT_NAME}」，默认发送「{HOMEWORK_DEFAULT_MESSAGE}」",
    )
    p.add_argument(
        "--agent-name",
        default=None,
        help="按名称精确查找 Agent（未传 agent_id 时）",
    )
    p.add_argument("--base-url", default=DEFAULT_BASE_URL, help="API 根地址")
    p.add_argument("--username", default="admin")
    p.add_argument("--password", default="123456")
    p.add_argument("--request-timeout", type=int, default=120, help="单次 HTTP 超时秒数")
    p.add_argument("--max-wait", type=int, default=300, help="轮询最长等待秒数")
    p.add_argument("--poll-interval", type=float, default=2.0)
    return p.parse_args()


if __name__ == "__main__":
    args = _parse_args()
    name: Optional[str] = None
    uid: Optional[str] = args.agent_id
    msg: str

    if args.homework and args.agent_name:
        print("[WARN] 同时指定 --homework 与 --agent-name，将使用 --agent-name 查找")

    if args.agent_name:
        name = args.agent_name
        msg = (
            args.user_input
            if args.user_input is not None
            else (
                HOMEWORK_DEFAULT_MESSAGE
                if args.homework
                else DEFAULT_ANDROID_PROMPT
            )
        )
    elif args.homework:
        name = HOMEWORK_AGENT_NAME
        msg = (
            args.user_input if args.user_input is not None else HOMEWORK_DEFAULT_MESSAGE
        )
    else:
        msg = (
            args.user_input
            if args.user_input is not None
            else DEFAULT_ANDROID_PROMPT
        )

    test_agent_execution(
        agent_id=uid,
        user_input=msg,
        base_url=args.base_url,
        username=args.username,
        password=args.password,
        agent_name=name if not uid else None,
        request_timeout=args.request_timeout,
        max_wait_time=args.max_wait,
        poll_interval=args.poll_interval,
    )