Let image_generate edit an image, or work from reference images

packages/redrob/src/tool/image_generate.ts

@@ -11,6 +11,7 @@ import { Watcher } from "@redrob-code/core/filesystem/watcher"
 import { FSUtil } from "@redrob-code/core/fs-util"
 import { CONSOLE_URL } from "@redrob-code/core/plugin/provider/redrob-constants"
 import { assertExternalDirectoryEffect } from "./external-directory"
+import { sniffAttachmentMime } from "@/util/media"
 
 /**
  * Image generation through the Redrob gateway, on the Redrob credential this engine already holds.
@@ -37,7 +38,10 @@ function gatewayUrl(): string {
 }
 
 const TIMEOUT_MS = 180_000
+/** A handful, which is what the image models take; each one is billed as image input. */
+const MAX_REFERENCES = 4
 const MAX_IMAGE_BYTES = 20 * 1024 * 1024
+const MAX_REFERENCE_BYTES = 10 * 1024 * 1024
 
 const EXTENSIONS: Record<string, string> = {
   "image/png": "png",
@@ -54,6 +58,9 @@ export const Parameters = Schema.Struct({
   model: Schema.optional(Schema.String).annotate({
     description: `Optional image model id served by the Redrob gateway. Defaults to ${DEFAULT_IMAGE_MODEL}.`,
   }),
+  references: Schema.optional(Schema.Array(Schema.String)).annotate({
+    description: `Optional paths of up to ${MAX_REFERENCES} images in the project to edit or to use as a reference (PNG, JPEG, WebP, or GIF). To edit an image, pass it here and describe the change in the prompt.`,
+  }),
 })
 
 const GatewayResponse = Schema.Struct({
@@ -103,10 +110,36 @@ export const ImageGenerateTool = Tool.define(
   "image_generate",
   Effect.gen(function* () {
     const http = yield* HttpClient.HttpClient
+
     const auth = yield* Auth.Service
     const fs = yield* FSUtil.Service
     const events = yield* EventV2Bridge.Service
 
+    /**
+     * The reference images, read and checked before anything is spent: in the project (or approved
+     * outside it, as for any read), an image of a type the models take, and not too large to send.
+     */
+    const readReferences = Effect.fn("ImageGenerateTool.readReferences")(function* (
+      paths: readonly string[],
+      directory: string,
+      ctx: Tool.Context,
+    ) {
+      if (paths.length > MAX_REFERENCES) throw new Error(`Pass at most ${MAX_REFERENCES} reference images.`)
+      const references: { path: string; url: string }[] = []
+      for (const given of paths) {
+        const filepath = path.isAbsolute(given) ? given : path.join(directory, given)
+        /* Asked before the file is opened, the same as any read outside the project. */
+        yield* assertExternalDirectoryEffect(ctx, filepath)
+        const bytes = yield* fs.readFile(filepath).pipe(Effect.orElseSucceed(() => undefined))
+        if (!bytes) throw new Error(`Reference image not found: ${given}`)
+        const mime = sniffAttachmentMime(bytes, "application/octet-stream")
+        if (!EXTENSIONS[mime]) throw new Error(`Not a PNG, JPEG, WebP, or GIF image: ${given}`)
+        if (bytes.byteLength > MAX_REFERENCE_BYTES) throw new Error(`Reference image is over 10 MB: ${given}`)
+        references.push({ path: filepath, url: `data:${mime};base64,${Buffer.from(bytes).toString("base64")}` })
+      }
+      return references
+    })
+
     return {
       description: DESCRIPTION,
       parameters: Parameters,
@@ -119,6 +152,7 @@ export const ImageGenerateTool = Tool.define(
           const model = params.model?.trim() || DEFAULT_IMAGE_MODEL
 
           const instance = yield* InstanceState.context
+          const references = yield* readReferences(params.references ?? [], instance.directory, ctx)
           const directory = path.join(instance.directory, "artifacts")
           const stem = slug(params.filename?.trim() || prompt)
 
@@ -145,7 +179,21 @@ export const ImageGenerateTool = Tool.define(
             HttpClientRequest.acceptJson,
             HttpClientRequest.bodyJson({
               model,
-              messages: [{ role: "user", content: prompt }],
+              /*
+               * Text first, then the images, which is the order every vendor documents. A plain string
+               * when there are none, so a text-only request is exactly what it was.
+               */
+              messages: [
+                {
+                  role: "user",
+                  content: references.length
+                    ? [
+                        { type: "text", text: prompt },
+                        ...references.map((reference) => ({ type: "image_url", image_url: { url: reference.url } })),
+                      ]
+                    : prompt,
+                },
+              ],
               modalities: ["image", "text"],
             }),
           )
@@ -187,7 +235,12 @@ export const ImageGenerateTool = Tool.define(
             output: [`Generated ${relative} with ${model}.`, caption ? `The model said: ${caption}` : undefined]
               .filter(Boolean)
               .join("\n"),
-            metadata: { filepath, model, ...(costUsd === undefined ? {} : { costUsd }) },
+            metadata: {
+              filepath,
+              model,
+              ...(references.length ? { references: references.map((reference) => reference.path) } : {}),
+              ...(costUsd === undefined ? {} : { costUsd }),
+            },
             attachments: [
               {
                 type: "file" as const,

packages/redrob/src/tool/image_generate.txt

@@ -6,4 +6,5 @@
 Usage notes:
   - Write the prompt as a full visual description: subject, setting, style, composition, colors, and any text that must appear in the image
   - `filename` is optional and has no extension; a name is derived from the prompt when it is omitted, and an existing file is never overwritten
+  - To edit an existing image, or to keep a person, product, or style consistent, pass the image's path in `references` and describe the change in the prompt; the original is left untouched and the result is saved as a new file
   - Leave `model` unset unless the user names an image model

packages/redrob/test/tool/image_generate.test.ts

@@ -155,6 +155,54 @@ describe("tool.image_generate", () => {
     }),
   )
 
+  it.instance("sends reference images after the prompt, to edit one, and leaves the original alone", () =>
+    Effect.gen(function* () {
+      const test = yield* TestInstance
+      const original = path.join(test.directory, "artifacts", "logo.png")
+      yield* Effect.promise(() => fs.mkdir(path.dirname(original), { recursive: true }))
+      yield* Effect.promise(() => fs.writeFile(original, Buffer.from(PIXEL_PNG, "base64")))
+      const { seen, server } = gateway(imageAnswer)
+      process.env.REDROB_CONSOLE_URL = server.url.toString()
+      try {
+        const result = yield* run({ prompt: "Make the logo blue", references: ["artifacts/logo.png"] })
+        expect(seen[0].body.messages).toEqual([
+          {
+            role: "user",
+            content: [
+              { type: "text", text: "Make the logo blue" },
+              { type: "image_url", image_url: { url: `data:image/png;base64,${PIXEL_PNG}` } },
+            ],
+          },
+        ])
+        expect(result.metadata.references).toEqual([original])
+        expect(result.title.endsWith(path.join("artifacts", "make-the-logo-blue.png"))).toBe(true)
+        expect(Buffer.from(yield* Effect.promise(() => fs.readFile(original))).toString("base64")).toBe(PIXEL_PNG)
+      } finally {
+        void server.stop(true)
+      }
+    }),
+  )
+
+  it.instance("refuses a reference that is missing, not an image, or one too many, before anything is sent", () =>
+    Effect.gen(function* () {
+      const test = yield* TestInstance
+      yield* Effect.promise(() => fs.writeFile(path.join(test.directory, "notes.txt"), "not an image"))
+      const { seen, server } = gateway(imageAnswer)
+      process.env.REDROB_CONSOLE_URL = server.url.toString()
+      try {
+        const missing = yield* run({ prompt: "x", references: ["nope.png"] }).pipe(Effect.exit)
+        expect(failure(missing)).toContain("Reference image not found: nope.png")
+        const text = yield* run({ prompt: "x", references: ["notes.txt"] }).pipe(Effect.exit)
+        expect(failure(text)).toContain("Not a PNG, JPEG, WebP, or GIF image: notes.txt")
+        const many = yield* run({ prompt: "x", references: ["a", "b", "c", "d", "e"] }).pipe(Effect.exit)
+        expect(failure(many)).toContain("at most 4 reference images")
+        expect(seen).toHaveLength(0)
+      } finally {
+        void server.stop(true)
+      }
+    }),
+  )
+
   it.instance("fails with what the model said when it returns no image", () =>
     Effect.gen(function* () {
       const { server } = gateway(() => Response.json({ choices: [{ message: { content: "I cannot draw that." } }] }))