190 lines
13 KiB
JavaScript
190 lines
13 KiB
JavaScript
import { readFile, writeFile, mkdir, rename } from 'node:fs/promises';
|
||
import { createReadStream } from 'node:fs';
|
||
import { Readable, addAbortSignal } from 'node:stream';
|
||
import { homedir } from 'node:os';
|
||
import { join, dirname, isAbsolute } from 'node:path';
|
||
import { fileURLToPath } from 'node:url';
|
||
import { createHash, randomUUID } from 'node:crypto';
|
||
import { validateBaseUrl, validateGeneration, generateImages, editImages, validateEditImage, listModels } from './engine.mjs';
|
||
|
||
export const name = 'imagegen-local';
|
||
export const inject = ['connection', 'tools', 'credentials', 'attachments', 'skills'];
|
||
const root = dirname(dirname(fileURLToPath(import.meta.url)));
|
||
const dataRoot = join(process.env.DSH_HOME || join(homedir(), '.dsh'), 'imagegen');
|
||
const settingsPath = join(dataRoot, 'settings.json');
|
||
const outputs = join(dataRoot, 'outputs');
|
||
const initial = { baseUrl: 'https://api.openai.com/v1', models: ['gpt-image-2', 'gpt-image-1.5'], defaultModel: '' };
|
||
// DSH CredentialKey requires <scope>/<id>, each matching [a-z][a-z0-9-]*.
|
||
// Use the registered plugin name as scope and prefix the digest: hex may start with a digit.
|
||
const keyFor = base => `${name}/endpoint-${createHash('sha256').update(base).digest('hex')}`;
|
||
const json = (value, status = 200) => new Response(JSON.stringify(value), { status, headers: { 'Content-Type': 'application/json', 'Cache-Control': 'no-store' } });
|
||
const schema = { type: 'object', additionalProperties: true };
|
||
const text = value => ({ type: 'text', text: JSON.stringify(value, null, 2) });
|
||
|
||
export async function apply(ctx) {
|
||
let config = { ...initial };
|
||
try { config = validateSettings(JSON.parse(await readFile(settingsPath, 'utf8'))); }
|
||
catch (error) { if (error.code !== 'ENOENT') throw new Error('生图配置无法读取，请检查 .dsh/imagegen/settings.json。'); }
|
||
let queue = Promise.resolve();
|
||
const lifetime = new AbortController();
|
||
ctx.on('dispose', () => lifetime.abort());
|
||
const signalFor = signal => signal ? AbortSignal.any([signal, lifetime.signal]) : lifetime.signal;
|
||
const readKey = async base => (await ctx.credentials.readRecord(keyFor(base)))?.key;
|
||
const view = async () => ({ ...config, configured: Boolean(await readKey(config.baseUrl)), outputDirectory: outputs });
|
||
async function save(input) {
|
||
const next = validateSettings(input);
|
||
if (input.apiKey !== undefined && (typeof input.apiKey !== 'string' || input.apiKey.length > 16384)) throw new Error('API Key 格式无效。');
|
||
if (input.clearKey && input.apiKey?.trim()) throw new Error('不能同时清除和保存 API Key。');
|
||
if (input.apiKey?.trim()) await ctx.credentials.modifyRecord(keyFor(next.baseUrl), async () => ({ kind: 'api-key', key: input.apiKey.trim() }));
|
||
if (input.clearKey === true) await ctx.credentials.deleteRecord(keyFor(next.baseUrl));
|
||
await mkdir(dataRoot, { recursive: true });
|
||
const temp = `${settingsPath}.${randomUUID()}.tmp`;
|
||
await writeFile(temp, JSON.stringify(next, null, 2), { mode: 0o600, flag: 'wx' });
|
||
await rename(temp, settingsPath);
|
||
config = next;
|
||
return view();
|
||
}
|
||
async function generate(input, signal, image) {
|
||
const args = validateGeneration(input);
|
||
const snapshot = { ...config };
|
||
const apiKey = await readKey(snapshot.baseUrl);
|
||
if (!apiKey) throw new Error('请先在 设置 → 生图 中为当前 Base URL 保存 API Key。');
|
||
const options = { baseUrl: snapshot.baseUrl, apiKey, input: args, signal: signalFor(signal) };
|
||
if (image !== undefined) await ctx.attachments.validateImage(validateEditImage(image));
|
||
const result = image === undefined ? await generateImages(options) : await editImages({ ...options, image });
|
||
await mkdir(outputs, { recursive: true });
|
||
const images = [];
|
||
for (const item of result.images) {
|
||
const extension = { 'image/png': 'png', 'image/jpeg': 'jpg', 'image/webp': 'webp', 'image/gif': 'gif' }[item.mediaType];
|
||
const filename = `${Date.now()}-${randomUUID()}.${extension}`;
|
||
const path = join(outputs, filename);
|
||
await writeFile(path, item.data, { flag: 'wx' });
|
||
// Keep original bytes on disk even when DSH normalizes the inline preview.
|
||
const attachment = await ctx.attachments.saveImage({ data: item.data, mediaType: item.mediaType, name: filename });
|
||
images.push({ path, filename, url: `/api/imagegen/file?name=${encodeURIComponent(filename)}`, attachment, ...(item.revisedPrompt ? { revisedPrompt: item.revisedPrompt } : {}) });
|
||
}
|
||
return { model: args.model, prompt: args.prompt, images };
|
||
}
|
||
function route(path, methods, handler, requestBody = 'buffered') {
|
||
ctx.connection.fetch.register({ path, methods, requestBody, fetch: async request => {
|
||
try { return await handler(request); }
|
||
catch (error) {
|
||
// Never include arbitrary provider responses or request headers in a browser error.
|
||
const message = String(error.message || '生图请求失败。');
|
||
return json({ error: message }, error.name === 'AbortError' ? 499 : 400);
|
||
}
|
||
} });
|
||
}
|
||
async function body(request) {
|
||
const raw = await request.text();
|
||
if (raw.length > 100000) throw new Error('请求内容过大。');
|
||
const parsed = JSON.parse(raw);
|
||
if (!parsed || typeof parsed !== 'object' || Array.isArray(parsed)) throw new Error('请求格式无效。');
|
||
return parsed;
|
||
}
|
||
route('/api/imagegen/settings', ['GET', 'POST'], async request => {
|
||
if (request.method === 'GET') return json(await view());
|
||
const input = await body(request);
|
||
const operation = queue.then(() => save(input));
|
||
queue = operation.catch(() => {});
|
||
return json(await operation);
|
||
});
|
||
route('/api/imagegen/models', ['POST'], async request => {
|
||
const snapshot = { ...config };
|
||
const apiKey = await readKey(snapshot.baseUrl);
|
||
if (!apiKey) throw new Error('请先保存 Base URL 和 API Key。');
|
||
return json({ models: await listModels({ baseUrl: snapshot.baseUrl, apiKey, signal: signalFor(request.signal) }) });
|
||
});
|
||
route('/api/imagegen/generate', ['POST'], async request => json(await generate(await body(request), request.signal)));
|
||
route('/api/imagegen/edit', ['POST'], async request => {
|
||
const limit = 21 * 1024 * 1024;
|
||
if (Number(request.headers.get('content-length')) > limit) throw new Error('编辑上传内容过大（图片最大 20 MiB）。');
|
||
if (!request.headers.get('content-type')?.startsWith('multipart/form-data;')) throw new Error('请使用图片上传表单。');
|
||
const bytes = await boundedBytes(request.body, limit, signalFor(request.signal));
|
||
let form;
|
||
try { form = await new Response(bytes, { headers: { 'Content-Type': request.headers.get('content-type') } }).formData(); }
|
||
catch { throw new Error('图片上传表单无效。'); }
|
||
if ([...form.keys()].some(key => !['image', 'input'].includes(key)) || form.getAll('image').length !== 1 || form.getAll('input').length !== 1) throw new Error('请只上传一张图片和一组编辑参数。');
|
||
const image = form.get('image'), raw = form.get('input');
|
||
if (!image || typeof image.arrayBuffer !== 'function' || image.size > 20 * 1024 * 1024) throw new Error('请选择不超过 20 MiB 的图片。');
|
||
if (typeof raw !== 'string' || raw.length > 100000) throw new Error('编辑参数无效。');
|
||
let input;
|
||
try { input = JSON.parse(raw); } catch { throw new Error('编辑参数必须为 JSON。'); }
|
||
return json(await generate(input, request.signal, Buffer.from(await image.arrayBuffer())));
|
||
}, 'streaming');
|
||
route('/api/imagegen/file', ['GET'], async request => {
|
||
const filename = new URL(request.url).searchParams.get('name') || '';
|
||
if (!/^\d+-[0-9a-f-]{36}\.(png|jpg|webp|gif)$/.test(filename)) return json({ error: '图片不存在。' }, 404);
|
||
let bytes;
|
||
try { bytes = await readFile(join(outputs, filename)); } catch (e) { if (e.code === 'ENOENT') return json({ error: '图片不存在。' }, 404); throw e; }
|
||
const ext = filename.split('.').pop();
|
||
return new Response(bytes, { headers: { 'Content-Type': { png: 'image/png', jpg: 'image/jpeg', webp: 'image/webp', gif: 'image/gif' }[ext], 'Cache-Control': 'private, max-age=31536000, immutable', 'X-Content-Type-Options': 'nosniff', 'Content-Disposition': `inline; filename="${filename}"` } });
|
||
});
|
||
ctx.tools.register({
|
||
name: 'imagegen_status',
|
||
description: 'Read image generation configuration (never reveals API keys). Returns Base URL, configured status, model choices and the user-selected default. Use before imagegen_generate; if the user has not specified a model and defaultModel is empty, ask them to choose.',
|
||
parameters: { type: 'object', properties: {}, additionalProperties: false },
|
||
output: { schema, render: (_, value) => [text(value)] },
|
||
execute: () => view(),
|
||
isConcurrencySafe: () => true,
|
||
});
|
||
const generationTool = {
|
||
name: 'imagegen_generate',
|
||
description: 'Generate images using the configured Images API. Requires an exact model explicitly chosen by the user in chat or Settings → 生图 (query imagegen_status for their default). Never silently switch models. Returns saved original file paths and inline image attachments. API calls may incur charges. Text-to-image only; no reference-image editing. Reuse dsh-imagegen skill prompt guidance.',
|
||
parameters: { type: 'object', properties: {
|
||
model: { type: 'string', description: 'Exact user-selected model ID. Required; no automatic fallback.' },
|
||
prompt: { type: 'string', description: 'Image prompt, preserving requested text and constraints.' },
|
||
n: { type: 'integer', description: 'Number of variations of this prompt (1–4), default 1.' },
|
||
size: { type: 'string', description: 'Optional auto or WIDTHxHEIGHT, if supported by this model.' },
|
||
quality: { type: 'string', enum: ['auto', 'low', 'medium', 'high', 'standard', 'hd'] },
|
||
background: { type: 'string', enum: ['auto', 'transparent', 'opaque'] },
|
||
output_format: { type: 'string', enum: ['png', 'jpeg', 'webp'] },
|
||
}, required: ['model', 'prompt'], additionalProperties: false },
|
||
timeoutMs: 240000,
|
||
output: { schema, render: (_, value) => [text({ model: value.model, prompt: value.prompt, images: value.images.map(({ path, revisedPrompt }) => ({ path, revisedPrompt })) }), ...value.images.map(image => ({ type: 'image', attachment: image.attachment }))] },
|
||
execute: (args, exec) => generate(args, exec.signal),
|
||
presentCall: args => ({ card: 'generic', title: `生图 · ${args.model}`, kind: 'fetch', rawInput: args }),
|
||
};
|
||
ctx.tools.register(generationTool);
|
||
ctx.tools.register({
|
||
...generationTool,
|
||
name: 'imagegen_edit',
|
||
description: 'Edit one user-provided image using POST /images/edits. Requires an absolute local image_path (uploaded attachment host path or existing image), exact user-selected model and edit prompt. Query imagegen_status first. PNG/JPEG/WebP up to 20 MiB; no masks or multiple inputs. Sends image bytes to the configured provider and may incur charges. Preserve requested unchanged features; never substitute text-only generation or silently change models.',
|
||
parameters: { ...generationTool.parameters, properties: { ...generationTool.parameters.properties, image_path: { type: 'string', description: 'Absolute host path of the user-authorized input image; use the attachment host path for a chat upload. Never guess a path.' } }, required: ['model', 'prompt', 'image_path'] },
|
||
execute: async (args, exec) => {
|
||
validateGeneration(args);
|
||
if (typeof args.image_path !== 'string' || !isAbsolute(args.image_path)) throw new Error('image_path 必须是图片的本机绝对路径。');
|
||
let image;
|
||
try { image = await boundedBytes(createReadStream(args.image_path), 20 * 1024 * 1024, signalFor(exec.signal)); }
|
||
catch { throw new Error('无法读取编辑图片，或图片超过 20 MiB，或读取已取消。'); }
|
||
return generate(args, exec.signal, image);
|
||
},
|
||
presentCall: args => ({ card: 'generic', title: `编辑图片 · ${args.model}`, kind: 'fetch', rawInput: args }),
|
||
});
|
||
ctx.skills.register({ name: 'dsh-imagegen', description: '使用生图插件生成图片，复用 imagegen 提示词规范，支持自定义 API 与显式模型选择。', source: 'runtime', path: join(root, 'skills', 'SKILL.md'), resourceBase: { kind: 'directory', path: join(root, 'skills') }, content: await readFile(join(root, 'skills', 'SKILL.md'), 'utf8') });
|
||
}
|
||
|
||
async function boundedBytes(source, limit, signal) {
|
||
if (!source) throw new Error('图片内容为空。');
|
||
const stream = typeof source.getReader === 'function' ? Readable.fromWeb(source) : source;
|
||
if (signal) addAbortSignal(signal, stream);
|
||
const chunks = [];
|
||
let length = 0;
|
||
for await (const chunk of stream) {
|
||
length += chunk.length;
|
||
if (length > limit) throw new Error('图片上传超过大小限制。');
|
||
chunks.push(chunk);
|
||
}
|
||
return Buffer.concat(chunks, length);
|
||
}
|
||
|
||
export function validateSettings(input) {
|
||
if (!input || typeof input !== 'object' || Array.isArray(input)) throw new Error('生图设置格式无效。');
|
||
const baseUrl = validateBaseUrl(input.baseUrl);
|
||
if (!Array.isArray(input.models) || input.models.length > 200 || input.models.some(m => typeof m !== 'string' || !m.trim() || m.length > 200)) throw new Error('模型列表无效，请每行填写一个模型 ID（最多 200 个）。');
|
||
const models = [...new Set(input.models.map(m => m.trim()))];
|
||
const defaultModel = typeof input.defaultModel === 'string' ? input.defaultModel.trim() : '';
|
||
if (defaultModel && !models.includes(defaultModel)) throw new Error('默认模型必须位于模型列表中。');
|
||
return { baseUrl, models, defaultModel };
|
||
}
|