@pi-unipi/image 2.2.5 → 2.4.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "@pi-unipi/image",
3
- "version": "2.2.5",
3
+ "version": "2.4.0",
4
4
  "description": "Image generation and image recognition tools for the Pi coding agent",
5
5
  "type": "module",
6
6
  "main": "src/index.ts",
@@ -34,7 +34,7 @@
34
34
  "access": "public"
35
35
  },
36
36
  "dependencies": {
37
- "@pi-unipi/core": "2.2.0"
37
+ "@pi-unipi/core": "2.4.0"
38
38
  },
39
39
  "peerDependencies": {
40
40
  "@earendil-works/pi-ai": "^0.80.0",
package/src/generate.ts CHANGED
@@ -131,6 +131,8 @@ export interface GenerateOptions {
131
131
  signal?: AbortSignal;
132
132
  /** Absolute directory for saved images; omit to skip saving. */
133
133
  outputDir?: string;
134
+ /** Source image; when set the request is an edit rather than a generation. */
135
+ inputImage?: { data: string; mimeType: string };
134
136
  now?: Date;
135
137
  /** Injected images collection, for tests. */
136
138
  images?: ImagesModelsLike;
@@ -141,7 +143,7 @@ export interface GenerateOptions {
141
143
  * @throws {Error} with an actionable message when generation fails.
142
144
  */
143
145
  export async function generateImage(options: GenerateOptions): Promise<GenerateResult> {
144
- const { prompt, model, signal, outputDir, now } = options;
146
+ const { prompt, model, signal, outputDir, now, inputImage } = options;
145
147
 
146
148
  if (!prompt.trim()) {
147
149
  throw new Error("A non-empty prompt is required.");
@@ -154,19 +156,19 @@ export async function generateImage(options: GenerateOptions): Promise<GenerateR
154
156
  );
155
157
  }
156
158
 
157
- // pi-ai's images collection carries its own provider set (currently only
158
- // `openrouter`) and is entirely separate from pi's chat model registry.
159
- // A chat provider registered by another extension can therefore list image
160
- // models that generation cannot actually drive pi-ai answers with a bare
161
- // "Unknown provider: x". Detect that here and say something useful.
159
+ // pi-ai's images collection has its own provider set, separate from pi's
160
+ // chat registry. `registerRegistryImageProviders()` bridges pi's providers
161
+ // in, but a model may still name a provider with no image route at all —
162
+ // pi-ai would answer with a bare "Unknown provider: x", so say something
163
+ // useful instead.
162
164
  if (!providerCanGenerate(imagesApi, model.provider)) {
163
165
  const supported = supportedProviders(imagesApi);
164
166
  throw new Error(
165
- `Provider "${model.provider}" cannot generate images.\n` +
166
- `→ Image generation is served by: ${supported.join(", ") || "openrouter"}.\n` +
167
- `→ "${model.provider}" is a chat provider; its image models are listed for ` +
168
- `recognition and reference, but generation must go through a supported provider.\n` +
169
- "→ Pick one with /unipi:image-settings.",
167
+ `Provider "${model.provider}" has no image-generation route.\n` +
168
+ `→ Available: ${supported.join(", ") || "openrouter"}.\n` +
169
+ " Providers are bridged from pi automatically; one without a baseUrl " +
170
+ "or an API key cannot be used.\n" +
171
+ "→ Pick another with /unipi:image-settings.",
170
172
  );
171
173
  }
172
174
 
@@ -174,7 +176,13 @@ export async function generateImage(options: GenerateOptions): Promise<GenerateR
174
176
  // a bare OPENROUTER_API_KEY still works.
175
177
  let apiKey: string | undefined;
176
178
  try {
177
- apiKey = (await imagesApi.getAuth(model))?.apiKey;
179
+ // pi-ai resolves to an `AuthResult`, i.e. `{ auth: { apiKey } }`. Older
180
+ // shapes put the key at the top level, so accept both — reading only one
181
+ // fails silently and looks like a missing credential.
182
+ const resolvedAuth = (await imagesApi.getAuth(model)) as
183
+ | { apiKey?: string; auth?: { apiKey?: string } }
184
+ | undefined;
185
+ apiKey = resolvedAuth?.auth?.apiKey ?? resolvedAuth?.apiKey;
178
186
  } catch {
179
187
  // Reported as a missing key below.
180
188
  }
@@ -183,14 +191,25 @@ export async function generateImage(options: GenerateOptions): Promise<GenerateR
183
191
  if (!apiKey) {
184
192
  throw new Error(
185
193
  `No API key for provider "${model.provider}".\n` +
186
- `→ Add one with /login, or set the provider's API key environment variable.\n` +
187
- `→ Image models are served through OpenRouter: https://openrouter.ai/keys`,
194
+ "→ Sign in with /login, or set the provider's API key environment variable.\n" +
195
+ `→ Expected environment variable: ` +
196
+ `${model.provider.toUpperCase().replace(/[^A-Z0-9]/g, "_")}_API_KEY`,
188
197
  );
189
198
  }
190
199
 
200
+ const input: Array<{ type: string; text?: string; data?: string; mimeType?: string }> =
201
+ [{ type: "text", text: prompt }];
202
+ if (inputImage) {
203
+ input.push({
204
+ type: "image",
205
+ data: inputImage.data,
206
+ mimeType: inputImage.mimeType,
207
+ });
208
+ }
209
+
191
210
  const result = (await imagesApi.generateImages(
192
211
  model,
193
- { input: [{ type: "text", text: prompt }] },
212
+ { input } as { input: Array<{ type: string; text?: string }> },
194
213
  { apiKey, ...(signal ? { signal } : {}) },
195
214
  )) as AssistantImagesLike;
196
215
 
package/src/index.ts CHANGED
@@ -21,6 +21,7 @@ import {
21
21
  import { registerImageCommands } from "./commands.js";
22
22
  import { registerImageTools } from "./tools.js";
23
23
  import { listImageGenModels, listVisionModels, type ChatModelRegistry } from "./models.js";
24
+ import { registerRegistryImageProviders } from "./register-providers.js";
24
25
  import { loadConfig } from "./settings.js";
25
26
 
26
27
  const VERSION = getPackageVersion(dirname(fileURLToPath(import.meta.url)));
@@ -43,6 +44,13 @@ export default function (pi: ExtensionAPI) {
43
44
  pi.on("session_start", async (_event, ctx) => {
44
45
  const config = loadConfig();
45
46
 
47
+ // Bridge pi's configured providers into pi-ai's images collection up front,
48
+ // so the settings picker and the info screen see them without a prior
49
+ // image_generate call. Best-effort: never block session start.
50
+ void registerRegistryImageProviders(
51
+ (ctx as unknown as { modelRegistry?: ChatModelRegistry }).modelRegistry,
52
+ ).catch(() => undefined);
53
+
46
54
  const tools: string[] = [];
47
55
  if (config.generate.enabled) tools.push(IMAGE_TOOLS.GENERATE);
48
56
  if (config.recognize.enabled) tools.push(IMAGE_TOOLS.RECOGNIZE);
package/src/models.ts CHANGED
@@ -211,12 +211,40 @@ export function listRegistryImageGenModels(
211
211
  provider: candidate.provider,
212
212
  name: candidate.name,
213
213
  api: candidate.api ?? "",
214
+ // Carry the endpoint through. The generic images adapter POSTs to
215
+ // `{baseUrl}/images/generations`, and this is the only place the
216
+ // registry's baseUrl is available — dropping it here surfaces later as
217
+ // "No baseUrl for image model ..." once generation is attempted.
218
+ ...(candidate.baseUrl ? { baseUrl: candidate.baseUrl } : {}),
214
219
  ...(candidate.output ? { output: candidate.output } : {}),
215
220
  });
216
221
  }
217
222
  return out;
218
223
  }
219
224
 
225
+ /**
226
+ * Find a provider's API endpoint in pi's registry.
227
+ *
228
+ * Needed because a model can reach generation without one: a user-typed
229
+ * "provider/model-id" is accepted at face value by `asExplicitModelRef`, and
230
+ * carries no baseUrl of its own.
231
+ */
232
+ export function findProviderBaseUrl(
233
+ registry: ChatModelRegistry | undefined,
234
+ provider: string,
235
+ ): string | undefined {
236
+ if (!registry) return undefined;
237
+ try {
238
+ const models = (registry.getAvailable?.() ?? registry.getAll()) as Array<{
239
+ provider?: string;
240
+ baseUrl?: string;
241
+ }>;
242
+ return models.find((m) => m?.provider === provider && m.baseUrl)?.baseUrl;
243
+ } catch {
244
+ return undefined;
245
+ }
246
+ }
247
+
220
248
  /**
221
249
  * Every selectable generation model: pi-ai's built-in catalog plus anything
222
250
  * contributed by registered providers, de-duplicated by "provider/id".
@@ -0,0 +1,282 @@
1
+ /**
2
+ * @pi-unipi/image — Generic OpenAI-compatible images adapter
3
+ *
4
+ * ONE adapter for every provider, rather than per-provider code. It speaks the
5
+ * OpenAI `POST {baseUrl}/images/generations` shape, which every gateway we have
6
+ * tested implements (OpenAI itself, OpenRouter, and OmniRoute's fan-out to
7
+ * openrouter/antigravity/codex/fal-ai backends).
8
+ *
9
+ * Why not pi-ai's built-in `api/openrouter-images`?
10
+ * Despite the name it drives `chat.completions` with `modalities:["image"]`.
11
+ * Gateways that do not implement that extension answer HTTP 200 with the model
12
+ * *narrating* the image ("Here's the image with the circle changed…") while
13
+ * silently dropping `message.images`. That is invisible data loss, so we use
14
+ * the dedicated images endpoint instead.
15
+ *
16
+ * Editing rides the same endpoint: `POST /images/generations` with an `image`
17
+ * array. `/images/edits` (multipart) is NOT used — gateways reject it for most
18
+ * providers ("Image edit is not supported for built-in provider ...").
19
+ */
20
+
21
+ import type { ImageGenModel } from "./models.js";
22
+
23
+ /** pi-ai's `ImagesContext` input parts. */
24
+ export interface ImagesInputPart {
25
+ type: string;
26
+ text?: string;
27
+ data?: string;
28
+ mimeType?: string;
29
+ }
30
+
31
+ export interface ImagesContextLike {
32
+ input: ImagesInputPart[];
33
+ }
34
+
35
+ export interface ImagesOptionsLike {
36
+ apiKey?: string;
37
+ signal?: AbortSignal;
38
+ headers?: Record<string, string | null>;
39
+ timeoutMs?: number;
40
+ /** Injectable fetch, for tests. */
41
+ fetchImpl?: typeof fetch;
42
+ }
43
+
44
+ /** pi-ai's `AssistantImages`. */
45
+ export interface AssistantImagesLike {
46
+ api: string;
47
+ provider: string;
48
+ model: string;
49
+ output: Array<{ type: string; text?: string; data?: string; mimeType?: string }>;
50
+ stopReason: "stop" | "error" | "aborted";
51
+ errorMessage?: string;
52
+ timestamp: number;
53
+ }
54
+
55
+ /** Images are slow — a minute is not unusual for a large model. */
56
+ const DEFAULT_TIMEOUT_MS = 240_000;
57
+
58
+ /**
59
+ * One returned image, normalized.
60
+ *
61
+ * Gateways disagree on the item shape; all three observed forms are accepted:
62
+ * - `{ b64_json, media_type }` — openrouter/* (note `media_type`, not `mimeType`)
63
+ * - `{ b64_json, revised_prompt }` — antigravity/*
64
+ * - `{ url: "data:image/png;base64,…" }` — codex/*
65
+ * A plain http(s) `url` is also tolerated and reported as text, since we cannot
66
+ * inline bytes we did not fetch.
67
+ */
68
+ interface RawImageItem {
69
+ b64_json?: unknown;
70
+ url?: unknown;
71
+ media_type?: unknown;
72
+ mime_type?: unknown;
73
+ mimeType?: unknown;
74
+ revised_prompt?: unknown;
75
+ }
76
+
77
+ function asString(value: unknown): string | undefined {
78
+ return typeof value === "string" && value.length > 0 ? value : undefined;
79
+ }
80
+
81
+ /** Pull `{ data, mimeType }` out of one response item, whatever its shape. */
82
+ export function normalizeImageItem(
83
+ item: RawImageItem,
84
+ ): { data: string; mimeType: string } | { text: string } | null {
85
+ const declared =
86
+ asString(item.media_type) ?? asString(item.mime_type) ?? asString(item.mimeType);
87
+
88
+ const b64 = asString(item.b64_json);
89
+ if (b64) return { data: b64, mimeType: declared ?? "image/png" };
90
+
91
+ const url = asString(item.url);
92
+ if (!url) return null;
93
+
94
+ // codex/* returns the bytes as a data: URL rather than b64_json.
95
+ const dataUrl = /^data:([^;,]+)(?:;[^,]*)*,(.*)$/s.exec(url);
96
+ if (dataUrl) {
97
+ const [, mime, payload] = dataUrl;
98
+ if (payload) return { data: payload, mimeType: declared ?? mime ?? "image/png" };
99
+ return null;
100
+ }
101
+
102
+ // A remote URL: surface it rather than silently dropping the result.
103
+ return { text: `Image available at: ${url}` };
104
+ }
105
+
106
+ /** Strip a trailing slash so `${base}/images/generations` is well-formed. */
107
+ function joinUrl(baseUrl: string, suffix: string): string {
108
+ return `${baseUrl.replace(/\/+$/, "")}/${suffix.replace(/^\/+/, "")}`;
109
+ }
110
+
111
+ /**
112
+ * Model ids are sent to the gateway verbatim.
113
+ *
114
+ * Do NOT try to "repair" a doubled-looking segment. OmniRoute genuinely serves
115
+ * `fal-ai/fal-ai/nano-banana-pro` (provider `fal-ai` + model `fal-ai/nano-...`),
116
+ * and rewriting it to `fal-ai/nano-banana-pro` yields a 404. Confusingly the
117
+ * gateway *also* advertises a `fal/...` alias in `/v1/models` that the images
118
+ * endpoint then rejects with "Invalid image model" — an upstream inconsistency
119
+ * we surface rather than guess around, because a wrong guess turns a clear
120
+ * error into a silently different model.
121
+ */
122
+ export function normalizeModelId(id: string): string {
123
+ return id;
124
+ }
125
+
126
+ /** Merge caller headers, dropping keys explicitly suppressed with null. */
127
+ function buildHeaders(
128
+ apiKey: string,
129
+ extra?: Record<string, string | null>,
130
+ ): Record<string, string> {
131
+ const headers: Record<string, string> = {
132
+ Authorization: `Bearer ${apiKey}`,
133
+ "Content-Type": "application/json",
134
+ };
135
+ for (const [key, value] of Object.entries(extra ?? {})) {
136
+ if (value === null) delete headers[key];
137
+ else headers[key] = value;
138
+ }
139
+ return headers;
140
+ }
141
+
142
+ /** Best-effort extraction of a provider error message. */
143
+ function describeError(status: number, statusText: string, body: string): string {
144
+ let detail = body.slice(0, 300);
145
+ try {
146
+ const parsed = JSON.parse(body) as { error?: { message?: string } | string };
147
+ if (typeof parsed.error === "string") detail = parsed.error;
148
+ else if (parsed.error?.message) detail = parsed.error.message;
149
+ } catch {
150
+ // Non-JSON body — the truncated text is the best we have.
151
+ }
152
+ return `${status} ${statusText}${detail ? `: ${detail}` : ""}`;
153
+ }
154
+
155
+ /** Add guidance for the gateway's confusing model-id errors. */
156
+ function annotateModelError(message: string, model: ImageGenModel): string {
157
+ if (/invalid image model|not found|unknown model/i.test(message)) {
158
+ return (
159
+ `${message}\n` +
160
+ `→ Model id sent: "${model.id}" (provider "${model.provider}").\n` +
161
+ "→ Some gateways list aliases they cannot serve. Try the id exactly as it " +
162
+ "appears in the provider's own catalog, or pick another with /unipi:image-settings."
163
+ );
164
+ }
165
+ return message;
166
+ }
167
+
168
+ /**
169
+ * Generate (or edit) images against an OpenAI-compatible endpoint.
170
+ *
171
+ * Satisfies pi-ai's `ProviderImages` interface, so the result is returned —
172
+ * never thrown — with `stopReason: "error"` on failure.
173
+ */
174
+ export async function generateImages(
175
+ model: ImageGenModel,
176
+ context: ImagesContextLike,
177
+ options?: ImagesOptionsLike,
178
+ ): Promise<AssistantImagesLike> {
179
+ const result: AssistantImagesLike = {
180
+ api: model.api || "openai-images",
181
+ provider: model.provider,
182
+ model: model.id,
183
+ output: [],
184
+ stopReason: "stop",
185
+ timestamp: Date.now(),
186
+ };
187
+
188
+ const fetchImpl = options?.fetchImpl ?? fetch;
189
+ const controller = new AbortController();
190
+ const timer = setTimeout(
191
+ () => controller.abort(),
192
+ options?.timeoutMs ?? DEFAULT_TIMEOUT_MS,
193
+ );
194
+ const onAbort = () => controller.abort();
195
+ options?.signal?.addEventListener("abort", onAbort, { once: true });
196
+
197
+ try {
198
+ const apiKey = options?.apiKey;
199
+ if (!apiKey) throw new Error(`No API key for provider: ${model.provider}`);
200
+ if (!model.baseUrl) {
201
+ throw new Error(`No baseUrl for image model ${model.provider}/${model.id}`);
202
+ }
203
+
204
+ // Text parts form the prompt; image parts switch the request into edit mode.
205
+ const prompt = context.input
206
+ .filter((part) => part.type === "text" && part.text)
207
+ .map((part) => part.text as string)
208
+ .join("\n")
209
+ .trim();
210
+
211
+ const images = context.input
212
+ .filter((part) => part.type === "image" && part.data)
213
+ .map((part) => `data:${part.mimeType || "image/png"};base64,${part.data}`);
214
+
215
+ if (!prompt) throw new Error("A non-empty prompt is required.");
216
+
217
+ const body: Record<string, unknown> = {
218
+ model: normalizeModelId(model.id),
219
+ prompt,
220
+ };
221
+ // Only send `image` for edits; some backends reject an empty array.
222
+ if (images.length > 0) body.image = images;
223
+
224
+ const response = await fetchImpl(joinUrl(model.baseUrl, "images/generations"), {
225
+ method: "POST",
226
+ headers: buildHeaders(apiKey, options?.headers),
227
+ body: JSON.stringify(body),
228
+ signal: controller.signal,
229
+ });
230
+
231
+ if (!response.ok) {
232
+ throw new Error(
233
+ describeError(response.status, response.statusText, await response.text()),
234
+ );
235
+ }
236
+
237
+ const payload = (await response.json()) as {
238
+ data?: RawImageItem[];
239
+ error?: { message?: string };
240
+ };
241
+
242
+ if (payload.error?.message) throw new Error(payload.error.message);
243
+
244
+ for (const item of payload.data ?? []) {
245
+ const normalized = normalizeImageItem(item);
246
+ if (!normalized) continue;
247
+ if ("text" in normalized) {
248
+ result.output.push({ type: "text", text: normalized.text });
249
+ } else {
250
+ result.output.push({
251
+ type: "image",
252
+ data: normalized.data,
253
+ mimeType: normalized.mimeType,
254
+ });
255
+ }
256
+ const revised = asString(item.revised_prompt);
257
+ if (revised && revised !== prompt) {
258
+ result.output.push({ type: "text", text: `Revised prompt: ${revised}` });
259
+ }
260
+ }
261
+
262
+ if (result.output.every((part) => part.type !== "image")) {
263
+ throw new Error("The provider returned no image data.");
264
+ }
265
+
266
+ return result;
267
+ } catch (error) {
268
+ const aborted = options?.signal?.aborted || controller.signal.aborted;
269
+ result.stopReason = options?.signal?.aborted ? "aborted" : "error";
270
+ result.errorMessage =
271
+ aborted && !options?.signal?.aborted
272
+ ? "Image request timed out."
273
+ : annotateModelError(
274
+ error instanceof Error ? error.message : String(error),
275
+ model,
276
+ );
277
+ return result;
278
+ } finally {
279
+ clearTimeout(timer);
280
+ options?.signal?.removeEventListener("abort", onAbort);
281
+ }
282
+ }
@@ -0,0 +1,220 @@
1
+ /**
2
+ * @pi-unipi/image — Bridge pi's chat providers into pi-ai's images collection
3
+ *
4
+ * pi-ai ships exactly one image provider (openrouter), so out of the box image
5
+ * generation demands an OpenRouter account even when the user has half a dozen
6
+ * other providers configured. pi's own registry knows those providers and their
7
+ * credentials, so we re-register each one as an *images* provider backed by the
8
+ * single generic OpenAI-compatible adapter.
9
+ *
10
+ * The result: any OpenAI-compatible provider the user configures in pi can
11
+ * generate and edit images with no image-specific setup, and no per-provider
12
+ * code here.
13
+ *
14
+ * ## Why capability detection stays heuristic
15
+ * pi's model registry cannot tell us which models emit images.
16
+ * `provider-composer.ts` builds each registered model from an explicit field
17
+ * list — `{id, name, api, provider, baseUrl, reasoning, input, cost,
18
+ * contextWindow, maxTokens, headers, compat}` — so an extension that attaches
19
+ * `output: ["image"]` has it silently dropped. `ProviderModelConfig` has no
20
+ * `output` field at all. Hence `looksLikeImageGenerator()` name-matching, plus
21
+ * explicit "provider/model-id" entry as the always-available escape hatch.
22
+ */
23
+
24
+ import * as imagesApi from "./openai-images-api.js";
25
+ import {
26
+ getImagesModels,
27
+ listRegistryImageGenModels,
28
+ type ChatModelRegistry,
29
+ type ImageGenModel,
30
+ } from "./models.js";
31
+
32
+ /** pi-ai's `createImagesProvider`, kept structural to avoid type coupling. */
33
+ interface CreateImagesProviderFn {
34
+ (input: {
35
+ id: string;
36
+ name?: string;
37
+ auth: unknown;
38
+ models: readonly ImageGenModel[];
39
+ api: unknown;
40
+ }): unknown;
41
+ }
42
+
43
+ /** The subset of a registry provider we need. */
44
+ interface ProviderLike {
45
+ id: string;
46
+ name?: string;
47
+ baseUrl?: string;
48
+ }
49
+
50
+ let registered = false;
51
+
52
+ /** Reset registration state. Test-only. */
53
+ export function __resetRegistrationForTests(): void {
54
+ registered = false;
55
+ }
56
+
57
+ /**
58
+ * Group discovered generator models by provider, attaching the provider's
59
+ * baseUrl so the adapter knows where to POST.
60
+ */
61
+ export function groupModelsByProvider(
62
+ models: ImageGenModel[],
63
+ baseUrlFor: (provider: string) => string | undefined,
64
+ ): Map<string, { baseUrl: string; models: ImageGenModel[] }> {
65
+ const grouped = new Map<string, { baseUrl: string; models: ImageGenModel[] }>();
66
+
67
+ for (const model of models) {
68
+ const baseUrl = model.baseUrl ?? baseUrlFor(model.provider);
69
+ // Without an endpoint the adapter cannot issue a request; skip rather than
70
+ // register a provider that is guaranteed to fail.
71
+ if (!baseUrl) continue;
72
+
73
+ let entry = grouped.get(model.provider);
74
+ if (!entry) {
75
+ entry = { baseUrl, models: [] };
76
+ grouped.set(model.provider, entry);
77
+ }
78
+ entry.models.push({ ...model, baseUrl });
79
+ }
80
+
81
+ return grouped;
82
+ }
83
+
84
+ /** Read a provider's baseUrl out of pi's registry. */
85
+ function providerBaseUrlLookup(
86
+ registry: ChatModelRegistry,
87
+ ): (provider: string) => string | undefined {
88
+ const cache = new Map<string, string | undefined>();
89
+
90
+ return (provider: string) => {
91
+ if (cache.has(provider)) return cache.get(provider);
92
+
93
+ let baseUrl: string | undefined;
94
+ try {
95
+ const models = (registry.getAvailable?.() ?? registry.getAll()) as Array<{
96
+ provider?: string;
97
+ baseUrl?: string;
98
+ }>;
99
+ baseUrl = models.find((m) => m?.provider === provider && m.baseUrl)?.baseUrl;
100
+ } catch {
101
+ baseUrl = undefined;
102
+ }
103
+
104
+ cache.set(provider, baseUrl);
105
+ return baseUrl;
106
+ };
107
+ }
108
+
109
+ /**
110
+ * Register every pi provider that looks capable of image generation into
111
+ * pi-ai's images collection.
112
+ *
113
+ * Idempotent and best-effort: a failure here must never break the extension,
114
+ * since generation still works for pi-ai's built-in providers.
115
+ *
116
+ * @returns the provider ids registered.
117
+ */
118
+ export async function registerRegistryImageProviders(
119
+ registry: ChatModelRegistry | undefined,
120
+ options?: { force?: boolean },
121
+ ): Promise<string[]> {
122
+ if (!registry) return [];
123
+ if (registered && !options?.force) return [];
124
+
125
+ const images = await getImagesModels();
126
+ if (!images) return [];
127
+
128
+ // `setProvider` is on MutableImagesModels; the built-in collection provides
129
+ // it, but guard in case a future pi-ai hands back an immutable one.
130
+ const mutable = images as unknown as {
131
+ setProvider?: (provider: unknown) => void;
132
+ getProvider?: (id: string) => unknown;
133
+ };
134
+ if (typeof mutable.setProvider !== "function") return [];
135
+
136
+ let createImagesProvider: CreateImagesProviderFn;
137
+ try {
138
+ const mod = (await import("@earendil-works/pi-ai")) as unknown as {
139
+ createImagesProvider?: CreateImagesProviderFn;
140
+ };
141
+ if (typeof mod.createImagesProvider !== "function") return [];
142
+ createImagesProvider = mod.createImagesProvider;
143
+ } catch {
144
+ return [];
145
+ }
146
+
147
+ const discovered = listRegistryImageGenModels(registry);
148
+ if (discovered.length === 0) {
149
+ registered = true;
150
+ return [];
151
+ }
152
+
153
+ const grouped = groupModelsByProvider(discovered, providerBaseUrlLookup(registry));
154
+ const added: string[] = [];
155
+
156
+ for (const [providerId, { models }] of grouped) {
157
+ // Never shadow a provider pi-ai serves natively — its own implementation
158
+ // is better informed than our generic adapter.
159
+ try {
160
+ if (mutable.getProvider?.(providerId)) continue;
161
+ } catch {
162
+ // Treat a lookup failure as "not present" and attempt registration.
163
+ }
164
+
165
+ try {
166
+ const provider = createImagesProvider({
167
+ id: providerId,
168
+ name: providerId,
169
+ models,
170
+ api: imagesApi,
171
+ auth: {
172
+ apiKey: {
173
+ name: `${providerId} API key`,
174
+ // Resolve through pi's own auth storage so the user never logs in
175
+ // twice. `resolve` MUST return an AuthResult (`{ auth: {...} }`);
176
+ // returning a bare key fails silently at request time.
177
+ resolve: async () => {
178
+ const key = await resolveProviderKey(registry, providerId);
179
+ return key ? { auth: { apiKey: key }, source: `pi:${providerId}` } : undefined;
180
+ },
181
+ },
182
+ },
183
+ });
184
+
185
+ mutable.setProvider(provider);
186
+ added.push(providerId);
187
+ } catch {
188
+ // One bad provider must not stop the rest.
189
+ }
190
+ }
191
+
192
+ registered = true;
193
+ return added;
194
+ }
195
+
196
+ /** Resolve a provider key from pi's auth storage, falling back to the env. */
197
+ async function resolveProviderKey(
198
+ registry: ChatModelRegistry,
199
+ provider: string,
200
+ ): Promise<string | undefined> {
201
+ try {
202
+ const key = await registry.getApiKeyForProvider?.(provider);
203
+ if (key) return key;
204
+ } catch {
205
+ // Fall through to the environment.
206
+ }
207
+ const envName = `${provider.toUpperCase().replace(/[^A-Z0-9]/g, "_")}_API_KEY`;
208
+ return process.env[envName] || undefined;
209
+ }
210
+
211
+ /** Provider ids pi-ai can currently generate with, after registration. */
212
+ export function registeredProviderIds(images: {
213
+ getProviders?: () => ReadonlyArray<{ id: string }>;
214
+ }): string[] {
215
+ try {
216
+ return (images.getProviders?.() ?? []).map((p) => p.id);
217
+ } catch {
218
+ return [];
219
+ }
220
+ }
package/src/tools.ts CHANGED
@@ -12,7 +12,9 @@ import { IMAGE_TOOLS } from "@pi-unipi/core";
12
12
 
13
13
  import { generateImage } from "./generate.js";
14
14
  import { loadImage } from "./image-source.js";
15
+ import { registerRegistryImageProviders } from "./register-providers.js";
15
16
  import {
17
+ findProviderBaseUrl,
16
18
  formatModelRef,
17
19
  listAllImageGenModels,
18
20
  resolveImageGenModel,
@@ -84,12 +86,15 @@ function registerGenerateTool(pi: ExtensionAPI): void {
84
86
  name: IMAGE_TOOLS.GENERATE,
85
87
  label: "Generate Image",
86
88
  description:
87
- "Generate an image from a text prompt using an image model. " +
88
- "The image is returned inline and, when enabled, saved to disk.",
89
+ "Generate an image from a text prompt, or edit an existing image by " +
90
+ "passing `image`. The result is returned inline and, when enabled, saved to disk.",
89
91
  promptSnippet: "Generate an image from a text prompt.",
90
92
  promptGuidelines: [
91
93
  "Use image_generate to create images from a text description.",
92
94
  "Write a detailed prompt — subject, style, composition and lighting all help.",
95
+ "Pass `image` to edit an existing image instead of generating a new one.",
96
+ "Editing regenerates the whole image, so unmentioned details may change.",
97
+ "Describe what you DO want; negation is unreliable in image models.",
93
98
  "Omit model to use the one configured in /unipi:image-settings.",
94
99
  "Generated images cost money per call; do not regenerate without being asked.",
95
100
  ],
@@ -97,6 +102,13 @@ function registerGenerateTool(pi: ExtensionAPI): void {
97
102
  prompt: Type.String({
98
103
  description: "Description of the image to generate. Be specific.",
99
104
  }),
105
+ image: Type.Optional(
106
+ Type.String({
107
+ description:
108
+ "Source image to edit: a local file path, data: URL, or base64 data. " +
109
+ "When set, the model edits this image instead of generating from scratch.",
110
+ }),
111
+ ),
100
112
  model: Type.Optional(
101
113
  Type.String({
102
114
  description:
@@ -109,22 +121,41 @@ function registerGenerateTool(pi: ExtensionAPI): void {
109
121
  try {
110
122
  const config = loadConfig();
111
123
  const registry = getRegistry(ctx);
124
+ // Bridge pi's own providers into pi-ai's images collection so the user
125
+ // is not forced onto OpenRouter. Idempotent and best-effort.
126
+ await registerRegistryImageProviders(registry);
112
127
  // Include image models contributed by registered providers, so the
113
128
  // tool can resolve anything the settings picker offers.
114
129
  const models = await listAllImageGenModels(registry);
115
130
 
116
131
  const requested = params.model?.trim() || config.generate.model;
117
- const resolved = resolveImageGenModel(requested, models);
118
- if (typeof resolved === "string") return errorResult(resolved);
132
+ const maybeResolved = resolveImageGenModel(requested, models);
133
+ if (typeof maybeResolved === "string") return errorResult(maybeResolved);
134
+
135
+ // A model may arrive without an endpoint — notably a user-typed
136
+ // "provider/model-id", accepted at face value. Fill it in from the
137
+ // registry so the adapter knows where to POST.
138
+ const registryBaseUrl = maybeResolved.baseUrl
139
+ ? undefined
140
+ : findProviderBaseUrl(registry, maybeResolved.provider);
141
+ const resolved = registryBaseUrl
142
+ ? { ...maybeResolved, baseUrl: registryBaseUrl }
143
+ : maybeResolved;
119
144
 
120
145
  // pi-ai resolves image auth from its own credential store; only fall
121
146
  // back to pi's chat-provider key when that comes up empty.
122
147
  const fallbackKey = await resolveApiKey(registry, resolved.provider);
123
148
 
149
+ // An input image switches the request into edit mode.
150
+ const sourceImage = params.image?.trim()
151
+ ? loadImage(params.image, ctx.cwd ?? process.cwd())
152
+ : undefined;
153
+
124
154
  const result = await generateImage({
125
155
  prompt: params.prompt,
126
156
  model: resolved,
127
157
  ...(fallbackKey ? { apiKey: fallbackKey } : {}),
158
+ ...(sourceImage ? { inputImage: sourceImage } : {}),
128
159
  signal,
129
160
  outputDir: config.generate.saveToDisk ? getOutputDir(config) : undefined,
130
161
  });
@@ -134,8 +165,8 @@ function registerGenerateTool(pi: ExtensionAPI): void {
134
165
  .filter((path): path is string => Boolean(path));
135
166
 
136
167
  const summary = [
137
- `Generated ${result.images.length} image${result.images.length === 1 ? "" : "s"} ` +
138
- `with ${formatModelRef(resolved)}.`,
168
+ `${sourceImage ? "Edited" : "Generated"} ${result.images.length} ` +
169
+ `image${result.images.length === 1 ? "" : "s"} with ${formatModelRef(resolved)}.`,
139
170
  saved.length > 0 ? `Saved to:\n${saved.map((p) => ` ${p}`).join("\n")}` : "",
140
171
  config.generate.saveToDisk && saved.length === 0
141
172
  ? "Could not write to the output directory — returning the image inline only."
@@ -23,6 +23,7 @@ import {
23
23
  type ChatModelRegistry,
24
24
  } from "../models.js";
25
25
  import { ImageModelSelectorOverlay, type SelectableModel } from "./model-selector.js";
26
+ import { registerRegistryImageProviders } from "../register-providers.js";
26
27
 
27
28
  const EXIT = "__exit__";
28
29
 
@@ -239,6 +240,10 @@ async function collectModels(
239
240
  .modelRegistry;
240
241
 
241
242
  if (kind === "generate") {
243
+ // Bridge pi's providers in first, so a model the user can actually run is
244
+ // not flagged "no image route" purely because we had not registered it yet.
245
+ await registerRegistryImageProviders(registry);
246
+
242
247
  // Include models from providers registered by other extensions, not just
243
248
  // pi-ai's built-in OpenRouter catalog.
244
249
  const models = await listAllImageGenModels(registry);
@@ -253,7 +258,7 @@ async function collectModels(
253
258
  // so flag them rather than letting the user pick a dead option.
254
259
  unavailable:
255
260
  generating.length > 0 && !generating.includes(m.provider)
256
- ? "cannot generate"
261
+ ? "no image route"
257
262
  : undefined,
258
263
  }));
259
264
  }