Houmeres/LLeMbas
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 3 Wiki Activity

Draw a picture, on a ComfyUI you are running

Browse Source 
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š
2026-08-05 14:13:19 +02:00
parent f4bf1bf670
commit 5e75948069
40 changed files with 4081 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CLAUDE.md
+97 -9
View File
@@ -20,7 +20,7 @@ lembas info # paths + counts, useful when confused
lembas secret-key # generate LEMBAS_SECRET_KEY
lembas create-admin # create or promote an admin
pytest # 1529 tests, ~93s
pytest # 1676 tests, ~102s
# PLAN.md tracks what is and is not built
ruff check . # lint (line length 100)
python scripts/build_artwork.py # regenerate artwork (SVG + PWA icons;
@@ -85,6 +85,7 @@ src/lembas/
admin_users.py users, groups, permissions
admin_audio.py speech-to-text and text-to-speech endpoints
admin_search.py web search provider and credentials
admin_images.py the ComfyUI, and the workflow templates on it
admin_prompts.py the prompt fragment editor and its preview
admin_suggestions.py the cards offered on the new-chat screen
admin_tools.py custom HTTP tools and MCP servers
@@ -104,6 +105,9 @@ src/lembas/
services/
llm/openai_client.py httpx streaming + model discovery
search/ ddgs, SearXNG and Firecrawl behind one shape
images/ drawing on a ComfyUI: comfy.py speaks HTTP,
workflow.py fills a template, tool.py ties them
to a chat and decides whether to keep the result
library/ documents, notes, memories, skills, FTS
mcp/ remote MCP servers: framing, transport, rows to tools
agent/ agent chats: the mode table, SSH, the six tools,
@@ -1938,6 +1942,91 @@ have to have been watching to understand. It is the one thing in a tool result
that is genuinely *not* untrusted: it is the reader's own words, so it is stated
as theirs and needs no fence.
**Image generation is a ComfyUI workflow with holes in it, and the holes are the
administrator's statement.** `services/images/` is three modules: `comfy.py`
speaks HTTP, `workflow.py` fills a template, `tool.py` ties them to a chat.
Which node holds the prompt is *declared* with `{{prompt}}` rather than sniffed
by node type — looking for the first `CLIPTextEncode` works on the shipped
workflow and on nothing else, and swaps positive for negative the first time
somebody reorders them.
**Substitution walks the parsed JSON, not 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. 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 hold arbitrary text. `seed` has no fixed default
— one would make every unspecified generation identical and make the retry loop
redraw the same rejected picture four times.
**One call is one finished image, and the retrying is inside the tool.**
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 — the admin's chosen vision model, else the chat's own if it has
vision, else nobody — is asked about *bytes* rather than about a row: 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. Rejected images are not stored — their verdicts are, in `event.text`.
**`task.image_review` is a `GROUP_TASKS` fragment**, so it is editable and
excluded from the harness, exactly like `task.title` and `task.compact` — and
clearing it switches reviewing off, the same way clearing `task.compact` switches
compaction off. It is biased hard towards KEEP on purpose: a reviewer that
retries on taste spends the GPU four times and usually ends up back at the first
image.
**Preserve VRAM unloads the chat's own connection and nothing else.**
`Connection.unload_url` is a column because the memory being freed belongs to one
machine: a local llama-swap answers `GET /unload`, and a box on the network has
no reason to be unloaded when ComfyUI wants memory *here*. Empty means "cannot be
unloaded", which is the honest default — there is no call that works everywhere.
The swap goes round the *review*, not round the tool: unload, generate, free
ComfyUI, ask the reviewer (which loads the LLM again), round again if it said no.
Two model loads per retry, which is why the two settings are independent and the
page says so when both are on. **Nothing loads the LLM back at the end** — the
reply's next request does, and llama-swap loads on demand; that step exists in
the description and not in the code, which is why the code says so.
**A generated image rides on the assistant message, so `message_payload` sends
images only on `user` turns.** No assistant message had ever carried one before,
so the distinction had never been drawn — and the moment one does, the
multimodal list form on an `assistant` turn is rejected by OpenAI and most local
runners, breaking not that turn but every later one in the chat. What follows and
is worth knowing: on a *later* turn the model cannot see the picture it made
(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.
**The runner writes the file; only the loop says which turn owns it.**
`event["attachment_id"]` is carried by `generation._run` exactly as
`event["canvas"]` and `event["plan"]` are, because `_persist` is the single
writer. `_bind_attachments` narrows on this chat and on rows still unbound, for
the reason `files.claim` does: the ids arrive on a dict a runner built.
**`files.store(keep_original=True)` skips the resize and the transcode, and
nothing else.** `_process_image` turns anything without alpha into JPEG q85 at
1400px, which is right for a phone photo and a visible loss on generated art.
Pillow still opens it, so a malformed file is still refused and the dimensions
are still measured rather than claimed.
**`/image` forces one tool for one round.** It sends the ordinary message with
`force_tool`, which becomes `tool_choice` — reusing the whole loop rather than
inventing a second generation path. `FORCEABLE_TOOLS` is an allow list because
this is read off a form, and `resolve_tools` still decides whether the tool
exists, so forcing one that was never offered does nothing. `payload.pop(
"tool_choice")` after the first round is load-bearing: left in place the reply
would draw a picture, be asked again, and draw another.
**`ToolContext` gained `chat_id`, and that fixed a tool nobody had ever run.**
`_run_scratch_write` read `context.chat_id` on a dataclass that had no such
field, so **every `scratch_write` call raised `AttributeError`** — swallowed by
`run_tool`'s blanket except into "the scratch_write tool failed", which reads
exactly like 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. There is one that runs it now.
**A control wired to a method its route does not serve fails silently.** The
agent-mode select posted with `hx-post` against a route that only answers
`PATCH`, so every change returned 405 and the mode never moved — for the whole
@@ -2058,14 +2147,13 @@ notes describe the machine.
## Not built yet
Image generation, and a nav entry marks where it goes. The tool loop in
`services/generation.py` is what a new capability plugs into — a tool is a
`ToolDef` reaching `tools.resolve_tools()` plus a permission and a capability
flag, not a new code path. Its guidance is the same shape: a
`prompts.register_source` yielding one `Fragment` per row puts it in the harness,
on the admin page and in the preview without touching the assembler, the save
handler or a template. Custom HTTP tools, MCP servers and agent chats are the
three worked examples.
Nothing large. The tool loop in `services/generation.py` is what a new
capability plugs into — a tool is a `ToolDef` reaching `tools.resolve_tools()`
plus a permission and a capability flag, not a new code path. Its guidance is
the same shape: a `prompts.register_source` yielding one `Fragment` per row puts
it in the harness, on the admin page and in the preview without touching the
assembler, the save handler or a template. Custom HTTP tools, MCP servers, agent
chats and image generation are the four worked examples.
**Nothing executes on this machine, and that is the design.** Agent chats run
their commands on a host reached over SSH. A local sandbox was designed in
PLAN.md
+26 -6
View File
@@ -7,8 +7,8 @@ that would be expensive to revisit. Kept current as work lands; the detail of
**Status:** usable daily. Streaming chat, attachments, reasoning, tool calling
with web search, custom HTTP tools and MCP servers, agent chats that work on a
machine over SSH, a knowledge library, notes, memory and skills, speech in and
out, users and groups, model administration, installable as an app. 1529 tests,
`ruff` clean.
out, image generation over ComfyUI, users and groups, model administration,
installable as an app. 1676 tests, `ruff` clean.
---
@@ -115,6 +115,30 @@ be a different project, not a refactor.
- [x] Local MCP over stdio is deliberately absent: spawning a subprocess would
run on this machine, which nothing here does
### Image generation
- [x] **Draws on a ComfyUI you are running**, as a tool the model chooses to
call and as an `/image` command that makes it call one. Never on this
machine, the same rule agent chats follow
- [x] **Multiple workflow templates** — a name, a description and a ComfyUI API
export with `{{prompt}}` and ten other placeholders where the values go.
The model picks between them by their descriptions, and by checkpoint,
falling back to the chat's usual and then the instance default when it
names neither
- [x] Model may set prompt, negative, seed, steps, cfg, width, height, sampler,
scheduler, denoise, checkpoint and template; **only the prompt is
required** and everything else has a default
- [x] **The result is checked before you see it** — optionally, a vision model
is shown the picture and the request and says keep or retry, up to a
configurable number of attempts. Only clearly wrong images are retried;
the last attempt is kept whatever it says, so a request always produces
something
- [x] **Preserve VRAM** — opt-in, for a machine that cannot hold both at once:
unload the chat's own language model, generate, free ComfyUI, and let the
next request load the model back. Per connection, so a box on the network
is never touched
- [x] Instance-wide extra instructions, injected into the harness beside the
tool's own guidance
### Agent chats
- [x] A chat is a **Chat** or an **Agent**, chosen when it starts and fixed
thereafter — a transcript whose earlier turns ran somewhere else is not
@@ -290,10 +314,6 @@ be a different project, not a refactor.
In the order they are likely to be worth doing.
### Image generation
Left until last from the start, as it needs heavy customisation. ComfyUI is
already running on this machine and is the obvious first target.
### Smaller things
- **OCR** for scanned PDFs
- **Conversation branching** — `Message.parent_id` exists unused; needs a UI for
src/lembas/__init__.py
+1 -1
View File
@@ -1,3 +1,3 @@
"""LLeMbas - a Middle-earth themed web UI for OpenAI-compatible LLM endpoints."""
__version__ = "0.6.5"
__version__ = "0.7.0"
src/lembas/api/admin.py
+8
View File
@@ -145,11 +145,19 @@ async def update_connection(
base_url: str = Form(...),
api_key: str = Form(""),
enabled: bool = Form(False),
unload_url: str = Form(""),
unload_method: str = Form("POST"),
) -> Response:
connection = _connection(db, connection_id)
connection.name = name.strip()[:120] or connection.name
connection.base_url = base_url.strip().rstrip("/")
connection.enabled = enabled
# How to ask this endpoint to drop its model, for image generation's
# Preserve VRAM. Empty means it cannot be unloaded, which is the honest
# answer for anything not running on the machine ComfyUI is on.
connection.unload_url = unload_url.strip()[:500]
method = unload_method.strip().upper()
connection.unload_method = method if method in ("GET", "POST") else "POST"
submitted = api_key.strip()
if submitted and submitted != UNCHANGED_SENTINEL:
src/lembas/api/admin_images.py
+381
View File
@@ -0,0 +1,381 @@
"""Image generation administration: the ComfyUI, and the workflows to run on it.
Two shapes on one nav entry, because they are two different kinds of thing. The
connection, the checkpoints and the switches are instance settings and get a
settings page. A workflow is an authored document with a name, a description and
a body, so the workflows are list-plus-detail -- the shape `CLAUDE.md` requires
of any admin list, and for the reason it gives: a page that renders a ten-line
JSON textarea per row is unusable at three rows.
Route order matters and is not alphabetical. `/admin/images/workflows/new` is
registered before `/admin/images/workflows/{workflow_id}`, or "new" is captured
as an id and 404s. That has already been a bug twice here.
"""
from __future__ import annotations
import json
import logging
import re
from datetime import UTC, datetime
from fastapi import APIRouter, Form, HTTPException, Request, Response, status
from fastapi.responses import RedirectResponse
from sqlalchemy import func, select
from lembas.api.deps import AdminUser, Db
from lembas.db.models import ImageWorkflow, Model
from lembas.services import settings_store
from lembas.services.crypto import UNCHANGED_SENTINEL, decrypt, keep_or_replace, mask
from lembas.services.images import comfy
from lembas.services.images import workflow as workflow_service
from lembas.services.llm.openai_client import LLMError
from lembas.web.templating import render
log = logging.getLogger(__name__)
router = APIRouter(prefix="/admin/images", tags=["admin-images"])
SLUG_PATTERN = re.compile(r"^[a-z0-9][a-z0-9_-]{0,47}$")
# The placeholders a workflow has to carry to be worth having. Without a prompt
# it draws the same picture whatever anybody types, which is the one failure
# somebody would not think to look for.
REQUIRED_PLACEHOLDERS = ("prompt",)
def _lines(text: str) -> list[str]:
"""One name per line, blanks dropped. The `admin_agents` pattern."""
seen: list[str] = []
for line in (text or "").splitlines():
name = line.strip()
if name and name not in seen:
seen.append(name)
return seen
def _config(db: Db) -> comfy.Config:
values = settings_store.images(db)
return comfy.Config(
base_url=str(values.get("base_url") or ""),
api_key=decrypt(str(values.get("api_key_encrypted") or "")),
timeout=30.0,
)
def _page(request: Request, db: Db, **extra) -> Response:
values = settings_store.images(db)
workflows = list(
db.scalars(select(ImageWorkflow).order_by(ImageWorkflow.position, ImageWorkflow.slug))
)
return render(
request,
"admin/images.html",
{
"values": values,
"workflows": workflows,
# Only models an administrator has marked as having vision can
# review, so the picker offers those and nothing else -- a list
# including text-only models would be a list of choices that
# silently do nothing.
"vision_models": list(
db.scalars(
select(Model)
.where(Model.enabled.is_(True))
.order_by(Model.position, Model.model_id)
)
),
"checkpoints_text": "\n".join(values.get("checkpoints") or []),
"masked": mask(decrypt(values.get("api_key_encrypted") or "")),
"unchanged": UNCHANGED_SENTINEL,
**extra,
},
)
@router.get("")
async def images_page(request: Request, db: Db, user: AdminUser, saved: str = "") -> Response:
return _page(request, db, saved=saved)
@router.post("")
async def save_images(
request: Request,
db: Db,
user: AdminUser,
enabled: bool = Form(False),
base_url: str = Form(""),
api_key: str = Form(""),
timeout: float = Form(600.0),
checkpoints: str = Form(""),
default_workflow_id: str = Form(""),
review_enabled: bool = Form(False),
review_model_id: str = Form(""),
max_tries: int = Form(4),
preserve_vram: bool = Form(False),
instructions: str = Form(""),
) -> Response:
"""Save the settings.
Every toggle defaults to False because an unticked checkbox is simply absent
from a form post -- that absence *is* the off signal, the rule
`admin_audio` states.
The discovered sampler and scheduler lists are deliberately not submitted
and not cleared here: they belong to whatever ComfyUI was tested, and a save
that only changed the instructions box has no opinion about them.
"""
current = settings_store.images(db)
settings_store.update(
db,
{
"enabled": enabled,
"base_url": base_url.strip().rstrip("/"),
"api_key_encrypted": keep_or_replace(
api_key, current.get("api_key_encrypted") or ""
),
"timeout": min(max(timeout, 10.0), 3600.0),
"checkpoints": _lines(checkpoints),
"default_workflow_id": default_workflow_id.strip(),
"review_enabled": review_enabled,
"review_model_id": review_model_id.strip(),
"max_tries": min(max(max_tries, 1), 10),
"preserve_vram": preserve_vram,
"instructions": instructions.strip()[:4000],
},
key=settings_store.IMAGES,
)
log.info("image generation %s by %s", "enabled" if enabled else "disabled", user.email)
return RedirectResponse(
"/admin/images?saved=Saved.", status_code=status.HTTP_303_SEE_OTHER
)
@router.post("/test")
async def test_images(request: Request, db: Db, user: AdminUser) -> Response:
"""Ask ComfyUI what it can do, and remember the answer.
Against the *saved* settings rather than the unsaved form, so what is tested
is what a chat would actually reach -- the same rule `/admin/search/test`
follows.
The lists are stored rather than only shown, because the request path may
never ask ComfyUI anything: `harness.context_variables` is synchronous and
the tool schema is built per request, so both read what this button wrote.
"""
config = _config(db)
if not config.configured:
return render(
request,
"admin/_images_result.html",
{"message": "Set a base URL first.", "message_kind": "error"},
)
try:
checkpoints, samplers, schedulers = await comfy.discover(config)
except LLMError as exc:
return render(
request,
"admin/_images_result.html",
{"message": exc.message, "message_kind": "error"},
)
stored = settings_store.images(db)
changes: dict = {"samplers": samplers, "schedulers": schedulers}
# The checkpoint list is filled in only when nobody has one yet, for the
# reason a refreshed connection does not overwrite a context length an
# administrator typed: they are usually narrowing it deliberately.
if not stored.get("checkpoints"):
changes["checkpoints"] = checkpoints
settings_store.update(db, changes, key=settings_store.IMAGES)
found = (
f"Found {len(checkpoints)} checkpoint{'' if len(checkpoints) == 1 else 's'}, "
f"{len(samplers)} samplers and {len(schedulers)} schedulers."
)
return render(
request,
"admin/_images_result.html",
{
"message": found,
"message_kind": "success",
"checkpoints": checkpoints,
"kept": bool(stored.get("checkpoints")),
},
)
# --- Workflows -----------------------------------------------------------------
def _workflow(db: Db, workflow_id: str) -> ImageWorkflow:
row = db.get(ImageWorkflow, workflow_id)
if row is None:
raise HTTPException(status.HTTP_404_NOT_FOUND, "That workflow no longer exists.")
return row
def _detail(request: Request, row: ImageWorkflow, *, is_new: bool, error: str = "", **extra):
return render(
request,
"admin/workflow_detail.html",
{
"workflow": row,
"is_new": is_new,
"error": error,
"placeholders": workflow_service.PLACEHOLDERS,
"workflow_text": extra.pop(
"workflow_text", json.dumps(row.workflow_json or {}, indent=2)
),
**extra,
},
)
def _populate(row: ImageWorkflow, form) -> None:
row.name = str(form.get("name") or "").strip()[:120]
row.description = str(form.get("description") or "").strip()[:2000]
row.enabled = "enabled" in form
def _problem(db: Db, row: ImageWorkflow, form, *, existing_id: str = "") -> str:
"""Why this cannot be saved, or an empty string.
A sentence rather than a 422, so a rejected save re-renders the form with
what was typed still in it -- losing forty lines of JSON to a validation
error is not a thing to do to somebody.
"""
if not row.name:
return "A workflow needs a name."
slug = str(form.get("slug") or "").strip().lower()
if not SLUG_PATTERN.match(slug):
return (
"The name the model uses must be lowercase letters, digits, "
"hyphens or underscores, and start with a letter or digit."
)
clash = db.scalar(select(ImageWorkflow).where(ImageWorkflow.slug == slug))
if clash is not None and clash.id != existing_id:
return f"There is already a workflow called “{slug}”."
row.slug = slug
raw = str(form.get("workflow") or "").strip()
if not raw:
return "Paste the workflow, in ComfyUI's API format."
try:
parsed = json.loads(raw)
except json.JSONDecodeError as exc:
return f"That is not valid JSON: {exc}"
if not isinstance(parsed, dict) or not parsed:
return (
"A ComfyUI API workflow is a JSON object keyed by node id. Use "
"“Export (API)” in ComfyUI rather than “Save”."
)
# The check worth having: a workflow with no {{prompt}} in it draws the same
# picture whatever anybody types, and would look like a broken model rather
# than an unparameterised template.
found = workflow_service.placeholders_in(parsed)
missing = [name for name in REQUIRED_PLACEHOLDERS if name not in found]
if missing:
return (
f"The workflow never uses {{{{{missing[0]}}}}}, so every image would be "
f"the same. Put it where the text prompt goes."
)
unknown = found - set(workflow_service.PLACEHOLDERS)
if unknown:
return f"Unknown placeholder {{{{{sorted(unknown)[0]}}}}}."
row.workflow_json = parsed
return ""
@router.get("/workflows/new")
async def new_workflow(request: Request, db: Db, user: AdminUser) -> Response:
"""A draft, never persisted -- the `admin_tools` shape.
Registered before `/workflows/{workflow_id}`: FastAPI matches in
registration order, and with the parameterised route first "new" is an id.
"""
from pathlib import Path
base = Path(__file__).resolve().parent.parent / "services/images/base_workflow.json"
draft = ImageWorkflow(
slug="",
name="",
description="",
workflow_json=json.loads(base.read_text(encoding="utf-8")),
enabled=True,
)
return _detail(request, draft, is_new=True)
@router.post("/workflows")
async def create_workflow(request: Request, db: Db, user: AdminUser) -> Response:
form = await request.form()
row = ImageWorkflow(workflow_json={})
_populate(row, form)
problem = _problem(db, row, form)
if problem:
return _detail(
request,
row,
is_new=True,
error=problem,
workflow_text=str(form.get("workflow") or ""),
)
row.position = (
db.scalar(select(func.coalesce(func.max(ImageWorkflow.position), -1))) or -1
) + 1
db.add(row)
db.commit()
log.info("%s added image workflow %s", user.email, row.slug)
return RedirectResponse(
f"/admin/images?saved=Added {row.name}.", status_code=status.HTTP_303_SEE_OTHER
)
@router.get("/workflows/{workflow_id}/edit")
async def edit_workflow(request: Request, db: Db, user: AdminUser, workflow_id: str) -> Response:
return _detail(request, _workflow(db, workflow_id), is_new=False)
@router.post("/workflows/{workflow_id}/delete")
async def delete_workflow(db: Db, user: AdminUser, workflow_id: str) -> Response:
row = _workflow(db, workflow_id)
name = row.name
db.delete(row)
db.commit()
log.info("%s deleted image workflow %s", user.email, name)
return RedirectResponse(
f"/admin/images?saved=Deleted {name}.", status_code=status.HTTP_303_SEE_OTHER
)
@router.post("/workflows/{workflow_id}")
async def update_workflow(request: Request, db: Db, user: AdminUser, workflow_id: str) -> Response:
row = _workflow(db, workflow_id)
form = await request.form()
# Validated against a draft, so a rejected save leaves the stored row alone
# and the form still holds what was typed.
draft = ImageWorkflow(workflow_json={}, position=row.position)
_populate(draft, form)
problem = _problem(db, draft, form, existing_id=row.id)
if problem:
draft.id = row.id
return _detail(
request,
draft,
is_new=False,
error=problem,
workflow_text=str(form.get("workflow") or ""),
)
_populate(row, form)
row.slug = draft.slug
row.workflow_json = draft.workflow_json
row.last_checked_at = datetime.now(UTC)
row.last_error = ""
db.commit()
log.info("%s updated image workflow %s", user.email, row.slug)
return RedirectResponse(
f"/admin/images?saved=Saved {row.name}.", status_code=status.HTTP_303_SEE_OTHER
)
src/lembas/api/admin_models.py
+1
View File
@@ -44,6 +44,7 @@ TOOL_CAPABILITIES = (
("tool_custom", "Custom tools"),
("tool_mcp", "MCP servers"),
("tool_ask", "Ask the reader"),
("tool_image", "Image generation"),
("tool_agent", "Agent execution"),
)
src/lembas/api/chats.py
+19 -2
View File
@@ -67,6 +67,13 @@ METRICS_INTERVAL = 1.0
# better than four hundred rows nobody meant to write.
MAX_QUEUED = 10
# Which tools a request may compel the model to call. An allow list rather than
# a passthrough: this becomes `tool_choice`, and a name read straight off a form
# would let anyone who can send a message decide what the model must do next.
# Being on this list is not permission to *use* the tool -- `resolve_tools` still
# decides that, and forcing one that was never offered simply does nothing.
FORCEABLE_TOOLS = frozenset({"image_generate"})
# How many things one chat may have switched off. There are a dozen families and
# sixty skills at most, so this is not a limit anybody reaches by hand -- it is
# there so a crafted POST cannot grow the column without bound.
@@ -816,12 +823,20 @@ async def post_message(
chat_id: str,
content: str = Form(""),
file_ids: list[str] = Form(default=[]),
force_tool: str = Form(""),
) -> Response:
"""Persist the user's turn and hand back the pair of bubbles.
The assistant bubble comes back empty, carrying the sse-connect attribute
that opens the stream below. Splitting it this way means the POST returns
immediately and the slow part is a separate, resumable connection.
`force_tool` is `/image` and nothing else. It is checked against a fixed
list rather than passed through: this ends up in `tool_choice`, and a name
taken from a form would let anybody who can send a message pick which tool
the model is compelled to call. Whether that tool is *offered* is still
decided by `resolve_tools`, so this can only ever narrow to something the
chat was already allowed.
"""
chat = _owned_chat(db, chat_id, user.id)
@@ -831,7 +846,8 @@ async def post_message(
if not content and not file_ids:
return Response(status_code=status.HTTP_204_NO_CONTENT)
return _send(request, db, chat, user, content, file_ids=file_ids)
forced = force_tool.strip() if force_tool.strip() in FORCEABLE_TOOLS else ""
return _send(request, db, chat, user, content, file_ids=file_ids, force_tool=forced)
def _reply_in_flight(db: DBSession, chat: Chat) -> bool:
@@ -871,6 +887,7 @@ def _send(
content: str,
*,
file_ids: list[str] | None = None,
force_tool: str = "",
) -> Response:
"""Write a turn, start the reply, and hand back the pair of bubbles.
@@ -927,7 +944,7 @@ def _send(
assistant_message = chat_service.create_message(
db, chat, ROLE_ASSISTANT, "", complete_=False, model_id=chat.model_id
)
generation_service.ensure(chat.id, assistant_message.id)
generation_service.ensure(chat.id, assistant_message.id, force_tool=force_tool)
# `user` is required by the shared message template, which renders both
# roles; without it the user bubble's initial blows up.
src/lembas/db/models/__init__.py
+2
View File
@@ -32,6 +32,7 @@ from lembas.db.models.chat import (
Message,
)
from lembas.db.models.connection import Connection, Model, model_groups
from lembas.db.models.image import ImageWorkflow
from lembas.db.models.library import (
AUTHOR_MODEL,
AUTHOR_USER,
@@ -119,6 +120,7 @@ __all__ = [
"Document",
"Folder",
"Group",
"ImageWorkflow",
"KnowledgeBase",
"McpServer",
"Memory",
src/lembas/db/models/chat.py
+13
View File
@@ -218,6 +218,19 @@ class Chat(UUIDPrimaryKey, Timestamps, Base):
# representations of "on" makes "why is this off?" unanswerable.
scope_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)
# What this chat generates pictures with when the model names neither. A
# preference rather than a constraint -- the model may still choose another
# template or checkpoint for a particular image, and the harness lists what
# is on offer -- so this is where "in this chat I am working in SDXL" is
# said once instead of in every prompt.
#
# Plain columns rather than keys in `scope_json`: that one narrows what a
# chat may *reach* and absent means on, which is the opposite of what an
# empty default here means. A workflow that has since been deleted reads
# back as no preference, so it is validated on use like `ssh_profile_id`.
image_workflow_id: Mapped[str | None] = mapped_column(String(32))
image_checkpoint: Mapped[str] = mapped_column(String(300), default="")
# Which files are open in the canvas panel, and which of them is in front.
# {"tabs": [{"key": "agent:/srv/app/main.py", "title": …, "source": …}],
# "active": "agent:/srv/app/main.py"}
src/lembas/db/models/connection.py
+12
View File
@@ -58,6 +58,18 @@ class Connection(UUIDPrimaryKey, Timestamps, Base):
# Extra headers merged into every request (e.g. OpenRouter's HTTP-Referer).
extra_headers_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)
# How to ask this endpoint to drop its model from memory, for the Preserve
# VRAM option in image generation. Per connection and not instance-wide,
# because the VRAM being freed is a particular machine's: llama-swap on this
# host answers `GET /unload`, while a remote vLLM has no such call and no
# reason to be unloaded when ComfyUI needs memory *here*.
#
# Empty means "this connection cannot be unloaded", which is the honest
# default -- there is no call that works everywhere, and guessing one would
# send an unexplained request to somebody's endpoint.
unload_url: Mapped[str] = mapped_column(String(500), default="")
unload_method: Mapped[str] = mapped_column(String(8), default="POST")
# Result of the most recent "Test & refresh", surfaced in the admin list.
last_checked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
last_error: Mapped[str] = mapped_column(Text, default="")
src/lembas/db/models/image.py
+60
View File
@@ -0,0 +1,60 @@
"""ComfyUI workflow templates an administrator saved.
A table rather than a list inside the settings group, for the reason
`McpServer.tools_json` is *not* a table: that one is a cache of somebody else's
document, replaced wholesale on every refresh, where each entry carries one
decision. These are the opposite -- authored by hand, individually named,
edited, reordered and deleted, and referenced by id from a chat. Everything a
table gives for free is exactly what is wanted.
Deliberately **no group access list**, unlike `CustomTool`. The whole feature is
already behind one capability flag and one permission; a second access system
covering which templates a person may pick would be a screen of checkboxes
nobody asked for, and the thing being restricted is the shape of a picture.
"""
from __future__ import annotations
from datetime import datetime
from typing import Any
from sqlalchemy import Boolean, DateTime, Integer, String, Text
from sqlalchemy.orm import Mapped, mapped_column
from lembas.db.base import Base, Timestamps, UUIDPrimaryKey
from lembas.db.types import JSONDict
class ImageWorkflow(UUIDPrimaryKey, Timestamps, Base):
"""One API-format ComfyUI workflow, with holes where the values go."""
__tablename__ = "image_workflows"
# What the *model* names when it picks this one, so it is short and
# lowercase for the same reason a tool's slug is: it lands in a schema enum
# and is generated by something that spells inconsistently.
slug: Mapped[str] = mapped_column(String(64), unique=True, nullable=False)
name: Mapped[str] = mapped_column(String(120), nullable=False)
# Sent to the model beside the slug, and the only thing it has to choose
# with. "Photographic, SDXL, slow" is a choice; "workflow 2" is not.
description: Mapped[str] = mapped_column(Text, default="")
# The workflow itself, in ComfyUI's API format, with `{{placeholders}}`
# where the parameters go. Stored parsed rather than as text so the admin
# form can only ever save something that is valid JSON -- a template that
# does not parse would fail at generation time, minutes later, in front of
# somebody who was not editing it.
workflow_json: Mapped[dict[str, Any]] = mapped_column(JSONDict, default=dict)
enabled: Mapped[bool] = mapped_column(Boolean, default=True, nullable=False)
position: Mapped[int] = mapped_column(Integer, default=0, nullable=False)
# The result of the last time somebody pressed Test, in the shape
# `CustomTool` and `McpServer` already use, so the row reads the same way in
# the list as theirs do.
last_checked_at: Mapped[datetime | None] = mapped_column(DateTime(timezone=True))
last_error: Mapped[str] = mapped_column(Text, default="")
def __repr__(self) -> str:
return f"<ImageWorkflow {self.slug}>"
src/lembas/main.py
+2
View File
@@ -16,6 +16,7 @@ from lembas.api import (
admin,
admin_agents,
admin_audio,
admin_images,
admin_models,
admin_prompts,
admin_search,
@@ -147,6 +148,7 @@ def create_app() -> FastAPI:
app.include_router(admin_models.router)
app.include_router(admin_audio.router)
app.include_router(admin_search.router)
app.include_router(admin_images.router)
app.include_router(admin_prompts.router)
app.include_router(admin_suggestions.router)
app.include_router(admin_tools.router)
src/lembas/security/permissions.py
+9
View File
@@ -95,6 +95,15 @@ PERMISSION_DEFS: tuple[PermissionDef, ...] = (
True,
"Chat",
),
PermissionDef(
"tools.image",
"Generate images",
"Let a model draw a picture and show it in the conversation. Only "
"offered when an image generator has been configured, and every "
"generation spends time on whatever machine is running it.",
True,
"Chat",
),
PermissionDef(
"tools.custom",
"Use custom tools",
src/lembas/services/chat.py
+24 -1
View File
@@ -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
src/lembas/services/files.py
+97 -5
View File
@@ -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}"
src/lembas/services/generation.py
+52 -3
View File
@@ -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
src/lembas/services/harness.py
+38
View File
@@ -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": "",
src/lembas/services/images/__init__.py
+13
View File
@@ -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
src/lembas/services/images/base_workflow.json
+52
View File
@@ -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" }
}
}
src/lembas/services/images/comfy.py
+305
View File
@@ -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 []
src/lembas/services/images/tool.py
+526
View File
@@ -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)
src/lembas/services/images/workflow.py
+178
View File
@@ -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
src/lembas/services/prompts.py
+105
View File
@@ -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",
src/lembas/services/settings_store.py
+72
View File
@@ -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"])
src/lembas/services/tool_labels.py
+8
View File
@@ -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",
}
src/lembas/services/tools.py
+73 -5
View File
@@ -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,
src/lembas/web/static/css/chat.css
+18
View File
@@ -452,6 +452,24 @@
overflow-y: auto;
}
/* What a generation drew, inside its own tool block.
Bounded by height rather than by width: a portrait at 1024x1536 sized to the
column would push everything under it off the screen, and this is the block's
supporting evidence rather than the thing itself -- the picture proper is an
attachment on the bubble above. `width: auto` with a max height is what keeps
a landscape and a portrait both looking deliberate. */
.tool-image { display: block; margin: 0; }
.tool-image img {
display: block;
max-width: 100%;
max-height: 20rem;
width: auto;
height: auto;
border-radius: var(--radius-sm);
background: var(--bg-sunken);
}
/* A diff, from a write or an update.
The same visual language as .tool-result__text above -- both answer "what did
src/lembas/web/static/js/commands.js
+42
View File
@@ -168,6 +168,48 @@
else window.lembas.toggleTheme();
}
},
{
name: "image",
summary: "Draw a picture",
argument: "what to draw",
/* Offered wherever there is a chat, not only where image generation is
switched on. `available()` filters `run()` as well as the menu, so a
command hidden here stops being a command and gets *sent as a message*
-- and "/image a red bicycle" arriving as prose is worse than being
told the feature is off. The server answers either way. */
when: function () { return !!chat(); },
run: function (rest) {
var wanted = (rest || "").trim();
if (!wanted) return note("Say what to draw: /image a red bicycle in the rain.");
var thread = el("#thread");
if (!thread) return;
var body = new FormData();
body.append("content", wanted);
/* The whole of what this command is. The turn goes through the ordinary
path -- same route, same bubbles, same stream -- and carries one extra
field that makes the first round call the image tool instead of
deciding whether to. The words still say what is wanted, so an
endpoint that ignores tool_choice steers on those alone. */
body.append("force_tool", "image_generate");
fetch("/api/chats/" + chat() + "/messages", {
method: "POST",
body: body,
credentials: "same-origin"
})
.then(function (r) { return r.text(); })
.then(function (html) {
thread.insertAdjacentHTML("beforeend", html);
/* Without this the assistant bubble's sse-connect is inert markup
and the reply never starts -- the same reason /compact processes
the thread it swapped in. */
if (window.htmx) window.htmx.process(thread);
if (window.lembas.scrollThread) window.lembas.scrollThread(true);
})
.catch(function () { note("Could not start the image.", "error"); });
}
},
{ name: "new", summary: "Start a new chat", run: function () { window.location = "/chat"; } },
{
name: "temp",
src/lembas/web/templates/admin/_connection_row.html
+19
View File
@@ -72,6 +72,25 @@
</p>
</div>
<div class="field">
<label class="field__label" for="unload-{{ connection.id }}">Unload URL</label>
<div class="btn-row">
<input class="input input--mono" id="unload-{{ connection.id }}" name="unload_url"
value="{{ connection.unload_url }}" placeholder="No unload call"
style="flex: 1; min-width: 0">
<select class="select" name="unload_method" style="flex: none">
<option value="POST" {{ 'selected' if connection.unload_method != 'GET' }}>POST</option>
<option value="GET" {{ 'selected' if connection.unload_method == 'GET' }}>GET</option>
</select>
</div>
<p class="field__hint">
Only used by image generation's Preserve VRAM, to free this endpoint's
memory while ComfyUI works. llama-swap answers <code>GET /unload</code>.
Leave it empty for anything on another machine — its memory is not the
memory being freed, and it would be an unexplained request.
</p>
</div>
<div class="field">
<label class="checkbox">
<input type="checkbox" name="enabled" value="true"
src/lembas/web/templates/admin/_images_result.html
+24
View File
@@ -0,0 +1,24 @@
{% from "_macros.html" import icon %}
{#
What ComfyUI said when it was asked what it has.
Everything here came off that machine and is escaped accordingly -- a
checkpoint filename is a filename somebody chose, not a value this
application generated.
#}
<div class="alert {{ 'alert--error' if message_kind == 'error' else 'alert--success' }}"
style="margin-top: var(--sp-3)">
{{ icon("warning" if message_kind == "error" else "check", "icon--sm") }}
<span>{{ message }}</span>
</div>
{% if checkpoints %}
<p class="field__hint">
{% if kept %}
Your list above was left alone. What ComfyUI has:
{% else %}
Filled in above:
{% endif %}
</p>
<pre class="tool-result__text">{{ checkpoints | join("\n") }}</pre>
{% endif %}
src/lembas/web/templates/admin/_layout.html
+4
View File
@@ -43,6 +43,10 @@
{{ icon("globe", "icon--sm") }}
<span class="nav-item__label">Web search</span>
</a>
<a class="nav-item {{ 'is-active' if section == 'images' }}" href="/admin/images">
{{ icon("image", "icon--sm") }}
<span class="nav-item__label">Image generation</span>
</a>
<a class="nav-item {{ 'is-active' if section == 'tools' }}" href="/admin/tools">
{{ icon("link", "icon--sm") }}
<span class="nav-item__label">Tools</span>
src/lembas/web/templates/admin/images.html
+274
View File
@@ -0,0 +1,274 @@
{% extends "admin/_layout.html" %}
{% from "_macros.html" import icon %}
{% set section = "images" %}
{% block title %}Image generation - LLeMbas{% endblock %}
{% block heading %}Image generation{% endblock %}
{% block admin_content %}
<p class="admin-lede">
Lets a model draw a picture and show it in the conversation, on a ComfyUI you
are running. It is offered as a tool the model chooses to call, so nothing
changes for a conversation that never asks for one — and it is only offered
once there is a ComfyUI, a workflow and at least one checkpoint, because a
tool that fails on its first call is worse than a tool nobody was given.
</p>
{% if saved %}
<div class="alert alert--success">{{ icon("check", "icon--sm") }} <span>{{ saved }}</span></div>
{% endif %}
<form method="post" action="/admin/images">
<section class="card">
<h2 class="card__title">
Image generation
{% if values.enabled %}<span class="badge badge--success">on</span>
{% else %}<span class="badge">off</span>{% endif %}
</h2>
<div class="field">
<label class="checkbox">
<input type="checkbox" name="enabled" value="true" {{ 'checked' if values.enabled }}>
<span>Offer image generation to models that support tools</span>
</label>
<p class="field__hint">
Who may use it is a permission — <code>tools.image</code> under
Groups &amp; permissions. Which models may is a checkbox on each model.
</p>
</div>
</section>
<section class="card">
<h2 class="card__title">ComfyUI</h2>
<div class="field">
<label class="field__label" for="base_url">Base URL</label>
<input class="input input--mono" id="base_url" name="base_url" type="url"
placeholder="http://127.0.0.1:8188" value="{{ values.base_url }}">
<p class="field__hint">
Where ComfyUI is listening. An address on this machine or this network is
fine here and is not checked against the request-forgery rules — you
typed it, unlike an address a model asks for.
</p>
</div>
<div class="field">
<label class="field__label" for="api_key">API key</label>
<input class="input input--mono" id="api_key" name="api_key" type="password"
autocomplete="off" placeholder="No key set"
value="{{ unchanged if values.api_key_encrypted else '' }}">
<p class="field__hint">
{% if values.api_key_encrypted %}
Currently <code>{{ masked }}</code>. Leave the dots alone to keep it, or
clear the field to remove it.
{% else %}
Only needed if something sits in front of ComfyUI. Encrypted at rest.
{% endif %}
</p>
</div>
<div class="field">
<label class="field__label" for="timeout">Timeout (seconds)</label>
<input class="input" id="timeout" name="timeout" type="number"
min="10" max="3600" step="10" value="{{ values.timeout | int }}">
<p class="field__hint">
How long to wait for one image, queue included. Far longer than any other
timeout here, because the thing being waited for genuinely takes that long.
</p>
</div>
</section>
<section class="card">
<h2 class="card__title">Checkpoints</h2>
<div class="field">
<label class="field__label" for="checkpoints">Available checkpoints</label>
<textarea class="textarea input--mono" id="checkpoints" name="checkpoints" rows="4"
spellcheck="false"
placeholder="sd_xl_base_1.0.safetensors">{{ checkpoints_text }}</textarea>
<p class="field__hint">
One filename per line, exactly as ComfyUI spells it. This is the list the
model chooses from, so leaving out a checkpoint is how you stop it being
used. Press Test below to read them off ComfyUI — that fills this in the
first time and never overwrites it afterwards.
</p>
</div>
<div class="btn-row">
<button class="btn" type="button"
hx-post="/admin/images/test" hx-target="#images-test-result">
{{ icon("refresh", "icon--sm") }} Test &amp; read what it has
</button>
</div>
<div id="images-test-result"></div>
</section>
<section class="card">
<h2 class="card__title">Checking the result</h2>
<div class="field">
<label class="checkbox">
<input type="checkbox" name="review_enabled" value="true"
{{ 'checked' if values.review_enabled }}>
<span>Look at each image before showing it, and try again if it is wrong</span>
</label>
<p class="field__hint">
A vision model is shown the picture and the request it came from, and says
keep or retry. Only clearly wrong images are retried — a missing subject, a
mangled picture — because taste is not a fault and the next attempt is not
promised to be better. The reader sees only the image that was kept.
</p>
</div>
<div class="field">
<label class="field__label" for="review_model_id">Model that reviews</label>
<select class="select" id="review_model_id" name="review_model_id">
<option value="">The chat's own model, when it has vision</option>
{% for model in vision_models %}
{% if model.capabilities_json.get("vision") %}
<option value="{{ model.id }}"
{{ 'selected' if values.review_model_id == model.id }}>
{{ model.label }}
</option>
{% endif %}
{% endfor %}
</select>
<p class="field__hint">
Only models marked as having vision are listed. If there is nothing to ask,
the first image is kept and nothing fails.
</p>
</div>
<div class="field">
<label class="field__label" for="max_tries">Attempts per image</label>
<input class="input" id="max_tries" name="max_tries" type="number"
min="1" max="10" value="{{ values.max_tries }}">
<p class="field__hint">
Including the first. The last attempt is kept whatever the review says, so
a request always produces a picture.
</p>
</div>
</section>
<section class="card">
<h2 class="card__title">Memory</h2>
<div class="field">
<label class="checkbox">
<input type="checkbox" name="preserve_vram" value="true"
{{ 'checked' if values.preserve_vram }}>
<span>Preserve VRAM: unload the language model while ComfyUI works</span>
</label>
<p class="field__hint">
For a machine that cannot hold both at once. Before generating, the chat's
own connection is asked to unload — set an unload URL on it under
Connections, or nothing happens. Afterwards ComfyUI is asked to free its
own models, and the language model loads again by itself on the next
request.
</p>
</div>
{% if values.preserve_vram and values.review_enabled %}
<div class="alert">
{{ icon("warning", "icon--sm") }}
<span>
Both are on, so every retry costs two model loads — one to review, one to
generate again. That is the slow combination; consider two attempts rather
than {{ values.max_tries }}.
</span>
</div>
{% endif %}
</section>
<section class="card">
<h2 class="card__title">Extra instructions</h2>
<div class="field">
<label class="field__label visually-hidden" for="instructions">Extra instructions</label>
<textarea class="textarea" id="instructions" name="instructions" rows="4"
placeholder="Prefer the SDXL template for anything photographic.">{{ values.instructions }}</textarea>
<p class="field__hint">
Added to what every model is told about image generation, on this instance.
Where “always put <em>text, watermark</em> in the negative prompt” lives.
Leave it empty and nothing is added at all.
</p>
</div>
</section>
<div class="btn-row">
<button class="btn btn--primary" type="submit">Save changes</button>
</div>
</form>
<section class="card">
<h2 class="card__title">Workflows</h2>
<p class="field__hint">
A workflow is a ComfyUI graph exported with <strong>Export (API)</strong>, with
<code>{{ '{{prompt}}' }}</code> and the other placeholders where the values go.
The model picks between them by their descriptions, so write the description
for a reader who cannot see the graph.
</p>
<div class="btn-row">
<a class="btn btn--primary" href="/admin/images/workflows/new">
{{ icon("plus", "icon--sm") }} Add a workflow
</a>
</div>
{% if not workflows %}
<div class="empty" style="padding: var(--sp-8) 0">
{{ icon("image", "empty__mark") }}
<p class="empty__text">
No workflows yet. One is needed before anything can be drawn; the default
one is filled in for you when you add the first.
</p>
</div>
{% else %}
<div class="model-rows">
{% for item in workflows %}
<div class="model-row {{ 'is-off' if not item.enabled }}">
<div class="model-row__main">
<span class="model-row__name">
{{ item.name }}
{% if not item.enabled %}<span class="badge badge--danger">disabled</span>{% endif %}
{% if values.default_workflow_id == item.id %}
<span class="badge">default</span>
{% endif %}
</span>
<code class="model-row__id">{{ item.slug }}</code>
{% if item.description %}
<p class="text-xs faint">{{ item.description }}</p>
{% endif %}
</div>
<div class="model-row__actions">
<a class="btn btn--sm" href="/admin/images/workflows/{{ item.id }}/edit">Edit</a>
</div>
</div>
{% endfor %}
</div>
<form method="post" action="/admin/images" class="field" style="margin-top: var(--sp-4)">
{# Every other value is resubmitted with it, because this posts to the same
handler as the form above and an absent field reads as cleared. #}
<input type="hidden" name="enabled" value="{{ 'true' if values.enabled }}">
<input type="hidden" name="base_url" value="{{ values.base_url }}">
<input type="hidden" name="api_key" value="{{ unchanged if values.api_key_encrypted }}">
<input type="hidden" name="timeout" value="{{ values.timeout }}">
<input type="hidden" name="checkpoints" value="{{ checkpoints_text }}">
<input type="hidden" name="review_enabled" value="{{ 'true' if values.review_enabled }}">
<input type="hidden" name="review_model_id" value="{{ values.review_model_id }}">
<input type="hidden" name="max_tries" value="{{ values.max_tries }}">
<input type="hidden" name="preserve_vram" value="{{ 'true' if values.preserve_vram }}">
<input type="hidden" name="instructions" value="{{ values.instructions }}">
<label class="field__label" for="default_workflow_id">Default workflow</label>
<div class="btn-row">
<select class="select" id="default_workflow_id" name="default_workflow_id">
<option value="">The first one</option>
{% for item in workflows %}
<option value="{{ item.id }}"
{{ 'selected' if values.default_workflow_id == item.id }}>{{ item.name }}</option>
{% endfor %}
</select>
<button class="btn" type="submit">Set default</button>
</div>
<p class="field__hint">
Used when the model names no template and the chat has no preference.
</p>
</form>
{% endif %}
</section>
{% endblock %}
src/lembas/web/templates/admin/workflow_detail.html
+101
View File
@@ -0,0 +1,101 @@
{% extends "admin/_layout.html" %}
{% from "_macros.html" import icon %}
{% set section = "images" %}
{% block title %}{{ workflow.name or "New workflow" }} - LLeMbas{% endblock %}
{% block heading %}{{ workflow.name or "New workflow" }}{% endblock %}
{% block admin_content %}
<nav class="crumbs">
<a class="crumbs__back" href="/admin/images">
{{ icon("chevron-right", "icon--sm crumbs__icon") }} Image generation
</a>
</nav>
{% if error %}
<div class="alert alert--error">{{ icon("warning", "icon--sm") }} <span>{{ error }}</span></div>
{% endif %}
<form method="post"
action="{{ '/admin/images/workflows' if is_new else '/admin/images/workflows/' ~ workflow.id }}"
class="form-grid">
<section class="card">
<h2 class="card__title">What it is</h2>
<div class="field">
<label class="field__label" for="name">Name</label>
<input class="input" id="name" name="name" required maxlength="120"
value="{{ workflow.name }}" placeholder="SDXL, photographic">
<p class="field__hint">Shown to you, in the list.</p>
</div>
<div class="field">
<label class="field__label" for="slug">Name the model uses</label>
<input class="input input--mono" id="slug" name="slug" required maxlength="48"
value="{{ workflow.slug }}" placeholder="sdxl-photo">
<p class="field__hint">
Lowercase letters, digits, hyphens and underscores. This is what the model
writes when it picks this workflow.
</p>
</div>
<div class="field">
<label class="field__label" for="description">Description</label>
<textarea class="textarea" id="description" name="description" rows="3"
placeholder="Photographic and slow. Best for people, interiors and product shots at 1024px.">{{ workflow.description }}</textarea>
<p class="field__hint">
The only thing the model has to choose with, so say what this is
<em>for</em> rather than what it contains. It never sees the graph.
</p>
</div>
<div class="field">
<label class="checkbox">
<input type="checkbox" name="enabled" value="true" {{ 'checked' if workflow.enabled }}>
<span>Offer this workflow</span>
</label>
</div>
</section>
<section class="card">
<h2 class="card__title">The workflow</h2>
<div class="field">
<label class="field__label" for="workflow">ComfyUI API format</label>
<textarea class="textarea input--mono" id="workflow" name="workflow" rows="20"
spellcheck="false">{{ workflow_text }}</textarea>
<p class="field__hint">
Export this from ComfyUI with <strong>Export (API)</strong>, not Save — the
two formats are different and only the API one can be submitted.
</p>
<p class="field__hint">
Then put a placeholder where each value goes. A placeholder that is the
<em>whole</em> value keeps its type, so
<code>"steps": {{ '{{steps}}' }}</code> sends the number 20 rather than the
text “20”; one inside a longer string is substituted as text, so
<code>"{{ '{{prompt}}' }}, masterpiece"</code> works. Anything you leave out
takes its default.
</p>
<p class="field__hint">
Available:
{% for name in placeholders %}<code>{{ '{{' ~ name ~ '}}' }}</code>{{ ", " if not loop.last }}{% endfor %}.
<code>{{ '{{prompt}}' }}</code> is required — without it every image would
be the same.
</p>
</div>
</section>
<div class="btn-row">
<button class="btn btn--primary" type="submit">
{{ "Add workflow" if is_new else "Save changes" }}
</button>
<a class="btn btn--ghost" href="/admin/images">Back</a>
{% if not is_new %}
<button class="btn btn--danger" type="submit" formnovalidate
formaction="/admin/images/workflows/{{ workflow.id }}/delete"
data-confirm-button="Delete the workflow “{{ workflow.name }}”? Chats that used it fall back to the default.">
Delete
</button>
{% endif %}
</div>
</form>
{% endblock %}
src/lembas/web/templates/chat/_tool_activity.html
+23 -1
View File
@@ -106,10 +106,32 @@
{% if event.error %}
<p class="tool-activity__error">{{ event.error }}</p>
{% elif not event.results and not event.text %}
{% elif not event.results and not event.text and not event.image %}
<p class="tool-activity__error">Nothing was found.</p>
{% endif %}
{% if event.image and event.image.id %}
{#
What was drawn, in the block that drew it. The picture is also an
attachment on this bubble, so the reader sees it without opening
anything; this copy is what makes the block make sense on its own once
it is expanded, beside the seed and the attempts.
The address is built here from an id and is never taken from the event.
That is the same rule the result links a few lines down follow, and it
matters more here: an event is a dict a runner wrote, and a `src` taken
from one would be a model-supplied URL fetched by the reader's browser.
`/api/files/{id}/content` is owner-checked on the way out.
#}
<a class="tool-image" href="/api/files/{{ event.image.id }}/content"
target="_blank" rel="noopener">
<img src="/api/files/{{ event.image.id }}/content" alt="The generated image"
loading="lazy"
{% if event.image.width %}width="{{ event.image.width }}"{% endif %}
{% if event.image.height %}height="{{ event.image.height }}"{% endif %}>
</a>
{% endif %}
{% if event.text %}
<pre class="tool-result__text">{{ event.text }}</pre>
{% endif %}
tests/test_admin_images.py
+272
View File
@@ -0,0 +1,272 @@
"""The image generation admin page, and the workflows behind it."""
from __future__ import annotations
import json
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import select
import lembas
from lembas.db.models import ImageWorkflow
from lembas.services import settings_store
from lembas.services.crypto import UNCHANGED_SENTINEL, decrypt
BASE_TEXT = (Path(lembas.__file__).parent / "services/images/base_workflow.json").read_text(
encoding="utf-8"
)
BASE = json.loads(BASE_TEXT)
@pytest.fixture
def admin(client: TestClient, db, registered):
from lembas.db.models import User
db.query(User).update({"role": "admin"})
db.commit()
return client
def _settings(client, **overrides):
data = {
"enabled": "true",
"base_url": "http://comfy.test:8188",
"api_key": "",
"timeout": "600",
"checkpoints": "sd.safetensors\nother.safetensors",
"default_workflow_id": "",
"review_enabled": "true",
"review_model_id": "",
"max_tries": "4",
"preserve_vram": "",
"instructions": "",
}
data.update(overrides)
return client.post("/admin/images", data=data, follow_redirects=False)
# --- The page ------------------------------------------------------------------
def test_the_page_renders_and_is_in_the_nav(admin):
page = admin.get("/admin/images").text
assert 'href="/admin/images"' in page
assert "ComfyUI" in page
def test_settings_are_saved(admin, db):
assert _settings(admin).status_code == 303
values = settings_store.images(db)
assert values["enabled"] is True
assert values["base_url"] == "http://comfy.test:8188"
assert values["checkpoints"] == ["sd.safetensors", "other.safetensors"]
assert values["review_enabled"] is True
def test_a_trailing_slash_is_stripped_from_the_url(admin, db):
_settings(admin, base_url="http://comfy.test:8188/")
assert settings_store.images(db)["base_url"] == "http://comfy.test:8188"
def test_an_absent_checkbox_is_off(admin, db):
"""An unticked box is simply missing from a form post -- that absence *is*
the off signal."""
data = {"base_url": "http://x", "timeout": "600", "max_tries": "4"}
admin.post("/admin/images", data=data, follow_redirects=False)
values = settings_store.images(db)
assert values["enabled"] is False
assert values["review_enabled"] is False
assert values["preserve_vram"] is False
def test_the_numbers_are_clamped(admin, db):
_settings(admin, max_tries="500", timeout="1")
values = settings_store.images(db)
assert values["max_tries"] == 10
assert values["timeout"] == 10.0
def test_max_tries_never_reaches_zero(admin, db):
"""Zero would mean the tool generates nothing and reports success. Unlike
the zeroes elsewhere in the settings, there is no reading of it anybody
wants."""
_settings(admin, max_tries="0")
assert settings_store.images(db)["max_tries"] == 1
def test_the_key_survives_a_save_that_did_not_touch_it(admin, db):
_settings(admin, api_key="secret-key")
assert decrypt(settings_store.images(db)["api_key_encrypted"]) == "secret-key"
_settings(admin, api_key=UNCHANGED_SENTINEL, instructions="something else")
assert decrypt(settings_store.images(db)["api_key_encrypted"]) == "secret-key"
def test_an_emptied_key_is_removed(admin, db):
_settings(admin, api_key="secret-key")
_settings(admin, api_key="")
assert settings_store.images(db)["api_key_encrypted"] == ""
def test_the_key_is_never_rendered(admin, db):
_settings(admin, api_key="secret-key")
assert "secret-key" not in admin.get("/admin/images").text
def test_testing_without_a_url_says_so(admin, db):
_settings(admin, base_url="")
assert "Set a base URL first" in admin.post("/admin/images/test").text
# --- Workflows -----------------------------------------------------------------
def test_new_is_not_read_as_an_id(admin):
"""FastAPI matches in registration order, so `/workflows/new` has to be
registered before `/workflows/{id}` or "new" is an id and 404s. This has
been a bug twice in this codebase."""
response = admin.get("/admin/images/workflows/new")
assert response.status_code == 200
assert "Export (API)" in response.text
def test_the_default_workflow_is_offered_to_start_from(admin):
"""Somebody setting this up for the first time should not have to find a
working API-format document before they can try anything."""
page = admin.get("/admin/images/workflows/new").text
assert "CheckpointLoaderSimple" in page
assert "{{prompt}}" in page
def test_a_workflow_is_created(admin, db):
response = admin.post(
"/admin/images/workflows",
data={
"slug": "sdxl",
"name": "SDXL",
"description": "photo",
"workflow": BASE_TEXT,
"enabled": "true",
},
follow_redirects=False,
)
assert response.status_code == 303
row = db.scalar(select(ImageWorkflow).where(ImageWorkflow.slug == "sdxl"))
assert row is not None
assert row.workflow_json["4"]["inputs"]["ckpt_name"] == "{{model}}"
def test_broken_json_is_refused_with_a_sentence(admin, db):
"""A sentence and not a 422: losing forty lines of pasted JSON to a
validation error is not a thing to do to somebody, so the form comes back
with what was typed still in it."""
response = admin.post(
"/admin/images/workflows",
data={"slug": "x", "name": "X", "workflow": "{not json"},
)
assert response.status_code == 200
assert "not valid JSON" in response.text
assert "{not json" in response.text, "and what was typed is still there"
assert db.scalar(select(ImageWorkflow)) is None
def test_a_workflow_with_no_prompt_placeholder_is_refused(admin, db):
"""It would draw the same picture whatever anybody typed, and would look
like a broken model rather than an unparameterised template."""
without = json.loads(BASE_TEXT)
without["6"]["inputs"]["text"] = "a fixed prompt"
response = admin.post(
"/admin/images/workflows",
data={"slug": "x", "name": "X", "workflow": json.dumps(without)},
)
assert "{{prompt}}" in response.text
assert "every image would be the same" in response.text
assert db.scalar(select(ImageWorkflow)) is None
def test_an_unknown_placeholder_is_refused(admin, db):
response = admin.post(
"/admin/images/workflows",
data={
"slug": "x",
"name": "X",
"workflow": json.dumps({"1": {"a": "{{prompt}} {{lora}}"}}),
},
)
assert "lora" in response.text
assert db.scalar(select(ImageWorkflow)) is None
def test_a_ui_format_export_is_refused_helpfully(admin, db):
"""The two ComfyUI formats look alike enough that this is the mistake
everybody makes first, and "invalid" would not tell them which button to
press instead."""
response = admin.post(
"/admin/images/workflows",
data={"slug": "x", "name": "X", "workflow": "[]"},
)
assert "Export (API)" in response.text
def test_a_duplicate_slug_is_refused(admin, db):
for _ in range(2):
response = admin.post(
"/admin/images/workflows",
data={"slug": "dup", "name": "Dup", "workflow": BASE_TEXT},
)
assert "already a workflow" in response.text
assert len(list(db.scalars(select(ImageWorkflow)))) == 1
def test_a_rejected_edit_leaves_the_stored_row_alone(admin, db):
admin.post(
"/admin/images/workflows",
data={"slug": "keep", "name": "Keep", "workflow": BASE_TEXT},
)
row = db.scalar(select(ImageWorkflow))
admin.post(
f"/admin/images/workflows/{row.id}",
data={"slug": "keep", "name": "Renamed", "workflow": "{broken"},
)
db.expire_all()
row = db.scalar(select(ImageWorkflow))
assert row.name == "Keep", "validated against a draft, so nothing was written"
def test_a_workflow_is_deleted(admin, db):
admin.post(
"/admin/images/workflows",
data={"slug": "gone", "name": "Gone", "workflow": BASE_TEXT},
)
row = db.scalar(select(ImageWorkflow))
admin.post(f"/admin/images/workflows/{row.id}/delete", follow_redirects=False)
db.expire_all()
assert db.scalar(select(ImageWorkflow)) is None
def test_only_an_administrator_can_reach_it(client: TestClient, db, registered):
"""The first account registered is an administrator, so this needs a second
one -- the `test_admin_tools` fixture's reason for existing."""
from lembas.db.models import User
client.post("/auth/logout")
client.post(
"/auth/register",
data={"name": "Sam", "email": "sam@shire.test", "password": "correct horse battery"},
follow_redirects=False,
)
user = db.scalar(select(User).where(User.email == "sam@shire.test"))
user.role = "user"
user.active = True
db.commit()
client.post(
"/auth/login",
data={"email": "sam@shire.test", "password": "correct horse battery"},
follow_redirects=False,
)
assert client.get("/admin/images").status_code in (302, 303, 403, 404)
assert client.post("/admin/images/workflows", data={}).status_code in (302, 303, 403, 404)
tests/test_canvas_stream.py
+35
View File
@@ -200,6 +200,41 @@ def test_scratch_write_opens_its_tab(db, user_id):
assert tool.risk == "read"
async def test_scratch_write_actually_runs(db, user_id):
"""It did not, for the whole life of the feature.
The runner read `context.chat_id` off a `ToolContext` that had no such
field, so every call raised `AttributeError` -- swallowed by `run_tool`'s
blanket except into "the scratch_write tool failed", which is
indistinguishable from a model calling it wrongly. Nothing exercised the
runner: the test above asserts the family and the risk, which are
attributes of the declaration rather than of the code.
"""
from lembas.db.models import Chat
from lembas.services.tools import REGISTRY, ToolContext
chat = Chat(user_id=user_id, model_id="m")
db.add(chat)
db.commit()
outcome = await REGISTRY["scratch_write"].run(
ToolContext(owner_id=user_id, chat_id=chat.id), {"text": "hello", "mode": "replace"}
)
assert outcome.event["status"] == "ok"
assert outcome.event.get("canvas", {}).get("source") == "scratch"
async def test_scratch_write_without_a_chat_says_so(db, user_id):
"""Which is what the `chat_id` check was always for."""
from lembas.services.tools import REGISTRY, ToolContext
outcome = await REGISTRY["scratch_write"].run(
ToolContext(owner_id=user_id), {"text": "hello"}
)
assert outcome.event["status"] == "error"
def test_the_event_survives_into_the_stored_transcript():
"""Harmless and mildly useful: `_tool_activity.html` reads named keys."""
event = {"name": "file_read", "canvas": {"key": "agent:/a.py"}}
tests/test_images_chat.py
+297
View File
@@ -0,0 +1,297 @@
"""Image generation where it meets the conversation.
Three seams, each of which fails silently if it is got wrong: how a generated
image is bound to the reply that made it, why it is never replayed on an
assistant turn, and what `/image` actually sends.
"""
from __future__ import annotations
from pathlib import Path
import pytest
from fastapi.testclient import TestClient
from sqlalchemy import select
import lembas
from lembas.db.models import ROLE_ASSISTANT, ROLE_USER, Attachment, Chat, Connection, Message, Model
from lembas.services import chat as chat_service
from lembas.services import files as files_service
from lembas.services import generation as generation_service
from lembas.services.crypto import encrypt
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 vision_chat(db, user_id):
connection = Connection(name="c", base_url="http://127.0.0.1:1", api_key_encrypted=encrypt(""))
db.add(connection)
db.commit()
db.add(
Model(
connection_id=connection.id,
model_id="m",
capabilities_json={"vision": True, "tools": True},
)
)
chat = Chat(user_id=user_id, model_id="m", connection_id=connection.id)
db.add(chat)
db.commit()
return chat
# --- Storing what was drawn ----------------------------------------------------
def test_a_generated_image_is_kept_as_it_arrived(db, user_id, vision_chat):
"""`_process_image` transcodes to JPEG q85 and downscales to 1400px, which
is right for a phone photo and a visible loss on the one output this feature
exists to produce."""
attachment = files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="out.png",
keep_original=True,
)
assert attachment.media_type == "image/png"
assert files_service.stored_path(attachment.stored_name).read_bytes() == PNG
def test_an_ordinary_upload_is_still_processed(db, user_id, vision_chat):
"""The flag is opt-in, and the protection it skips still applies to
everything that arrives from outside."""
attachment = files_service.store(
db, user_id=user_id, chat_id=vision_chat.id, payload=PNG, filename="photo.png"
)
assert attachment.media_type == "image/jpeg"
def test_a_corrupt_image_is_still_refused(db, user_id, vision_chat):
"""What `keep_original` skips is the resize and the transcode, not the
check that this is an image at all."""
with pytest.raises(files_service.FileError):
files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=b"\x89PNG\r\n\x1a\n" + b"rubbish",
filename="broken.png",
keep_original=True,
)
def test_the_loop_binds_the_image_to_the_reply(db, user_id, vision_chat):
"""A runner cannot write the message row -- `_persist` is the single writer
-- so the runner makes the attachment and the loop says which turn owns it,
exactly as it already does for a canvas tab."""
message = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "here it is")
attachment = files_service.store(
db, user_id=user_id, chat_id=vision_chat.id, payload=PNG, filename="a.png"
)
generation_service._bind_attachments(db, vision_chat, message, [attachment.id])
db.commit()
db.refresh(attachment)
assert attachment.message_id == message.id
def test_binding_refuses_a_row_from_another_chat(db, user_id, vision_chat):
"""The ids arrive on a tool event, which is a dict a runner built. Without
the narrowing a forged one would pull somebody else's file into this
conversation -- the reason `files.claim` is scoped the same way."""
other = Chat(user_id=user_id, model_id="m")
db.add(other)
db.commit()
elsewhere = files_service.store(
db, user_id=user_id, chat_id=other.id, payload=PNG, filename="a.png"
)
message = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "x")
generation_service._bind_attachments(db, vision_chat, message, [elsewhere.id])
db.commit()
db.refresh(elsewhere)
assert elsewhere.message_id is None
def test_binding_refuses_a_row_already_bound(db, user_id, vision_chat):
first = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "one")
second = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "two")
attachment = files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="a.png",
message_id=first.id,
)
generation_service._bind_attachments(db, vision_chat, second, [attachment.id])
db.commit()
db.refresh(attachment)
assert attachment.message_id == first.id
# --- What reaches the model ----------------------------------------------------
def test_an_image_on_an_assistant_turn_is_never_replayed(db, user_id, vision_chat):
"""The multimodal list form on an `assistant` turn is rejected outright by
OpenAI and by most local runners -- and it would break not that turn but
every later one in the chat."""
reply = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "here it is")
files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="a.png",
message_id=reply.id,
)
db.refresh(reply)
payload = chat_service.message_payload(reply, vision=True)
assert isinstance(payload["content"], str), "no image_url parts on an assistant turn"
assert payload["content"] == "here it is"
def test_an_image_a_person_sent_still_reaches_the_model(db, user_id, vision_chat):
"""The rule narrows assistant turns and nothing else."""
turn = chat_service.create_message(db, vision_chat, ROLE_USER, "what is this?")
files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="a.png",
message_id=turn.id,
)
db.refresh(turn)
payload = chat_service.message_payload(turn, vision=True)
assert isinstance(payload["content"], list)
assert any(part["type"] == "image_url" for part in payload["content"])
def test_a_generated_image_still_renders_in_the_bubble(
client: TestClient, db, user_id, vision_chat
):
"""It is an attachment on the assistant message, and `_message.html` renders
attachments for either role -- so the reader sees it without the template
learning anything new."""
reply = chat_service.create_message(db, vision_chat, ROLE_ASSISTANT, "here it is")
attachment = files_service.store(
db,
user_id=user_id,
chat_id=vision_chat.id,
payload=PNG,
filename="a.png",
message_id=reply.id,
)
page = client.get(f"/chat/{vision_chat.id}").text
assert f"/api/files/{attachment.id}/content" in page
# --- Forcing the tool ----------------------------------------------------------
def test_the_forced_tool_reaches_the_request(db, vision_chat):
tools = [{"type": "function", "function": {"name": "image_generate", "parameters": {}}}]
body = chat_service.build_request(db, vision_chat, tools=tools, force_tool="image_generate")
assert body["tool_choice"] == {
"type": "function",
"function": {"name": "image_generate"},
}
def test_nothing_is_forced_by_default(db, vision_chat):
"""A provider strict about unknown parameters must see exactly the request
it always did until somebody types a slash command."""
tools = [{"type": "function", "function": {"name": "image_generate", "parameters": {}}}]
assert "tool_choice" not in chat_service.build_request(db, vision_chat, tools=tools)
def test_a_tool_that_was_not_offered_cannot_be_forced(db, vision_chat):
"""`resolve_tools` still decides what exists. Forcing something absent from
the array is a request most endpoints reject outright."""
tools = [{"type": "function", "function": {"name": "web_search", "parameters": {}}}]
body = chat_service.build_request(db, vision_chat, tools=tools, force_tool="image_generate")
assert "tool_choice" not in body
def test_forcing_needs_a_tools_array_at_all(db, vision_chat):
assert "tool_choice" not in chat_service.build_request(
db, vision_chat, force_tool="image_generate"
)
def test_the_endpoint_only_accepts_names_from_the_allow_list(
client: TestClient, db, registered, vision_chat, monkeypatch
):
"""This becomes `tool_choice`, so a name read straight off a form would let
anyone who can send a message decide what the model must do next."""
from lembas.api import chats as chats_api
seen: dict = {}
monkeypatch.setattr(
chats_api.generation_service,
"ensure",
lambda chat_id, message_id, *, force_tool="": seen.update(force_tool=force_tool),
)
client.post(
f"/api/chats/{vision_chat.id}/messages",
data={"content": "hello", "force_tool": "shell_run"},
)
assert seen["force_tool"] == "", "not on the list, so not forced"
# That first turn left an unfinished assistant row behind, and a second
# message while one is in flight is *queued* rather than sent -- so it would
# never reach `ensure` at all. Finish it first.
db.query(Message).filter(Message.complete.is_(False)).update({"complete": True})
db.commit()
client.post(
f"/api/chats/{vision_chat.id}/messages",
data={"content": "a bicycle", "force_tool": "image_generate"},
)
assert seen["force_tool"] == "image_generate"
def test_the_image_command_is_in_the_table():
"""`/help` reads this list, so a command missing from it is one nobody can
discover -- the direction this actually rots."""
source = (Path(lembas.__file__).parent / "web/static/js/commands.js").read_text(
encoding="utf-8"
)
assert 'name: "image"' in source
assert 'body.append("force_tool", "image_generate")' in source
assert "htmx.process" in source, "or the reply's sse-connect is inert markup"
def test_a_forced_round_does_not_force_the_next_one(db, vision_chat):
"""Leaving `tool_choice` in place would make every round call the tool
again: draw a picture, be asked again, draw another."""
source = (Path(lembas.__file__).parent / "services/generation.py").read_text(encoding="utf-8")
assert 'payload.pop("tool_choice", None)' in source
def test_a_queued_turn_is_not_lost_when_it_was_forced(db, user_id, vision_chat):
"""`/image` typed while a reply is streaming queues like anything else. The
forcing is on the generation, so a queued turn simply arrives unforced --
which is right: by then the model has the words and the context."""
db.add(Message(chat_id=vision_chat.id, role=ROLE_ASSISTANT, content="", complete=False))
db.commit()
from lembas.api import chats as chats_api
assert chats_api._reply_in_flight(db, vision_chat) is True
assert db.scalar(select(Attachment)) is None
tests/test_images_comfy.py
+281
View File
@@ -0,0 +1,281 @@
"""The ComfyUI client, against recorded shapes.
Every response body here was read off a real ComfyUI 0.27.0 while this was being
written, rather than reconstructed from documentation -- including the one that
matters most, `/history` answering with an empty object while the job is still
queued.
"""
from __future__ import annotations
import httpx
import pytest
from lembas.services.images import comfy
CONFIG = comfy.Config(base_url="http://comfy.test:8188", timeout=5.0)
HISTORY_DONE = {
"p1": {
"status": {"status_str": "success", "completed": True, "messages": []},
"outputs": {
"9": {"images": [{"filename": "LLeMbas_00001_.png", "subfolder": "", "type": "output"}]}
},
}
}
def _handler(routes):
"""Answer by path, and record what was asked."""
def handle(request: httpx.Request) -> httpx.Response:
for path, responder in routes.items():
if request.url.path.startswith(path):
return responder(request)
return httpx.Response(404, json={"error": "no route"})
return handle
# --- Submitting ----------------------------------------------------------------
async def test_a_workflow_is_queued_and_its_id_returned(mock_http):
seen: dict = {}
def prompt(request):
seen["body"] = request.read().decode()
return httpx.Response(200, json={"prompt_id": "p1", "number": 3, "node_errors": {}})
mock_http(_handler({"/prompt": prompt}))
assert await comfy.submit(CONFIG, {"3": {"class_type": "KSampler"}}) == "p1"
assert "KSampler" in seen["body"]
assert "client_id" in seen["body"], "ComfyUI keys its progress socket on this"
async def test_a_graph_error_names_the_node(mock_http):
""" "Invalid prompt" against a twelve-node document says nothing. The
commonest cause by far is a checkpoint that does not exist on that machine,
and the node id is what points at it."""
mock_http(
_handler(
{
"/prompt": lambda r: httpx.Response(
400,
json={
"error": {"type": "prompt_outputs_failed_validation"},
"node_errors": {
"4": {"errors": [{"message": "value not in list: ckpt_name"}]}
},
},
)
}
)
)
with pytest.raises(comfy.ComfyError) as caught:
await comfy.submit(CONFIG, {})
assert "node 4" in caught.value.message
assert "ckpt_name" in caught.value.message
async def test_an_accepted_workflow_with_no_id_is_refused(mock_http):
mock_http(_handler({"/prompt": lambda r: httpx.Response(200, json={"number": 1})}))
with pytest.raises(comfy.ComfyError):
await comfy.submit(CONFIG, {})
async def test_an_unreachable_comfyui_says_so(mock_http):
def refuse(request):
raise httpx.ConnectError("nope", request=request)
mock_http(_handler({"/prompt": refuse}))
with pytest.raises(comfy.ComfyError) as caught:
await comfy.submit(CONFIG, {})
assert "Could not reach ComfyUI" in caught.value.message
# --- Waiting -------------------------------------------------------------------
async def test_an_empty_history_means_not_yet(mock_http, monkeypatch):
"""`/history/{id}` is empty while the job is queued and gains the whole
record when it ends -- so empty is "not yet", not "nothing". Reading it the
other way makes every generation fail the instant it is submitted."""
monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01)
calls = {"n": 0}
def history(request):
calls["n"] += 1
return httpx.Response(200, json=HISTORY_DONE if calls["n"] >= 3 else {})
mock_http(_handler({"/history": history}))
refs = await comfy.await_images(CONFIG, "p1")
assert calls["n"] == 3, "it kept asking"
assert [ref.filename for ref in refs] == ["LLeMbas_00001_.png"]
async def test_waiting_gives_up_eventually(mock_http, monkeypatch):
monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01)
mock_http(_handler({"/history": lambda r: httpx.Response(200, json={})}))
with pytest.raises(comfy.ComfyError) as caught:
await comfy.await_images(comfy.Config(base_url=CONFIG.base_url, timeout=0.05), "p1")
assert "did not finish" in caught.value.message
async def test_a_failed_workflow_is_not_an_empty_one(mock_http, monkeypatch):
monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01)
mock_http(
_handler(
{
"/history": lambda r: httpx.Response(
200,
json={
"p1": {"status": {"status_str": "error", "completed": True}, "outputs": {}}
},
)
}
)
)
with pytest.raises(comfy.ComfyError):
await comfy.await_images(CONFIG, "p1")
async def test_images_are_read_from_every_node(mock_http, monkeypatch):
"""A template is somebody else's document and may save from a node called
anything, or from two of them. Looking for `SaveImage` by name works on the
shipped workflow and on nothing else."""
monkeypatch.setattr(comfy, "POLL_INTERVAL", 0.01)
mock_http(
_handler(
{
"/history": lambda r: httpx.Response(
200,
json={
"p1": {
"status": {"status_str": "success", "completed": True},
"outputs": {
"12": {
"images": [
{"filename": "a.png", "subfolder": "s", "type": "output"}
]
},
"40": {"images": [{"filename": "b.png"}]},
},
}
},
)
}
)
)
refs = await comfy.await_images(CONFIG, "p1")
assert [r.filename for r in refs] == ["a.png", "b.png"]
assert refs[0].subfolder == "s"
# --- Fetching ------------------------------------------------------------------
async def test_the_image_comes_back_as_bytes(mock_http):
seen: dict = {}
def view(request):
seen["params"] = dict(request.url.params)
return httpx.Response(200, content=b"\x89PNG\r\n\x1a\n" + b"x" * 100)
mock_http(_handler({"/view": view}))
payload = await comfy.fetch_image(CONFIG, comfy.Ref("a.png", "sub", "output"))
assert payload.startswith(b"\x89PNG")
assert seen["params"] == {"filename": "a.png", "subfolder": "sub", "type": "output"}
async def test_an_oversized_image_is_refused(mock_http, monkeypatch):
"""The one place an external service hands back raw bytes that are written
to disk. Neither `audio.speak` nor `openai_client` has a cap to copy."""
monkeypatch.setattr(comfy, "MAX_IMAGE_BYTES", 64)
mock_http(_handler({"/view": lambda r: httpx.Response(200, content=b"x" * 4096)}))
with pytest.raises(comfy.ComfyError):
await comfy.fetch_image(CONFIG, comfy.Ref("a.png"))
async def test_an_empty_file_is_refused(mock_http):
mock_http(_handler({"/view": lambda r: httpx.Response(200, content=b"")}))
with pytest.raises(comfy.ComfyError):
await comfy.fetch_image(CONFIG, comfy.Ref("a.png"))
# --- Freeing -------------------------------------------------------------------
async def test_freeing_never_raises(mock_http):
"""It runs on the way out of a generation that has already produced its
image. Failing the whole tool because a memory hint was refused would be
turning a tidy-up into an error."""
def boom(request):
raise httpx.ConnectError("gone", request=request)
mock_http(_handler({"/free": boom}))
await comfy.free(CONFIG) # must not raise
async def test_freeing_asks_for_both(mock_http):
seen: dict = {}
def free(request):
seen["body"] = request.read().decode()
return httpx.Response(200, json={})
mock_http(_handler({"/free": free}))
await comfy.free(CONFIG)
assert "unload_models" in seen["body"] and "free_memory" in seen["body"]
# --- Discovery -----------------------------------------------------------------
async def test_discovery_reads_the_option_lists(mock_http):
"""The shape is `{node: {input: {required: {field: [[...values], {meta}]}}}}`
-- a list whose first element is the options."""
mock_http(
_handler(
{
"/object_info/CheckpointLoaderSimple": lambda r: httpx.Response(
200,
json={
"CheckpointLoaderSimple": {
"input": {
"required": {"ckpt_name": [["a.safetensors", "b.safetensors"]]}
}
}
},
),
"/object_info/KSampler": lambda r: httpx.Response(
200,
json={
"KSampler": {
"input": {
"required": {
"sampler_name": [["euler", "dpmpp_2m"]],
"scheduler": [["normal", "karras"]],
}
}
}
},
),
}
)
)
checkpoints, samplers, schedulers = await comfy.discover(CONFIG)
assert checkpoints == ["a.safetensors", "b.safetensors"]
assert samplers == ["euler", "dpmpp_2m"]
assert schedulers == ["normal", "karras"]
async def test_a_custom_node_pack_that_changes_the_shape_is_survived(mock_http):
"""It is somebody else's schema. An empty list is a page that says "none
found"; an exception is a page that says nothing at all."""
mock_http(
_handler(
{
"/object_info": lambda r: httpx.Response(200, json={"Whatever": {"input": {}}}),
}
)
)
assert await comfy.discover(CONFIG) == ([], [], [])
tests/test_images_tool.py
+395
View File
@@ -0,0 +1,395 @@
"""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"
tests/test_images_workflow.py
+122
View File
@@ -0,0 +1,122 @@
"""Filling a ComfyUI template.
The thing worth pinning here is types. ComfyUI validates its inputs, so a
workflow whose `steps` arrives as the string "20" is refused -- and the refusal
happens on somebody's machine a minute later rather than in a test.
"""
from __future__ import annotations
import json
from pathlib import Path
import pytest
import lembas
from lembas.services.images import workflow as wf
BASE = json.loads(
(Path(lembas.__file__).parent / "services/images/base_workflow.json").read_text(
encoding="utf-8"
)
)
def _filled(**given):
return wf.fill(BASE, wf.resolve({"prompt": "a bicycle", **given}))
# --- Substitution --------------------------------------------------------------
def test_a_whole_placeholder_keeps_its_type():
"""`"steps": "{{steps}}"` has to become the number 20, not the text "20".
Substituting textually is the obvious implementation and produces a document
ComfyUI refuses."""
sampler = _filled(steps=25, cfg=7.5, seed=42)["3"]["inputs"]
assert sampler["steps"] == 25
assert isinstance(sampler["steps"], int)
assert sampler["cfg"] == 7.5
assert isinstance(sampler["cfg"], float)
assert sampler["seed"] == 42
assert isinstance(sampler["seed"], int)
def test_a_placeholder_inside_a_string_is_text():
"""Which is what makes `"{{prompt}}, masterpiece"` a usable template."""
filled = wf.fill(
{"n": {"inputs": {"text": "{{prompt}}, masterpiece"}}}, wf.resolve({"prompt": "a cat"})
)
assert filled["n"]["inputs"]["text"] == "a cat, masterpiece"
def test_the_prompt_and_the_negative_go_to_different_nodes():
filled = _filled(negative="blurry")
assert filled["6"]["inputs"]["text"] == "a bicycle"
assert filled["7"]["inputs"]["text"] == "blurry"
def test_the_links_between_nodes_survive():
"""A workflow is a graph, and `["4", 0]` is an edge rather than a value. A
filler that walked only dicts would flatten every one of them."""
filled = _filled()
assert filled["3"]["inputs"]["model"] == ["4", 0]
assert filled["8"]["inputs"]["vae"] == ["4", 2]
assert filled["5"]["inputs"]["batch_size"] == 1
def test_an_unknown_placeholder_is_left_alone():
"""`prompts.substitute`'s rule. A literal `{{x}}` is not a feature, but
silently deleting one is worse than leaving it where somebody can see it."""
filled = wf.fill({"n": {"inputs": {"lora": "{{lora_name}}"}}}, wf.resolve({"prompt": "x"}))
assert filled["n"]["inputs"]["lora"] == "{{lora_name}}"
def test_the_shipped_template_uses_every_placeholder_it_should():
assert wf.placeholders_in(BASE) == set(wf.PLACEHOLDERS)
# --- Resolution ----------------------------------------------------------------
def test_what_is_not_asked_for_takes_its_default():
values = wf.resolve({"prompt": "a bicycle"})
assert values["steps"] == 20
assert values["cfg"] == 8.0
assert values["width"] == 512
assert values["negative"] == "text, watermark"
assert values["sampler"] == "euler"
def test_the_seed_is_random_when_nobody_names_one():
"""A fixed default would make every unspecified generation identical -- and
would make the retry loop produce the same rejected image four times."""
seeds = {wf.resolve({"prompt": "x"})["seed"] for _ in range(8)}
assert len(seeds) == 8
def test_a_named_seed_is_kept():
assert wf.resolve({"prompt": "x", "seed": 1234})["seed"] == 1234
@pytest.mark.parametrize("empty", [None, ""])
def test_null_and_empty_mean_no_opinion(empty):
"""A model that emits `"steps": null` rather than omitting the key is common
enough that reading it as a request for zero steps would be a bug nobody
could see."""
assert wf.resolve({"prompt": "x", "steps": empty})["steps"] == 20
def test_out_of_range_numbers_are_clamped_rather_than_refused():
"""A model asking for 300 steps has misjudged rather than misbehaved, and a
clarifying round costs more than doing the sensible thing."""
values = wf.resolve({"prompt": "x", "steps": 5000, "cfg": -4, "denoise": 12})
assert values["steps"] == 150
assert values["cfg"] == 0.0
assert values["denoise"] == 1.0
def test_nonsense_falls_back_instead_of_raising():
"""Arguments come from a model, so "twenty" is a thing that will arrive."""
assert wf.resolve({"prompt": "x", "steps": "twenty"})["steps"] == 20
def test_a_seed_larger_than_comfyui_allows_is_wrapped():
assert 0 <= wf.resolve({"prompt": "x", "seed": wf.MAX_SEED + 5})["seed"] <= wf.MAX_SEED