mirror of
https://github.com/pewdiepie-archdaemon/odysseus.git
synced 2026-10-16 11:52:25 +02:00
Server-side photo editor: natural-language editor agent route
(/api/editor/agent), photo_editor tool family + style guide, image-op
provider proxy with progress, email server hardening, gallery/hwfit
routes, and graceful shutdown.
Source: vendor/photocraft (ad1cff44) editor-server files, merged onto
upstream/dev (a8c147b2). Conflict resolution documented in
pr_work/B_conflicts.md.
Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
1493 lines
62 KiB
Python
1493 lines
62 KiB
Python
"""Tests for routes/editor_agent_routes.py (NL editor planning + bridge).
|
|
|
|
The route wires an injected planner callable (messages, owner) -> raw plan and always
|
|
falls back to a deterministic keyword planner, so the editor gets a usable plan even
|
|
without a model. Tests use a minimal FastAPI app (no full app import) and inject fake
|
|
planners to assert parse/vision/check-mode/bridge-validation behavior.
|
|
"""
|
|
import json
|
|
|
|
import pytest
|
|
from fastapi import FastAPI
|
|
from fastapi.testclient import TestClient
|
|
|
|
from routes import editor_agent_routes as eaar
|
|
|
|
|
|
@pytest.fixture(autouse=True)
|
|
def _reset_planner():
|
|
eaar.set_editor_planner(None)
|
|
yield
|
|
eaar.set_editor_planner(None)
|
|
|
|
|
|
@pytest.fixture()
|
|
def client():
|
|
app = FastAPI()
|
|
app.include_router(eaar.setup_editor_agent_routes())
|
|
return TestClient(app)
|
|
|
|
|
|
def test_plan_uses_injected_planner_and_parses_steps(client):
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: json.dumps(
|
|
{"reply": "ok", "steps": [{"op": "adjustments.apply", "params": {"type": "vibrance"}}]}
|
|
)
|
|
)
|
|
r = client.post("/api/editor/plan", json={"prompt": "make it vivid", "state": {}})
|
|
assert r.status_code == 200
|
|
data = r.json()
|
|
assert data["planner"] == "llm"
|
|
assert data["steps"] == [{"op": "adjustments.apply", "params": {"type": "vibrance"}}]
|
|
|
|
|
|
def test_plan_sends_screenshot_as_vision_part(client):
|
|
seen = {}
|
|
|
|
def fake(messages, owner):
|
|
seen["messages"] = messages
|
|
return json.dumps({"reply": "ok", "steps": [{"op": "filters.apply", "params": {"id": "sharpen"}}]})
|
|
|
|
eaar.set_editor_planner(fake)
|
|
r = client.post(
|
|
"/api/editor/plan",
|
|
json={"prompt": "sharpen it", "state": {}, "screenshot": "data:image/png;base64,AAAA"},
|
|
)
|
|
assert r.status_code == 200
|
|
content = seen["messages"][1]["content"]
|
|
assert any(p.get("type") == "image_url" for p in content if isinstance(p, dict))
|
|
assert any(
|
|
p.get("type") == "image_url" and p["image_url"]["url"].startswith("data:image/")
|
|
for p in content if isinstance(p, dict)
|
|
)
|
|
|
|
|
|
def test_plan_without_screenshot_sends_plain_text(client):
|
|
seen = {}
|
|
|
|
def fake(messages, owner):
|
|
seen["messages"] = messages
|
|
return json.dumps({"reply": "ok", "steps": [{"op": "filters.apply", "params": {"id": "sharpen"}}]})
|
|
|
|
eaar.set_editor_planner(fake)
|
|
client.post("/api/editor/plan", json={"prompt": "sharpen it", "state": {}})
|
|
assert isinstance(seen["messages"][1]["content"], str)
|
|
|
|
|
|
def test_plan_falls_back_when_planner_yields_nothing(client):
|
|
eaar.set_editor_planner(lambda messages, owner: None)
|
|
r = client.post("/api/editor/plan", json={"prompt": "make it black and white", "state": {}})
|
|
data = r.json()
|
|
assert data["planner"] == "fallback"
|
|
assert any(op["op"] == "adjustments.apply" for op in data["steps"])
|
|
|
|
|
|
def test_plan_parses_dict_return_and_filters_unknown_ops(client):
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: {
|
|
"reply": "x",
|
|
"steps": [{"op": "evil.op"}, {"op": "text.add", "params": {"text": "hi"}}],
|
|
}
|
|
)
|
|
data = client.post("/api/editor/plan", json={"prompt": "add text", "state": {}}).json()
|
|
assert [s["op"] for s in data["steps"]] == ["text.add"]
|
|
|
|
|
|
def test_plan_awaits_async_planner(client):
|
|
async def fake(messages, owner):
|
|
return json.dumps({"reply": "ok", "steps": [{"op": "history.mark", "params": {}}]})
|
|
|
|
eaar.set_editor_planner(fake)
|
|
data = client.post("/api/editor/plan", json={"prompt": "mark", "state": {}}).json()
|
|
assert data["steps"][0]["op"] == "history.mark"
|
|
|
|
|
|
def test_check_mode_reports_done_when_model_completes(client):
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps({"done": True, "reply": "looks good"}))
|
|
data = client.post(
|
|
"/api/editor/plan",
|
|
json={"prompt": "make it warm", "state": {"iteration": 2}, "mode": "check",
|
|
"screenshot": "data:image/png;base64,AAAA"},
|
|
).json()
|
|
assert data["done"] is True
|
|
assert data["steps"] == []
|
|
|
|
|
|
def test_check_mode_returns_remaining_steps(client):
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: json.dumps(
|
|
{"done": False, "reply": "more contrast", "steps": [{"op": "adjustments.apply", "params": {}}]}
|
|
)
|
|
)
|
|
data = client.post(
|
|
"/api/editor/plan",
|
|
json={"prompt": "make it vivid", "state": {"iteration": 1}, "mode": "check"},
|
|
).json()
|
|
assert data.get("done") is not True
|
|
assert data["steps"][0]["op"] == "adjustments.apply"
|
|
|
|
|
|
def test_check_mode_with_unparseable_planner_stops_loop(client):
|
|
eaar.set_editor_planner(lambda messages, owner: "not json at all")
|
|
data = client.post(
|
|
"/api/editor/plan", json={"prompt": "x", "state": {}, "mode": "check"}
|
|
).json()
|
|
assert data["done"] is True
|
|
assert data["steps"] == []
|
|
|
|
|
|
def test_normalize_steps_drops_ops_outside_catalog():
|
|
assert eaar._normalize_steps({"op": "evil.op"}) == []
|
|
assert eaar._normalize_steps(
|
|
{"steps": [{"op": "text.add"}, {"op": "nope"}, {"not": "a step"}]}
|
|
) == [{"op": "text.add", "params": {}}]
|
|
|
|
|
|
def test_bridge_command_only_queues_known_ops(client):
|
|
good = client.post("/api/editor/bridge/command", json={"session": "s", "op": "screenshot"}).json()
|
|
assert good["queued"]
|
|
bad = client.post("/api/editor/bridge/command", json={"session": "s", "op": "evil.op"}).json()
|
|
assert bad["queued"] == []
|
|
|
|
|
|
def test_bridge_command_accepts_selection_export_mask(client):
|
|
# The generative tools fetch the current selection via this op; if it is not
|
|
# in the catalog the bridge drops the command and the mask comes back empty.
|
|
assert "selection.exportMask" in eaar._OP_SET
|
|
ok = client.post(
|
|
"/api/editor/bridge/command",
|
|
json={"session": "s", "op": "selection.exportMask", "wait": False},
|
|
).json()
|
|
assert ok["queued"], ok
|
|
|
|
|
|
def test_set_editor_planner_is_used_by_setup_when_set():
|
|
sentinel = lambda messages, owner: None # noqa: E731
|
|
eaar.set_editor_planner(sentinel)
|
|
app = FastAPI()
|
|
app.include_router(eaar.setup_editor_agent_routes())
|
|
assert eaar._EDITOR_PLANNER is sentinel
|
|
|
|
|
|
def _recipe_ops(steps):
|
|
out = []
|
|
for step in steps:
|
|
params = step.get("params") or {}
|
|
if step.get("op") in ("adjustments.apply", "adjustments.applyMasked"):
|
|
out.append((step.get("op"), params.get("type")))
|
|
elif step.get("op") in ("filters.apply", "filters.applyMasked"):
|
|
out.append((step.get("op"), params.get("id")))
|
|
else:
|
|
out.append((step.get("op"), None))
|
|
return out
|
|
|
|
|
|
def test_agent_path_enforces_v3_recipes_over_a_timid_model_plan(client, monkeypatch):
|
|
"""The /api/editor/agent path (not just /plan) must follow the style guide.
|
|
|
|
The sampling model returns one mild brightness-contrast step for every
|
|
generic intent; the guide is authoritative, so the calibrated recipe ops
|
|
must appear in the agent response steps.
|
|
"""
|
|
weak = [{"op": "adjustments.apply",
|
|
"params": {"type": "brightness-contrast",
|
|
"params": {"brightness": 10, "contrast": 10}}}]
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps({"reply": "ok", "steps": weak}))
|
|
|
|
async def _fake_wait(fut, timeout=None, **kwargs): # noqa: ANN001
|
|
return {"ok": True, "result": {"width": 100, "height": 100}}
|
|
|
|
monkeypatch.setattr(eaar.asyncio, "wait_for", _fake_wait)
|
|
|
|
cases = {
|
|
"restore this old photo": {("photo.autoTone", None),
|
|
("adjustments.apply", "white-balance"),
|
|
("adjustments.apply", "vibrance")},
|
|
"make it pretty": {("photo.autoTone", None),
|
|
("adjustments.apply", "vibrance"),
|
|
("filters.apply", "clarity")},
|
|
"make it pop": {("photo.autoTone", None),
|
|
("adjustments.apply", "vibrance"),
|
|
("filters.apply", "clarity")},
|
|
"make it look like summer": {("adjustments.apply", "white-balance")},
|
|
}
|
|
for prompt, required in cases.items():
|
|
resp = client.post("/api/editor/agent",
|
|
json={"prompt": prompt, "state": {"width": 100, "height": 100}})
|
|
assert resp.status_code == 200, (prompt, resp.text)
|
|
data = resp.json()
|
|
ops = set(_recipe_ops(data["steps"]))
|
|
assert required <= ops, (prompt, ops)
|
|
if prompt != "make it look like summer":
|
|
assert len(data["steps"]) >= 3, (prompt, data["steps"])
|
|
|
|
|
|
def test_agent_path_low_key_lighting_uses_auto_tone_plus_one_shadow_lift(client, monkeypatch):
|
|
"""Round-6 case 5: 'fix the lighting' on a night photo.
|
|
|
|
The model returned a single adjustments.apply step while claiming an
|
|
auto-tone pass; the guide must supply photo.autoTone + one shadows lift,
|
|
and the executed steps must contain the auto tone.
|
|
"""
|
|
weak = [{"op": "adjustments.apply",
|
|
"params": {"type": "shadows-highlights",
|
|
"params": {"shadows": 20, "highlights": -10}}}]
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps(
|
|
{"reply": "A single auto-tone pass.", "steps": weak}))
|
|
|
|
async def _fake_wait(fut, timeout=None, **kwargs): # noqa: ANN001
|
|
return {"ok": True, "result": {"width": 900, "height": 600}}
|
|
|
|
monkeypatch.setattr(eaar.asyncio, "wait_for", _fake_wait)
|
|
|
|
resp = client.post("/api/editor/agent",
|
|
json={"prompt": "fix the lighting", "state": {"width": 900, "height": 600}})
|
|
assert resp.status_code == 200, resp.text
|
|
data = resp.json()
|
|
ops = _recipe_ops(data["steps"])
|
|
assert ("photo.autoTone", None) in ops, ops
|
|
shadow = [s for s in data["steps"] if s.get("params", {}).get("type") == "shadows-highlights"]
|
|
assert shadow and shadow[0]["params"]["params"]["highlights"] <= 0, data["steps"]
|
|
assert len(data["steps"]) == 2, data["steps"]
|
|
|
|
|
|
def test_agent_path_reply_describes_executed_ops_not_the_plan(client, monkeypatch):
|
|
"""A reply must not claim an edit (auto tone / white balance) that the
|
|
executed steps do not contain; it is rewritten from the real steps."""
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps({
|
|
"reply": "Applied a single auto-tone pass and white balance.",
|
|
"steps": [{"op": "adjustments.apply",
|
|
"params": {"type": "vibrance", "params": {"vibrance": 20}}}],
|
|
}))
|
|
|
|
async def _fake_wait(fut, timeout=None, **kwargs): # noqa: ANN001
|
|
return {"ok": True, "result": {"width": 100, "height": 100}}
|
|
|
|
monkeypatch.setattr(eaar.asyncio, "wait_for", _fake_wait)
|
|
|
|
resp = client.post("/api/editor/agent",
|
|
json={"prompt": "boost the saturation", "state": {"width": 100, "height": 100}})
|
|
assert resp.status_code == 200, resp.text
|
|
data = resp.json()
|
|
reply = data["reply"].lower()
|
|
assert "auto" not in reply, data["reply"]
|
|
assert "white balance" not in reply, data["reply"]
|
|
assert "vibrance" in reply, data["reply"]
|
|
|
|
|
|
def _install_provider_and_fake_tools(monkeypatch, *, failing=False):
|
|
"""Provider is available and the AI content tools resolve instantly."""
|
|
monkeypatch.setattr(eaar, "_image_provider_available", lambda owner: True)
|
|
|
|
class _FakeTool:
|
|
async def execute(self, content, ctx): # noqa: ANN001
|
|
if failing:
|
|
return ("ai_edit_image", {"exit_code": 1, "error": "provider exploded"})
|
|
return (
|
|
"ai_edit_image",
|
|
{
|
|
"exit_code": 0,
|
|
"image_id": "gen-1",
|
|
"image_url": "/api/generated-image/gen-1.png",
|
|
"placed": True,
|
|
"output": "applied",
|
|
},
|
|
)
|
|
|
|
from src.agent_tools import editor_tools
|
|
|
|
monkeypatch.setattr(editor_tools, "AiEditImageTool", _FakeTool)
|
|
monkeypatch.setattr(editor_tools, "GenerateImageLayerTool", _FakeTool)
|
|
|
|
async def _fake_wait(fut, timeout=None, **kwargs): # noqa: ANN001
|
|
return {"ok": True, "result": {"width": 100, "height": 100}}
|
|
|
|
monkeypatch.setattr(eaar.asyncio, "wait_for", _fake_wait)
|
|
|
|
|
|
def test_agent_path_routes_generative_intents_when_provider_present(client, monkeypatch):
|
|
"""Gen-6 bug 4: with an image provider, object-removal / background-change
|
|
prompts must use ai_edit_image, and plain 'blur the background' stays a blur."""
|
|
_install_provider_and_fake_tools(monkeypatch)
|
|
|
|
def _ai_ops(data):
|
|
return {s.get("op") for s in data["steps"]}
|
|
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "remove the distractions", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
assert "ai_edit_image" in _ai_ops(r.json()), r.json()["steps"]
|
|
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "make the background nicer", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
assert "ai_edit_image" in _ai_ops(r.json()), r.json()["steps"]
|
|
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "add a full moon in the sky", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
assert "ai_edit_image" in _ai_ops(r.json()), r.json()["steps"]
|
|
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "remove the person on the left", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
assert "ai_edit_image" in _ai_ops(r.json()), r.json()["steps"]
|
|
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "create a new image of a full moon", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
assert "generate_image" in _ai_ops(r.json()), r.json()["steps"]
|
|
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "blur the background", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
assert _ai_ops(r.json()) & {"ai_edit_image", "generate_image"} == set(), r.json()["steps"]
|
|
|
|
|
|
def test_agent_path_recolor_is_exactly_object_mask_plus_masked_hue(client, monkeypatch):
|
|
"""#125/#127: on the agent path a recolour must be EXACTLY
|
|
[selection.samText, adjustments.applyMasked] — the VLM planner used to
|
|
append whole-image adjustments.apply after the masked hue shift."""
|
|
_install_provider_and_fake_tools(monkeypatch)
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: json.dumps({
|
|
"reply": "recolor",
|
|
"steps": [
|
|
{"op": "selection.samText", "params": {"text": "red car"}},
|
|
{"op": "adjustments.applyMasked",
|
|
"params": {"mode": "hue-saturation", "hue": -120, "sourceHue": 0}},
|
|
{"op": "adjustments.apply", "params": {"type": "vibrance", "params": {"vibrance": 20}}},
|
|
{"op": "adjustments.apply", "params": {"type": "brightness-contrast", "params": {"brightness": 15}}},
|
|
],
|
|
})
|
|
)
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "make only the red car blue",
|
|
"session": "s1", "editor_session": "s1",
|
|
"state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
ops = [s["op"] for s in r.json()["steps"]]
|
|
assert ops == ["selection.samText", "adjustments.applyMasked"], ops
|
|
assert "selection.fill" not in ops
|
|
|
|
|
|
def test_agent_path_strips_global_grades_after_a_masked_recolor(client, monkeypatch):
|
|
"""Even when the deterministic recolour special does not fire, a planner's
|
|
masked-recolour plan must not keep trailing whole-image adjustments."""
|
|
_install_provider_and_fake_tools(monkeypatch)
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: json.dumps({
|
|
"reply": "tweak",
|
|
"steps": [
|
|
{"op": "selection.samText", "params": {"text": "the object"}},
|
|
{"op": "adjustments.applyMasked",
|
|
"params": {"mode": "hue-saturation", "hue": 90, "sourceHue": 200}},
|
|
{"op": "adjustments.apply", "params": {"type": "vibrance", "params": {"vibrance": 20}}},
|
|
{"op": "photo.autoTone", "params": {}},
|
|
],
|
|
})
|
|
)
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "do a quick tweak",
|
|
"session": "s2", "editor_session": "s2",
|
|
"state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
ops = [s["op"] for s in r.json()["steps"]]
|
|
assert ops == ["selection.samText", "adjustments.applyMasked"], ops
|
|
|
|
|
|
def test_agent_path_reports_failure_when_a_generative_step_fails(client, monkeypatch):
|
|
"""Gen-6 bug 5: a failed generative step must return ok:false with the real
|
|
error, never ok:true."""
|
|
_install_provider_and_fake_tools(monkeypatch, failing=True)
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "remove the distractions", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert data["ok"] is False, data
|
|
assert "provider exploded" in data["reply"], data["reply"]
|
|
assert any(
|
|
r.get("op") in ("ai_edit_image", "generate_image") and r.get("ok") is False
|
|
for r in data["results"]
|
|
), data["results"]
|
|
|
|
|
|
def test_bridge_default_timeout_allows_slow_generative_steps():
|
|
"""Gen-6 bug 2: the bridge/agent waits must exceed the provider latency."""
|
|
assert eaar._BRIDGE_DEFAULT_TIMEOUT >= 300
|
|
|
|
|
|
def test_inpaint_proxy_decodes_data_url_base64():
|
|
"""Gen-6 bug 1: the image endpoints must accept both raw base64 and a
|
|
'data:...;base64,' data URL (the editor tools send data URLs)."""
|
|
import base64
|
|
|
|
from routes.gallery import gallery_routes as gr
|
|
|
|
raw = b"\x89PNG\r\n\x1a\n payload"
|
|
data_url = "data:image/png;base64," + base64.b64encode(raw).decode("ascii")
|
|
assert gr._b64_bytes(data_url) == raw
|
|
assert gr._b64_bytes(base64.b64encode(raw).decode("ascii")) == raw
|
|
from fastapi import HTTPException
|
|
|
|
with pytest.raises(HTTPException):
|
|
gr._b64_bytes("")
|
|
|
|
|
|
def test_generative_routing_uses_instruction_vs_masked():
|
|
"""Qwen-Image-Edit instruction approach: most generative edits are whole-image
|
|
instruction edits (no mask); masked inpaint only for a clear regional object
|
|
removal; generate_image only for a brand-new image."""
|
|
from routes.editor_agent_routes import _generative_mode, _generative_steps
|
|
|
|
bg = _generative_steps("replace", "make the background nicer")
|
|
assert len(bg) == 1 and bg[0]["op"] == "ai_edit_image"
|
|
assert bg[0]["params"]["mode"] == "instruction"
|
|
assert "keep" in bg[0]["params"]["prompt"].lower() # names what must stay
|
|
|
|
add = _generative_steps("add", "add a full moon in the sky")
|
|
assert add[-1]["op"] == "ai_edit_image" and add[-1]["params"]["mode"] == "instruction"
|
|
|
|
cleanup = _generative_steps("remove", "remove the distractions")
|
|
assert cleanup[-1]["op"] == "ai_edit_image" and cleanup[-1]["params"]["mode"] == "instruction"
|
|
|
|
masked = _generative_steps("remove", "remove the person on the left")
|
|
assert masked[0]["op"] == "selection.subject" # tight mask first
|
|
assert masked[-1]["op"] == "ai_edit_image" and masked[-1]["params"]["mode"] == "masked"
|
|
assert _generative_mode("remove", "remove the person on the left") == "masked"
|
|
|
|
fresh = _generative_steps("add", "create a new image of a full moon")
|
|
assert fresh[-1]["op"] == "generate_image"
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"prompt",
|
|
[
|
|
"add wildflowers",
|
|
"make it snowy",
|
|
"make it autumn",
|
|
"add sunglasses to her",
|
|
"add a rainbow",
|
|
"generative fill the sky",
|
|
],
|
|
)
|
|
def test_generative_content_intents_never_route_to_adjustments(prompt):
|
|
"""Bench round 4: adding elements, season/weather changes and an explicit
|
|
"generative fill" must be a whole-image instruction edit (steps 8, keep
|
|
clause), never a colour adjustment/filter."""
|
|
deterministic = [
|
|
{"op": "adjustments.apply", "params": {"type": "brightness-contrast", "params": {"brightness": 10}}}
|
|
]
|
|
kind = eaar._generative_kind(prompt)
|
|
assert kind in ("add", "replace"), (prompt, kind)
|
|
steps = eaar._ensure_generative(prompt, list(deterministic))
|
|
assert not any(s.get("op") == "adjustments.apply" for s in steps), steps
|
|
ai = [s for s in steps if s.get("op") == "ai_edit_image"]
|
|
assert ai, steps
|
|
assert ai[-1]["params"]["mode"] == "instruction"
|
|
assert ai[-1]["params"].get("steps") == 8
|
|
assert "keep" in ai[-1]["params"]["prompt"].lower()
|
|
|
|
|
|
@pytest.mark.parametrize(
|
|
"prompt", ["make it black and white", "blur the background", "add a vignette", "make it look like summer"]
|
|
)
|
|
def test_tone_intents_still_do_not_route_to_generative(prompt):
|
|
assert eaar._generative_kind(prompt) is None
|
|
|
|
|
|
def test_season_prompt_names_the_season_and_keeps_subject():
|
|
params = eaar._generative_step("replace", "make it snowy")["params"]
|
|
low = params["prompt"].lower()
|
|
assert "snow" in low and "keep" in low and params["mode"] == "instruction" and params["steps"] == 8
|
|
|
|
|
|
def test_removal_and_expand_use_eight_steps():
|
|
"""Gen-6 (C/D): removal and expand edits ask for 8 sampler steps (the server
|
|
default of 4 is too weak) and naming what to remove / where to fill."""
|
|
from routes.editor_agent_routes import _generative_step, _plan_format_steps
|
|
|
|
rm = _generative_step("remove", "remove the distractions")["params"]
|
|
assert rm["steps"] == 8, rm
|
|
assert "landscape" in rm["prompt"].lower(), rm["prompt"]
|
|
|
|
masked = _generative_step("remove", "remove the person on the left")["params"]
|
|
assert masked["steps"] == 8 and masked["mode"] == "masked", masked
|
|
|
|
ext = _plan_format_steps("extend the canvas to 16:9", {"width": 900, "height": 600}, True)[0]["params"]
|
|
assert ext["steps"] == 8 and ext["feather"] is True, ext
|
|
|
|
|
|
def test_extend_plan_uses_instruction_expand_with_feather():
|
|
from routes.editor_agent_routes import _plan_format_steps
|
|
|
|
steps = _plan_format_steps("extend the canvas to 16:9", {"width": 900, "height": 600}, True)
|
|
assert steps and steps[0]["op"] == "ai.generativeExpand"
|
|
p = steps[0]["params"]
|
|
assert p.get("mode") == "instruction" and p.get("feather") is True
|
|
assert "border" in p.get("prompt", "").lower()
|
|
|
|
|
|
def test_ai_edit_image_schema_exposes_mode():
|
|
from src.tool_schemas import FUNCTION_TOOL_SCHEMAS
|
|
|
|
schema = next(
|
|
s["function"] for s in FUNCTION_TOOL_SCHEMAS if s["function"]["name"] == "ai_edit_image"
|
|
)
|
|
mode = schema["parameters"]["properties"]["mode"]
|
|
assert set(mode["enum"]) == {"instruction", "masked"}
|
|
|
|
|
|
def test_ai_edit_image_tool_instruction_needs_no_mask_masked_does(monkeypatch):
|
|
"""Instruction mode sends a full-white mask (whole-image edit); masked mode
|
|
still requires a real selection and errors when none is available."""
|
|
import asyncio
|
|
import base64
|
|
import io
|
|
|
|
from PIL import Image
|
|
|
|
import src.agent_tools.editor_tools as et
|
|
|
|
buf = io.BytesIO()
|
|
Image.new("RGB", (32, 32), (10, 20, 30)).save(buf, "PNG")
|
|
img = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
|
|
|
|
captured = {}
|
|
|
|
async def _fake_post(path, payload, owner, timeout=120):
|
|
captured["path"] = path
|
|
captured["payload"] = payload
|
|
return {"image": "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()}
|
|
|
|
async def _no_mask(owner, session):
|
|
return None
|
|
|
|
async def _no_place(*a, **k):
|
|
return True
|
|
|
|
monkeypatch.setattr(et, "_post", _fake_post)
|
|
monkeypatch.setattr(et, "_fetch_selection_mask", _no_mask)
|
|
monkeypatch.setattr(et, "_place_layer", _no_place)
|
|
monkeypatch.setattr(et, "_persist_png", lambda *a, **k: {})
|
|
monkeypatch.setattr(et, "_provider_available", lambda owner: True)
|
|
|
|
tool = et.AiEditImageTool()
|
|
# Instruction mode: no mask supplied, still succeeds and sends a full mask.
|
|
desc, res = asyncio.run(tool.execute(
|
|
json.dumps({"op": "generative", "mode": "instruction", "prompt": "nice beach", "image": img}),
|
|
{"owner": "u"},
|
|
))
|
|
assert res.get("exit_code") == 0, res
|
|
assert captured["payload"].get("mode") == "instruction"
|
|
assert str(captured["payload"].get("mask", "")).startswith("data:image/png;base64,")
|
|
|
|
# Masked mode without any selection and no grounding box: fall back to
|
|
# whole-image instruction mode instead of placing an empty/transparent layer.
|
|
captured.clear()
|
|
desc, res = asyncio.run(tool.execute(
|
|
json.dumps({"op": "inpaint", "mode": "masked", "prompt": "remove it", "image": img}),
|
|
{"owner": "u"},
|
|
))
|
|
assert res.get("exit_code") == 0, res
|
|
assert captured["payload"].get("mode") == "instruction"
|
|
|
|
# Masked mode with a grounding box (OWL-ViT) builds a dilated box mask and
|
|
# stays in masked mode.
|
|
async def _ground(owner, image, label):
|
|
return [0.0, 0.0, 8.0, 8.0]
|
|
|
|
monkeypatch.setattr(et, "_ground_box", _ground)
|
|
captured.clear()
|
|
desc, res = asyncio.run(tool.execute(
|
|
json.dumps({"op": "inpaint", "mode": "masked", "prompt": "remove the person on the left", "image": img}),
|
|
{"owner": "u"},
|
|
))
|
|
assert res.get("exit_code") == 0, res
|
|
assert captured["payload"].get("mode") == "masked"
|
|
assert str(captured["payload"].get("mask", "")).startswith("data:image/png;base64,")
|
|
|
|
|
|
def test_place_layer_forwards_edit_mode(monkeypatch):
|
|
"""Instruction-mode results must tell the client to build a difference mask."""
|
|
import asyncio
|
|
|
|
import src.agent_tools.editor_tools as et
|
|
|
|
seen = {}
|
|
|
|
async def _fake_post(path, payload, owner, timeout=120):
|
|
seen["payload"] = payload
|
|
return {"ok": True}
|
|
|
|
monkeypatch.setattr(et, "_post", _fake_post)
|
|
|
|
asyncio.run(et._place_layer("data:image/png;base64,AA", "AI inpaint", "u", "s1", mode="instruction"))
|
|
params = seen["payload"]["params"]
|
|
assert params["mode"] == "instruction"
|
|
assert params["differenceMask"] is True
|
|
|
|
asyncio.run(et._place_layer("data:image/png;base64,AA", "AI inpaint", "u", "s1", masked=True, mode="masked"))
|
|
params = seen["payload"]["params"]
|
|
assert params["mode"] == "masked"
|
|
assert params["differenceMask"] is False
|
|
|
|
# No mode (e.g. generate_image) => keys omitted so the client keeps its behavior.
|
|
asyncio.run(et._place_layer("data:image/png;base64,AA", "Generated image", "u", "s1"))
|
|
params = seen["payload"]["params"]
|
|
assert "mode" not in params and "differenceMask" not in params
|
|
|
|
|
|
def test_plan_route_forces_generative_over_a_deterministic_model_plan(client):
|
|
"""/plan must agree with /agent: a generative intent is never left on a
|
|
deterministic substitute even if the model proposes one."""
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: {"reply": "ok", "steps": [{"op": "history.mark", "params": {}}]}
|
|
)
|
|
try:
|
|
for prompt in ("remove the person on the left", "add a full moon in the sky",
|
|
"make the background nicer"):
|
|
r = client.post("/api/editor/plan", json={"prompt": prompt, "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
ops = [s["op"] for s in r.json()["steps"]]
|
|
assert "ai_edit_image" in ops, (prompt, r.json()["steps"])
|
|
r = client.post("/api/editor/plan", json={"prompt": "create a new image of a full moon",
|
|
"state": {"width": 100, "height": 100}})
|
|
assert "generate_image" in [s["op"] for s in r.json()["steps"]], r.json()["steps"]
|
|
finally:
|
|
eaar.set_editor_planner(None)
|
|
|
|
|
|
def test_scene_replacement_prompt_names_the_scene_and_uses_eight_steps():
|
|
"""Gen-6 case 2: a background/scene replacement must name concrete scene
|
|
details and ask for 8 sampler steps, not a bare 'a beach'."""
|
|
from routes.editor_agent_routes import _generative_steps
|
|
|
|
steps = _generative_steps("replace", "put her on a beach")
|
|
ai = steps[-1]
|
|
assert ai["op"] == "ai_edit_image", steps
|
|
params = ai["params"]
|
|
assert params["mode"] == "instruction", params
|
|
assert params.get("steps") == 8, params
|
|
low = params["prompt"].lower()
|
|
assert "beach" in low and ("sand" in low or "sea" in low), params["prompt"]
|
|
assert "keep" in low, params["prompt"]
|
|
# Deterministic template, passed to the provider unmodified.
|
|
assert params["prompt"].startswith(
|
|
"Replace the entire background behind the person with"), params["prompt"]
|
|
assert params["prompt"].endswith("Keep the person exactly the same."), params["prompt"]
|
|
|
|
generic = _generative_steps("replace", "make the background nicer")[-1]["params"]
|
|
assert generic.get("steps") == 8, generic
|
|
assert "background" in generic["prompt"].lower(), generic
|
|
|
|
|
|
def test_scene_prompt_overrides_a_vague_model_ai_step():
|
|
"""A model-written ai_edit_image prompt must not drop the scene details: the
|
|
background-replacement intent forces the templated scene prompt."""
|
|
from routes.editor_agent_routes import _ensure_generative
|
|
|
|
model_steps = [{"op": "ai_edit_image",
|
|
"params": {"op": "generative", "mode": "instruction",
|
|
"prompt": "put her outside somewhere"}}]
|
|
out = _ensure_generative("put her on a beach", model_steps)
|
|
ai = [s for s in out if s.get("op") == "ai_edit_image"][-1]
|
|
prompt = ai["params"]["prompt"]
|
|
assert prompt.startswith("Replace the entire background behind the person with")
|
|
assert prompt.endswith("Keep the person exactly the same.")
|
|
assert "beach" in prompt.lower() and "sand" in prompt.lower()
|
|
assert ai["params"]["mode"] == "instruction" and ai["params"].get("mask") is None
|
|
|
|
|
|
def test_agent_path_masked_ai_step_receives_canvas_screenshot(client, monkeypatch):
|
|
"""Gen-6 case 1: after a selection step, the masked ai_edit_image must still
|
|
receive the canvas image (from the request screenshot), not fail with
|
|
'needs an image'."""
|
|
import base64
|
|
|
|
from src.agent_tools import editor_tools
|
|
|
|
seen = {}
|
|
|
|
class _FakeAI:
|
|
async def execute(self, content, ctx): # noqa: ANN001
|
|
seen["args"] = json.loads(content)
|
|
return ("ai_edit_image", {"exit_code": 0, "image_id": "g", "placed": True, "output": "ok"})
|
|
|
|
monkeypatch.setattr(editor_tools, "AiEditImageTool", _FakeAI)
|
|
monkeypatch.setattr(eaar, "_image_provider_available", lambda owner: True)
|
|
|
|
async def _fake_wait(fut, timeout=None, **kwargs): # noqa: ANN001
|
|
return {"ok": True, "result": {"width": 100, "height": 100}}
|
|
|
|
monkeypatch.setattr(eaar.asyncio, "wait_for", _fake_wait)
|
|
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: json.dumps(
|
|
{
|
|
"reply": "ok",
|
|
"steps": [
|
|
{"op": "selection.subject", "params": {"label": "person"}},
|
|
{
|
|
"op": "ai_edit_image",
|
|
"params": {"op": "inpaint", "mode": "masked", "prompt": "remove the person on the left"},
|
|
},
|
|
],
|
|
}
|
|
)
|
|
)
|
|
shot = "data:image/png;base64," + base64.b64encode(b"canvas").decode("ascii")
|
|
r = client.post(
|
|
"/api/editor/agent",
|
|
json={
|
|
"prompt": "remove the person on the left",
|
|
"session": "s1",
|
|
"editor_session": "s1",
|
|
"state": {"width": 100, "height": 100},
|
|
"screenshot": shot,
|
|
},
|
|
)
|
|
assert r.status_code == 200, r.text
|
|
assert seen.get("args"), r.json()
|
|
assert seen["args"].get("image") == shot, seen["args"]
|
|
assert seen["args"].get("mode") == "masked", seen["args"]
|
|
|
|
|
|
def _bridge_app():
|
|
app = FastAPI()
|
|
app.include_router(eaar.setup_editor_agent_routes())
|
|
return app
|
|
|
|
|
|
async def _asgi_post(app, path, payload):
|
|
sent = []
|
|
scope = {"type": "http", "method": "POST", "path": path, "query_string": b"",
|
|
"headers": [(b"content-type", b"application/json")]}
|
|
data = json.dumps(payload).encode()
|
|
recv = [{"type": "http.request", "body": data, "more_body": False}]
|
|
|
|
async def receive():
|
|
return recv.pop(0) if recv else {"type": "http.disconnect"}
|
|
|
|
async def send(msg):
|
|
if msg["type"] == "http.response.body" and msg.get("body"):
|
|
sent.append(msg["body"])
|
|
|
|
await app(scope, receive, send)
|
|
return b"".join(sent)
|
|
|
|
|
|
def test_bridge_roundtrip_subscribe_command_result():
|
|
"""URGENT regression: a real subscribe -> command -> result round trip must
|
|
deliver the command to the SSE subscriber and resolve the waiting /command."""
|
|
import asyncio
|
|
|
|
app = _bridge_app()
|
|
|
|
async def scenario():
|
|
out, stop = [], asyncio.Event()
|
|
|
|
async def receive():
|
|
await stop.wait()
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(msg):
|
|
if msg["type"] == "http.response.body" and msg.get("body"):
|
|
out.append(msg["body"])
|
|
|
|
sse = asyncio.create_task(app(
|
|
{"type": "http", "method": "GET", "path": "/api/editor/bridge/events",
|
|
"query_string": b"session=S", "headers": []}, receive, send))
|
|
for _ in range(200):
|
|
if any(b": connected" in b for b in out):
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
|
|
cmd_result = {}
|
|
|
|
async def issue():
|
|
cmd_result["body"] = await _asgi_post(
|
|
app, "/api/editor/bridge/command",
|
|
{"session": "S", "steps": [{"op": "document.info"}], "wait": True, "timeout": 5})
|
|
|
|
cmd_task = asyncio.create_task(issue())
|
|
cmd_id = None
|
|
for _ in range(300):
|
|
for b in out:
|
|
for line in b.decode().splitlines():
|
|
if line.startswith("data:"):
|
|
cmd_id = json.loads(line[5:].strip())["id"]
|
|
break
|
|
if cmd_id:
|
|
break
|
|
if cmd_id:
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
assert cmd_id, out
|
|
await _asgi_post(app, "/api/editor/bridge/result",
|
|
{"id": cmd_id, "ok": True, "result": {"width": 10, "height": 20}})
|
|
await asyncio.wait_for(cmd_task, 5)
|
|
body = json.loads(cmd_result["body"].decode())
|
|
assert body.get("ok") is True and body.get("result") == {"width": 10, "height": 20}, body
|
|
stop.set()
|
|
try:
|
|
await asyncio.wait_for(sse, 2)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_bridge_broadcasts_to_every_subscriber():
|
|
"""URGENT regression: with two live subscribers on the same session (EventSource
|
|
reconnects overlap), a command must reach BOTH, never be eaten by one."""
|
|
import asyncio
|
|
|
|
app = _bridge_app()
|
|
|
|
async def scenario():
|
|
outs = {"A": [], "B": []}
|
|
stops = {"A": asyncio.Event(), "B": asyncio.Event()}
|
|
|
|
def subscriber(name):
|
|
async def receive():
|
|
await stops[name].wait()
|
|
return {"type": "http.disconnect"}
|
|
|
|
async def send(msg):
|
|
if msg["type"] == "http.response.body" and msg.get("body"):
|
|
outs[name].append(msg["body"])
|
|
|
|
return asyncio.create_task(app(
|
|
{"type": "http", "method": "GET", "path": "/api/editor/bridge/events",
|
|
"query_string": b"session=D", "headers": []}, receive, send))
|
|
|
|
tasks = [subscriber("A"), subscriber("B")]
|
|
for _ in range(200):
|
|
if all(any(b": connected" in b for b in outs[n]) for n in ("A", "B")):
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
|
|
await _asgi_post(app, "/api/editor/bridge/command",
|
|
{"session": "D", "steps": [{"op": "screenshot"}], "wait": False})
|
|
for _ in range(300):
|
|
if all(any(b"data:" in b for b in outs[n]) for n in ("A", "B")):
|
|
break
|
|
await asyncio.sleep(0.01)
|
|
|
|
for name in ("A", "B"):
|
|
assert any(b'"op": "screenshot"' in b for b in outs[name]), (name, outs)
|
|
for e in stops.values():
|
|
e.set()
|
|
for t in tasks:
|
|
try:
|
|
await asyncio.wait_for(t, 2)
|
|
except (asyncio.TimeoutError, asyncio.CancelledError):
|
|
pass
|
|
|
|
asyncio.run(scenario())
|
|
|
|
|
|
def test_editor_image_models_route_shape():
|
|
"""The editor model selector reads the available image models from here."""
|
|
from fastapi import FastAPI
|
|
from starlette.testclient import TestClient
|
|
|
|
app = FastAPI()
|
|
app.include_router(eaar.setup_editor_agent_routes())
|
|
c = TestClient(app)
|
|
r = c.get("/api/editor/image-models")
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert isinstance(data.get("models"), list)
|
|
for m in data["models"]:
|
|
assert set(("id", "label", "kind", "provider")).issubset(m)
|
|
assert m["kind"] in ("edit", "generate")
|
|
|
|
|
|
def test_ai_edit_image_tool_forwards_and_reports_model(monkeypatch):
|
|
"""A picked model is forwarded to the provider (_model) and reported back."""
|
|
import asyncio
|
|
import base64
|
|
import io
|
|
|
|
from PIL import Image
|
|
|
|
import src.agent_tools.editor_tools as et
|
|
|
|
buf = io.BytesIO()
|
|
Image.new("RGB", (32, 32), (10, 20, 30)).save(buf, "PNG")
|
|
img = "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()
|
|
|
|
captured = {}
|
|
|
|
async def _fake_post(path, payload, owner, timeout=120):
|
|
captured["payload"] = payload
|
|
return {"image": img, "model": "qwen-image-edit-2511"}
|
|
|
|
async def _no_place(*a, **k):
|
|
return True
|
|
|
|
monkeypatch.setattr(et, "_post", _fake_post)
|
|
monkeypatch.setattr(et, "_place_layer", _no_place)
|
|
monkeypatch.setattr(et, "_persist_png", lambda *a, **k: {})
|
|
monkeypatch.setattr(et, "_provider_available", lambda owner: True)
|
|
|
|
_desc, res = asyncio.run(et.AiEditImageTool().execute(
|
|
json.dumps({"op": "generative", "mode": "instruction", "prompt": "beach",
|
|
"image": img, "model": "qwen-image-edit-2511", "place": False}),
|
|
{"owner": "u"},
|
|
))
|
|
assert res.get("exit_code") == 0, res
|
|
assert captured["payload"].get("_model") == "qwen-image-edit-2511", captured["payload"]
|
|
assert res.get("model") == "qwen-image-edit-2511", res
|
|
|
|
|
|
def test_agent_path_injects_selected_model_into_ai_step(client, monkeypatch):
|
|
"""The /api/editor/agent body `model` reaches the AI step tool payload."""
|
|
monkeypatch.setattr(eaar, "_image_provider_available", lambda owner: True)
|
|
captured = {}
|
|
|
|
class _CapturingTool:
|
|
async def execute(self, content, ctx): # noqa: ANN001
|
|
captured["content"] = json.loads(content)
|
|
return ("ai_edit_image", {"exit_code": 0, "image_id": "g1", "placed": True,
|
|
"output": "ok", "model": "qwen-image-edit-2511"})
|
|
|
|
from src.agent_tools import editor_tools
|
|
|
|
monkeypatch.setattr(editor_tools, "AiEditImageTool", _CapturingTool)
|
|
monkeypatch.setattr(editor_tools, "GenerateImageLayerTool", _CapturingTool)
|
|
|
|
async def _fake_wait(fut, timeout=None, **kwargs): # noqa: ANN001
|
|
return {"ok": True, "result": {"width": 100, "height": 100}}
|
|
|
|
monkeypatch.setattr(eaar.asyncio, "wait_for", _fake_wait)
|
|
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "put her on a beach", "state": {"width": 100, "height": 100},
|
|
"model": "qwen-image-edit-2511"})
|
|
assert r.status_code == 200, r.text
|
|
assert captured.get("content", {}).get("model") == "qwen-image-edit-2511", captured
|
|
data = r.json()
|
|
assert data["results"][-1]["result"].get("model") == "qwen-image-edit-2511", data
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Editor agent/planner model (local Qwen3.5-9B) resolution + reporting.
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def test_editor_agent_models_route_shape(client):
|
|
"""The agent-model picker reads the local planner model + hosted fallback."""
|
|
r = client.get("/api/editor/agent-models")
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert isinstance(data.get("models"), list)
|
|
assert isinstance(data.get("probe"), dict)
|
|
for m in data["models"]:
|
|
assert set(("id", "label", "source", "vision")).issubset(m)
|
|
assert m["source"] in ("local", "hosted")
|
|
|
|
|
|
def test_list_editor_agent_models_includes_base_url():
|
|
import asyncio
|
|
|
|
import src.editor_agent_model as eam
|
|
|
|
models = asyncio.run(eam.list_editor_agent_models(None))
|
|
assert any(m["source"] == "local" for m in models), models
|
|
for m in models:
|
|
assert "base_url" in m
|
|
|
|
|
|
def test_plan_reports_agent_model(client):
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: json.dumps(
|
|
{"reply": "ok", "steps": [{"op": "filters.apply", "params": {"id": "sharpen"}}]}
|
|
)
|
|
)
|
|
r = client.post("/api/editor/plan",
|
|
json={"prompt": "sharpen it", "state": {}, "agent_model": "MYPLANNER"})
|
|
assert r.status_code == 200, r.text
|
|
assert r.json().get("agent_model") == "MYPLANNER", r.json()
|
|
|
|
|
|
def test_agent_reports_agent_model(client, monkeypatch):
|
|
monkeypatch.setattr(eaar, "_image_provider_available", lambda owner: True)
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: json.dumps(
|
|
{"reply": "ok", "steps": [{"op": "adjustments.apply", "params": {"type": "vibrance"}}]}
|
|
)
|
|
)
|
|
|
|
async def _fake_wait(fut, timeout=None, **kwargs): # noqa: ANN001
|
|
try:
|
|
fut.close()
|
|
except Exception:
|
|
pass
|
|
return {"ok": True, "result": {"width": 100, "height": 100}}
|
|
|
|
monkeypatch.setattr(eaar.asyncio, "wait_for", _fake_wait)
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "make it vivid", "state": {"width": 100, "height": 100},
|
|
"agent_model": "MYPLANNER"})
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert data.get("agent_model") == "MYPLANNER", data
|
|
assert data["trace"][-1].get("agent_model") == "MYPLANNER" or data["trace"][1].get("agent_model") == "MYPLANNER", data
|
|
|
|
|
|
def test_editor_agent_model_override_is_first_candidate(monkeypatch):
|
|
import src.editor_agent_model as eam
|
|
|
|
monkeypatch.setattr(eam, "local_bases", lambda: ["http://local1"])
|
|
monkeypatch.setattr(eam, "local_model", lambda: "LMODEL")
|
|
cands = eam._candidates(None, "http://ovr::OVERRIDE")
|
|
assert cands[0] == {"base": "http://ovr", "model": "OVERRIDE", "headers": None,
|
|
"source": "override"}
|
|
assert any(c["source"] == "local" and c["base"] == "http://local1" for c in cands)
|
|
|
|
|
|
def test_editor_agent_model_uses_vision_first(monkeypatch):
|
|
import asyncio
|
|
|
|
import src.editor_agent_model as eam
|
|
|
|
calls = []
|
|
|
|
async def fake_post(base, model, messages, headers, timeout, max_tokens=1024): # noqa: ANN001
|
|
calls.append((base, eam._has_image(messages)))
|
|
return "PLAN" if base == "http://local1" and eam._has_image(messages) else None
|
|
|
|
monkeypatch.setattr(eam, "_candidates", lambda owner, override: [
|
|
{"base": "http://local1", "model": "m1", "headers": None, "source": "local"}])
|
|
monkeypatch.setattr(eam, "_post_chat", fake_post)
|
|
|
|
msgs = [{"role": "user", "content": [
|
|
{"type": "text", "text": "x"},
|
|
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}]
|
|
res = asyncio.run(eam.plan_with_editor_agent(msgs, None))
|
|
assert res == ("PLAN", "m1", "local"), res
|
|
assert calls == [("http://local1", True)]
|
|
|
|
|
|
def test_editor_agent_model_caption_fallback(monkeypatch):
|
|
import asyncio
|
|
|
|
import src.editor_agent_model as eam
|
|
|
|
async def fake_post(base, model, messages, headers, timeout, max_tokens=1024): # noqa: ANN001
|
|
return None if eam._has_image(messages) else "PLAN"
|
|
|
|
async def fake_caption(url, owner=None, timeout=90): # noqa: ANN001
|
|
return "a beach at sunset"
|
|
|
|
monkeypatch.setattr(eam, "_candidates", lambda owner, override: [
|
|
{"base": "http://local1", "model": "m1", "headers": None, "source": "local"}])
|
|
monkeypatch.setattr(eam, "_post_chat", fake_post)
|
|
monkeypatch.setattr(eam, "caption_image", fake_caption)
|
|
|
|
msgs = [{"role": "user", "content": [
|
|
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}]
|
|
res = asyncio.run(eam.plan_with_editor_agent(msgs, None))
|
|
assert res == ("PLAN", "m1", "local+caption"), res
|
|
|
|
|
|
def test_editor_agent_model_returns_none_when_all_fail(monkeypatch):
|
|
import asyncio
|
|
|
|
import src.editor_agent_model as eam
|
|
|
|
async def fake_post(*a, **k): # noqa: ANN001
|
|
return None
|
|
|
|
async def fake_caption(*a, **k): # noqa: ANN001
|
|
return None
|
|
|
|
monkeypatch.setattr(eam, "_candidates", lambda owner, override: [
|
|
{"base": "b", "model": "m", "headers": None, "source": "local"}])
|
|
monkeypatch.setattr(eam, "_post_chat", fake_post)
|
|
monkeypatch.setattr(eam, "caption_image", fake_caption)
|
|
|
|
msgs = [{"role": "user", "content": [
|
|
{"type": "image_url", "image_url": {"url": "data:image/png;base64,AA"}}]}]
|
|
assert asyncio.run(eam.plan_with_editor_agent(msgs, None)) is None
|
|
|
|
|
|
def test_plan_special_recolor_overrides_a_model_plan(client):
|
|
# A deterministic recolour must win over any model-written plan
|
|
# (no LLM variance / no samText+placeImage substitution).
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps({
|
|
"reply": "ok",
|
|
"steps": [{"op": "adjustments.apply", "params": {"type": "vibrance"}}]}))
|
|
r = client.post("/api/editor/plan", json={
|
|
"prompt": "make only the red car blue",
|
|
"state": {"width": 900, "height": 600}})
|
|
assert r.status_code == 200
|
|
data = r.json()
|
|
assert data["planner"] == "deterministic"
|
|
assert [s["op"] for s in data["steps"]] == [
|
|
"selection.samText", "adjustments.applyMasked"]
|
|
assert "selection.fill" not in [s["op"] for s in data["steps"]]
|
|
|
|
|
|
def test_self_check_endpoint_parses_retry(client, monkeypatch):
|
|
import src.editor_agent_model as eam
|
|
|
|
async def _retry(*a, **k): # noqa: ANN001
|
|
return "RETRY: brighten the subject"
|
|
|
|
monkeypatch.setattr(eam, "_post_chat", _retry)
|
|
r = client.post("/api/editor/self-check", json={
|
|
"before": "data:image/png;base64,AA",
|
|
"after": "data:image/png;base64,AA",
|
|
"prompt": "did it get brighter?"})
|
|
assert r.status_code == 200
|
|
data = r.json()
|
|
assert data["done"] is False
|
|
assert "brighten" in data["retry_hint"]
|
|
assert "RETRY" in data["text"]
|
|
|
|
|
|
def test_self_check_endpoint_parses_done(client, monkeypatch):
|
|
import src.editor_agent_model as eam
|
|
|
|
async def _done(*a, **k): # noqa: ANN001
|
|
return "DONE"
|
|
|
|
monkeypatch.setattr(eam, "_post_chat", _done)
|
|
r = client.post("/api/editor/self-check", json={
|
|
"before": "data:image/png;base64,AA",
|
|
"after": "data:image/png;base64,AA",
|
|
"prompt": "ok?"})
|
|
assert r.status_code == 200
|
|
assert r.json()["done"] is True
|
|
|
|
|
|
def test_self_check_endpoint_unavailable_when_vlm_down(client, monkeypatch):
|
|
import src.editor_agent_model as eam
|
|
|
|
async def _none(*a, **k): # noqa: ANN001
|
|
return None
|
|
|
|
monkeypatch.setattr(eam, "_post_chat", _none)
|
|
r = client.post("/api/editor/self-check", json={
|
|
"before": "data:image/png;base64,AA",
|
|
"after": "data:image/png;base64,AA",
|
|
"prompt": "ok?"})
|
|
assert r.status_code == 200
|
|
data = r.json()
|
|
assert data["done"] is True
|
|
assert data.get("reason") == "unavailable"
|
|
|
|
|
|
def test_self_check_endpoint_unavailable_without_images(client):
|
|
r = client.post("/api/editor/self-check", json={"prompt": "ok?"})
|
|
assert r.status_code == 200
|
|
data = r.json()
|
|
assert data["done"] is True
|
|
assert data.get("reason") == "unavailable"
|
|
|
|
|
|
# ── Round: planner-unsure asks; whole-image restyle is a model edit ──────────
|
|
def test_plan_asks_when_no_intent_matches(client):
|
|
"""A prompt mapped to no concrete intent must ASK (2-3 concrete choices),
|
|
never silently apply a default photo.autoTone/adjustments/filters stack."""
|
|
eaar.set_editor_planner(lambda messages, owner: None)
|
|
r = client.post("/api/editor/plan", json={"prompt": "give it some oomph", "state": {}})
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert data["steps"] == [], data
|
|
ask = data.get("ask_user")
|
|
assert isinstance(ask, dict), data
|
|
assert 2 <= len(ask["options"]) <= 3, ask
|
|
for s in data["steps"]:
|
|
assert s.get("op") not in {"adjustments.apply", "filters.apply", "photo.autoTone"}
|
|
|
|
|
|
def test_agent_path_asks_when_no_intent_matches(client):
|
|
eaar.set_editor_planner(lambda messages, owner: None)
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "give it some oomph", "session": "s1",
|
|
"editor_session": "s1", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert data.get("ask_user_pending") is True, data
|
|
assert data["steps"] == [], data
|
|
assert data["results"] == [], data
|
|
ask = data.get("ask_user")
|
|
assert isinstance(ask, dict), data
|
|
assert 2 <= len(ask["options"]) <= 3, ask
|
|
|
|
|
|
@pytest.mark.parametrize("prompt", [
|
|
"make it look like an oil painting",
|
|
"turn this into an anime",
|
|
"make it a watercolour",
|
|
"pixel art this",
|
|
"make it surreal",
|
|
"restyle as a comic",
|
|
])
|
|
def test_restyle_prompts_classified_as_model_edit(prompt):
|
|
assert eaar._generative_kind(prompt) == "restyle"
|
|
|
|
|
|
@pytest.mark.parametrize("prompt", [
|
|
"make it look like an oil painting",
|
|
"turn this into an anime",
|
|
"make it a watercolour",
|
|
])
|
|
def test_plan_routes_restyle_to_ai_edit_image(client, prompt):
|
|
eaar.set_editor_planner(lambda messages, owner: None)
|
|
data = client.post("/api/editor/plan", json={"prompt": prompt, "state": {}}).json()
|
|
ops = [s["op"] for s in data["steps"]]
|
|
assert ops == ["ai_edit_image"], data
|
|
params = data["steps"][0]["params"]
|
|
assert params["mode"] == "instruction", params
|
|
assert "exactly the same" in params["prompt"], params
|
|
|
|
|
|
def test_restyle_does_not_hijack_tone_or_new_image(client):
|
|
"""Restyle must not grab tone prompts, nor a brand-new image request."""
|
|
for prompt in ("make it black and white", "blur the background", "add a vignette",
|
|
"make it look snowy", "make it warmer"):
|
|
assert eaar._generative_kind(prompt) != "restyle", prompt
|
|
assert eaar._generative_kind("create an illustration of a cat") == "add"
|
|
|
|
|
|
def test_agent_path_forces_restyle_over_a_local_adjustment_plan(client, monkeypatch):
|
|
_install_provider_and_fake_tools(monkeypatch)
|
|
eaar.set_editor_planner(
|
|
lambda messages, owner: json.dumps({
|
|
"reply": "vibrance",
|
|
"steps": [{"op": "adjustments.apply",
|
|
"params": {"type": "vibrance", "params": {"vibrance": 20}}}],
|
|
})
|
|
)
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "make it look like an oil painting", "session": "s1",
|
|
"editor_session": "s1", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
ops = [s["op"] for s in r.json()["steps"]]
|
|
assert "ai_edit_image" in ops, r.json()["steps"]
|
|
assert "adjustments.apply" not in ops, r.json()["steps"]
|
|
|
|
|
|
def test_plan_asks_when_model_returns_only_a_generic_tone_plan(client):
|
|
"""Even when the model returns a plan, a bare whole-image tone stack for an
|
|
unspecified prompt is a guess -> ask (2-3 concrete choices)."""
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps({
|
|
"reply": "punchy",
|
|
"steps": [
|
|
{"op": "photo.autoTone", "params": {}},
|
|
{"op": "adjustments.apply", "params": {"type": "vibrance", "params": {"vibrance": 25}}},
|
|
{"op": "filters.apply", "params": {"id": "clarity"}},
|
|
],
|
|
}))
|
|
data = client.post("/api/editor/plan", json={"prompt": "give it some oomph", "state": {}}).json()
|
|
assert data["steps"] == [], data
|
|
ask = data.get("ask_user")
|
|
assert isinstance(ask, dict), data
|
|
assert 2 <= len(ask["options"]) <= 3, ask
|
|
|
|
|
|
def test_plan_keeps_a_generic_tone_plan_when_a_concrete_intent_matched(client):
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps({
|
|
"reply": "punchy",
|
|
"steps": [
|
|
{"op": "photo.autoTone", "params": {}},
|
|
{"op": "adjustments.apply", "params": {"type": "vibrance", "params": {"vibrance": 25}}},
|
|
],
|
|
}))
|
|
data = client.post("/api/editor/plan", json={"prompt": "make it pop", "state": {}}).json()
|
|
assert data["steps"], data
|
|
assert not data.get("ask_user"), data
|
|
|
|
|
|
def test_agent_path_asks_when_model_returns_only_a_generic_tone_plan(client):
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps({
|
|
"reply": "punchy", "steps": [{"op": "photo.autoTone", "params": {}}]}))
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "give it some oomph", "session": "s1",
|
|
"editor_session": "s1", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert data.get("ask_user_pending") is True, data
|
|
assert data["steps"] == [], data
|
|
|
|
|
|
# ── Strict ops schema (structured outputs) ──────────────────────────────────
|
|
|
|
def test_plan_schema_is_a_strict_ops_enum():
|
|
props = eaar.PLAN_SCHEMA["properties"]
|
|
enum = props["steps"]["items"]["properties"]["op"]["enum"]
|
|
# New model-facing ops plus the deterministic catalog.
|
|
for op in ("generate_edit", "select_object", "ask_user", "adjustments.apply",
|
|
"filters.apply", "selection.samText"):
|
|
assert op in enum, op
|
|
assert props["steps"]["items"]["required"] == ["op"]
|
|
assert eaar.PLAN_RESPONSE_FORMAT["type"] == "json_schema"
|
|
assert eaar.PLAN_RESPONSE_FORMAT["json_schema"]["schema"] is eaar.PLAN_SCHEMA
|
|
|
|
|
|
def test_validate_plan_translates_model_ops_to_execution_ops():
|
|
gen = eaar._validate_plan(
|
|
{"reply": "r", "steps": [{"op": "generate_edit",
|
|
"params": {"prompt": "make it snowy",
|
|
"mode": "instruction"}}]})
|
|
assert gen["steps"][0]["op"] == "ai_edit_image"
|
|
sel = eaar._validate_plan(
|
|
{"reply": "r", "steps": [{"op": "select_object", "params": {"text": "the red car"}}]})
|
|
assert sel["steps"][0]["op"] == "selection.samText"
|
|
assert sel["steps"][0]["params"]["text"] == "the red car"
|
|
|
|
|
|
def test_validate_plan_normalizes_ai_edit_image_inner_op():
|
|
# Regression: the :8097 model authors a restyle as ai_edit_image with NO inner
|
|
# op; it must be repaired to op "generative" (instruction mode).
|
|
no_inner = eaar._validate_plan({"reply": "r", "steps": [
|
|
{"op": "ai_edit_image", "params": {"mode": "instruction", "prompt": "oil painting"}}]})
|
|
assert no_inner["steps"][0]["params"]["op"] == "generative"
|
|
assert no_inner["steps"][0]["params"]["mode"] == "instruction"
|
|
# An invented inner op ("restyle") is repaired too.
|
|
bad = eaar._validate_plan({"reply": "r", "steps": [
|
|
{"op": "ai_edit_image", "params": {"op": "restyle", "mode": "instruction", "prompt": "anime"}}]})
|
|
assert bad["steps"][0]["params"]["op"] == "generative"
|
|
# generate_edit is translated AND normalized.
|
|
ge = eaar._validate_plan({"reply": "r", "steps": [
|
|
{"op": "generate_edit", "params": {"prompt": "restyle"}}]})
|
|
assert ge["steps"][0]["op"] == "ai_edit_image"
|
|
assert ge["steps"][0]["params"]["op"] == "generative"
|
|
# A masked edit with no inner op becomes inpaint.
|
|
masked = eaar._validate_plan({"reply": "r", "steps": [
|
|
{"op": "ai_edit_image", "params": {"mode": "masked", "prompt": "remove the sign"}}]})
|
|
assert masked["steps"][0]["params"]["op"] == "inpaint"
|
|
assert masked["steps"][0]["params"]["mode"] == "masked"
|
|
# A valid inner op is left untouched.
|
|
keep = eaar._validate_plan({"reply": "r", "steps": [
|
|
{"op": "ai_edit_image", "params": {"op": "remove-bg"}}]})
|
|
assert keep["steps"][0]["params"]["op"] == "remove-bg"
|
|
|
|
|
|
def test_validate_plan_drops_a_single_bad_step_and_empty_plans():
|
|
# A single unknown/hallucinated op is DROPPED; the valid steps survive.
|
|
mixed = eaar._validate_plan(
|
|
{"reply": "r", "steps": [{"op": "adjustments.apply", "params": {}}, {"op": "evil.op"}]})
|
|
assert [s["op"] for s in mixed["steps"]] == ["adjustments.apply"]
|
|
# A plan whose only step is invalid has nothing left -> None (ask_user upstream).
|
|
assert eaar._validate_plan({"reply": "r", "steps": [{"op": "convert_to_grayscale"}]}) is None
|
|
assert eaar._validate_plan({"reply": "r", "steps": []}) is None
|
|
assert eaar._validate_plan("not a dict") is None
|
|
|
|
|
|
def test_validate_plan_ask_user_needs_two_to_three_choices():
|
|
ok = eaar._validate_plan({"reply": "q", "steps": [],
|
|
"ask_user": {"question": "Which?", "options": ["A", "B"]}})
|
|
assert ok["ask_user"]["options"] == [{"label": "A", "description": ""},
|
|
{"label": "B", "description": ""}]
|
|
trimmed = eaar._validate_plan({"reply": "q", "steps": [],
|
|
"ask_user": {"question": "Which?", "options": ["A", "B", "C", "D"]}})
|
|
assert len(trimmed["ask_user"]["options"]) == 3
|
|
# Fewer than two choices is not a valid ask_user.
|
|
assert eaar._validate_plan({"reply": "q", "steps": [],
|
|
"ask_user": {"question": "Which?", "options": ["only"]}}) is None
|
|
|
|
|
|
def test_plan_route_invalid_model_plan_asks_not_a_tone_default(client):
|
|
# A hallucinated op invalidates the plan; the route must ask, never tweak.
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps(
|
|
{"reply": "made it gray", "steps": [{"op": "convert_to_grayscale", "params": {}}]}))
|
|
data = client.post("/api/editor/plan", json={"prompt": "do the thing", "state": {}}).json()
|
|
assert data.get("ask_user"), data
|
|
assert data["steps"] == [], data
|
|
|
|
|
|
def test_plan_route_translates_generate_edit(client):
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps(
|
|
{"reply": "restyling", "steps": [{"op": "generate_edit",
|
|
"params": {"prompt": "make it a watercolour",
|
|
"mode": "instruction"}}]}))
|
|
data = client.post("/api/editor/plan",
|
|
json={"prompt": "do the thing", "state": {"width": 100, "height": 100}}).json()
|
|
assert [s["op"] for s in data["steps"]] == ["ai_edit_image"], data
|
|
|
|
|
|
def test_restyle_keeps_a_model_chosen_generate_edit(client, monkeypatch):
|
|
# With the strict schema + generate_edit available, the planner itself picks
|
|
# the AI op for a restyle; the _RESTYLE_RE keyword list is only a fast path,
|
|
# so a valid model generate_edit is NOT overwritten by the templated step.
|
|
_install_provider_and_fake_tools(monkeypatch)
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps(
|
|
{"reply": "oil painting", "steps": [{"op": "generate_edit",
|
|
"params": {"prompt": "render as an impressionist oil painting",
|
|
"mode": "instruction"}}]}))
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "make it look like an oil painting", "session": "s1",
|
|
"editor_session": "s1", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert [s["op"] for s in data["steps"]] == ["ai_edit_image"], data["steps"]
|
|
|
|
|
|
def test_agent_restyle_with_model_shape_repairs_inner_op(client, monkeypatch):
|
|
"""Regression (m2334): in AGENT mode the :8097 model authors a restyle as an
|
|
ai_edit_image step WITHOUT a valid inner op, which used to be blocked. It must
|
|
be repaired to op 'generative' and run (ok)."""
|
|
_install_provider_and_fake_tools(monkeypatch)
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps(
|
|
{"reply": "painting", "steps": [{"op": "ai_edit_image",
|
|
"params": {"mode": "instruction",
|
|
"prompt": "restyle as an oil painting"}}]}))
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "make it look like an oil painting", "session": "s1",
|
|
"editor_session": "s1", "state": {"width": 100, "height": 100}})
|
|
assert r.status_code == 200, r.text
|
|
data = r.json()
|
|
assert data["ok"] is True, data
|
|
ai = [s for s in data["steps"] if s["op"] == "ai_edit_image"]
|
|
assert ai and ai[0]["params"].get("op") == "generative", data["steps"]
|
|
|
|
|
|
def test_restyle_keyword_is_still_a_fast_path_when_model_returns_no_ai_op(client, monkeypatch):
|
|
# Control: a model adjustment plan for a restyle is still forced to the
|
|
# deterministic restyle step (the fast path), so the keyword list stays a
|
|
# safety net.
|
|
_install_provider_and_fake_tools(monkeypatch)
|
|
eaar.set_editor_planner(lambda messages, owner: json.dumps(
|
|
{"reply": "vivid", "steps": [{"op": "adjustments.apply",
|
|
"params": {"type": "vibrance"}}]}))
|
|
r = client.post("/api/editor/agent",
|
|
json={"prompt": "make it look like an oil painting", "session": "s1",
|
|
"editor_session": "s1", "state": {"width": 100, "height": 100}})
|
|
data = r.json()
|
|
assert [s["op"] for s in data["steps"]] == ["ai_edit_image"], data["steps"]
|
|
|
|
|
|
def test_editor_agent_plan_response_format_only_sent_to_local(monkeypatch):
|
|
from src import editor_agent_model as eam
|
|
|
|
calls = []
|
|
|
|
async def _fake_post(base, model, messages, headers, timeout, **kwargs):
|
|
calls.append((base, kwargs.get("response_format")))
|
|
return json.dumps({"reply": "r", "steps": [{"op": "adjustments.apply",
|
|
"params": {"type": "vibrance"}}]})
|
|
|
|
monkeypatch.setattr(eam, "_post_chat", _fake_post)
|
|
monkeypatch.setattr(eam, "_candidates", lambda owner, override: [
|
|
{"base": "http://100.69.120.65:8097/v1", "model": "Qwen/Qwen3.5-9B",
|
|
"headers": None, "source": "local"},
|
|
{"base": "https://api.deepseek.com/v1", "model": "deepseek-v4-flash",
|
|
"headers": None, "source": "hosted"},
|
|
])
|
|
import asyncio
|
|
fmt = {"type": "json_schema", "json_schema": {"name": "x", "schema": {}}}
|
|
out = asyncio.run(eam.plan_with_editor_agent(
|
|
[{"role": "user", "content": "x"}], response_format=fmt))
|
|
assert out is not None
|
|
# The local (first) candidate got the structured-output constraint; the
|
|
# request short-circuits on the first success, so only the local call ran.
|
|
assert calls and calls[0][0] == "http://100.69.120.65:8097/v1"
|
|
assert calls[0][1] is fmt
|
|
assert all(base != "https://api.deepseek.com/v1" for base, _ in calls)
|