diff --git a/CLAUDE.md b/CLAUDE.md index 0a2a259..6c3296b 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -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 diff --git a/PLAN.md b/PLAN.md index b6c95dd..bf08369 100644 --- a/PLAN.md +++ b/PLAN.md @@ -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 diff --git a/src/lembas/__init__.py b/src/lembas/__init__.py index 62307b4..b2fe638 100644 --- a/src/lembas/__init__.py +++ b/src/lembas/__init__.py @@ -1,3 +1,3 @@ """LLeMbas - a Middle-earth themed web UI for OpenAI-compatible LLM endpoints.""" -__version__ = "0.6.5" +__version__ = "0.7.0" diff --git a/src/lembas/api/admin.py b/src/lembas/api/admin.py index 9a6a3ee..351aa2d 100644 --- a/src/lembas/api/admin.py +++ b/src/lembas/api/admin.py @@ -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: diff --git a/src/lembas/api/admin_images.py b/src/lembas/api/admin_images.py new file mode 100644 index 0000000..92b6446 --- /dev/null +++ b/src/lembas/api/admin_images.py @@ -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 + ) diff --git a/src/lembas/api/admin_models.py b/src/lembas/api/admin_models.py index 9604f3e..6f2d80a 100644 --- a/src/lembas/api/admin_models.py +++ b/src/lembas/api/admin_models.py @@ -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"), ) diff --git a/src/lembas/api/chats.py b/src/lembas/api/chats.py index 6fa7b53..f6b11a3 100644 --- a/src/lembas/api/chats.py +++ b/src/lembas/api/chats.py @@ -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. diff --git a/src/lembas/db/models/__init__.py b/src/lembas/db/models/__init__.py index df9c91e..b59f167 100644 --- a/src/lembas/db/models/__init__.py +++ b/src/lembas/db/models/__init__.py @@ -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", diff --git a/src/lembas/db/models/chat.py b/src/lembas/db/models/chat.py index badb928..54b0098 100644 --- a/src/lembas/db/models/chat.py +++ b/src/lembas/db/models/chat.py @@ -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"} diff --git a/src/lembas/db/models/connection.py b/src/lembas/db/models/connection.py index 0e0b7ab..c91c8f9 100644 --- a/src/lembas/db/models/connection.py +++ b/src/lembas/db/models/connection.py @@ -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="") diff --git a/src/lembas/db/models/image.py b/src/lembas/db/models/image.py new file mode 100644 index 0000000..8f65a16 --- /dev/null +++ b/src/lembas/db/models/image.py @@ -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"" diff --git a/src/lembas/main.py b/src/lembas/main.py index 54fac6f..40e1754 100644 --- a/src/lembas/main.py +++ b/src/lembas/main.py @@ -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) diff --git a/src/lembas/security/permissions.py b/src/lembas/security/permissions.py index b9b4fba..6e87328 100644 --- a/src/lembas/security/permissions.py +++ b/src/lembas/security/permissions.py @@ -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", diff --git a/src/lembas/services/chat.py b/src/lembas/services/chat.py index c4d67bb..d9400b6 100644 --- a/src/lembas/services/chat.py +++ b/src/lembas/services/chat.py @@ -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 diff --git a/src/lembas/services/files.py b/src/lembas/services/files.py index 6ada44b..8eb095c 100644 --- a/src/lembas/services/files.py +++ b/src/lembas/services/files.py @@ -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}" diff --git a/src/lembas/services/generation.py b/src/lembas/services/generation.py index 6d79ef1..78051e6 100644 --- a/src/lembas/services/generation.py +++ b/src/lembas/services/generation.py @@ -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 diff --git a/src/lembas/services/harness.py b/src/lembas/services/harness.py index 8021b79..7ec9f11 100644 --- a/src/lembas/services/harness.py +++ b/src/lembas/services/harness.py @@ -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": "", diff --git a/src/lembas/services/images/__init__.py b/src/lembas/services/images/__init__.py new file mode 100644 index 0000000..980dd1a --- /dev/null +++ b/src/lembas/services/images/__init__.py @@ -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 diff --git a/src/lembas/services/images/base_workflow.json b/src/lembas/services/images/base_workflow.json new file mode 100644 index 0000000..5a12e6e --- /dev/null +++ b/src/lembas/services/images/base_workflow.json @@ -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" } + } +} diff --git a/src/lembas/services/images/comfy.py b/src/lembas/services/images/comfy.py new file mode 100644 index 0000000..959aab8 --- /dev/null +++ b/src/lembas/services/images/comfy.py @@ -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 [] diff --git a/src/lembas/services/images/tool.py b/src/lembas/services/images/tool.py new file mode 100644 index 0000000..59f0c92 --- /dev/null +++ b/src/lembas/services/images/tool.py @@ -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) diff --git a/src/lembas/services/images/workflow.py b/src/lembas/services/images/workflow.py new file mode 100644 index 0000000..9f8cb89 --- /dev/null +++ b/src/lembas/services/images/workflow.py @@ -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 diff --git a/src/lembas/services/prompts.py b/src/lembas/services/prompts.py index faf2447..976d7e0 100644 --- a/src/lembas/services/prompts.py +++ b/src/lembas/services/prompts.py @@ -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", diff --git a/src/lembas/services/settings_store.py b/src/lembas/services/settings_store.py index 3813fac..a690277 100644 --- a/src/lembas/services/settings_store.py +++ b/src/lembas/services/settings_store.py @@ -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"]) diff --git a/src/lembas/services/tool_labels.py b/src/lembas/services/tool_labels.py index c20f461..97f3474 100644 --- a/src/lembas/services/tool_labels.py +++ b/src/lembas/services/tool_labels.py @@ -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", } diff --git a/src/lembas/services/tools.py b/src/lembas/services/tools.py index 4be73d5..726e604 100644 --- a/src/lembas/services/tools.py +++ b/src/lembas/services/tools.py @@ -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, diff --git a/src/lembas/web/static/css/chat.css b/src/lembas/web/static/css/chat.css index 0e43f7b..671a35d 100644 --- a/src/lembas/web/static/css/chat.css +++ b/src/lembas/web/static/css/chat.css @@ -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 diff --git a/src/lembas/web/static/js/commands.js b/src/lembas/web/static/js/commands.js index e76720b..2694ad7 100644 --- a/src/lembas/web/static/js/commands.js +++ b/src/lembas/web/static/js/commands.js @@ -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", diff --git a/src/lembas/web/templates/admin/_connection_row.html b/src/lembas/web/templates/admin/_connection_row.html index b402d61..66e8669 100644 --- a/src/lembas/web/templates/admin/_connection_row.html +++ b/src/lembas/web/templates/admin/_connection_row.html @@ -72,6 +72,25 @@

+
+ +
+ + +
+

+ Only used by image generation's Preserve VRAM, to free this endpoint's + memory while ComfyUI works. llama-swap answers GET /unload. + Leave it empty for anything on another machine — its memory is not the + memory being freed, and it would be an unexplained request. +

+
+
+ +{% if checkpoints %} +

+ {% if kept %} + Your list above was left alone. What ComfyUI has: + {% else %} + Filled in above: + {% endif %} +

+
{{ checkpoints | join("\n") }}
+{% endif %} diff --git a/src/lembas/web/templates/admin/_layout.html b/src/lembas/web/templates/admin/_layout.html index 3a02485..34d9bc4 100644 --- a/src/lembas/web/templates/admin/_layout.html +++ b/src/lembas/web/templates/admin/_layout.html @@ -43,6 +43,10 @@ {{ icon("globe", "icon--sm") }} Web search + + {{ icon("image", "icon--sm") }} + Image generation + {{ icon("link", "icon--sm") }} Tools diff --git a/src/lembas/web/templates/admin/images.html b/src/lembas/web/templates/admin/images.html new file mode 100644 index 0000000..c37f007 --- /dev/null +++ b/src/lembas/web/templates/admin/images.html @@ -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 %} +

+ 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. +

+ +{% if saved %} +
{{ icon("check", "icon--sm") }} {{ saved }}
+{% endif %} + +
+
+

+ Image generation + {% if values.enabled %}on + {% else %}off{% endif %} +

+
+ +

+ Who may use it is a permission — tools.image under + Groups & permissions. Which models may is a checkbox on each model. +

+
+
+ +
+

ComfyUI

+
+ + +

+ 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. +

+
+ +
+ + +

+ {% if values.api_key_encrypted %} + Currently {{ masked }}. 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 %} +

+
+ +
+ + +

+ 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. +

+
+
+ +
+

Checkpoints

+
+ + +

+ 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. +

+
+ +
+ +
+
+
+ +
+

Checking the result

+
+ +

+ 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. +

+
+ +
+ + +

+ Only models marked as having vision are listed. If there is nothing to ask, + the first image is kept and nothing fails. +

+
+ +
+ + +

+ Including the first. The last attempt is kept whatever the review says, so + a request always produces a picture. +

+
+
+ +
+

Memory

+
+ +

+ 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. +

+
+ + {% if values.preserve_vram and values.review_enabled %} +
+ {{ icon("warning", "icon--sm") }} + + 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 }}. + +
+ {% endif %} +
+ +
+

Extra instructions

+
+ + +

+ Added to what every model is told about image generation, on this instance. + Where “always put text, watermark in the negative prompt” lives. + Leave it empty and nothing is added at all. +

+
+
+ +
+ +
+
+ +
+

Workflows

+

+ A workflow is a ComfyUI graph exported with Export (API), with + {{ '{{prompt}}' }} 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. +

+ +
+ + {{ icon("plus", "icon--sm") }} Add a workflow + +
+ + {% if not workflows %} +
+ {{ icon("image", "empty__mark") }} +

+ No workflows yet. One is needed before anything can be drawn; the default + one is filled in for you when you add the first. +

+
+ {% else %} +
+ {% for item in workflows %} +
+
+ + {{ item.name }} + {% if not item.enabled %}disabled{% endif %} + {% if values.default_workflow_id == item.id %} + default + {% endif %} + + {{ item.slug }} + {% if item.description %} +

{{ item.description }}

+ {% endif %} +
+
+ Edit +
+
+ {% endfor %} +
+ +
+ {# 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. #} + + + + + + + + + + + +
+ + +
+

+ Used when the model names no template and the chat has no preference. +

+
+ {% endif %} +
+{% endblock %} diff --git a/src/lembas/web/templates/admin/workflow_detail.html b/src/lembas/web/templates/admin/workflow_detail.html new file mode 100644 index 0000000..5b440da --- /dev/null +++ b/src/lembas/web/templates/admin/workflow_detail.html @@ -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 %} + + +{% if error %} +
{{ icon("warning", "icon--sm") }} {{ error }}
+{% endif %} + +
+
+

What it is

+ +
+ + +

Shown to you, in the list.

+
+ +
+ + +

+ Lowercase letters, digits, hyphens and underscores. This is what the model + writes when it picks this workflow. +

+
+ +
+ + +

+ The only thing the model has to choose with, so say what this is + for rather than what it contains. It never sees the graph. +

+
+ +
+ +
+
+ +
+

The workflow

+
+ + +

+ Export this from ComfyUI with Export (API), not Save — the + two formats are different and only the API one can be submitted. +

+

+ Then put a placeholder where each value goes. A placeholder that is the + whole value keeps its type, so + "steps": {{ '{{steps}}' }} sends the number 20 rather than the + text “20”; one inside a longer string is substituted as text, so + "{{ '{{prompt}}' }}, masterpiece" works. Anything you leave out + takes its default. +

+

+ Available: + {% for name in placeholders %}{{ '{{' ~ name ~ '}}' }}{{ ", " if not loop.last }}{% endfor %}. + {{ '{{prompt}}' }} is required — without it every image would + be the same. +

+
+
+ +
+ + Back + {% if not is_new %} + + {% endif %} +
+
+{% endblock %} diff --git a/src/lembas/web/templates/chat/_tool_activity.html b/src/lembas/web/templates/chat/_tool_activity.html index c85e3d9..ef45f71 100644 --- a/src/lembas/web/templates/chat/_tool_activity.html +++ b/src/lembas/web/templates/chat/_tool_activity.html @@ -106,10 +106,32 @@ {% if event.error %}

{{ event.error }}

- {% elif not event.results and not event.text %} + {% elif not event.results and not event.text and not event.image %}

Nothing was found.

{% 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. + #} + + The generated image + + {% endif %} + {% if event.text %} 
{{ event.text }}
{% endif %} diff --git a/tests/test_admin_images.py b/tests/test_admin_images.py new file mode 100644 index 0000000..549e4c9 --- /dev/null +++ b/tests/test_admin_images.py @@ -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) diff --git a/tests/test_canvas_stream.py b/tests/test_canvas_stream.py index 0905989..2383349 100644 --- a/tests/test_canvas_stream.py +++ b/tests/test_canvas_stream.py @@ -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"}} diff --git a/tests/test_images_chat.py b/tests/test_images_chat.py new file mode 100644 index 0000000..75350c1 --- /dev/null +++ b/tests/test_images_chat.py @@ -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 diff --git a/tests/test_images_comfy.py b/tests/test_images_comfy.py new file mode 100644 index 0000000..4810a0c --- /dev/null +++ b/tests/test_images_comfy.py @@ -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) == ([], [], []) diff --git a/tests/test_images_tool.py b/tests/test_images_tool.py new file mode 100644 index 0000000..b80bb5b --- /dev/null +++ b/tests/test_images_tool.py @@ -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" diff --git a/tests/test_images_workflow.py b/tests/test_images_workflow.py new file mode 100644 index 0000000..ed83efa --- /dev/null +++ b/tests/test_images_workflow.py @@ -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