Agent/examples/test_nanobanana/run.py

"""
NanoBanana 工具测试脚本

测试场景：
1. 纯文本生成图像
2. 基于第一个测试生成的图像，转换为不同风格
"""

import asyncio
import sys
from pathlib import Path
from dotenv import load_dotenv

# 添加项目根目录到 Python 路径
project_root = Path(__file__).parent.parent.parent
sys.path.insert(0, str(project_root))

# 从 Agent 根目录加载 .env 文件
env_path = project_root / ".env"
if env_path.exists():
    load_dotenv(env_path)
    print(f"✅ 已加载环境变量: {env_path}")
else:
    print(f"⚠️  未找到 .env 文件: {env_path}")

from agent.core.runner import AgentRunner, RunConfig
from agent.trace.store import FileSystemTraceStore
from agent.trace.models import Trace, Message
from agent.llm.openrouter import create_openrouter_llm_call


async def test_text_to_image(output_dir: Path) -> Path:
    """测试1: 纯文本生成图像

    Returns:
        生成的图片路径
    """
    print("\n" + "="*60)
    print("测试1: 纯文本生成图像")
    print("="*60)

    # 导入自定义工具
    import examples.test_nanobanana.tool

    store = FileSystemTraceStore(base_path=".trace")
    llm_call = create_openrouter_llm_call(model="anthropic/claude-sonnet-4.5")

    runner = AgentRunner(
        trace_store=store,
        llm_call=llm_call,
    )

    config = RunConfig(
        model="anthropic/claude-sonnet-4.5",
        temperature=0.3,
        max_iterations=10,
    )

    # 使用绝对路径
    output_path = output_dir / "cat.png"

    task = f"""
请使用 nanobanana 工具生成一张图片：
- 内容：一只可爱的橙色小猫，坐在窗台上看着外面的雨
- 风格：水彩画风格
- 保存到：{output_path}

注意：路径必须使用绝对路径。
"""

    print(f"\n任务: 生成小猫图片")
    print(f"输出路径: {output_path}\n")

    # 将任务转换为消息格式
    messages = [{"role": "user", "content": task}]

    generated_image = None
    try:
        async for item in runner.run(messages=messages, config=config):
            # 处理 Trace 对象（整体状态变化）
            if isinstance(item, Trace):
                if item.status == "running":
                    print(f"[Trace] 开始: {item.trace_id[:8]}...")
                elif item.status == "completed":
                    print(f"\n[Trace] ✅ 完成")
                    print(f"  - Total messages: {item.total_messages}")
                    print(f"  - Total tokens: {item.total_tokens}")
                    print(f"  - Total cost: ${item.total_cost:.4f}")
                    if output_path.exists():
                        generated_image = output_path
                        print(f"  - 生成的图片: {generated_image}")
                elif item.status == "failed":
                    print(f"\n[Trace] ❌ 失败: {item.error_message}")
                elif item.status == "stopped":
                    print(f"\n[Trace] ⏸️ 已停止")

            # 处理 Message 对象（执行过程）
            elif isinstance(item, Message):
                if item.role == "assistant":
                    content = item.content
                    if isinstance(content, dict):
                        text = content.get("text", "")
                        tool_calls = content.get("tool_calls")

                        if text and not tool_calls:
                            # 纯文本回复（最终响应）
                            print(f"\n[Response] {text}")
                        elif tool_calls:
                            # 工具调用
                            tool_names = [tc.get("function", {}).get("name") for tc in tool_calls]
                            print(f"[Tool Call] {', '.join(tool_names)}")

                elif item.role == "tool":
                    # 工具结果
                    if isinstance(item.content, dict):
                        tool_name = item.content.get("tool_name", "unknown")
                        print(f"[Tool Result] {tool_name}")
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()

    return generated_image


async def test_image_to_image(input_image: Path, output_dir: Path):
    """测试2: 基于参考图像生成新图像

    Args:
        input_image: 参考图像路径（来自测试1）
        output_dir: 输出目录
    """
    print("\n" + "="*60)
    print("测试2: 基于参考图像生成新图像")
    print("="*60)

    if not input_image or not input_image.exists():
        print("⚠️  跳过测试2: 测试1没有成功生成图片")
        return

    print(f"参考图像: {input_image}")

    # 导入自定义工具
    import examples.test_nanobanana.tool

    store = FileSystemTraceStore(base_path=".trace")
    llm_call = create_openrouter_llm_call(model="anthropic/claude-sonnet-4.5")

    runner = AgentRunner(
        trace_store=store,
        llm_call=llm_call,
    )

    config = RunConfig(
        model="anthropic/claude-sonnet-4.5",
        temperature=0.3,
        max_iterations=10,
    )

    # 使用绝对路径
    output_path = output_dir / "cat_oil_painting.png"

    task = f"""
请使用 nanobanana 工具，基于参考图像生成一张新图片：
- 参考图像：{input_image}
- 要求：保持小猫和窗台的元素，但改变风格为油画风格
- 保存到：{output_path}

注意：路径必须使用绝对路径。
"""

    print(f"\n任务: 将小猫图片转换为油画风格")
    print(f"输出路径: {output_path}\n")

    # 将任务转换为消息格式
    messages = [{"role": "user", "content": task}]

    try:
        async for item in runner.run(messages=messages, config=config):
            # 处理 Trace 对象（整体状态变化）
            if isinstance(item, Trace):
                if item.status == "running":
                    print(f"[Trace] 开始: {item.trace_id[:8]}...")
                elif item.status == "completed":
                    print(f"\n[Trace] ✅ 完成")
                    print(f"  - Total messages: {item.total_messages}")
                    print(f"  - Total tokens: {item.total_tokens}")
                    print(f"  - Total cost: ${item.total_cost:.4f}")
                    if output_path.exists():
                        print(f"  - 生成的图片: {output_path}")
                elif item.status == "failed":
                    print(f"\n[Trace] ❌ 失败: {item.error_message}")
                elif item.status == "stopped":
                    print(f"\n[Trace] ⏸️ 已停止")

            # 处理 Message 对象（执行过程）
            elif isinstance(item, Message):
                if item.role == "assistant":
                    content = item.content
                    if isinstance(content, dict):
                        text = content.get("text", "")
                        tool_calls = content.get("tool_calls")

                        if text and not tool_calls:
                            # 纯文本回复（最终响应）
                            print(f"\n[Response] {text}")
                        elif tool_calls:
                            # 工具调用
                            tool_names = [tc.get("function", {}).get("name") for tc in tool_calls]
                            print(f"[Tool Call] {', '.join(tool_names)}")

                elif item.role == "tool":
                    # 工具结果
                    if isinstance(item.content, dict):
                        tool_name = item.content.get("tool_name", "unknown")
                        print(f"[Tool Result] {tool_name}")
    except Exception as e:
        print(f"\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()


async def main():
    """运行所有测试"""
    print("\n" + "="*60)
    print("NanoBanana 工具测试")
    print("="*60)

    # 创建输出目录（使用绝对路径）
    output_dir = Path(__file__).parent / "output"
    output_dir.mkdir(exist_ok=True)
    print(f"输出目录: {output_dir.absolute()}\n")

    try:
        # 测试1: 纯文本生成
        generated_image = await test_text_to_image(output_dir)

        # 测试2: 图像转换（基于测试1的结果）
        if generated_image:
            await test_image_to_image(generated_image, output_dir)
        else:
            print("\n⚠️  跳过测试2: 测试1未成功生成图片")

        print("\n" + "="*60)
        print("测试完成！")
        print("="*60)
        print(f"\n查看生成的图片:")
        if generated_image and generated_image.exists():
            print(f"  1. {generated_image}")
        oil_painting = output_dir / "cat_oil_painting.png"
        if oil_painting.exists():
            print(f"  2. {oil_painting}")

    except KeyboardInterrupt:
        print("\n\n⚠️  测试被用户中断")
    except Exception as e:
        print(f"\n\n❌ 测试失败: {e}")
        import traceback
        traceback.print_exc()


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