comfyui_o1key/tests/test_prompt_optimizer.py

"""Offline tests for the image and video AI-writing helpers."""

from __future__ import annotations

import asyncio
import base64
from io import BytesIO
import json
from pathlib import Path
import sys
import tempfile
import unittest

from PIL import Image


ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT))

from utils import chat_support as CHAT  # noqa: E402


class _FakeResponse:
    def __init__(self, status, payload):
        self.status = status
        self.payload = payload

    async def __aenter__(self):
        return self

    async def __aexit__(self, _exc_type, _exc, _tb):
        return False

    async def json(self, **_kwargs):
        return self.payload


class _FakeSession:
    def __init__(self, response):
        self.response = response
        self.calls = []

    def post(self, url, **kwargs):
        self.calls.append((url, kwargs))
        return self.response


class PromptOptimizerTests(unittest.TestCase):
    def test_payload_uses_fixed_model_high_reasoning_and_ordered_images(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            Image.new("RGB", (80, 120), "red").save(root / "first.png")
            Image.new("RGB", (120, 80), "blue").save(root / "second.png")

            payload = CHAT.build_prompt_optimization_payload(
                "保留人物身份，改成雨夜街景",
                [
                    {"name": "first.png", "subfolder": "", "type": "input"},
                    {"name": "second.png", "subfolder": "", "type": "input"},
                ],
                directory,
            )

        self.assertEqual(payload["model"], "gpt-5.6-sol")
        self.assertEqual(payload["reasoning_effort"], "high")
        self.assertFalse(payload["stream"])
        self.assertIn("视觉元素绑定优先", payload["messages"][0]["content"])
        self.assertIn("必须保持不变", payload["messages"][0]["content"])

        content = payload["messages"][1]["content"]
        self.assertEqual(content[1]["text"], "参考图 1（上传顺序第 1 张）")
        self.assertEqual(content[3]["text"], "参考图 2（上传顺序第 2 张）")
        image_items = [item for item in content if item["type"] == "image_url"]
        self.assertEqual(len(image_items), 2)
        for item in image_items:
            url = item["image_url"]["url"]
            self.assertTrue(url.startswith("data:image/jpeg;base64,"))
            with Image.open(BytesIO(base64.b64decode(url.split(",", 1)[1]))) as image:
                self.assertEqual(image.format, "JPEG")
        body_size = len(
            json.dumps(payload, ensure_ascii=False, separators=(",", ":")).encode("utf-8")
        )
        self.assertLessEqual(body_size, CHAT.PROMPT_OPTIMIZER_BODY_LIMIT_BYTES)

    def test_rejects_reference_path_outside_input_directory(self):
        with tempfile.TemporaryDirectory() as directory:
            with self.assertRaisesRegex(ValueError, "路径不安全"):
                CHAT.build_prompt_optimization_payload(
                    "测试提示词",
                    [{"name": "outside.png", "subfolder": "..", "type": "input"}],
                    directory,
                )

    def test_non_streaming_request_extracts_only_the_final_prompt(self):
        session = _FakeSession(_FakeResponse(200, {
            "choices": [{"message": {"content": "```text\n最终视觉提示词\n```"}}],
        }))
        with tempfile.TemporaryDirectory() as directory:
            result = asyncio.run(CHAT.optimize_image_prompt(
                session,
                "https://example.invalid",
                "offline-test-token",
                "原始提示词",
                [],
                directory,
            ))

        self.assertEqual(result, "最终视觉提示词")
        self.assertEqual(len(session.calls), 1)
        url, kwargs = session.calls[0]
        self.assertEqual(url, "https://example.invalid/v1/chat/completions")
        self.assertEqual(kwargs["json"]["model"], "gpt-5.6-sol")
        self.assertEqual(kwargs["json"]["reasoning_effort"], "high")
        self.assertFalse(kwargs["json"]["stream"])

    def test_video_payload_uses_dedicated_preset_and_ordered_frame_roles(self):
        with tempfile.TemporaryDirectory() as directory:
            root = Path(directory)
            Image.new("RGB", (96, 128), "green").save(root / "start.png")
            Image.new("RGB", (128, 96), "yellow").save(root / "end.png")
            payload = CHAT.build_video_prompt_writing_payload(
                "人物从门口走到窗前",
                [
                    {"name": "start.png", "subfolder": "", "type": "input"},
                    {"name": "end.png", "subfolder": "", "type": "input"},
                ],
                directory,
                {
                    "generation_mode": "first_last_frame",
                    "duration": "8",
                    "aspect_ratio": "16:9",
                    "generate_audio": True,
                    "reference_video_count": 1,
                    "reference_audio_count": 2,
                },
            )

        self.assertEqual(payload["model"], "gpt-5.6-sol")
        self.assertEqual(payload["reasoning_effort"], "high")
        self.assertFalse(payload["stream"])
        system_prompt = payload["messages"][0]["content"]
        self.assertIn("连续的时间过程", system_prompt)
        self.assertIn("首尾帧生视频", system_prompt)
        content = payload["messages"][1]["content"]
        self.assertIn("生成模式：首尾帧生视频", content[0]["text"])
        self.assertIn("目标时长：8", content[0]["text"])
        self.assertIn("生成音频：开启", content[0]["text"])
        self.assertIn("参考视频 1 个、参考音频 2 个", content[0]["text"])
        self.assertEqual(content[1]["text"], "首帧（第 1 张分析图）")
        self.assertEqual(content[3]["text"], "尾帧（第 2 张分析图）")
        self.assertEqual(len([item for item in content if item["type"] == "image_url"]), 2)

    def test_video_request_uses_non_streaming_completion(self):
        session = _FakeSession(_FakeResponse(200, {
            "choices": [{"message": {"content": "连续镜头中的最终视频提示词"}}],
        }))
        with tempfile.TemporaryDirectory() as directory:
            result = asyncio.run(CHAT.write_video_prompt(
                session,
                "https://example.invalid",
                "offline-test-token",
                "原始视频创意",
                [],
                directory,
                {"generation_mode": "text", "duration": "5"},
            ))

        self.assertEqual(result, "连续镜头中的最终视频提示词")
        self.assertEqual(session.calls[0][0], "https://example.invalid/v1/chat/completions")
        request = session.calls[0][1]["json"]
        self.assertIn("AI 视频生成提示词编写助手", request["messages"][0]["content"])
        self.assertFalse(request["stream"])


if __name__ == "__main__":
    unittest.main(verbosity=2)