aiagent/scripts/test_coding_agent.py

"""
测试代码编程助手 Agent — 验证流式和非流式对话是否正常工作。
"""
import json
import urllib.request
import urllib.parse
import time
import sys

BASE = "http://localhost:8037"
AGENT_ID = "010c0813-d45c-4c97-b3fc-21cedc6d4f9d"


def req(method, path, headers=None, body=None, raw_body=None, timeout=15):
    hdrs = {"Content-Type": "application/json"}
    if headers:
        hdrs.update(headers)
    data = raw_body if raw_body else (json.dumps(body).encode() if body else None)
    r = urllib.request.Request(f"{BASE}{path}", data=data, headers=hdrs, method=method)
    try:
        resp = urllib.request.urlopen(r, timeout=timeout)
        return resp.status, json.loads(resp.read())
    except urllib.request.HTTPError as e:
        return e.code, json.loads(e.read())
    except Exception as e:
        return 0, {"error": str(e)}


def login():
    _, _ = req("POST", "/api/v1/auth/register", body={
        "username": "codingbot", "email": "coding@test.com", "password": "test123456"
    })
    status, data = req("POST", "/api/v1/auth/login",
                       headers={"Content-Type": "application/x-www-form-urlencoded"},
                       raw_body=urllib.parse.urlencode(
                           {"username": "codingbot", "password": "test123456"}).encode())
    if status != 200:
        print(f"[FAIL] Login: {data}")
        sys.exit(1)
    token = data["access_token"]
    print(f"[OK] Login, token: {token[:20]}...")
    return {"Authorization": f"Bearer {token}", "Content-Type": "application/json"}


def test_non_streaming(auth, message, timeout=120):
    """测试非流式对话 POST /api/v1/agent-chat/{agent_id}"""
    body = json.dumps({"message": message, "temperature": 0.3}).encode()
    r = urllib.request.Request(
        f"{BASE}/api/v1/agent-chat/{AGENT_ID}",
        data=body, headers=auth, method="POST"
    )
    start = time.time()
    try:
        resp = urllib.request.urlopen(r, timeout=timeout)
        elapsed = time.time() - start
        result = json.loads(resp.read())
        content = result.get("content", "")
        print(f"[OK] 非流式 | {elapsed:6.1f}s | 内容={len(content)}字 | "
              f"迭代={result.get('iterations_used')} | "
              f"工具={result.get('tool_calls_made')} | "
              f"截断={result.get('truncated')}")
        print(f"     前100字: {content[:100]}")
        return True
    except urllib.request.HTTPError as e:
        elapsed = time.time() - start
        print(f"[FAIL] 非流式 | {elapsed:6.1f}s | HTTP {e.code}: {e.read().decode()[:200]}")
        return False
    except Exception as e:
        elapsed = time.time() - start
        print(f"[FAIL] 非流式 | {elapsed:6.1f}s | {e}")
        return False


def test_streaming(auth, message, timeout=120):
    """测试流式对话 POST /api/v1/agent-chat/{agent_id}/stream"""
    body = json.dumps({"message": message, "temperature": 0.3}).encode()
    r = urllib.request.Request(
        f"{BASE}/api/v1/agent-chat/{AGENT_ID}/stream",
        data=body, headers=auth, method="POST"
    )
    start = time.time()
    try:
        resp = urllib.request.urlopen(r, timeout=timeout)
        data = resp.read().decode()
        elapsed = time.time() - start

        # 解析 SSE 事件
        events = []
        for part in data.split("\n\n"):
            part = part.strip()
            if not part:
                continue
            lines = part.split("\n")
            event_type = ""
            event_data = {}
            for line in lines:
                if line.startswith("event: "):
                    event_type = line[7:]
                elif line.startswith("data: "):
                    try:
                        event_data = json.loads(line[6:])
                    except json.JSONDecodeError:
                        event_data = {"raw": line[6:]}
            if event_type:
                events.append({"type": event_type, "data": event_data})

        # 分析
        event_types = [e["type"] for e in events]
        content = ""
        for e in events:
            if e["type"] == "final":
                content = e["data"].get("content", "")

        print(f"[OK] 流式     | {elapsed:6.1f}s | {len(events)}个事件 | "
              f"内容={len(content)}字")
        print(f"     事件序列: {event_types}")
        print(f"     前100字: {content[:100]}")

        # 验证
        assert "final" in event_types, "缺少 final 事件"
        assert events[-1]["type"] == "final", "最后一个事件不是 final"
        print(f"     验证通过: final事件为末, 含内容")
        return True
    except urllib.request.HTTPError as e:
        elapsed = time.time() - start
        print(f"[FAIL] 流式     | {elapsed:6.1f}s | HTTP {e.code}: {e.read().decode()[:200]}")
        return False
    except Exception as e:
        elapsed = time.time() - start
        print(f"[FAIL] 流式     | {elapsed:6.1f}s | {e}")
        return False


def main():
    print("=" * 60)
    print("  代码编程助手 - Agent 对话测试")
    print("=" * 60)

    auth = login()

    test_cases = [
        ("问候", "你好"),
        ("代码", "写一个Python函数判断素数"),
        ("搜索", "grep_search工具怎么用？"),
        ("文件", "帮我读一下README.md的第一行"),
    ]

    passed = 0
    failed = 0

    for name, msg in test_cases:
        print(f"\n--- 测试: {name} ---")
        ok = test_non_streaming(auth, msg)
        if ok:
            passed += 1
        else:
            failed += 1

        ok = test_streaming(auth, msg)
        if ok:
            passed += 1
        else:
            failed += 1

    print("\n" + "=" * 60)
    print(f"  结果: {passed} 通过, {failed} 失败, 共 {passed + failed} 测试")
    print("=" * 60)

    if failed > 0:
        sys.exit(1)


if __name__ == "__main__":
    main()