test_run.py

from __future__ import annotations

import asyncio
import logging
import os
from pathlib import Path

from agents.cognitive_annotator import CognitiveAnnotatorAgent


PROJECT_ROOT = Path(__file__).resolve().parent
IMAGE_DIR = PROJECT_ROOT / "clearn_base_data"
OUTPUT_PATH = PROJECT_ROOT / "test_sft_data.jsonl"
IMAGE_SUFFIXES = {".jpg", ".jpeg", ".png", ".bmp", ".webp"}

QWEN_API_KEY = os.getenv("QWEN_API_KEY")
QWEN_BASE_URL = os.getenv("QWEN_BASE_URL", "https://dashscope.aliyuncs.com/compatible-mode/v1")
QWEN_MODEL = os.getenv("QWEN_MODEL", "qwen3.5-plus")


async def main() -> None:
    """运行 CognitiveAnnotatorAgent 的最小异步测试。"""
    logging.basicConfig(
        level=logging.INFO,
        format="%(asctime)s | %(levelname)s | %(name)s | %(message)s",
    )

    if not QWEN_API_KEY:
        raise ValueError("未读取到 QWEN_API_KEY，请先设置环境变量后再运行测试脚本。")

    if not IMAGE_DIR.exists():
        raise FileNotFoundError(f"图片目录不存在: {IMAGE_DIR}")

    image_paths = sorted(
        str(path)
        for path in IMAGE_DIR.iterdir()
        if path.is_file() and path.suffix.lower() in IMAGE_SUFFIXES
    )[:3]

    if not image_paths:
        raise FileNotFoundError(f"目录中未找到可用图片: {IMAGE_DIR}")

    print(f"本次测试图片目录: {IMAGE_DIR}")
    print(f"本次测试输出文件: {OUTPUT_PATH}")
    print(f"Qwen Base URL: {QWEN_BASE_URL}")
    print(f"Qwen Model: {QWEN_MODEL}")
    print("本次仅测试前 3 张图片:")
    for image_path in image_paths:
        print(f"- {image_path}")

    agent = CognitiveAnnotatorAgent(
        config={
            "api_key": QWEN_API_KEY,
            "base_url": QWEN_BASE_URL,
            "model": QWEN_MODEL,
            "max_concurrency": 3,
            "max_retries": 3,
            "request_timeout": 120,
            "output_path": str(OUTPUT_PATH),
        }
    )

    try:
        result = await agent.annotate_images_async(image_paths)
    finally:
        await agent.aclose()

    print("\nAgent 返回结果:")
    print(result)

    if OUTPUT_PATH.exists():
        print(f"\nJSONL 文件已生成: {OUTPUT_PATH}")
        with OUTPUT_PATH.open("r", encoding="utf-8") as file:
            for index, line in enumerate(file):
                if index >= 3:
                    break
                print(line.rstrip())
    else:
        print("\n未检测到输出 JSONL 文件，请检查日志与 API 配置。")


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