mirror of
https://github.com/bytedance/deer-flow.git
synced 2026-09-15 09:08:38 +00:00
a22c6169b3
* feat(skills): support OpenAI-compatible image generation * fix(skills): address image provider review feedback
427 lines
15 KiB
Python
427 lines
15 KiB
Python
import base64
|
|
import sys
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
sys.path.insert(0, str(Path(__file__).resolve().parent))
|
|
from skill_loader import FakeResp, load # noqa: E402
|
|
|
|
img = load("image-generation")
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def clean_env(monkeypatch):
|
|
for k in [
|
|
"GEMINI_API_KEY",
|
|
"MINIMAX_API_KEY",
|
|
"IMAGE_GENERATION_PROVIDER",
|
|
"MINIMAX_API_HOST",
|
|
"MINIMAX_IMAGE_MODEL",
|
|
"IMAGE_GENERATION_API_KEY",
|
|
"IMAGE_GENERATION_BASE_URL",
|
|
"IMAGE_GENERATION_MODEL",
|
|
"IMAGE_GENERATION_SIZE",
|
|
]:
|
|
monkeypatch.delenv(k, raising=False)
|
|
|
|
|
|
def test_resolve_prefers_gemini(monkeypatch):
|
|
monkeypatch.setenv("GEMINI_API_KEY", "g")
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
assert (
|
|
img._resolve_provider("IMAGE_GENERATION_PROVIDER", "gemini", True) == "gemini"
|
|
)
|
|
|
|
|
|
def test_resolve_falls_back_to_minimax(monkeypatch):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
assert (
|
|
img._resolve_provider("IMAGE_GENERATION_PROVIDER", "gemini", False) == "minimax"
|
|
)
|
|
|
|
|
|
def test_resolve_falls_back_to_openai_after_minimax(monkeypatch):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
assert (
|
|
img._resolve_provider("IMAGE_GENERATION_PROVIDER", "gemini", False) == "minimax"
|
|
)
|
|
|
|
|
|
def test_resolve_falls_back_to_openai(monkeypatch):
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
assert (
|
|
img._resolve_provider("IMAGE_GENERATION_PROVIDER", "gemini", False) == "openai"
|
|
)
|
|
|
|
|
|
def test_resolve_override_wins(monkeypatch):
|
|
monkeypatch.setenv("GEMINI_API_KEY", "g")
|
|
monkeypatch.setenv("IMAGE_GENERATION_PROVIDER", "MiniMax")
|
|
assert (
|
|
img._resolve_provider("IMAGE_GENERATION_PROVIDER", "gemini", True) == "minimax"
|
|
)
|
|
|
|
|
|
def test_resolve_errors_when_none(monkeypatch):
|
|
with pytest.raises(ValueError, match="IMAGE_GENERATION_API_KEY"):
|
|
img._resolve_provider("IMAGE_GENERATION_PROVIDER", "gemini", False)
|
|
|
|
|
|
def test_minimax_builds_payload_and_writes(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
raw = b"PNGBYTES"
|
|
captured = {}
|
|
|
|
def fake_post(url, headers=None, json=None, **kw):
|
|
captured["url"] = url
|
|
captured["headers"] = headers
|
|
captured["json"] = json
|
|
return FakeResp(
|
|
{
|
|
"data": {"image_base64": [base64.b64encode(raw).decode()]},
|
|
"base_resp": {"status_code": 0, "status_msg": "success"},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
out = tmp_path / "o.jpg"
|
|
prompt_file = tmp_path / "p.json"
|
|
prompt_file.write_text("a red apple", encoding="utf-8")
|
|
msg = img.generate_image(str(prompt_file), [], str(out), "16:9")
|
|
|
|
assert out.read_bytes() == raw
|
|
assert captured["url"].endswith("/v1/image_generation")
|
|
assert captured["headers"]["Authorization"] == "Bearer m"
|
|
assert captured["json"]["model"] == "image-01"
|
|
assert captured["json"]["response_format"] == "base64"
|
|
assert captured["json"]["aspect_ratio"] == "16:9"
|
|
assert captured["json"]["n"] == 1
|
|
assert captured["json"]["prompt_optimizer"] is True
|
|
assert "Successfully generated image" in msg
|
|
|
|
|
|
def test_minimax_reference_image_as_data_url(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
captured = {}
|
|
|
|
def fake_post(url, headers=None, json=None, **kw):
|
|
captured["json"] = json
|
|
return FakeResp(
|
|
{
|
|
"data": {"image_base64": [base64.b64encode(b"x").decode()]},
|
|
"base_resp": {"status_code": 0},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
ref = tmp_path / "ref.jpg"
|
|
ref.write_bytes(b"\xff\xd8refbytes")
|
|
prompt_file = tmp_path / "p.json"
|
|
prompt_file.write_text("scene", encoding="utf-8")
|
|
img.generate_image(str(prompt_file), [str(ref)], str(tmp_path / "o.jpg"), "1:1")
|
|
|
|
subj = captured["json"]["subject_reference"]
|
|
assert subj[0]["type"] == "character"
|
|
assert subj[0]["image_file"].startswith("data:image/jpeg;base64,")
|
|
import base64 as _b64
|
|
|
|
encoded = subj[0]["image_file"].split(",", 1)[1]
|
|
assert _b64.b64decode(encoded) == b"\xff\xd8refbytes"
|
|
|
|
|
|
def test_minimax_raises_on_base_resp_error(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
|
|
def fake_post(url, headers=None, json=None, **kw):
|
|
return FakeResp(
|
|
{"base_resp": {"status_code": 1004, "status_msg": "auth failed"}}
|
|
)
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
prompt_file = tmp_path / "p.json"
|
|
prompt_file.write_text("x", encoding="utf-8")
|
|
with pytest.raises(Exception) as e:
|
|
img.generate_image(str(prompt_file), [], str(tmp_path / "o.jpg"), "1:1")
|
|
assert "1004" in str(e.value)
|
|
|
|
|
|
def test_minimax_extracts_json_prompt_field(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
captured = {}
|
|
|
|
def fake_post(url, headers=None, json=None, **kw):
|
|
captured["json"] = json
|
|
return FakeResp(
|
|
{
|
|
"data": {"image_base64": [base64.b64encode(b"x").decode()]},
|
|
"base_resp": {"status_code": 0},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
prompt_file = tmp_path / "p.json"
|
|
prompt_file.write_text(
|
|
'{"prompt": "a red barn at dawn", "style": "watercolor", '
|
|
'"composition": "rule of thirds", "negative_prompt": "blurry"}',
|
|
encoding="utf-8",
|
|
)
|
|
img.generate_image(str(prompt_file), [], str(tmp_path / "o.jpg"), "16:9")
|
|
|
|
# Only the JSON `prompt` field reaches MiniMax — no other fields, no JSON syntax.
|
|
assert captured["json"]["prompt"] == "a red barn at dawn"
|
|
assert captured["json"]["prompt_optimizer"] is True
|
|
|
|
|
|
def test_minimax_plaintext_prompt_passes_through(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
captured = {}
|
|
|
|
def fake_post(url, headers=None, json=None, **kw):
|
|
captured["json"] = json
|
|
return FakeResp(
|
|
{
|
|
"data": {"image_base64": [base64.b64encode(b"x").decode()]},
|
|
"base_resp": {"status_code": 0},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
prompt_file = tmp_path / "p.txt"
|
|
prompt_file.write_text("a red apple on a table", encoding="utf-8")
|
|
img.generate_image(str(prompt_file), [], str(tmp_path / "o.jpg"), "1:1")
|
|
|
|
assert captured["json"]["prompt"] == "a red apple on a table"
|
|
|
|
|
|
def test_minimax_rejects_overlong_prompt_without_calling_api(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
|
|
def fake_post(url, headers=None, json=None, **kw): # pragma: no cover
|
|
raise AssertionError("must not call the API when the prompt is over the limit")
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
prompt_file = tmp_path / "p.json"
|
|
prompt_file.write_text('{"prompt": "' + "x" * 1600 + '"}', encoding="utf-8")
|
|
out = tmp_path / "o.jpg"
|
|
msg = img.generate_image(str(prompt_file), [], str(out), "16:9")
|
|
|
|
assert "1500" in msg
|
|
assert "character" in msg.lower()
|
|
assert not out.exists()
|
|
|
|
|
|
def test_minimax_creates_nested_output_dir(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("MINIMAX_API_KEY", "m")
|
|
|
|
def fake_post(url, headers=None, json=None, **kw):
|
|
return FakeResp(
|
|
{
|
|
"data": {"image_base64": [base64.b64encode(b"img").decode()]},
|
|
"base_resp": {"status_code": 0},
|
|
}
|
|
)
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
prompt_file = tmp_path / "p.txt"
|
|
prompt_file.write_text("a cat", encoding="utf-8")
|
|
out = tmp_path / "nested" / "dir" / "o.jpg"
|
|
img.generate_image(str(prompt_file), [], str(out), "1:1")
|
|
|
|
assert out.read_bytes() == b"img"
|
|
|
|
|
|
def test_unknown_provider_raises(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("IMAGE_GENERATION_PROVIDER", "unknown")
|
|
monkeypatch.setenv("GEMINI_API_KEY", "g")
|
|
pf = tmp_path / "p.json"
|
|
pf.write_text("x", encoding="utf-8")
|
|
with pytest.raises(ValueError):
|
|
img.generate_image(str(pf), [], str(tmp_path / "o.jpg"), "1:1")
|
|
|
|
|
|
def test_openai_compatible_generation_writes_base64(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("IMAGE_GENERATION_PROVIDER", "openai")
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
monkeypatch.setenv("IMAGE_GENERATION_BASE_URL", "https://models.example/v1/")
|
|
captured = {}
|
|
|
|
def fake_post(url, headers=None, json=None, **kwargs):
|
|
captured.update(url=url, headers=headers, json=json, kwargs=kwargs)
|
|
return FakeResp({"data": [{"b64_json": base64.b64encode(b"image").decode()}]})
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
prompt_file = tmp_path / "prompt.json"
|
|
prompt_file.write_text('{"prompt": "a red deer"}', encoding="utf-8")
|
|
output_file = tmp_path / "nested" / "image.png"
|
|
|
|
result = img.generate_image(str(prompt_file), [], str(output_file), "16:9")
|
|
|
|
assert captured["url"] == "https://models.example/v1/images/generations"
|
|
assert captured["headers"]["Authorization"] == "Bearer image-key"
|
|
assert captured["json"]["model"] == "gpt-image-2.5-flare"
|
|
assert captured["json"]["size"] == "1536x1024"
|
|
assert captured["json"]["output_format"] == "png"
|
|
assert "response_format" not in captured["json"]
|
|
assert output_file.read_bytes() == b"image"
|
|
assert "Successfully generated image" in result
|
|
|
|
|
|
def test_openai_compatible_generation_downloads_url(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("IMAGE_GENERATION_PROVIDER", "openai-compatible")
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
|
|
monkeypatch.setattr(
|
|
img.requests,
|
|
"post",
|
|
lambda *args, **kwargs: FakeResp(
|
|
{"data": [{"url": "https://cdn.example/image.png"}]}
|
|
),
|
|
)
|
|
monkeypatch.setattr(
|
|
img.requests,
|
|
"get",
|
|
lambda *args, **kwargs: FakeResp(content=b"downloaded-image"),
|
|
)
|
|
prompt_file = tmp_path / "prompt.txt"
|
|
prompt_file.write_text("a red deer", encoding="utf-8")
|
|
output_file = tmp_path / "image.png"
|
|
|
|
img.generate_image(str(prompt_file), [], str(output_file), "1:1")
|
|
|
|
assert output_file.read_bytes() == b"downloaded-image"
|
|
|
|
|
|
def test_openai_compatible_dall_e_uses_response_format(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("IMAGE_GENERATION_PROVIDER", "openai")
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
monkeypatch.setenv("IMAGE_GENERATION_MODEL", "dall-e-3")
|
|
captured = {}
|
|
|
|
def fake_post(url, headers=None, json=None, **kwargs):
|
|
captured["json"] = json
|
|
return FakeResp({"data": [{"b64_json": base64.b64encode(b"image").decode()}]})
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
prompt_file = tmp_path / "prompt.txt"
|
|
prompt_file.write_text("a red deer", encoding="utf-8")
|
|
|
|
img.generate_image(str(prompt_file), [], str(tmp_path / "image.png"), "1:1")
|
|
|
|
assert captured["json"]["response_format"] == "b64_json"
|
|
assert "output_format" not in captured["json"]
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("model", "aspect_ratio", "expected"),
|
|
[
|
|
("gpt-image-2.5-flare", "16:9", "1536x1024"),
|
|
("gpt-image-2.5-flare", "9:16", "1024x1536"),
|
|
("dall-e-3", "16:9", "1792x1024"),
|
|
("dall-e-3", "9:16", "1024x1792"),
|
|
("dall-e-2", "16:9", "1024x1024"),
|
|
],
|
|
)
|
|
def test_openai_size_matches_model(model, aspect_ratio, expected):
|
|
assert img._openai_size(aspect_ratio, model) == expected
|
|
|
|
|
|
def test_openai_size_override_wins(monkeypatch):
|
|
monkeypatch.setenv("IMAGE_GENERATION_SIZE", "2048x1024")
|
|
|
|
assert img._openai_size("1:1", "gpt-image-2.5-flare") == "2048x1024"
|
|
|
|
|
|
def test_openai_dall_e_rejects_invalid_size_override(monkeypatch):
|
|
monkeypatch.setenv("IMAGE_GENERATION_SIZE", "1536x1024")
|
|
|
|
with pytest.raises(ValueError, match="dall-e-3 size must be one of"):
|
|
img._openai_size("16:9", "dall-e-3")
|
|
|
|
|
|
@pytest.mark.parametrize("extension", [".jpg", ".webp", ".unknown"])
|
|
def test_openai_dall_e_rejects_non_png_output(monkeypatch, tmp_path, extension):
|
|
monkeypatch.setenv("IMAGE_GENERATION_PROVIDER", "openai")
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
monkeypatch.setenv("IMAGE_GENERATION_MODEL", "dall-e-3")
|
|
monkeypatch.setattr(
|
|
img.requests,
|
|
"post",
|
|
lambda *args, **kwargs: pytest.fail("request must not be sent"),
|
|
)
|
|
prompt_file = tmp_path / "prompt.txt"
|
|
prompt_file.write_text("a red deer", encoding="utf-8")
|
|
|
|
with pytest.raises(ValueError, match=r"\.png extension"):
|
|
img.generate_image(
|
|
str(prompt_file), [], str(tmp_path / f"image{extension}"), "1:1"
|
|
)
|
|
|
|
|
|
@pytest.mark.parametrize("model", ["dall-e-2", "dall-e-3"])
|
|
def test_openai_dall_e_rejects_reference_images(monkeypatch, tmp_path, model):
|
|
monkeypatch.setenv("IMAGE_GENERATION_PROVIDER", "openai")
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
monkeypatch.setenv("IMAGE_GENERATION_MODEL", model)
|
|
monkeypatch.setattr(
|
|
img.requests,
|
|
"post",
|
|
lambda *args, **kwargs: pytest.fail("request must not be sent"),
|
|
)
|
|
prompt_file = tmp_path / "prompt.txt"
|
|
prompt_file.write_text("add snow", encoding="utf-8")
|
|
reference = tmp_path / "reference.png"
|
|
reference.write_bytes(b"png")
|
|
|
|
with pytest.raises(ValueError, match="reference-image editing is not supported"):
|
|
img.generate_image(
|
|
str(prompt_file), [str(reference)], str(tmp_path / "image.png"), "1:1"
|
|
)
|
|
|
|
|
|
def test_openai_compatible_reference_images_use_edits(monkeypatch, tmp_path):
|
|
monkeypatch.setenv("IMAGE_GENERATION_PROVIDER", "openai")
|
|
monkeypatch.setenv("IMAGE_GENERATION_API_KEY", "image-key")
|
|
captured = {}
|
|
|
|
def fake_post(url, headers=None, data=None, files=None, **kwargs):
|
|
captured.update(url=url, data=data, files=files)
|
|
return FakeResp({"data": [{"b64_json": base64.b64encode(b"edited").decode()}]})
|
|
|
|
monkeypatch.setattr(img.requests, "post", fake_post)
|
|
prompt_file = tmp_path / "prompt.txt"
|
|
prompt_file.write_text("add snow", encoding="utf-8")
|
|
reference = tmp_path / "reference.png"
|
|
reference.write_bytes(b"png")
|
|
output_file = tmp_path / "image.png"
|
|
|
|
img.generate_image(str(prompt_file), [str(reference)], str(output_file), "9:16")
|
|
|
|
assert captured["url"].endswith("/images/edits")
|
|
assert captured["data"]["size"] == "1024x1536"
|
|
assert captured["files"][0][0] == "image[]"
|
|
assert output_file.read_bytes() == b"edited"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
("filename", "expected"),
|
|
[
|
|
("image.png", "png"),
|
|
("image.jpg", "jpeg"),
|
|
("image.jpeg", "jpeg"),
|
|
("image.webp", "webp"),
|
|
("image.unknown", "png"),
|
|
],
|
|
)
|
|
def test_openai_output_format_matches_filename(filename, expected):
|
|
assert img._openai_output_format(filename) == expected
|
|
|
|
|
|
def test_guess_mime_by_extension():
|
|
assert img._guess_mime("/a/b.png") == "image/png"
|
|
assert img._guess_mime("/a/b.webp") == "image/webp"
|
|
assert img._guess_mime("/a/b.jpg") == "image/jpeg"
|
|
assert img._guess_mime("/a/b.unknown") == "image/jpeg"
|