Tool_Agent/tools/local/gpt_image_2/tests/test_generate.py

import os
import sys
import base64
import re
import requests
from pathlib import Path
from dotenv import load_dotenv

load_dotenv()

# 添加父目录到路径
sys.path.insert(0, str(Path(__file__).parent.parent))

from openai import OpenAI

api_key = os.getenv("APIYI_KEY")
if not api_key:
    raise ValueError("APIYI_KEY environment variable is required")

client = OpenAI(
    api_key=api_key,
    base_url="https://api.apiyi.com/v1"
)

def extract_url_from_markdown(text):
    """从 markdown 格式中提取 URL"""
    # 匹配 ![image](url) 格式
    match = re.search(r'!\[.*?\]\((https?://[^\)]+)\)', text)
    if match:
        return match.group(1)
    # 如果不是 markdown 格式，直接返回
    return text.strip()

def test_generate_image():
    print("Testing GPT Image 2 generation with gpt-image-2-all model...")
    
    # 使用 chat/completions 接口
    response = client.chat.completions.create(
        model="gpt-image-2-all",
        messages=[
            {"role": "user", "content": "1024x1024 square image, a cute cat sitting on a sunny windowsill"}
        ],
        extra_body={"response_format": {"type": "url"}}
    )
    
    assert response.choices and len(response.choices) > 0, "No response returned"
    content = response.choices[0].message.content
    assert content, "Empty content"
    
    # 提取真实的 URL
    image_url = extract_url_from_markdown(content)
    print(f"[SUCCESS] Image generated successfully!")
    print(f"   Raw content: {content}")
    print(f"   Image URL: {image_url}")
    
    # 下载图像保存到本地
    output_dir = Path(__file__).parent / "output"
    output_dir.mkdir(exist_ok=True)
    
    img_response = requests.get(image_url, timeout=30)
    output_path = output_dir / "test_cat.png"
    output_path.write_bytes(img_response.content)
    
    print(f"   Saved to: {output_path}")
    print(f"   File size: {len(img_response.content)} bytes")
    return True

if __name__ == "__main__":
    test_generate_image()