Draw a picture, on a ComfyUI you are running
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>
This commit is contained in:
Jaroslav Beneš
@@ -135,7 +135,16 @@ def message_payload(message: Message, *, vision: bool) -> dict[str, Any]:
|
||||
# in view, which is how these models are trained to read a prompt.
|
||||
text = f"{documents}\n\n{text}" if text else documents
|
||||
|
||||
images = message.images if vision else []
|
||||
# Images ride on a *user* turn and nowhere else. Until image generation
|
||||
# existed no assistant message had ever carried one, so this was never a
|
||||
# distinction worth drawing -- and the moment one does, the multimodal list
|
||||
# form on an `assistant` turn is rejected outright by OpenAI and by most
|
||||
# local runners, which would break not that turn but every later one in the
|
||||
# chat. What follows from it, and is worth knowing rather than discovering:
|
||||
# a model cannot see the picture it made on a *subsequent* turn (tool
|
||||
# results are not replayed either), so "make it bluer" regenerates rather
|
||||
# than edits. Honest for a text-to-image workflow with no img2img path.
|
||||
images = message.images if (vision and message.role == ROLE_USER) else []
|
||||
if not images:
|
||||
return {"role": message.role, "content": text}
|
||||
|
||||
@@ -304,6 +313,7 @@ def build_request(
|
||||
upto: Message | None = None,
|
||||
tools: list[dict[str, Any]] | None = None,
|
||||
user=None,
|
||||
force_tool: str = "",
|
||||
) -> dict[str, Any]:
|
||||
"""The whole request body, tools and harness included.
|
||||
|
||||
@@ -347,6 +357,19 @@ def build_request(
|
||||
}
|
||||
if tools:
|
||||
body["tools"] = tools
|
||||
# Making the model call one particular tool, for `/image` -- the whole
|
||||
# of what that command is. Only ever sent alongside a tools array and
|
||||
# only when something asked for it, so a provider strict about unknown
|
||||
# parameters sees exactly the request it always did until somebody types
|
||||
# a slash command.
|
||||
#
|
||||
# An endpoint that ignores `tool_choice` is not a failure here: the turn
|
||||
# still carries the instruction in words, so the model is being steered
|
||||
# twice and the weaker half is the one that can be dropped.
|
||||
if force_tool and any(
|
||||
(tool.get("function") or {}).get("name") == force_tool for tool in tools
|
||||
):
|
||||
body["tool_choice"] = {"type": "function", "function": {"name": force_tool}}
|
||||
|
||||
apply_effort(body, (chat.params_json or {}).get("reasoning_effort"))
|
||||
return body
|
||||
|
||||
@@ -263,15 +263,59 @@ def _process_text(payload: bytes, filename: str) -> Prepared:
|
||||
)
|
||||
|
||||
|
||||
def prepare(payload: bytes, filename: str) -> Prepared:
|
||||
"""Inspect an upload, decide what it is, and process it accordingly."""
|
||||
def _keep_image(payload: bytes) -> Prepared:
|
||||
"""An image stored as it arrived, measured but not re-encoded.
|
||||
|
||||
`_process_image` exists to protect the window from a phone camera: eight
|
||||
megapixels of JPEG become 1400px of JPEG at quality 85, and for something
|
||||
somebody photographed that is all upside. For an image *this application
|
||||
asked a diffusion model to make*, at a size somebody chose, it is a visible
|
||||
loss on the one output the feature exists to produce -- soft detail and
|
||||
ringing on exactly the fine texture the prompt was about.
|
||||
|
||||
Still opened by Pillow, so a malformed file is still refused and the
|
||||
dimensions are still real rather than claimed; still bounded by
|
||||
`MAX_UPLOAD_BYTES` in `prepare`. What is skipped is only the resize and the
|
||||
transcode.
|
||||
"""
|
||||
detected = _detect_image(payload)
|
||||
if detected is None:
|
||||
raise FileError("That is not an image.")
|
||||
media_type, extension = detected
|
||||
try:
|
||||
with Image.open(io.BytesIO(payload)) as image:
|
||||
image.load()
|
||||
width, height = image.size
|
||||
except Image.DecompressionBombError as exc:
|
||||
raise FileError("That image's dimensions are implausibly large.") from exc
|
||||
except (UnidentifiedImageError, OSError, ValueError) as exc:
|
||||
raise FileError("That image could not be read. Is it corrupt?") from exc
|
||||
|
||||
return Prepared(
|
||||
payload=payload,
|
||||
kind=KIND_IMAGE,
|
||||
media_type=media_type,
|
||||
extension=extension,
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
|
||||
|
||||
def prepare(payload: bytes, filename: str, *, keep_original: bool = False) -> Prepared:
|
||||
"""Inspect an upload, decide what it is, and process it accordingly.
|
||||
|
||||
`keep_original` is for an image the application produced rather than one
|
||||
somebody sent: see `_keep_image`. It applies to images only -- there is no
|
||||
argument for keeping an unparsed PDF, and the text path stores its bytes
|
||||
verbatim already.
|
||||
"""
|
||||
if not payload:
|
||||
raise FileError("That file is empty.")
|
||||
if len(payload) > MAX_UPLOAD_BYTES:
|
||||
raise FileError(f"Files must be under {MAX_UPLOAD_BYTES // (1024 * 1024)} MB.")
|
||||
|
||||
if _detect_image(payload) is not None:
|
||||
return _process_image(payload)
|
||||
return _keep_image(payload) if keep_original else _process_image(payload)
|
||||
if _looks_like_pdf(payload):
|
||||
return _process_pdf(payload)
|
||||
return _process_text(payload, filename)
|
||||
@@ -291,9 +335,19 @@ def store(
|
||||
chat_id: str | None,
|
||||
payload: bytes,
|
||||
filename: str,
|
||||
keep_original: bool = False,
|
||||
source_path: str = "",
|
||||
source_label: str = "",
|
||||
message_id: str | None = None,
|
||||
) -> Attachment:
|
||||
"""Process and persist an upload. Raises FileError if it is unusable."""
|
||||
prepared = prepare(payload, filename)
|
||||
"""Process and persist an upload. Raises FileError if it is unusable.
|
||||
|
||||
`message_id` is normally left null -- an upload is bound to a turn by
|
||||
`claim()` when the message is sent. A generated image is the mirror image of
|
||||
that: it exists *because* a reply is being written, so it says which turn it
|
||||
belongs to at the moment it is made.
|
||||
"""
|
||||
prepared = prepare(payload, filename, keep_original=keep_original)
|
||||
|
||||
stored_name = f"{secrets.token_hex(16)}{prepared.extension}"
|
||||
(attachments_dir() / stored_name).write_bytes(prepared.payload)
|
||||
@@ -301,6 +355,7 @@ def store(
|
||||
attachment = Attachment(
|
||||
user_id=user_id,
|
||||
chat_id=chat_id,
|
||||
message_id=message_id,
|
||||
filename=safe_display_name(filename),
|
||||
stored_name=stored_name,
|
||||
media_type=prepared.media_type,
|
||||
@@ -312,6 +367,8 @@ def store(
|
||||
pages=prepared.pages,
|
||||
truncated=prepared.truncated,
|
||||
extraction_error=prepared.extraction_error,
|
||||
source_path=source_path[:1000],
|
||||
source_label=source_label[:200],
|
||||
)
|
||||
db.add(attachment)
|
||||
db.commit()
|
||||
@@ -549,3 +606,38 @@ def data_uri(attachment: Attachment) -> str | None:
|
||||
return None
|
||||
encoded = base64.b64encode(path.read_bytes()).decode("ascii")
|
||||
return f"data:{attachment.media_type};base64,{encoded}"
|
||||
|
||||
|
||||
def preview_data_uri(payload: bytes, *, max_edge: int = MAX_IMAGE_EDGE) -> str | None:
|
||||
"""The same thing for bytes in hand, downscaled, for a model to look at.
|
||||
|
||||
Fidelity and weight are two different jobs. What is stored is what ComfyUI
|
||||
produced, because that is the artefact somebody keeps; what is *shown to a
|
||||
model to be judged* wants to be small, because a 400KB PNG is 550KB of
|
||||
base64 in a request that exists only to answer one question.
|
||||
|
||||
Takes bytes rather than an Attachment: the reviewer looks at an image that
|
||||
may be about to be thrown away, and writing a row for something rejected
|
||||
seconds later is work with nothing to show for it.
|
||||
"""
|
||||
import base64
|
||||
|
||||
try:
|
||||
with Image.open(io.BytesIO(payload)) as image:
|
||||
image.load()
|
||||
frame = image.convert("RGB")
|
||||
longest = max(frame.size)
|
||||
if longest > max_edge:
|
||||
scale = max_edge / longest
|
||||
frame = frame.resize(
|
||||
(max(1, int(frame.width * scale)), max(1, int(frame.height * scale))),
|
||||
Image.LANCZOS,
|
||||
)
|
||||
buffer = io.BytesIO()
|
||||
frame.save(buffer, format="JPEG", quality=JPEG_QUALITY, optimize=True)
|
||||
except (Image.DecompressionBombError, UnidentifiedImageError, OSError, ValueError):
|
||||
log.warning("could not build a preview of a generated image", exc_info=True)
|
||||
return None
|
||||
|
||||
encoded = base64.b64encode(buffer.getvalue()).decode("ascii")
|
||||
return f"data:image/jpeg;base64,{encoded}"
|
||||
|
||||
@@ -218,10 +218,21 @@ class Generation:
|
||||
# -- the one frame that reaches a browser after a reply is over.
|
||||
drained: bool = False
|
||||
injected_ids: list[str] = field(default_factory=list)
|
||||
# Images this reply produced, waiting to be bound to its message row. The
|
||||
# runner writes the file and the `Attachment`; only `_persist` may say which
|
||||
# turn it belongs to, which is the same division of labour `canvas` above
|
||||
# follows and for the same reason.
|
||||
attachment_ids: list[str] = field(default_factory=list)
|
||||
# How many times *in a row* this reply has ended with plan tasks still open
|
||||
# and been told to carry on. Reset the moment it calls a tool again, so the
|
||||
# count is of consecutive stops rather than of stops in total.
|
||||
nudges: int = 0
|
||||
# A tool this reply must call, set by `/image` and by nothing else. It goes
|
||||
# into the *first* request only -- `_run` rebuilds the payload's messages
|
||||
# per round but keeps this body, and `tool_choice` left in place would make
|
||||
# every later round call the tool again, which is a loop rather than a
|
||||
# command. Cleared once the first round has gone out.
|
||||
force_tool: str = ""
|
||||
|
||||
def touch(self) -> None:
|
||||
self.version += 1
|
||||
@@ -362,7 +373,7 @@ def _prune() -> None:
|
||||
_TASKS.pop(message_id, None)
|
||||
|
||||
|
||||
def ensure(chat_id: str, message_id: str) -> Generation:
|
||||
def ensure(chat_id: str, message_id: str, *, force_tool: str = "") -> Generation:
|
||||
"""Start generating this reply if it is not already under way.
|
||||
|
||||
Idempotent, because more than one thing can ask for it: the route that
|
||||
@@ -377,7 +388,7 @@ def ensure(chat_id: str, message_id: str) -> Generation:
|
||||
if existing is not None:
|
||||
return existing
|
||||
|
||||
generation = Generation(chat_id=chat_id, message_id=message_id)
|
||||
generation = Generation(chat_id=chat_id, message_id=message_id, force_tool=force_tool)
|
||||
_RUNNING[message_id] = generation
|
||||
_TASKS[message_id] = asyncio.create_task(_run(generation))
|
||||
return generation
|
||||
@@ -471,7 +482,7 @@ async def _run(generation: Generation) -> None:
|
||||
toolset = tools_service.resolve_tools(db, chat, owner)
|
||||
offered = toolset.schemas
|
||||
payload = chat_service.build_request(
|
||||
db, chat, upto=message, tools=offered, user=owner
|
||||
db, chat, upto=message, tools=offered, user=owner, force_tool=generation.force_tool
|
||||
)
|
||||
question = _question_from(payload)
|
||||
needs_title = not chat.title_generated
|
||||
@@ -777,6 +788,12 @@ async def _run(generation: Generation) -> None:
|
||||
# somebody through all of them -- or away from a file they
|
||||
# are editing -- is what makes a panel like this unusable.
|
||||
canvas_service.open_tab(generation.canvas, opened, activate=False)
|
||||
if attachment_id := outcome.event.get("attachment_id"):
|
||||
# A generated image. The runner wrote the row and the bytes;
|
||||
# binding it to this reply is the loop's job for the reason
|
||||
# the canvas tab above is -- a runner cannot write the
|
||||
# message row, and `_persist` is the single writer.
|
||||
generation.attachment_ids.append(str(attachment_id))
|
||||
if outcome.event.get("plan"):
|
||||
generation.plan = outcome.event["plan"]
|
||||
# Only `plan_submit` sets this. `plan_update` writes the
|
||||
@@ -809,6 +826,10 @@ async def _run(generation: Generation) -> None:
|
||||
messages.append(added)
|
||||
|
||||
payload = {**payload, "messages": messages}
|
||||
# `/image` forces the first round to call the tool. Leaving it set
|
||||
# would force *every* round to, so the reply could never finish --
|
||||
# it would draw a picture, be asked again, and draw another.
|
||||
payload.pop("tool_choice", None)
|
||||
|
||||
# A plan ends the turn. One more request so the model can say what
|
||||
# it proposed and why -- a bubble containing only a card reads as
|
||||
@@ -1902,6 +1923,28 @@ def _inject(generation: Generation, chat_id: str, vision: bool) -> dict | None:
|
||||
return entry
|
||||
|
||||
|
||||
def _bind_attachments(db, chat, message, ids: list[str]) -> None:
|
||||
"""Bind images this reply produced to the bubble that produced them.
|
||||
|
||||
The narrowing is the point. The ids arrive on a tool event, and an event is
|
||||
a dict a runner built -- so the query names this chat and refuses a row that
|
||||
is already bound, exactly as `files.claim` does for an upload, and for the
|
||||
identical reason: without it a forged id would attach somebody else's file
|
||||
to this conversation.
|
||||
"""
|
||||
from lembas.db.models import Attachment
|
||||
|
||||
rows = db.scalars(
|
||||
select(Attachment).where(
|
||||
Attachment.id.in_(ids),
|
||||
Attachment.chat_id == chat.id,
|
||||
Attachment.message_id.is_(None),
|
||||
)
|
||||
)
|
||||
for attachment in rows:
|
||||
attachment.message_id = message.id
|
||||
|
||||
|
||||
def _persist(generation: Generation, title: str, elapsed: float) -> None:
|
||||
"""Write the finished reply, name the chat, and set the unread flag.
|
||||
|
||||
@@ -1944,6 +1987,12 @@ def _persist(generation: Generation, title: str, elapsed: float) -> None:
|
||||
# the snapshot above was seeded when the reply began, and
|
||||
# somebody may have opened a tab by hand since.
|
||||
chat.canvas_json = canvas_service.merge(chat.canvas_json, generation.canvas)
|
||||
if generation.attachment_ids:
|
||||
# Images this reply made, bound to it here because this is the
|
||||
# only writer. Scoped to rows this chat owns and still unbound,
|
||||
# for the reason `files.claim` is scoped: an id that came back
|
||||
# on an event must not be able to pull in somebody else's file.
|
||||
_bind_attachments(db, chat, message, generation.attachment_ids)
|
||||
if generation.plan:
|
||||
# This bubble now carries the plan in force, and the chat points
|
||||
# at it so the harness can find it with one primary-key lookup
|
||||
|
||||
@@ -115,6 +115,32 @@ def _tool_names(tools: list[dict[str, Any]]) -> str:
|
||||
)
|
||||
|
||||
|
||||
def _image_templates(db: DBSession) -> str:
|
||||
"""One line per workflow, name and description.
|
||||
|
||||
The description is the load-bearing half, the same way it is for a skill:
|
||||
it is the only thing the model has to choose with, and "workflow-2" is not
|
||||
a choice. Capped, because a list of thirty costs the window on every
|
||||
request forever.
|
||||
"""
|
||||
from lembas.db.models import ImageWorkflow
|
||||
|
||||
rows = list(
|
||||
db.scalars(
|
||||
select(ImageWorkflow)
|
||||
.where(ImageWorkflow.enabled.is_(True))
|
||||
.order_by(ImageWorkflow.position, ImageWorkflow.slug)
|
||||
.limit(12)
|
||||
)
|
||||
)
|
||||
return "\n".join(f"- {row.slug}: {row.description or row.name}" for row in rows)
|
||||
|
||||
|
||||
def _image_models(db: DBSession) -> str:
|
||||
"""The checkpoints an administrator has listed, comma separated."""
|
||||
return ", ".join(settings_store.images(db).get("checkpoints") or [])
|
||||
|
||||
|
||||
def _document_names(db: DBSession, chat) -> str:
|
||||
"""The names of the non-image files attached anywhere in this chat."""
|
||||
from lembas.db.models import Attachment
|
||||
@@ -181,6 +207,18 @@ def context_variables(
|
||||
if "skills" in families
|
||||
else ""
|
||||
),
|
||||
# What can be drawn, and with what. Guarded by family for the reason the
|
||||
# memory block is: an instance with no ComfyUI must not pay a settings
|
||||
# read and a table scan to tell a model about a tool it was not offered.
|
||||
# Database reads only -- `context_variables` is synchronous and on the
|
||||
# request path, so asking ComfyUI itself what it has would hold the
|
||||
# request open while somebody's box thought about it. The admin page
|
||||
# discovers; this reads what it stored.
|
||||
"image_templates": _image_templates(db) if "image" in families else "",
|
||||
"image_models": _image_models(db) if "image" in families else "",
|
||||
"image_instructions": (
|
||||
str(settings_store.images(db).get("instructions") or "") if "image" in families else ""
|
||||
),
|
||||
"knowledge_bases": "",
|
||||
"document_names": "",
|
||||
"agent_target": "",
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
"""Making pictures, on a ComfyUI somebody else is running.
|
||||
|
||||
Three modules, split along the same seam the rest of the codebase uses:
|
||||
`comfy.py` speaks HTTP and knows nothing about chats, `workflow.py` turns a
|
||||
stored template plus a model's arguments into the document ComfyUI wants, and
|
||||
`tool.py` is the `ToolDef` that ties them to a conversation.
|
||||
|
||||
Nothing here executes anything locally. That is the same rule agent chats
|
||||
follow: the work happens on a service reached over HTTP, chosen and configured
|
||||
by an administrator, and the security of it is the security of that service.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
@@ -0,0 +1,52 @@
|
||||
{
|
||||
"3": {
|
||||
"inputs": {
|
||||
"seed": "{{seed}}",
|
||||
"steps": "{{steps}}",
|
||||
"cfg": "{{cfg}}",
|
||||
"sampler_name": "{{sampler}}",
|
||||
"scheduler": "{{scheduler}}",
|
||||
"denoise": "{{denoise}}",
|
||||
"model": ["4", 0],
|
||||
"positive": ["6", 0],
|
||||
"negative": ["7", 0],
|
||||
"latent_image": ["5", 0]
|
||||
},
|
||||
"class_type": "KSampler",
|
||||
"_meta": { "title": "KSampler" }
|
||||
},
|
||||
"4": {
|
||||
"inputs": { "ckpt_name": "{{model}}" },
|
||||
"class_type": "CheckpointLoaderSimple",
|
||||
"_meta": { "title": "Load Checkpoint" }
|
||||
},
|
||||
"5": {
|
||||
"inputs": {
|
||||
"width": "{{width}}",
|
||||
"height": "{{height}}",
|
||||
"batch_size": 1
|
||||
},
|
||||
"class_type": "EmptyLatentImage",
|
||||
"_meta": { "title": "Empty Latent Image" }
|
||||
},
|
||||
"6": {
|
||||
"inputs": { "text": "{{prompt}}", "clip": ["4", 1] },
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": { "title": "CLIP Text Encode (Prompt)" }
|
||||
},
|
||||
"7": {
|
||||
"inputs": { "text": "{{negative}}", "clip": ["4", 1] },
|
||||
"class_type": "CLIPTextEncode",
|
||||
"_meta": { "title": "CLIP Text Encode (Negative)" }
|
||||
},
|
||||
"8": {
|
||||
"inputs": { "samples": ["3", 0], "vae": ["4", 2] },
|
||||
"class_type": "VAEDecode",
|
||||
"_meta": { "title": "VAE Decode" }
|
||||
},
|
||||
"9": {
|
||||
"inputs": { "filename_prefix": "LLeMbas", "images": ["8", 0] },
|
||||
"class_type": "SaveImage",
|
||||
"_meta": { "title": "Save Image" }
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,305 @@
|
||||
"""Talking to ComfyUI.
|
||||
|
||||
Four calls and a discovery one, all plain httpx. `fetch.fetch` cannot be reused
|
||||
for the same reasons `custom_tools` gives -- it is GET-only, bodyless, and
|
||||
refuses every content type that is not HTML or text, which is both the JSON here
|
||||
and the PNG at the end of it.
|
||||
|
||||
**The base URL is exempt from the SSRF guard, and that is deliberate rather than
|
||||
forgotten.** `fetch.check_url` exists to stop a *model or a reader* pointing the
|
||||
application at something on the private network; this address was typed by an
|
||||
administrator into the admin page, exactly like `Connection.base_url` and the two
|
||||
audio endpoints, none of which are checked either. Saying so here because the
|
||||
default value is `127.0.0.1:8188`, which is precisely the shape the guard exists
|
||||
to refuse and therefore looks like a hole rather than a decision.
|
||||
|
||||
Progress is **polled, not streamed**. ComfyUI offers a WebSocket for it, and
|
||||
holding one open for the length of a generation is the live-connection state the
|
||||
whole `agent/ssh.py` design forbids; polling `/history` is self-healing across a
|
||||
restart of either side, and the thing being waited for takes tens of seconds, so
|
||||
a poll costs nothing anybody can measure.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import time
|
||||
import uuid
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
|
||||
from lembas.services.llm.openai_client import (
|
||||
LLMError,
|
||||
describe_http_error,
|
||||
)
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# What one generated image may weigh. A cap is required rather than tidy: this is
|
||||
# the only place in the codebase where an external service hands back raw bytes
|
||||
# that are then written to disk, and neither `audio.speak` nor `openai_client`
|
||||
# has one to copy. Generous, because a 2048px PNG is a legitimate several
|
||||
# megabytes and refusing it would be refusing the feature.
|
||||
MAX_IMAGE_BYTES = 32 * 1024 * 1024
|
||||
|
||||
# How often to ask whether it has finished, and how long to keep asking. The
|
||||
# interval is not adaptive: unlike a background job, which may run for hours,
|
||||
# a generation is over in tens of seconds and the whole reply is parked on it.
|
||||
POLL_INTERVAL = 1.0
|
||||
# How long to wait for the queue *before* our own job starts running. A busy
|
||||
# ComfyUI with somebody else's batch in front of us is not an error.
|
||||
DEFAULT_TIMEOUT = 600.0
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Config:
|
||||
"""Everything a call needs, lifted out of the settings group.
|
||||
|
||||
A snapshot rather than a session, for the reason `ToolContext` is one: a
|
||||
generation outlives the request that resolved it.
|
||||
"""
|
||||
|
||||
base_url: str
|
||||
api_key: str = ""
|
||||
timeout: float = DEFAULT_TIMEOUT
|
||||
|
||||
@property
|
||||
def configured(self) -> bool:
|
||||
return bool(self.base_url)
|
||||
|
||||
def url(self, path: str) -> str:
|
||||
return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}"
|
||||
|
||||
def headers(self) -> dict[str, str]:
|
||||
# ComfyUI itself has no auth; a key is only ever for something in front
|
||||
# of it, so an empty one must not become `Authorization: Bearer `.
|
||||
return {"Authorization": f"Bearer {self.api_key}"} if self.api_key else {}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Ref:
|
||||
"""Where a finished image lives on the far side."""
|
||||
|
||||
filename: str
|
||||
subfolder: str = ""
|
||||
kind: str = "output"
|
||||
|
||||
|
||||
class ComfyError(LLMError):
|
||||
"""Anything that stopped a generation, in words worth showing somebody."""
|
||||
|
||||
|
||||
def _transport_error(exc: httpx.RequestError, config: Config) -> ComfyError:
|
||||
"""The `wrap_transport_error` shape, said about ComfyUI rather than an LLM.
|
||||
|
||||
Not reused directly: that one names the request timeout from the deployment
|
||||
settings, which is not the timeout in force here.
|
||||
"""
|
||||
if isinstance(exc, httpx.ConnectError):
|
||||
return ComfyError(
|
||||
f"Could not reach ComfyUI at {config.base_url}. Is it running and the URL correct?"
|
||||
)
|
||||
if isinstance(exc, httpx.TimeoutException):
|
||||
return ComfyError(f"ComfyUI at {config.base_url} did not respond in time.")
|
||||
return ComfyError(f"Could not reach ComfyUI at {config.base_url}: {exc}")
|
||||
|
||||
|
||||
async def _get_json(config: Config, path: str, *, timeout: float = 30.0) -> Any:
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=timeout) as client:
|
||||
response = await client.get(config.url(path), headers=config.headers())
|
||||
response.raise_for_status()
|
||||
return response.json()
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise ComfyError(describe_http_error(exc), status_code=exc.response.status_code) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise _transport_error(exc, config) from exc
|
||||
except (ValueError, json.JSONDecodeError) as exc:
|
||||
raise ComfyError(f"ComfyUI sent something that is not JSON: {exc}") from exc
|
||||
|
||||
|
||||
async def submit(config: Config, workflow: dict[str, Any]) -> str:
|
||||
"""Queue a workflow, and answer with the id it was given.
|
||||
|
||||
A `node_errors` block is a refusal rather than a failure: the workflow was
|
||||
accepted as JSON and rejected as a graph, usually because a checkpoint name
|
||||
does not exist on that machine. It is reported with the node named, because
|
||||
"invalid prompt" against a twelve-node document says nothing.
|
||||
"""
|
||||
body = {"prompt": workflow, "client_id": uuid.uuid4().hex}
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=60.0) as client:
|
||||
response = await client.post(
|
||||
config.url("prompt"), headers=config.headers(), json=body
|
||||
)
|
||||
if response.status_code >= 400:
|
||||
raise ComfyError(_refusal(response))
|
||||
data = response.json()
|
||||
except ComfyError:
|
||||
raise
|
||||
except httpx.RequestError as exc:
|
||||
raise _transport_error(exc, config) from exc
|
||||
except (ValueError, json.JSONDecodeError) as exc:
|
||||
raise ComfyError(f"ComfyUI sent something that is not JSON: {exc}") from exc
|
||||
|
||||
if errors := (data.get("node_errors") or {}):
|
||||
raise ComfyError(_describe_nodes(errors))
|
||||
prompt_id = str(data.get("prompt_id") or "")
|
||||
if not prompt_id:
|
||||
raise ComfyError("ComfyUI accepted the workflow but did not say what to call it.")
|
||||
return prompt_id
|
||||
|
||||
|
||||
def _refusal(response: httpx.Response) -> str:
|
||||
"""Why ComfyUI would not take a workflow, in one sentence."""
|
||||
try:
|
||||
payload = response.json()
|
||||
except (ValueError, json.JSONDecodeError):
|
||||
return f"ComfyUI refused the workflow (HTTP {response.status_code})."
|
||||
if isinstance(payload, dict):
|
||||
if errors := (payload.get("node_errors") or {}):
|
||||
return _describe_nodes(errors)
|
||||
if message := payload.get("error"):
|
||||
if isinstance(message, dict):
|
||||
message = message.get("message") or message.get("type") or ""
|
||||
return f"ComfyUI refused the workflow: {message}"
|
||||
return f"ComfyUI refused the workflow (HTTP {response.status_code})."
|
||||
|
||||
|
||||
def _describe_nodes(errors: dict[str, Any]) -> str:
|
||||
parts: list[str] = []
|
||||
for node, detail in list(errors.items())[:4]:
|
||||
messages = detail.get("errors") if isinstance(detail, dict) else None
|
||||
first = ""
|
||||
if isinstance(messages, list) and messages:
|
||||
entry = messages[0]
|
||||
first = entry.get("message", "") if isinstance(entry, dict) else str(entry)
|
||||
parts.append(f"node {node}: {first}" if first else f"node {node}")
|
||||
return "ComfyUI refused the workflow — " + "; ".join(parts)
|
||||
|
||||
|
||||
async def await_images(config: Config, prompt_id: str) -> list[Ref]:
|
||||
"""Wait for one queued workflow and answer with what it saved.
|
||||
|
||||
`/history/{id}` is empty while the job is queued or running and gains the
|
||||
whole record when it ends, so an empty answer is "not yet" rather than
|
||||
"nothing" -- which is why the deadline is the only thing that ends this.
|
||||
"""
|
||||
deadline = time.monotonic() + config.timeout
|
||||
while True:
|
||||
record = (await _get_json(config, f"history/{prompt_id}")).get(prompt_id)
|
||||
if isinstance(record, dict) and (record.get("status") or {}).get("completed"):
|
||||
status = record.get("status") or {}
|
||||
if status.get("status_str") not in (None, "success"):
|
||||
raise ComfyError(
|
||||
f"ComfyUI could not finish the workflow ({status.get('status_str')})."
|
||||
)
|
||||
return _refs_in(record.get("outputs") or {})
|
||||
if time.monotonic() > deadline:
|
||||
raise ComfyError(
|
||||
f"ComfyUI did not finish within {config.timeout:.0f}s. "
|
||||
"It may still be working; the queue is on its own page."
|
||||
)
|
||||
await asyncio.sleep(POLL_INTERVAL)
|
||||
|
||||
|
||||
def _refs_in(outputs: dict[str, Any]) -> list[Ref]:
|
||||
"""Every image any node saved, in node order.
|
||||
|
||||
Every node is read rather than a `SaveImage` being looked for by name: a
|
||||
template is somebody else's document and may save from a node called
|
||||
anything, or from two of them.
|
||||
"""
|
||||
refs: list[Ref] = []
|
||||
for node in outputs.values():
|
||||
for image in (node or {}).get("images") or []:
|
||||
if filename := str(image.get("filename") or ""):
|
||||
refs.append(
|
||||
Ref(
|
||||
filename=filename,
|
||||
subfolder=str(image.get("subfolder") or ""),
|
||||
kind=str(image.get("type") or "output"),
|
||||
)
|
||||
)
|
||||
return refs
|
||||
|
||||
|
||||
async def fetch_image(config: Config, ref: Ref) -> bytes:
|
||||
"""The bytes of one finished image."""
|
||||
params = {"filename": ref.filename, "subfolder": ref.subfolder, "type": ref.kind}
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
response = await client.get(
|
||||
config.url("view"), headers=config.headers(), params=params
|
||||
)
|
||||
response.raise_for_status()
|
||||
payload = response.content
|
||||
except httpx.HTTPStatusError as exc:
|
||||
raise ComfyError(describe_http_error(exc), status_code=exc.response.status_code) from exc
|
||||
except httpx.RequestError as exc:
|
||||
raise _transport_error(exc, config) from exc
|
||||
|
||||
if not payload:
|
||||
raise ComfyError(f"ComfyUI returned an empty file for {ref.filename}.")
|
||||
if len(payload) > MAX_IMAGE_BYTES:
|
||||
raise ComfyError(
|
||||
f"{ref.filename} is {len(payload) // (1024 * 1024)}MB, over the "
|
||||
f"{MAX_IMAGE_BYTES // (1024 * 1024)}MB limit."
|
||||
)
|
||||
return payload
|
||||
|
||||
|
||||
async def free(config: Config) -> None:
|
||||
"""Ask ComfyUI to drop its models from memory.
|
||||
|
||||
Best-effort by design and never raised into the caller: this runs on the way
|
||||
out of a generation that has already produced its image, and failing the
|
||||
whole tool because a memory hint was refused would be turning a tidy-up into
|
||||
an error. The consequence of it silently not working is VRAM staying used,
|
||||
which is the state Preserve VRAM was already in before it was switched on.
|
||||
"""
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=30.0) as client:
|
||||
await client.post(
|
||||
config.url("free"),
|
||||
headers=config.headers(),
|
||||
json={"unload_models": True, "free_memory": True},
|
||||
)
|
||||
except Exception: # noqa: BLE001 - a hint that failed is not a failed generation
|
||||
log.debug("could not free ComfyUI at %s", config.base_url, exc_info=True)
|
||||
|
||||
|
||||
async def discover(config: Config) -> tuple[list[str], list[str], list[str]]:
|
||||
"""What this ComfyUI can actually do: checkpoints, samplers, schedulers.
|
||||
|
||||
For the admin page only. Never called from the request path -- the tool
|
||||
reads the stored lists, exactly as the project listing is read from a cache
|
||||
rather than walked, because a keystroke must not wait on a machine.
|
||||
"""
|
||||
checkpoints = _options(
|
||||
await _get_json(config, "object_info/CheckpointLoaderSimple"),
|
||||
"CheckpointLoaderSimple",
|
||||
"ckpt_name",
|
||||
)
|
||||
sampler_info = await _get_json(config, "object_info/KSampler")
|
||||
samplers = _options(sampler_info, "KSampler", "sampler_name")
|
||||
schedulers = _options(sampler_info, "KSampler", "scheduler")
|
||||
return checkpoints, samplers, schedulers
|
||||
|
||||
|
||||
def _options(payload: Any, node: str, field: str) -> list[str]:
|
||||
"""The allowed values of one input, out of an `/object_info` document.
|
||||
|
||||
The shape is `{node: {input: {required: {field: [[...values], {...meta}]}}}}`
|
||||
-- a list whose first element is the list of options. Read defensively: this
|
||||
is somebody else's schema and a custom node pack can change it.
|
||||
"""
|
||||
try:
|
||||
spec = payload[node]["input"]["required"][field][0]
|
||||
except (KeyError, IndexError, TypeError):
|
||||
return []
|
||||
return [str(value) for value in spec] if isinstance(spec, list) else []
|
||||
@@ -0,0 +1,526 @@
|
||||
"""The tool that makes a picture, and the loop that decides to keep it.
|
||||
|
||||
One call is one finished image. The alternative -- return every attempt to the
|
||||
conversation and let the model decide whether to call again -- costs a full
|
||||
round per retry, makes the ceiling advisory rather than enforced, and shows the
|
||||
reader every reject on the way past. So the retrying happens here, and what
|
||||
comes back is the image that was kept.
|
||||
|
||||
**Three things are ordered rather than incidental.**
|
||||
|
||||
*The reviewer is asked about bytes, not about a row.* An attempt that is going
|
||||
to be thrown away should not leave an `Attachment` behind, so the judge is shown
|
||||
a downscaled preview built in memory and only the kept image is ever written.
|
||||
|
||||
*Preserve VRAM swaps around the review, not around the tool.* The sequence is
|
||||
unload the LLM, generate, free ComfyUI, ask the reviewer (which loads the LLM
|
||||
again), and round once more if it said no. Two model loads per retry, which is
|
||||
why the two settings are independent and the admin page says so.
|
||||
|
||||
*Nothing loads the LLM back at the end.* The reply's next request does it, and
|
||||
llama-swap -- or Ollama, or anything else worth pointing this at -- loads on
|
||||
demand. A step that exists in the description and not in the code looks like an
|
||||
omission, so it is said here instead.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import logging
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import httpx
|
||||
from sqlalchemy import select
|
||||
|
||||
from lembas.services.images import comfy, workflow
|
||||
from lembas.services.llm.openai_client import Endpoint, LLMError, complete
|
||||
from lembas.services.tools import RISK_WRITE, ToolContext, ToolDef, ToolOutcome
|
||||
|
||||
log = logging.getLogger(__name__)
|
||||
|
||||
# What the reviewer is allowed to write back. It is one verdict and one line of
|
||||
# reason, and a model that writes an essay about a picture is a model whose
|
||||
# answer nobody reads.
|
||||
MAX_VERDICT_TOKENS = 200
|
||||
|
||||
# How long to wait for a connection to admit it has unloaded. Short: this is a
|
||||
# hint before a slow operation, and a machine that will not answer it is one
|
||||
# where the generation should go ahead anyway rather than fail.
|
||||
UNLOAD_TIMEOUT = 30.0
|
||||
|
||||
SCHEMA: dict[str, Any] = {
|
||||
"type": "object",
|
||||
"properties": {
|
||||
# First, and the only required one, because `tools.parse_arguments`
|
||||
# puts the whole raw string into the first required parameter when a
|
||||
# model emits arguments that are not valid JSON. That failure is common
|
||||
# with small models, and this way it degrades into a prompt rather than
|
||||
# into a seed.
|
||||
"prompt": {
|
||||
"type": "string",
|
||||
"description": "What to draw. Describe the subject, the setting and the style.",
|
||||
},
|
||||
"negative": {
|
||||
"type": "string",
|
||||
"description": "What to keep out of the picture. Defaults to 'text, watermark'.",
|
||||
},
|
||||
"template": {
|
||||
"type": "string",
|
||||
"description": "Which workflow to use. Omit to use this chat's usual one.",
|
||||
},
|
||||
"model": {
|
||||
"type": "string",
|
||||
"description": "Which checkpoint to draw with. Omit to use this chat's usual one.",
|
||||
},
|
||||
"seed": {
|
||||
"type": "integer",
|
||||
"description": "Omit for a new random image; repeat one to get the same image again.",
|
||||
},
|
||||
"steps": {"type": "integer", "description": "Sampling steps. Default 20."},
|
||||
"cfg": {"type": "number", "description": "Prompt adherence. Default 8."},
|
||||
"width": {"type": "integer", "description": "Pixels. Default 512."},
|
||||
"height": {"type": "integer", "description": "Pixels. Default 512."},
|
||||
"sampler": {"type": "string", "description": "Sampler name. Default euler."},
|
||||
"scheduler": {"type": "string", "description": "Scheduler name. Default normal."},
|
||||
"denoise": {"type": "number", "description": "0 to 1. Default 1."},
|
||||
},
|
||||
"required": ["prompt"],
|
||||
}
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Attempt:
|
||||
"""One generated image and what was decided about it."""
|
||||
|
||||
number: int
|
||||
seed: int
|
||||
kept: bool
|
||||
verdict: str = ""
|
||||
|
||||
|
||||
def config_of(context: ToolContext) -> comfy.Config:
|
||||
"""The client snapshot, with the key decrypted at the last moment."""
|
||||
from lembas.services.crypto import decrypt
|
||||
|
||||
values = context.image_config or {}
|
||||
return comfy.Config(
|
||||
base_url=str(values.get("base_url") or ""),
|
||||
api_key=decrypt(str(values.get("api_key_encrypted") or "")),
|
||||
timeout=float(values.get("timeout") or comfy.DEFAULT_TIMEOUT),
|
||||
)
|
||||
|
||||
|
||||
def _choices(db, values: dict[str, Any]) -> tuple[list[Any], list[str]]:
|
||||
"""The templates and checkpoints on offer, for the schema and the harness."""
|
||||
from lembas.db.models import ImageWorkflow
|
||||
|
||||
rows = list(
|
||||
db.scalars(
|
||||
select(ImageWorkflow)
|
||||
.where(ImageWorkflow.enabled.is_(True))
|
||||
.order_by(ImageWorkflow.position, ImageWorkflow.slug)
|
||||
)
|
||||
)
|
||||
return rows, [str(name) for name in (values.get("checkpoints") or [])]
|
||||
|
||||
|
||||
def schema_for(db, values: dict[str, Any]) -> dict[str, Any]:
|
||||
"""The parameter schema, with this instance's own choices in it.
|
||||
|
||||
`template` and `model` become enums because a name that does not exist is a
|
||||
refusal from ComfyUI and a wasted round; `sampler` and `scheduler` stay
|
||||
plain strings because there are forty-four and nine of them, and an enum
|
||||
that size costs tokens on every request forever to prevent a mistake worth
|
||||
one sentence of correction.
|
||||
"""
|
||||
rows, checkpoints = _choices(db, values)
|
||||
schema = json.loads(json.dumps(SCHEMA))
|
||||
if rows:
|
||||
schema["properties"]["template"]["enum"] = [row.slug for row in rows]
|
||||
schema["properties"]["template"]["description"] = "Which workflow to use. " + "; ".join(
|
||||
f"{row.slug}: {row.description or row.name}" for row in rows[:12]
|
||||
)
|
||||
if checkpoints:
|
||||
schema["properties"]["model"]["enum"] = checkpoints
|
||||
return schema
|
||||
|
||||
|
||||
def tool_def(db, values: dict[str, Any]) -> ToolDef:
|
||||
return ToolDef(
|
||||
name="image_generate",
|
||||
family="image",
|
||||
description=(
|
||||
"Draw a picture from a description and show it to the person you are "
|
||||
"talking to. Returns once the image has been made and is on screen."
|
||||
),
|
||||
parameters=schema_for(db, values),
|
||||
run=run,
|
||||
# Not RISK_READ: it spends somebody's GPU for a minute and puts a new
|
||||
# artefact in the conversation. In an agent chat that means the mode
|
||||
# decides whether to ask first, which is the right answer for a call
|
||||
# that cannot be undone by reading something again.
|
||||
risk=RISK_WRITE,
|
||||
)
|
||||
|
||||
|
||||
# --- Preserve VRAM -------------------------------------------------------------
|
||||
async def _unload_llm(context: ToolContext) -> bool:
|
||||
"""Ask this chat's own endpoint to drop its model. Best-effort.
|
||||
|
||||
*This chat's own* is the whole of the design. The unload hook is a column on
|
||||
`Connection`, so a chat talking to a local llama-swap unloads that and a
|
||||
chat talking to a box on the network unloads nothing -- its VRAM is not the
|
||||
VRAM ComfyUI is about to want.
|
||||
"""
|
||||
from lembas.db.models import Connection
|
||||
from lembas.db.session import session_scope
|
||||
|
||||
url = ""
|
||||
method = "POST"
|
||||
try:
|
||||
with session_scope() as db:
|
||||
connection = db.get(Connection, context.connection_id)
|
||||
if connection is not None:
|
||||
url = (connection.unload_url or "").strip()
|
||||
method = (connection.unload_method or "POST").upper()
|
||||
except Exception: # noqa: BLE001 - a hint that could not be looked up is not a failure
|
||||
log.debug("could not read the unload hook", exc_info=True)
|
||||
return False
|
||||
|
||||
if not url:
|
||||
return False
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=UNLOAD_TIMEOUT) as client:
|
||||
await client.request(method, url)
|
||||
return True
|
||||
except Exception: # noqa: BLE001 - see the module docstring: a hint, not a step
|
||||
log.info("could not unload the model at %s", url, exc_info=True)
|
||||
return False
|
||||
|
||||
|
||||
# --- The reviewer --------------------------------------------------------------
|
||||
def _reviewer(context: ToolContext) -> tuple[Endpoint, str] | None:
|
||||
"""The model that judges an image, or None if there is nobody to ask.
|
||||
|
||||
The admin's choice first, then the chat's own model when it has vision. A
|
||||
chat on a text-only model with no reviewer configured simply keeps the first
|
||||
image, which is the behaviour with review switched off -- said here rather
|
||||
than failing, because "you asked for a picture and got an error about
|
||||
vision" is a worse answer than a picture.
|
||||
"""
|
||||
from lembas.db.models import Connection, Model
|
||||
from lembas.db.session import session_scope
|
||||
|
||||
values = context.image_config or {}
|
||||
if not values.get("review_enabled"):
|
||||
return None
|
||||
|
||||
wanted = str(values.get("review_model_id") or "")
|
||||
try:
|
||||
with session_scope() as db:
|
||||
model = None
|
||||
if wanted:
|
||||
model = db.get(Model, wanted)
|
||||
if model is None and context.model_id:
|
||||
model = db.scalar(
|
||||
select(Model).where(
|
||||
Model.model_id == context.model_id,
|
||||
Model.connection_id == context.connection_id,
|
||||
)
|
||||
)
|
||||
if model is None or not (model.capabilities_json or {}).get("vision"):
|
||||
return None
|
||||
connection = db.get(Connection, model.connection_id)
|
||||
if connection is None or not connection.enabled:
|
||||
return None
|
||||
return Endpoint.from_connection(connection), model.model_id
|
||||
except Exception: # noqa: BLE001 - no reviewer is a degraded mode, not an error
|
||||
log.warning("could not resolve an image reviewer", exc_info=True)
|
||||
return None
|
||||
|
||||
|
||||
async def _review(
|
||||
context: ToolContext, endpoint: Endpoint, model_id: str, prompt: str, payload: bytes
|
||||
) -> tuple[bool, str]:
|
||||
"""Show the reviewer the image and ask whether to keep it.
|
||||
|
||||
Answers `(keep, reason)`. **Anything that goes wrong is a keep**: the
|
||||
reviewer is a second opinion on a picture that already exists, and losing an
|
||||
image because a judging request timed out would be the check destroying the
|
||||
thing it was checking.
|
||||
"""
|
||||
from lembas.db.session import session_scope
|
||||
from lembas.services import files as files_service
|
||||
from lembas.services import prompts as prompts_service
|
||||
|
||||
preview = files_service.preview_data_uri(payload, max_edge=768)
|
||||
if preview is None:
|
||||
return True, ""
|
||||
|
||||
with session_scope() as db:
|
||||
instruction = prompts_service.resolve(db, "task.image_review")
|
||||
# An administrator who cleared the fragment has switched reviewing off, the
|
||||
# same way clearing `task.compact` switches compaction off. Nothing is asked
|
||||
# of anyone and the image is kept.
|
||||
if not instruction.strip():
|
||||
return True, ""
|
||||
|
||||
body = {
|
||||
"model": model_id,
|
||||
"messages": [
|
||||
{"role": "system", "content": instruction},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": f"The request was: {prompt}"},
|
||||
{"type": "image_url", "image_url": {"url": preview}},
|
||||
],
|
||||
},
|
||||
],
|
||||
"max_tokens": MAX_VERDICT_TOKENS,
|
||||
"temperature": 0,
|
||||
}
|
||||
try:
|
||||
answer = (await complete(endpoint, body)).strip()
|
||||
except LLMError as exc:
|
||||
log.info("could not review a generated image: %s", exc.message)
|
||||
return True, ""
|
||||
|
||||
verdict, _, reason = answer.partition("\n")
|
||||
keep = not verdict.strip().upper().startswith("RETRY")
|
||||
return keep, (reason or verdict).strip()[:300]
|
||||
|
||||
|
||||
# --- The runner ----------------------------------------------------------------
|
||||
async def run(context: ToolContext, args: dict[str, Any]) -> ToolOutcome:
|
||||
"""Generate one image, review it if there is anybody to ask, and keep one."""
|
||||
from lembas.db.session import session_scope
|
||||
from lembas.services import files as files_service
|
||||
|
||||
event: dict[str, Any] = {
|
||||
"name": "image_generate",
|
||||
"kind": "image",
|
||||
"query": str(args.get("prompt") or "")[:200],
|
||||
"results": [],
|
||||
}
|
||||
|
||||
prompt = str(args.get("prompt") or "").strip()
|
||||
if not prompt:
|
||||
return ToolOutcome(
|
||||
"No prompt was given, so nothing was drawn. Say what the picture should show.",
|
||||
{**event, "status": "error", "error": "No prompt."},
|
||||
)
|
||||
if not context.chat_id:
|
||||
return ToolOutcome(
|
||||
"Images can only be generated inside a chat.",
|
||||
{**event, "status": "error", "error": "No chat."},
|
||||
)
|
||||
|
||||
values = context.image_config or {}
|
||||
config = config_of(context)
|
||||
if not config.configured:
|
||||
return ToolOutcome(
|
||||
"No image generator is configured on this instance.",
|
||||
{**event, "status": "error", "error": "No ComfyUI configured."},
|
||||
)
|
||||
|
||||
# Resolve the template and the checkpoint: what the model asked for, then
|
||||
# this chat's usual, then the instance default. Every rung is a preference
|
||||
# and none of them is a constraint, which is what lets a model that only
|
||||
# wrote a prompt still get a picture.
|
||||
try:
|
||||
with session_scope() as db:
|
||||
rows, checkpoints = _choices(db, values)
|
||||
wanted = str(args.get("template") or "")
|
||||
chosen = _pick(rows, wanted, context.image_workflow_id, values)
|
||||
if chosen is None:
|
||||
return ToolOutcome(
|
||||
"No image workflow has been set up on this instance.",
|
||||
{**event, "status": "error", "error": "No workflow."},
|
||||
)
|
||||
template = json.loads(json.dumps(chosen.workflow_json or {}))
|
||||
template_slug, template_name = chosen.slug, chosen.name
|
||||
except ToolOutcome: # pragma: no cover - defensive
|
||||
raise
|
||||
except Exception as exc: # noqa: BLE001
|
||||
log.exception("could not resolve an image workflow")
|
||||
return ToolOutcome(
|
||||
f"The image workflow could not be read: {exc}",
|
||||
{**event, "status": "error", "error": str(exc)},
|
||||
)
|
||||
|
||||
checkpoint = _checkpoint(str(args.get("model") or ""), context.image_checkpoint, checkpoints)
|
||||
if checkpoint is None:
|
||||
return ToolOutcome(
|
||||
"No checkpoint is available. An administrator has to list them on the "
|
||||
"image generation page.",
|
||||
{**event, "status": "error", "error": "No checkpoint."},
|
||||
)
|
||||
|
||||
given = {name: args.get(name) for name in workflow.PLACEHOLDERS if name in args}
|
||||
given["model"] = checkpoint
|
||||
given["prompt"] = prompt
|
||||
|
||||
reviewer = _reviewer(context)
|
||||
tries = int(values.get("max_tries") or 1) if reviewer else 1
|
||||
preserve = bool(values.get("preserve_vram"))
|
||||
|
||||
attempts: list[Attempt] = []
|
||||
kept: tuple[bytes, dict[str, Any]] | None = None
|
||||
|
||||
try:
|
||||
for number in range(1, tries + 1):
|
||||
if preserve:
|
||||
await _unload_llm(context)
|
||||
|
||||
params = workflow.resolve({**given, "seed": args.get("seed") if number == 1 else None})
|
||||
refs = await comfy.await_images(
|
||||
config, await comfy.submit(config, workflow.fill(template, params))
|
||||
)
|
||||
if not refs:
|
||||
raise comfy.ComfyError("ComfyUI finished but saved no image.")
|
||||
payload = await comfy.fetch_image(config, refs[0])
|
||||
|
||||
if preserve:
|
||||
await comfy.free(config)
|
||||
|
||||
if reviewer is None:
|
||||
attempts.append(Attempt(number, params["seed"], kept=True))
|
||||
kept = (payload, params)
|
||||
break
|
||||
|
||||
endpoint, model_id = reviewer
|
||||
keep, reason = await _review(context, endpoint, model_id, prompt, payload)
|
||||
last = number == tries
|
||||
attempts.append(Attempt(number, params["seed"], kept=keep or last, verdict=reason))
|
||||
if keep or last:
|
||||
kept = (payload, params)
|
||||
break
|
||||
except comfy.ComfyError as exc:
|
||||
if preserve:
|
||||
await comfy.free(config)
|
||||
return ToolOutcome(
|
||||
f"The image could not be generated: {exc.message}",
|
||||
{**event, "status": "error", "error": exc.message},
|
||||
)
|
||||
|
||||
if preserve:
|
||||
await comfy.free(config)
|
||||
if kept is None: # pragma: no cover - the loop always keeps its last attempt
|
||||
return ToolOutcome(
|
||||
"Nothing was generated.", {**event, "status": "error", "error": "No image."}
|
||||
)
|
||||
|
||||
payload, params = kept
|
||||
try:
|
||||
with session_scope() as db:
|
||||
attachment = files_service.store(
|
||||
db,
|
||||
user_id=context.owner_id,
|
||||
chat_id=context.chat_id,
|
||||
payload=payload,
|
||||
filename=f"{template_slug}-{params['seed']}.png",
|
||||
# What ComfyUI made, at the size it made it. See `_keep_image`.
|
||||
keep_original=True,
|
||||
source_label="Image generation",
|
||||
source_path=f"{checkpoint} · seed {params['seed']}",
|
||||
)
|
||||
attachment_id = attachment.id
|
||||
width, height = attachment.width, attachment.height
|
||||
except Exception as exc: # noqa: BLE001
|
||||
log.exception("could not store a generated image")
|
||||
return ToolOutcome(
|
||||
f"The image was generated but could not be saved: {exc}",
|
||||
{**event, "status": "error", "error": str(exc)},
|
||||
)
|
||||
|
||||
return ToolOutcome(
|
||||
_describe(prompt, template_name, checkpoint, params, attempts),
|
||||
{
|
||||
**event,
|
||||
"status": "ok",
|
||||
"detail": f"{template_name} · {checkpoint}",
|
||||
"text": _transcript(params, attempts),
|
||||
# Bound to the reply by `generation._persist`, the single writer. A
|
||||
# runner may create the row; only the loop may say which turn owns
|
||||
# it.
|
||||
"attachment_id": attachment_id,
|
||||
"image": {"id": attachment_id, "width": width, "height": height},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _pick(rows: list[Any], wanted: str, chat_default: str, values: dict[str, Any]) -> Any:
|
||||
"""The workflow to use: asked for, then the chat's, then the instance's."""
|
||||
by_slug = {row.slug: row for row in rows}
|
||||
if wanted and wanted in by_slug:
|
||||
return by_slug[wanted]
|
||||
by_id = {row.id: row for row in rows}
|
||||
if chat_default and chat_default in by_id:
|
||||
return by_id[chat_default]
|
||||
fallback = str(values.get("default_workflow_id") or "")
|
||||
if fallback and fallback in by_id:
|
||||
return by_id[fallback]
|
||||
return rows[0] if rows else None
|
||||
|
||||
|
||||
def _checkpoint(wanted: str, chat_default: str, available: list[str]) -> str | None:
|
||||
"""The checkpoint to draw with, on the same ladder.
|
||||
|
||||
A name the instance does not have is ignored rather than passed through: it
|
||||
would reach ComfyUI, be refused, and cost a round to discover -- and the
|
||||
model was shown the list it may choose from.
|
||||
"""
|
||||
if wanted and wanted in available:
|
||||
return wanted
|
||||
if chat_default and chat_default in available:
|
||||
return chat_default
|
||||
return available[0] if available else None
|
||||
|
||||
|
||||
def _describe(
|
||||
prompt: str, template: str, checkpoint: str, params: dict[str, Any], attempts: list[Attempt]
|
||||
) -> str:
|
||||
"""What the model reads back.
|
||||
|
||||
It is told the image is already on screen, because otherwise the commonest
|
||||
next thing it does is offer to show it -- and there is nothing it could do
|
||||
to comply.
|
||||
"""
|
||||
lines = [
|
||||
"The image has been generated and is shown to them. It is not a link and "
|
||||
"needs no further action.",
|
||||
f"Prompt: {prompt}",
|
||||
f"Template {template}, checkpoint {checkpoint}, "
|
||||
f"{params['width']}x{params['height']}, seed {params['seed']}, "
|
||||
f"{params['steps']} steps, cfg {params['cfg']}.",
|
||||
]
|
||||
if len(attempts) > 1:
|
||||
rejected = [a for a in attempts if not a.kept]
|
||||
lines.append(
|
||||
f"It took {len(attempts)} attempts; the earlier ones were rejected on review "
|
||||
f"({'; '.join(a.verdict for a in rejected if a.verdict) or 'no reason given'})."
|
||||
)
|
||||
return "\n".join(lines)
|
||||
|
||||
|
||||
def _transcript(params: dict[str, Any], attempts: list[Attempt]) -> str:
|
||||
"""What the reader sees when they open the tool block.
|
||||
|
||||
The rejected attempts are recorded here and their images are not kept. A
|
||||
transcript full of pictures somebody's model decided against is noise, and
|
||||
the disk they would occupy buys nothing -- what is worth knowing is that it
|
||||
took three goes and why the first two did not do.
|
||||
"""
|
||||
lines = [
|
||||
f"seed {params['seed']} · {params['steps']} steps · cfg {params['cfg']} · "
|
||||
f"{params['sampler']}/{params['scheduler']} · denoise {params['denoise']}"
|
||||
]
|
||||
if len(attempts) > 1:
|
||||
lines.append("")
|
||||
for attempt in attempts:
|
||||
state = "kept" if attempt.kept else "rejected"
|
||||
reason = f" — {attempt.verdict}" if attempt.verdict else ""
|
||||
lines.append(f"Attempt {attempt.number} (seed {attempt.seed}): {state}{reason}")
|
||||
return "\n".join(lines)
|
||||
@@ -0,0 +1,178 @@
|
||||
"""Turning a stored template and a model's arguments into a ComfyUI workflow.
|
||||
|
||||
A template is an API-format workflow with `{{placeholders}}` where the values
|
||||
go. Which node holds the prompt is therefore the administrator's statement
|
||||
rather than something guessed from node types -- sniffing for the first
|
||||
`CLIPTextEncode` works on the shipped template and on nothing else, and gets
|
||||
positive and negative the wrong way round the first time somebody reorders them.
|
||||
|
||||
**Substitution walks the parsed JSON, not the text of it.** A value that is
|
||||
*exactly* `"{{steps}}"` is replaced by the number 20, not by the string "20";
|
||||
ComfyUI validates types and refuses the second. A placeholder inside a longer
|
||||
string still substitutes as text, which is what makes
|
||||
`"{{prompt}}, masterpiece"` work. Doing it textually would also mean a prompt
|
||||
containing a quotation mark produced a document that no longer parses, on the
|
||||
one input guaranteed to contain arbitrary text.
|
||||
|
||||
The names are the tool's parameter names, so there is one vocabulary: what a
|
||||
model may set, what the admin page documents and what a template may reference
|
||||
cannot drift apart.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import re
|
||||
import secrets
|
||||
from typing import Any
|
||||
|
||||
# Every hole a template may carry. A name outside this set is left alone, the
|
||||
# same rule `prompts.substitute` follows -- a literal `{{x}}` is not a feature,
|
||||
# but silently deleting one is worse than leaving it visible.
|
||||
PLACEHOLDERS = (
|
||||
"model",
|
||||
"prompt",
|
||||
"negative",
|
||||
"seed",
|
||||
"steps",
|
||||
"cfg",
|
||||
"width",
|
||||
"height",
|
||||
"sampler",
|
||||
"scheduler",
|
||||
"denoise",
|
||||
)
|
||||
|
||||
# The defaults, taken from the base template. `seed` is deliberately absent: it
|
||||
# has no fixed default, because one would make every generation that did not
|
||||
# name a seed identical -- and would make the retry loop produce the same
|
||||
# rejected image four times over.
|
||||
DEFAULTS: dict[str, Any] = {
|
||||
"negative": "text, watermark",
|
||||
"steps": 20,
|
||||
"cfg": 8.0,
|
||||
"width": 512,
|
||||
"height": 512,
|
||||
"sampler": "euler",
|
||||
"scheduler": "normal",
|
||||
"denoise": 1.0,
|
||||
}
|
||||
|
||||
# ComfyUI's own ranges, read off `/object_info`. Clamped rather than refused: a
|
||||
# model that asks for 300 steps has misjudged rather than misbehaved, and one
|
||||
# clarifying round to say so is worse than doing the sensible thing.
|
||||
LIMITS: dict[str, tuple[float, float]] = {
|
||||
"steps": (1, 150),
|
||||
"cfg": (0.0, 30.0),
|
||||
"width": (64, 2048),
|
||||
"height": (64, 2048),
|
||||
"denoise": (0.0, 1.0),
|
||||
}
|
||||
|
||||
# ComfyUI's seed is a uint64. Generated here rather than left to the far side
|
||||
# so the value can be reported back -- "it looked like this and here is how to
|
||||
# get it again" is most of what a seed is for.
|
||||
MAX_SEED = 2**64 - 1
|
||||
|
||||
_PLACEHOLDER = re.compile(r"\{\{\s*([a-z][a-z0-9_]*)\s*\}\}")
|
||||
|
||||
|
||||
def random_seed() -> int:
|
||||
return secrets.randbelow(MAX_SEED)
|
||||
|
||||
|
||||
def resolve(given: dict[str, Any]) -> dict[str, Any]:
|
||||
"""The full parameter set: what was asked for, over the defaults.
|
||||
|
||||
Absent and null are both "no opinion". A model that emits `"seed": null`
|
||||
rather than omitting the key is common enough that treating it as a request
|
||||
for seed zero would be a bug nobody could see.
|
||||
"""
|
||||
values: dict[str, Any] = {**DEFAULTS}
|
||||
for name, value in (given or {}).items():
|
||||
if name in PLACEHOLDERS and value is not None and value != "":
|
||||
values[name] = value
|
||||
|
||||
values["seed"] = _whole(values.get("seed"), default=random_seed()) % (MAX_SEED + 1)
|
||||
for name in ("steps", "width", "height"):
|
||||
values[name] = _clamp(_whole(values.get(name), DEFAULTS[name]), name)
|
||||
for name in ("cfg", "denoise"):
|
||||
values[name] = _clamp(_decimal(values.get(name), DEFAULTS[name]), name)
|
||||
for name in ("prompt", "negative", "sampler", "scheduler", "model"):
|
||||
values[name] = str(values.get(name) or "")
|
||||
return values
|
||||
|
||||
|
||||
def _whole(value: Any, default: int) -> int:
|
||||
try:
|
||||
return int(float(value))
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _decimal(value: Any, default: float) -> float:
|
||||
try:
|
||||
return float(value)
|
||||
except (TypeError, ValueError):
|
||||
return default
|
||||
|
||||
|
||||
def _clamp(value: Any, name: str) -> Any:
|
||||
low, high = LIMITS.get(name, (None, None))
|
||||
if low is None:
|
||||
return value
|
||||
clamped = min(max(value, low), high)
|
||||
return int(clamped) if isinstance(value, int) else clamped
|
||||
|
||||
|
||||
def fill(template: Any, values: dict[str, Any]) -> Any:
|
||||
"""A copy of the template with its placeholders replaced.
|
||||
|
||||
Recursive over dicts and lists, because a workflow is nested and a
|
||||
placeholder can be anywhere in it -- including inside a node's `_meta`,
|
||||
which is harmless and should not be treated specially.
|
||||
"""
|
||||
if isinstance(template, dict):
|
||||
return {key: fill(value, values) for key, value in template.items()}
|
||||
if isinstance(template, list):
|
||||
return [fill(item, values) for item in template]
|
||||
if isinstance(template, str):
|
||||
return _fill_string(template, values)
|
||||
return template
|
||||
|
||||
|
||||
def _fill_string(text: str, values: dict[str, Any]) -> Any:
|
||||
"""One string, which may *become* a number.
|
||||
|
||||
The whole-value case is what keeps types right: `"{{steps}}"` is the number
|
||||
and not a string that looks like one. Anything else is ordinary text
|
||||
substitution, so `"{{prompt}}, masterpiece"` reads as a sentence.
|
||||
"""
|
||||
whole = _PLACEHOLDER.fullmatch(text.strip())
|
||||
if whole is not None:
|
||||
return values.get(whole.group(1), text)
|
||||
|
||||
def swap(match: re.Match[str]) -> str:
|
||||
name = match.group(1)
|
||||
return str(values[name]) if name in values else match.group(0)
|
||||
|
||||
return _PLACEHOLDER.sub(swap, text)
|
||||
|
||||
|
||||
def placeholders_in(template: Any) -> set[str]:
|
||||
"""Every `{{name}}` a template uses, for the admin page to report.
|
||||
|
||||
A template that mentions none of them is almost certainly a workflow pasted
|
||||
straight out of ComfyUI without being parameterised, which would generate
|
||||
the same picture whatever anybody typed. Worth saying at save time rather
|
||||
than leaving somebody to discover it.
|
||||
"""
|
||||
found: set[str] = set()
|
||||
if isinstance(template, dict):
|
||||
for value in template.values():
|
||||
found |= placeholders_in(value)
|
||||
elif isinstance(template, list):
|
||||
for item in template:
|
||||
found |= placeholders_in(item)
|
||||
elif isinstance(template, str):
|
||||
found |= {match.group(1) for match in _PLACEHOLDER.finditer(template)}
|
||||
return found
|
||||
@@ -218,6 +218,24 @@ VARIABLES: tuple[Variable, ...] = (
|
||||
"Skill index",
|
||||
"Each available skill's name and when to use it, one per line.",
|
||||
),
|
||||
Variable(
|
||||
"image_templates",
|
||||
"Image templates",
|
||||
"Each enabled image workflow's name and what it is for, one per line. "
|
||||
"Empty when none has been set up.",
|
||||
),
|
||||
Variable(
|
||||
"image_models",
|
||||
"Image checkpoints",
|
||||
"The checkpoints an administrator has listed on the image generation "
|
||||
"page, comma separated.",
|
||||
),
|
||||
Variable(
|
||||
"image_instructions",
|
||||
"Image house rules",
|
||||
"Whatever an administrator wrote in the Extra instructions box on the "
|
||||
"image generation page. Empty when they wrote nothing.",
|
||||
),
|
||||
Variable(
|
||||
"knowledge_bases",
|
||||
"Knowledge bases",
|
||||
@@ -1001,6 +1019,65 @@ BUILTIN: tuple[Fragment, ...] = (
|
||||
"since that is all you will see next time."
|
||||
),
|
||||
),
|
||||
Fragment(
|
||||
key="tool.image",
|
||||
label="Generating an image",
|
||||
group=GROUP_TOOLS,
|
||||
order=243,
|
||||
families=("image",),
|
||||
hint="Appears when image generation is offered. The sentence about the "
|
||||
"picture already being on screen is the one that earns its place: "
|
||||
"without it the commonest thing a model does next is offer to show you "
|
||||
"the image, which it has no way of doing and which has already "
|
||||
"happened.",
|
||||
default=(
|
||||
"- You can draw a picture with image_generate. Describe what you want in "
|
||||
"the prompt as fully as you can — subject, setting, lighting, style — "
|
||||
"because the prompt is the whole of what the picture is made from.\n"
|
||||
"- The picture appears in the conversation as soon as the tool returns. "
|
||||
"It is already on screen: do not offer to show it, link to it or "
|
||||
"describe how to open it.\n"
|
||||
"- Only the prompt is required. Everything else has a sensible default, "
|
||||
"so set a parameter when you have a reason to and leave it out "
|
||||
"otherwise. Repeat a seed to get the same picture again."
|
||||
),
|
||||
),
|
||||
Fragment(
|
||||
key="tool.image_choices",
|
||||
label="Image models and templates",
|
||||
group=GROUP_TOOLS,
|
||||
order=244,
|
||||
families=("image",),
|
||||
requires=("image_templates",),
|
||||
variables=("image_templates", "image_models"),
|
||||
hint="Only once there is at least one workflow to choose between. Split "
|
||||
"from the fragment above for the reason `tool.skills` is split from "
|
||||
"`tool.skills_write`: an instance with one template should not be told "
|
||||
"to weigh up its options, and a list that is not there is worse than no "
|
||||
"sentence about it.",
|
||||
default=(
|
||||
"- The templates you can draw with, and what each is for:\n"
|
||||
"{{image_templates}}\n"
|
||||
"- The checkpoints you can name: {{image_models}}\n"
|
||||
"- Choose the template and checkpoint that suit what is being asked for. "
|
||||
"If none obviously fits, leave both out and the usual ones are used."
|
||||
),
|
||||
),
|
||||
Fragment(
|
||||
key="tool.image_instructions",
|
||||
label="Image generation: house rules",
|
||||
group=GROUP_TOOLS,
|
||||
order=245,
|
||||
families=("image",),
|
||||
requires=("image_instructions",),
|
||||
variables=("image_instructions",),
|
||||
hint="Whatever an administrator wrote in the Extra instructions box on "
|
||||
"the image generation page. Absent entirely when that box is empty, "
|
||||
"which is why this is a fragment of its own rather than a paragraph in "
|
||||
"the one above -- an empty heading saying nothing is worse than no "
|
||||
"heading.",
|
||||
default="{{image_instructions}}",
|
||||
),
|
||||
Fragment(
|
||||
key="tool.scratch",
|
||||
label="The scratch document",
|
||||
@@ -1325,6 +1402,34 @@ BUILTIN: tuple[Fragment, ...] = (
|
||||
"Assistant: {{answer}}"
|
||||
),
|
||||
),
|
||||
Fragment(
|
||||
key="task.image_review",
|
||||
label="Reviewing a generated image",
|
||||
group=GROUP_TASKS,
|
||||
order=405,
|
||||
hint="A separate one-message request carrying the picture that was just "
|
||||
"made, asked of a vision model before the reader is shown anything. "
|
||||
"Clear it to stop reviewing: the first image is then kept, which is "
|
||||
"what happens anyway when nothing on the instance has vision. The bias "
|
||||
"towards KEEP is deliberate — a reviewer that retries on taste rather "
|
||||
"than on faults spends somebody's GPU four times over and usually ends "
|
||||
"up back at the first image.",
|
||||
default=(
|
||||
"You are checking a picture that was just generated against the request "
|
||||
"it was generated from. Judge only whether it is a competent attempt at "
|
||||
"what was asked for.\n"
|
||||
"\n"
|
||||
"Answer on the first line with one word: KEEP or RETRY. If RETRY, put "
|
||||
"one short sentence on the second line saying what is wrong.\n"
|
||||
"\n"
|
||||
"Say RETRY only for something clearly wrong: the subject that was asked "
|
||||
"for is missing, the image is mangled or unreadable, or it shows "
|
||||
"something quite different from the request. Say KEEP for anything that "
|
||||
"answers the request, including work you would have composed "
|
||||
"differently. Taste is not a fault, and there is no guarantee the next "
|
||||
"attempt will be better."
|
||||
),
|
||||
),
|
||||
Fragment(
|
||||
key="task.compact",
|
||||
label="Compaction summary",
|
||||
|
||||
@@ -29,6 +29,7 @@ DEFAULT_CHAT_ROUNDS = 0
|
||||
SEARCH = "search"
|
||||
PROMPTS = "prompts"
|
||||
AGENTS = "agents"
|
||||
IMAGES = "images"
|
||||
|
||||
|
||||
def _general_defaults() -> dict[str, Any]:
|
||||
@@ -232,12 +233,52 @@ def _prompts_defaults() -> dict[str, Any]:
|
||||
}
|
||||
|
||||
|
||||
def _images_defaults() -> dict[str, Any]:
|
||||
"""Generating pictures on a ComfyUI somebody else is running."""
|
||||
return {
|
||||
"enabled": False,
|
||||
"base_url": "",
|
||||
"api_key_encrypted": "",
|
||||
# A generation is tens of seconds and a queue in front of it can be
|
||||
# minutes. Far longer than any other timeout here, because the thing
|
||||
# being waited for genuinely takes that long.
|
||||
"timeout": 600.0,
|
||||
# What this ComfyUI advertises, discovered by the Test button and stored
|
||||
# so the request path never has to ask. The checkpoints are also the
|
||||
# enum a model chooses from, which is why an empty list means the tool
|
||||
# is not offered: a model naming a checkpoint that does not exist gets a
|
||||
# refusal from ComfyUI and spends a round finding out.
|
||||
"checkpoints": [],
|
||||
"samplers": [],
|
||||
"schedulers": [],
|
||||
"default_workflow_id": "",
|
||||
# Whether a vision model looks at what came back and says whether to
|
||||
# keep it. Deliberately independent of `preserve_vram` below: on a
|
||||
# machine that can hold both models this costs nothing, and on one that
|
||||
# cannot it costs two model loads per retry, which is a judgement only
|
||||
# the person running it can make.
|
||||
"review_enabled": False,
|
||||
# Which model judges. Empty means the chat's own, when it has vision.
|
||||
"review_model_id": "",
|
||||
"max_tries": 4,
|
||||
# Unload the chat's own LLM while ComfyUI works, and free ComfyUI
|
||||
# afterwards. For a machine that cannot hold both at once. Off by
|
||||
# default: it makes every generation slower, and most people have the
|
||||
# memory.
|
||||
"preserve_vram": False,
|
||||
# Instance-wide guidance, injected into the harness beside the tool's
|
||||
# own. Where "always add these words to the negative prompt" lives.
|
||||
"instructions": "",
|
||||
}
|
||||
|
||||
|
||||
_DEFAULTS: dict[str, Any] = {
|
||||
GENERAL: _general_defaults,
|
||||
AUDIO: _audio_defaults,
|
||||
SEARCH: _search_defaults,
|
||||
PROMPTS: _prompts_defaults,
|
||||
AGENTS: _agents_defaults,
|
||||
IMAGES: _images_defaults,
|
||||
}
|
||||
|
||||
|
||||
@@ -350,3 +391,34 @@ def agents(db: DBSession) -> dict[str, Any]:
|
||||
)
|
||||
values["background_max_jobs"] = min(max(int(values.get("background_max_jobs") or 0), 1), 100)
|
||||
return values
|
||||
|
||||
|
||||
def images(db: DBSession) -> dict[str, Any]:
|
||||
"""Image generation settings, with the numbers clamped.
|
||||
|
||||
Clamped on read rather than on save, for the reason `agents` gives: a value
|
||||
stored by an earlier version cannot bite either. `max_tries` has a floor of
|
||||
one because zero would mean the tool generates nothing at all and reports
|
||||
success -- there is no reading of "no tries" that anybody wants, unlike the
|
||||
zeroes above, which each mean something.
|
||||
"""
|
||||
values = get_group(db, IMAGES)
|
||||
values["timeout"] = min(max(float(values.get("timeout") or 0), 10.0), 3600.0)
|
||||
values["max_tries"] = min(max(int(values.get("max_tries") or 1), 1), 10)
|
||||
for name in ("checkpoints", "samplers", "schedulers"):
|
||||
stored = values.get(name)
|
||||
values[name] = [str(item) for item in stored] if isinstance(stored, list) else []
|
||||
return values
|
||||
|
||||
|
||||
def images_ready(db: DBSession) -> bool:
|
||||
"""Whether image generation can actually happen.
|
||||
|
||||
Three things, and the checkpoint list is the one worth stating: without it a
|
||||
model has nothing to name, and ComfyUI refuses a workflow whose checkpoint
|
||||
does not exist -- so offering the tool would be offering a round that ends
|
||||
in a refusal. Read by the tool gate, which is why it lives here beside the
|
||||
values rather than in `tools.py` with the other gates.
|
||||
"""
|
||||
values = images(db)
|
||||
return bool(values["enabled"] and values["base_url"] and values["checkpoints"])
|
||||
|
||||
@@ -62,6 +62,7 @@ LABELS: dict[str, str] = {
|
||||
"notes_edit": "Note updated",
|
||||
"notes_delete": "Note deleted",
|
||||
"scratch_write": "Canvas written",
|
||||
"image_generate": "Image",
|
||||
"memory_add": "Memory saved",
|
||||
"memory_forget": "Memory removed",
|
||||
"skill_get": "Skill read",
|
||||
@@ -94,6 +95,7 @@ ICONS: dict[str, str] = {
|
||||
"notes_edit": "pencil",
|
||||
"notes_delete": "trash",
|
||||
"scratch_write": "file-text",
|
||||
"image_generate": "image",
|
||||
"memory_add": "star",
|
||||
"memory_forget": "trash",
|
||||
"skill_get": "sparkle",
|
||||
@@ -109,6 +111,7 @@ KIND_ICONS: dict[str, str] = {
|
||||
"fetch": "link",
|
||||
"custom": "link",
|
||||
"mcp": "server",
|
||||
"image": "image",
|
||||
}
|
||||
FALLBACK_ICON = "sparkle"
|
||||
|
||||
@@ -130,6 +133,7 @@ ACTIONS: dict[str, str] = {
|
||||
"notes_edit": "Change a note",
|
||||
"notes_delete": "Delete a note",
|
||||
"scratch_write": "Write in the canvas",
|
||||
"image_generate": "Generate an image",
|
||||
"memory_add": "Remember something",
|
||||
"memory_forget": "Forget something",
|
||||
"skill_get": "Read a skill",
|
||||
@@ -152,6 +156,10 @@ DETAIL_KEYS: dict[str, str] = {
|
||||
"knowledge_search": "query",
|
||||
"notes_search": "query",
|
||||
"job_stop": "id",
|
||||
# The thing being agreed to is what will be drawn, not which sampler draws
|
||||
# it. Also what makes the box on the card editable: a prompt corrected
|
||||
# before it runs is the commonest useful edit this feature will see.
|
||||
"image_generate": "prompt",
|
||||
}
|
||||
|
||||
|
||||
|
||||
@@ -113,6 +113,12 @@ FAMILY_SCRATCH = "scratch"
|
||||
# services/agent/session.py:resolve, which answers all three at once.
|
||||
FAMILY_AGENT = "agent"
|
||||
|
||||
# Drawing a picture on a ComfyUI an administrator configured. A family of its
|
||||
# own for the reason `fetch` is one: an instance may reasonably want a model
|
||||
# that can look things up but not spend a minute of GPU on every request, and
|
||||
# the whole cost of this one is somewhere else.
|
||||
FAMILY_IMAGE = "image"
|
||||
|
||||
# The built-in families, in the order they are offered.
|
||||
FAMILIES = (
|
||||
FAMILY_SEARCH,
|
||||
@@ -123,6 +129,7 @@ FAMILIES = (
|
||||
FAMILY_SKILLS,
|
||||
FAMILY_SCRATCH,
|
||||
FAMILY_ASK,
|
||||
FAMILY_IMAGE,
|
||||
FAMILY_AGENT,
|
||||
)
|
||||
|
||||
@@ -164,6 +171,11 @@ class ToolContext:
|
||||
"""
|
||||
|
||||
owner_id: str
|
||||
# Which conversation this call belongs to. Needed by anything that writes
|
||||
# something the chat owns rather than something the *reader* owns -- the
|
||||
# scratch document, a generated image -- and empty for a call with no chat
|
||||
# behind it, which is what those runners check first.
|
||||
chat_id: str = ""
|
||||
search_config: dict[str, Any] = field(default_factory=dict)
|
||||
# Which knowledge bases this chat is scoped to. Empty means "everything the
|
||||
# owner can see", which is what a chat with none attached should do.
|
||||
@@ -186,6 +198,17 @@ class ToolContext:
|
||||
# a model can name a skill it was never shown and the runner would fetch it
|
||||
# anyway. Same rule as "what may be run is what was offered".
|
||||
skills_off: frozenset[str] = field(default_factory=frozenset)
|
||||
# Image generation, snapshotted like everything else here. `image_config` is
|
||||
# the instance settings group; the two below are this chat's preferences,
|
||||
# used when the model names neither. `model_id` and `connection_id` are what
|
||||
# the reviewer and the Preserve VRAM unload need to find the chat's own
|
||||
# endpoint -- its own, and no other, because the VRAM being freed belongs to
|
||||
# one machine.
|
||||
image_config: dict[str, Any] = field(default_factory=dict)
|
||||
image_workflow_id: str = ""
|
||||
image_checkpoint: str = ""
|
||||
model_id: str = ""
|
||||
connection_id: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
@@ -1097,7 +1120,7 @@ REGISTRY: dict[str, ToolDef] = {
|
||||
|
||||
|
||||
def _family_allowed(
|
||||
family: str, *, config: dict, capabilities: dict, allowed: dict
|
||||
family: str, *, config: dict, capabilities: dict, allowed: dict, images: bool = False
|
||||
) -> bool:
|
||||
"""Whether one family is on for this chat.
|
||||
|
||||
@@ -1122,6 +1145,14 @@ def _family_allowed(
|
||||
# attach path keeps working, because that one is a person's instruction
|
||||
# rather than a model's choice.
|
||||
return bool(allowed.get("tools.fetch") and config.get("fetch_enabled"))
|
||||
if gate == FAMILY_IMAGE:
|
||||
# Its own branch rather than a name in the tuple below, and the second
|
||||
# half is why: an instance with no ComfyUI, or one with no checkpoints
|
||||
# listed, must not offer this at all. A model that calls it there spends
|
||||
# a round to be told the thing it was offered does not work, which is
|
||||
# the shape `resolve_tools` already refuses for `skill_get` with an
|
||||
# empty library. `settings_store.images_ready` answers all three.
|
||||
return bool(allowed.get("tools.image") and images)
|
||||
if gate in (FAMILY_CUSTOM, FAMILY_MCP, FAMILY_ASK, FAMILY_AGENT, FAMILY_SCRATCH):
|
||||
# Deliberately without `library.use`: an HTTP endpoint an administrator
|
||||
# wrote has nothing to do with this person's own documents and notes,
|
||||
@@ -1167,6 +1198,20 @@ def _agent_defs(db: DBSession, chat: Chat | None, user: User | None) -> list[Too
|
||||
return agent_tools.tool_defs(context)
|
||||
|
||||
|
||||
def _image_defs(db: DBSession, values: dict | None = None) -> list[ToolDef]:
|
||||
"""The image tool, whose schema carries this instance's own choices.
|
||||
|
||||
Built per request rather than at import, because the templates a model may
|
||||
name and the checkpoints it may draw with are rows and settings. That is the
|
||||
same reason a custom tool cannot live in `REGISTRY`, and it is why this has
|
||||
to be listed in `registry(db)` below as well -- a name that resolves to no
|
||||
family is a tool whose guidance never reaches the model.
|
||||
"""
|
||||
from lembas.services.images import tool as image_tool
|
||||
|
||||
return [image_tool.tool_def(db, values if values is not None else settings_store.images(db))]
|
||||
|
||||
|
||||
def _book(defs: list[ToolDef]) -> dict[str, ToolDef]:
|
||||
"""Keyed by name, first claim winning.
|
||||
|
||||
@@ -1196,7 +1241,9 @@ def registry(db: DBSession) -> dict[str, ToolDef]:
|
||||
"""
|
||||
from lembas.services.agent import tools as agent_tools
|
||||
|
||||
return _book([*_row_defs(db, None, everything=True), *agent_tools.tool_defs()])
|
||||
return _book(
|
||||
[*_row_defs(db, None, everything=True), *agent_tools.tool_defs(), *_image_defs(db)]
|
||||
)
|
||||
|
||||
|
||||
def families(db: DBSession) -> tuple[str, ...]:
|
||||
@@ -1220,10 +1267,21 @@ def resolve_tools(db: DBSession, chat: Chat, user: User | None) -> ToolSet:
|
||||
|
||||
allowed = permissions.resolve(db, user)
|
||||
config = settings_store.search(db)
|
||||
image_values = settings_store.images(db)
|
||||
images_ready = settings_store.images_ready(db)
|
||||
|
||||
# Resolved against what this reader may see, not against everything that
|
||||
# exists: a tool restricted to a group is not offered outside it.
|
||||
book = _book([*_row_defs(db, user), *_agent_defs(db, chat, user)])
|
||||
# exists: a tool restricted to a group is not offered outside it. The image
|
||||
# tool is built only when it could be offered, because building its schema
|
||||
# reads the workflow table and there is no sense doing that for an instance
|
||||
# with no ComfyUI.
|
||||
book = _book(
|
||||
[
|
||||
*_row_defs(db, user),
|
||||
*_agent_defs(db, chat, user),
|
||||
*(_image_defs(db, image_values) if images_ready else []),
|
||||
]
|
||||
)
|
||||
|
||||
# What this chat has switched off, applied AFTER the gates and never
|
||||
# instead of them. A chat can only ever *narrow* what the model's
|
||||
@@ -1240,7 +1298,11 @@ def resolve_tools(db: DBSession, chat: Chat, user: User | None) -> ToolSet:
|
||||
tool
|
||||
for tool in book.values()
|
||||
if _family_allowed(
|
||||
tool.family, config=config, capabilities=capabilities, allowed=allowed
|
||||
tool.family,
|
||||
config=config,
|
||||
capabilities=capabilities,
|
||||
allowed=allowed,
|
||||
images=images_ready,
|
||||
)
|
||||
and gate_of(tool.family) not in off
|
||||
# Nothing to read and nothing to improve. Offering `skill_get` with
|
||||
@@ -1327,7 +1389,13 @@ def context_for(
|
||||
return ToolContext(
|
||||
agent=agent_session.resolve(db, chat, user) if chat is not None else None,
|
||||
owner_id=user.id if user else "",
|
||||
chat_id=chat.id if chat is not None else "",
|
||||
search_config=settings_store.search(db),
|
||||
image_config=settings_store.images(db),
|
||||
image_workflow_id=(chat.image_workflow_id or "") if chat is not None else "",
|
||||
image_checkpoint=(chat.image_checkpoint or "") if chat is not None else "",
|
||||
model_id=(chat.model_id or "") if chat is not None else "",
|
||||
connection_id=(chat.connection_id or "") if chat is not None else "",
|
||||
base_ids=[base.id for base in chat.knowledge_bases] if chat is not None else [],
|
||||
skills_off=scoped_skills_off(chat),
|
||||
tools=tools.by_name if tools is not None else None,
|
||||
|
||||
Reference in New Issue
Block a user