Jaroslav Beneš
5e75948069
The last unbuilt capability, and built the way CLAUDE.md said it had to be: a
ToolDef reaching resolve_tools plus a permission and a capability flag, not a new
code path. The only genuinely new UI is one branch in the transcript.
services/images/ is three modules. comfy.py speaks HTTP -- submit, poll /history,
fetch the PNG, /free, and an /object_info discovery for the admin page only.
Polled and not socketed, because holding a connection open for the length of a
generation is the live-connection state the whole ssh.py design forbids, and the
thing being waited for takes tens of seconds anyway. The base URL is exempt from
the SSRF guard by construction, exactly as Connection.base_url and the audio
endpoints are -- said out loud in the docstring, because a default of
127.0.0.1:8188 is precisely the shape that guard exists to refuse and therefore
reads as a hole rather than a decision.
workflow.py fills a template, and the one thing that matters is that it walks the
parsed JSON rather than the text of it. A value that is exactly "{{steps}}"
becomes the number 20; ComfyUI validates types and refuses the string. A
placeholder inside a longer string is still text, which is what makes
"{{prompt}}, masterpiece" work -- and text substitution would additionally mean a
prompt containing a quotation mark produced a document that no longer parses, on
the one input guaranteed to hold arbitrary text. Which node holds the prompt is
the administrator's statement rather than a guess from node types: sniffing for
the first CLIPTextEncode works on the shipped workflow and on nothing else, and
swaps positive for negative the first time somebody reorders them. seed has no
fixed default, because one would make every unspecified generation identical and
make the retry loop redraw the same rejected picture four times.
tool.py is one call, one finished image. Returning every attempt to the
conversation would cost a round each, make the ceiling advisory rather than
enforced, and walk the reader past every reject -- so the reviewer lives inside
the tool and is asked about *bytes*: an attempt about to be discarded should not
leave an Attachment behind, so it sees a downscaled preview built in memory and
only the kept image is written. Anything that goes wrong in review is a keep;
losing a picture because a judging request timed out would be the check
destroying the thing it was checking. The last attempt is kept whatever the
verdict, so a request always produces something. Rejects are recorded, not
stored.
Preserve VRAM unloads the chat's own connection and nothing else, because the
memory being freed belongs to one machine: local llama-swap answers GET /unload,
and a box on the network has no reason to be unloaded when ComfyUI wants memory
here. The swap goes round the review rather than round the tool, which costs two
model loads per retry -- so the two settings are independent and the page warns
when both are on. Nothing loads the LLM back: the reply's next request does, and
that step exists in the description and not in the code, so the code says so.
Two rules elsewhere had to be drawn for the first time. message_payload sends
images only on user turns -- no assistant message had ever carried one, and the
moment one does the multimodal list form on an assistant turn is rejected by
OpenAI and most local runners, breaking every later turn in the chat. And
files.store gained keep_original, because _process_image turns anything without
alpha into JPEG q85 at 1400px: right for a phone photo, a visible loss on the one
output this feature exists to produce.
/image sends the ordinary message with force_tool, which becomes tool_choice for
the first round only -- left in place the reply would draw a picture, be asked
again, and draw another. FORCEABLE_TOOLS is an allow list because the name is
read off a form.
ToolContext gained chat_id, and that fixed a tool nobody had ever successfully
run: _run_scratch_write read context.chat_id on a dataclass with no such field,
so every call raised AttributeError, swallowed by run_tool's blanket except into
"the scratch_write tool failed" -- indistinguishable from a model calling it
wrongly. The test that existed asserted the family and the risk, which are
properties of the declaration rather than of the code.
Verified against the real ComfyUI 0.27.0 on this machine rather than against
documentation: every endpoint shape here was read off it, a generation ran end to
end through the client, the reviewer was shown a matching and a mismatched prompt
and answered KEEP and RETRY correctly, and the unload hook fired for the local
llama-swap and not for the remote box.
Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
396 lines
15 KiB
Python
396 lines
15 KiB
Python
"""Generating an image: what is offered, and what the loop does.
|
|
|
|
The gate tests matter more than they look. Image generation is the first tool
|
|
whose availability depends on an instance being *configured* as well as
|
|
permitted, and the branch it needs in `_family_allowed` is easy to leave out --
|
|
without it the family falls through to the last line and silently requires
|
|
`library.use`, which has nothing to do with drawing a picture.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
|
|
import lembas
|
|
from lembas.db.models import Chat, Connection, ImageWorkflow, Model, User
|
|
from lembas.services import settings_store
|
|
from lembas.services import tools as tools_service
|
|
from lembas.services.crypto import encrypt
|
|
from lembas.services.images import comfy
|
|
from lembas.services.images import tool as image_tool
|
|
from lembas.services.tools import ToolContext
|
|
|
|
BASE = json.loads(
|
|
(Path(lembas.__file__).parent / "services/images/base_workflow.json").read_text(
|
|
encoding="utf-8"
|
|
)
|
|
)
|
|
PNG = (
|
|
b"\x89PNG\r\n\x1a\n\x00\x00\x00\rIHDR\x00\x00\x00\x01\x00\x00\x00\x01\x08\x02\x00\x00\x00"
|
|
b"\x90wS\xde\x00\x00\x00\x0cIDATx\x9cc```\x00\x00\x00\x04\x00\x01\xf6\x178U\x00\x00\x00"
|
|
b"\x00IEND\xaeB`\x82"
|
|
)
|
|
|
|
|
|
@pytest.fixture
|
|
def configured(db, user_id):
|
|
"""An instance that can draw: a ComfyUI, a checkpoint and a workflow."""
|
|
connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt(""))
|
|
db.add(connection)
|
|
db.commit()
|
|
model = Model(
|
|
connection_id=connection.id,
|
|
model_id="m",
|
|
capabilities_json={"tools": True, "tool_image": True},
|
|
)
|
|
db.add(model)
|
|
db.add(ImageWorkflow(slug="base", name="Base", description="general", workflow_json=BASE))
|
|
db.commit()
|
|
chat = Chat(user_id=user_id, model_id="m", connection_id=connection.id)
|
|
db.add(chat)
|
|
settings_store.update(
|
|
db,
|
|
{
|
|
"enabled": True,
|
|
"base_url": "http://comfy.test:8188",
|
|
"checkpoints": ["sd.safetensors", "other.safetensors"],
|
|
"review_enabled": False,
|
|
},
|
|
key=settings_store.IMAGES,
|
|
)
|
|
db.commit()
|
|
return chat
|
|
|
|
|
|
def _context(db, user_id, chat) -> ToolContext:
|
|
return tools_service.context_for(db, db.get(User, user_id), chat)
|
|
|
|
|
|
# --- Whether it is offered at all ----------------------------------------------
|
|
def test_it_is_offered_once_the_instance_can_draw(db, user_id, configured):
|
|
offered = tools_service.resolve_tools(db, configured, db.get(User, user_id))
|
|
assert "image_generate" in offered.by_name
|
|
|
|
|
|
def test_it_is_withheld_without_a_comfyui(db, user_id, configured):
|
|
settings_store.update(db, {"base_url": ""}, key=settings_store.IMAGES)
|
|
assert (
|
|
"image_generate"
|
|
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
|
|
)
|
|
|
|
|
|
def test_it_is_withheld_with_no_checkpoints(db, user_id, configured):
|
|
"""A model naming a checkpoint that does not exist gets a refusal from
|
|
ComfyUI and spends a round finding out, so an instance that has listed none
|
|
should not offer the tool -- the same rule `skill_get` follows for an empty
|
|
library."""
|
|
settings_store.update(db, {"checkpoints": []}, key=settings_store.IMAGES)
|
|
assert (
|
|
"image_generate"
|
|
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
|
|
)
|
|
|
|
|
|
def test_it_is_withheld_when_the_feature_is_off(db, user_id, configured):
|
|
settings_store.update(db, {"enabled": False}, key=settings_store.IMAGES)
|
|
assert (
|
|
"image_generate"
|
|
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
|
|
)
|
|
|
|
|
|
def test_it_is_withheld_without_the_permission(db, user_id, configured):
|
|
user = db.get(User, user_id)
|
|
user.role = "user" # administrators are given every permission
|
|
settings_store.update(db, {"default_permissions": {"tools.image": False}})
|
|
db.commit()
|
|
assert "image_generate" not in tools_service.resolve_tools(db, configured, user).by_name
|
|
|
|
|
|
def test_it_does_not_need_the_library_permission(db, user_id, configured):
|
|
"""The trap the explicit branch in `_family_allowed` exists to avoid: with
|
|
no branch the family falls to the last line and requires `library.use`,
|
|
which is a coincidence of naming rather than a rule."""
|
|
user = db.get(User, user_id)
|
|
user.role = "user"
|
|
settings_store.update(db, {"default_permissions": {"library.use": False}})
|
|
db.commit()
|
|
assert "image_generate" in tools_service.resolve_tools(db, configured, user).by_name
|
|
|
|
|
|
def test_the_model_capability_can_withhold_it(db, user_id, configured):
|
|
model = db.scalar(Model.__table__.select().where(Model.model_id == "m"))
|
|
db.query(Model).filter(Model.model_id == "m").update(
|
|
{"capabilities_json": {"tools": True, "tool_image": False}}
|
|
)
|
|
db.commit()
|
|
assert model is not None
|
|
assert (
|
|
"image_generate"
|
|
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
|
|
)
|
|
|
|
|
|
def test_a_chat_can_switch_it_off(db, user_id, configured):
|
|
configured.scope_json = {"families": {"image": False}}
|
|
db.commit()
|
|
assert (
|
|
"image_generate"
|
|
not in tools_service.resolve_tools(db, configured, db.get(User, user_id)).by_name
|
|
)
|
|
|
|
|
|
def test_the_registry_knows_it_so_the_guidance_applies(db):
|
|
"""`harness._families` maps an offered name back to a family through
|
|
`registry(db)`. A tool missing from there resolves to no family, and its
|
|
fragment never appears -- the omission that cost custom tools their guidance
|
|
once already."""
|
|
assert tools_service.registry(db)["image_generate"].family == "image"
|
|
assert "image" in tools_service.families(db)
|
|
|
|
|
|
# --- The schema ----------------------------------------------------------------
|
|
def test_the_schema_offers_this_instance_s_own_choices(db, user_id, configured):
|
|
schema = image_tool.schema_for(db, settings_store.images(db))
|
|
assert schema["properties"]["model"]["enum"] == ["sd.safetensors", "other.safetensors"]
|
|
assert schema["properties"]["template"]["enum"] == ["base"]
|
|
assert "general" in schema["properties"]["template"]["description"]
|
|
|
|
|
|
def test_only_the_prompt_is_required_and_it_is_first(db, user_id, configured):
|
|
"""`tools.parse_arguments` puts the raw string into the first *required*
|
|
parameter when a model emits arguments that are not valid JSON -- common
|
|
with small models. This way that degrades into a prompt rather than a seed.
|
|
"""
|
|
schema = image_tool.schema_for(db, settings_store.images(db))
|
|
assert schema["required"] == ["prompt"]
|
|
assert next(iter(schema["properties"])) == "prompt"
|
|
|
|
|
|
def test_the_samplers_are_not_an_enum(db, user_id, configured):
|
|
"""Forty-four of them, on every request, forever, to prevent a mistake worth
|
|
one sentence of correction."""
|
|
schema = image_tool.schema_for(db, settings_store.images(db))
|
|
assert "enum" not in schema["properties"]["sampler"]
|
|
|
|
|
|
# --- The loop ------------------------------------------------------------------
|
|
class FakeComfy:
|
|
"""Stands in for the far side, counting what it was asked to do."""
|
|
|
|
def __init__(self):
|
|
self.submits: list[dict] = []
|
|
self.frees = 0
|
|
|
|
async def submit(self, config, wf):
|
|
self.submits.append(wf)
|
|
return f"p{len(self.submits)}"
|
|
|
|
async def await_images(self, config, prompt_id):
|
|
return [comfy.Ref(f"{prompt_id}.png")]
|
|
|
|
async def fetch_image(self, config, ref):
|
|
return PNG
|
|
|
|
async def free(self, config):
|
|
self.frees += 1
|
|
|
|
|
|
@pytest.fixture
|
|
def fake(monkeypatch):
|
|
stub = FakeComfy()
|
|
for name in ("submit", "await_images", "fetch_image", "free"):
|
|
monkeypatch.setattr(comfy, name, getattr(stub, name))
|
|
return stub
|
|
|
|
|
|
async def test_one_call_produces_one_stored_image(db, user_id, configured, fake):
|
|
from lembas.db.models import Attachment
|
|
|
|
outcome = await image_tool.run(
|
|
_context(db, user_id, configured), {"prompt": "a bicycle", "model": "sd.safetensors"}
|
|
)
|
|
|
|
assert outcome.event["status"] == "ok"
|
|
attachment = db.get(Attachment, outcome.event["attachment_id"])
|
|
assert attachment is not None
|
|
assert attachment.media_type == "image/png", "kept as ComfyUI made it, not transcoded"
|
|
assert attachment.message_id is None, "the loop binds it, not the runner"
|
|
assert attachment.source_label == "Image generation"
|
|
assert len(fake.submits) == 1
|
|
|
|
|
|
async def test_the_model_is_told_it_is_already_on_screen(db, user_id, configured, fake):
|
|
"""Without this the commonest next thing a model does is offer to show you
|
|
the image, which it has no way of doing and which has already happened."""
|
|
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "a bicycle"})
|
|
assert "already" in outcome.content or "is shown to them" in outcome.content
|
|
|
|
|
|
async def test_a_prompt_is_required(db, user_id, configured, fake):
|
|
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": " "})
|
|
assert outcome.event["status"] == "error"
|
|
assert not fake.submits
|
|
|
|
|
|
async def test_an_unknown_checkpoint_falls_back_rather_than_failing(db, user_id, configured, fake):
|
|
"""It would reach ComfyUI, be refused, and cost a round to discover -- and
|
|
the model was shown the list it may choose from."""
|
|
await image_tool.run(
|
|
_context(db, user_id, configured), {"prompt": "x", "model": "invented.safetensors"}
|
|
)
|
|
assert fake.submits[0]["4"]["inputs"]["ckpt_name"] == "sd.safetensors"
|
|
|
|
|
|
async def test_the_chat_s_own_preference_is_used_when_the_model_names_none(
|
|
db, user_id, configured, fake
|
|
):
|
|
configured.image_checkpoint = "other.safetensors"
|
|
db.commit()
|
|
await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert fake.submits[0]["4"]["inputs"]["ckpt_name"] == "other.safetensors"
|
|
|
|
|
|
async def test_a_failure_out_there_is_reported_not_raised(
|
|
db, user_id, configured, fake, monkeypatch
|
|
):
|
|
async def refuse(config, wf):
|
|
raise comfy.ComfyError("ComfyUI refused the workflow — node 4")
|
|
|
|
monkeypatch.setattr(comfy, "submit", refuse)
|
|
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert outcome.event["status"] == "error"
|
|
assert "node 4" in outcome.event["error"]
|
|
|
|
|
|
# --- Reviewing -----------------------------------------------------------------
|
|
def _with_review(db, chat, *, tries=3, vision=True):
|
|
db.query(Model).filter(Model.model_id == "m").update(
|
|
{"capabilities_json": {"tools": True, "tool_image": True, "vision": vision}}
|
|
)
|
|
settings_store.update(
|
|
db, {"review_enabled": True, "max_tries": tries}, key=settings_store.IMAGES
|
|
)
|
|
db.commit()
|
|
|
|
|
|
async def test_a_kept_image_stops_the_loop(db, user_id, configured, fake, monkeypatch):
|
|
_with_review(db, configured)
|
|
monkeypatch.setattr(image_tool, "_review", _verdicts([(True, "")]))
|
|
|
|
await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert len(fake.submits) == 1
|
|
|
|
|
|
async def test_a_rejected_image_is_tried_again(db, user_id, configured, fake, monkeypatch):
|
|
_with_review(db, configured)
|
|
monkeypatch.setattr(image_tool, "_review", _verdicts([(False, "no bicycle"), (True, "")]))
|
|
|
|
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert len(fake.submits) == 2
|
|
assert "no bicycle" in outcome.event["text"]
|
|
assert "2 attempts" in outcome.content
|
|
|
|
|
|
async def test_the_ceiling_is_enforced_and_the_last_one_is_kept(
|
|
db, user_id, configured, fake, monkeypatch
|
|
):
|
|
"""A request always produces a picture. Returning nothing after four
|
|
rejections would be spending a minute of somebody's GPU to say no."""
|
|
_with_review(db, configured, tries=3)
|
|
monkeypatch.setattr(image_tool, "_review", _verdicts([(False, "nope")] * 5))
|
|
|
|
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert len(fake.submits) == 3
|
|
assert outcome.event["status"] == "ok"
|
|
assert outcome.event["attachment_id"]
|
|
|
|
|
|
async def test_each_attempt_gets_a_new_seed(db, user_id, configured, fake, monkeypatch):
|
|
"""Retrying with the same seed regenerates the same rejected image."""
|
|
_with_review(db, configured, tries=3)
|
|
monkeypatch.setattr(
|
|
image_tool, "_review", _verdicts([(False, "no"), (False, "no"), (True, "")])
|
|
)
|
|
|
|
await image_tool.run(_context(db, user_id, configured), {"prompt": "x", "seed": 99})
|
|
seeds = [wf["3"]["inputs"]["seed"] for wf in fake.submits]
|
|
assert seeds[0] == 99, "the seed that was asked for is honoured"
|
|
assert len(set(seeds)) == 3, "and the retries are not the same picture again"
|
|
|
|
|
|
async def test_nothing_is_reviewed_without_a_vision_model(db, user_id, configured, fake):
|
|
""" "You asked for a picture and got an error about vision" is a worse answer
|
|
than a picture."""
|
|
_with_review(db, configured, vision=False)
|
|
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert len(fake.submits) == 1
|
|
assert outcome.event["status"] == "ok"
|
|
|
|
|
|
def _verdicts(sequence):
|
|
calls = {"n": 0}
|
|
|
|
async def review(context, endpoint, model_id, prompt, payload):
|
|
index = min(calls["n"], len(sequence) - 1)
|
|
calls["n"] += 1
|
|
return sequence[index]
|
|
|
|
return review
|
|
|
|
|
|
# --- Preserve VRAM -------------------------------------------------------------
|
|
async def test_nothing_is_unloaded_by_default(db, user_id, configured, fake, monkeypatch):
|
|
calls = {"n": 0}
|
|
|
|
async def unload(context):
|
|
calls["n"] += 1
|
|
return True
|
|
|
|
monkeypatch.setattr(image_tool, "_unload_llm", unload)
|
|
await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert calls["n"] == 0
|
|
assert fake.frees == 0
|
|
|
|
|
|
async def test_preserve_vram_unloads_and_frees(db, user_id, configured, fake, monkeypatch):
|
|
settings_store.update(db, {"preserve_vram": True}, key=settings_store.IMAGES)
|
|
db.commit()
|
|
calls = {"n": 0}
|
|
|
|
async def unload(context):
|
|
calls["n"] += 1
|
|
return True
|
|
|
|
monkeypatch.setattr(image_tool, "_unload_llm", unload)
|
|
await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert calls["n"] == 1, "once, before generating"
|
|
assert fake.frees >= 1, "and ComfyUI is asked to let go afterwards"
|
|
|
|
|
|
async def test_a_connection_with_no_unload_url_is_never_called(db, user_id, configured):
|
|
"""The whole of the per-connection design: a chat talking to a box on the
|
|
network must not have that box unloaded, because its memory is not the
|
|
memory ComfyUI wants."""
|
|
context = _context(db, user_id, configured)
|
|
assert await image_tool._unload_llm(context) is False
|
|
|
|
|
|
async def test_an_unload_that_fails_does_not_stop_the_generation(
|
|
db, user_id, configured, fake, monkeypatch
|
|
):
|
|
"""It is a hint before a slow operation. A machine that will not answer it is
|
|
one where the picture should still be drawn."""
|
|
connection = db.scalar(Connection.__table__.select())
|
|
db.query(Connection).update({"unload_url": "http://127.0.0.1:9/unload"})
|
|
settings_store.update(db, {"preserve_vram": True}, key=settings_store.IMAGES)
|
|
db.commit()
|
|
assert connection is not None
|
|
|
|
outcome = await image_tool.run(_context(db, user_id, configured), {"prompt": "x"})
|
|
assert outcome.event["status"] == "ok"
|