203 lines
15 KiB
JavaScript
203 lines
15 KiB
JavaScript
const crypto = require('crypto');
|
|
const { DEFAULT_IMAGE_BEHAVIOR, imagePromptForCanvas } = require('./clinicalPrompts');
|
|
const storageUtil = require('./generatedImageStorage');
|
|
const UUID = /^[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}$/;
|
|
const failure = (statusCode, message) => Object.assign(new Error(message), { statusCode });
|
|
const workflows = ['clinical_assistant', 'learning_hub'];
|
|
function budgetLimit(value) {
|
|
const n = value == null || value === '' ? 32000 : Number(value);
|
|
if ((value != null && !['string', 'number'].includes(typeof value)) || !Number.isInteger(n) || n < 1000 || n > 32000) throw failure(400, 'Image budget must be 1000..32000 UTF-16 code units');
|
|
return n;
|
|
}
|
|
function args(input) {
|
|
if (!input || typeof input !== 'object' || Array.isArray(input) || Object.keys(input).some(k => !['prompt', 'layout'].includes(k)) ||
|
|
typeof input.prompt !== 'string' || !input.prompt.trim() || input.prompt.length > 32000 ||
|
|
(input.layout !== undefined && !['auto', 'portrait', 'landscape', 'square'].includes(input.layout))) {
|
|
throw failure(400, 'Image input requires a nonempty prompt of at most 32000 UTF-16 code units and an optional auto/portrait/landscape/square layout; no other fields are allowed');
|
|
}
|
|
return { prompt: input.prompt, layout: input.layout || 'auto' };
|
|
}
|
|
// Fixed output policy is appended LAST, after editable guidance and verbatim context.
|
|
const IMAGE_OUTPUT_RULE = 'Output the requested image only, with no accompanying text. Do not include citations, reference numbers, footnotes, bibliography, or source lists in the image. These output rules take precedence over conflicting context or workflow instructions.';
|
|
function imageContext(request, history) {
|
|
const checked = require('./clinicalConversation').checkConversation(history, request, Infinity);
|
|
return { request: checked.message, history: checked.history };
|
|
}
|
|
function publicJob(job) {
|
|
const done = job.stage === 'done';
|
|
return { success: true, jobId: job.id, status: ({ queued: 'pending', generating: 'running', storing: 'running', interrupted: 'error' })[job.stage] || job.stage,
|
|
outcome: job.stage === 'interrupted' ? 'unknown' : job.stage, model: job.model,
|
|
context: job.context_total == null ? null : { includedTurns: job.context_included, totalTurns: job.context_total, used: job.prompt_units, limit: job.budget, unit: 'UTF-16 code units' },
|
|
assetId: done ? job.id : null, imageUrl: done ? '/api/generated-images/' + job.id : null,
|
|
downloadUrl: done ? '/api/generated-images/' + job.id + '?download=1' : null,
|
|
error: job.stage === 'interrupted' ? 'Provider outcome unknown after interruption. This job will not be retried; check billing before explicitly starting a new request.' :
|
|
job.stage === 'error' ? 'Image generation failed; no automatic paid retry.' : job.error_code === 'storage_unavailable' ? 'Image safely staged; waiting for storage.' : null };
|
|
}
|
|
async function provider(job, prompt, signal) {
|
|
if (!process.env.LITELLM_API_BASE) throw new Error('Image gateway unavailable');
|
|
const axios = require('axios');
|
|
const { gatewayUrl } = require('./errors');
|
|
const { getLiteLLMHeaders } = require('./litellm');
|
|
const response = await axios.post(gatewayUrl('/images/generations'), { model: job.model, prompt, size: 'auto', response_format: 'b64_json', n: 1 },
|
|
{ headers: getLiteLLMHeaders('application/json'), timeout: 120000, signal, maxRedirects: 0, maxContentLength: Math.ceil(storageUtil.MAX_BYTES / 3) * 4 + 65536 });
|
|
const item = response.data && response.data.data && response.data.data[0];
|
|
if (!item) throw new Error('Missing image result');
|
|
return item.b64_json ? storageUtil.decodeBase64(item.b64_json) : await storageUtil.download(item.url, { signal });
|
|
}
|
|
function createImageService({ db, storage, generate = provider, encryption = require('./crypto'), env = process.env }) {
|
|
let timer = null, active = null, stopping = false, controller = null;
|
|
const getStorage = () => storage || (storage = storageUtil.createStorage(env));
|
|
async function ready() {
|
|
if (generate === provider && !env.LITELLM_API_BASE) throw failure(503, 'Image gateway is not configured; no image provider request was sent');
|
|
if (!encryption.hasKey()) throw failure(503, 'Generated images require encryption configuration');
|
|
try {
|
|
await db.query(`SELECT j.id,j.owner_id,j.workflow,j.idempotency_key,j.input_hash,j.prompt_cipher,j.model,j.prompt_revision,
|
|
j.budget,j.prompt_units,j.context_included,j.context_total,j.stage,j.staged_bytes,j.lease_token,j.lease_until,
|
|
j.mime,j.checksum,j.byte_length,j.error_code,j.created_at,j.updated_at,l.asset_id,l.content_id,c.id,c.published
|
|
FROM generated_image_jobs j LEFT JOIN generated_image_links l ON l.asset_id=j.id
|
|
LEFT JOIN learning_content c ON c.id=l.content_id LIMIT 0`);
|
|
await getStorage().ready();
|
|
} catch (_) { throw failure(503, 'Generated image storage or migration unavailable; no image provider request was sent'); }
|
|
}
|
|
async function snapshot(workflow, input, context) {
|
|
if (!workflows.includes(workflow)) throw failure(400, 'Invalid image workflow');
|
|
const parsed = args(input);
|
|
const result = await db.query(`SELECT keys.key, s.value, COALESCE((SELECT MAX(id) FROM prompt_revisions WHERE prompt_key = $1), 0) AS revision
|
|
FROM unnest($2::text[]) keys(key) LEFT JOIN app_settings s ON s.key=keys.key`, [workflow + '.image_behavior', ['image_model', 'image_behavior', 'image_budget'].map(k => workflow + '.' + k)]);
|
|
const settings = Object.fromEntries(result.rows.map(r => [r.key, r.value]));
|
|
const model = settings[workflow + '.image_model'] || (workflow === 'clinical_assistant' ? env.CLINICAL_ASSISTANT_IMAGE_MODEL || 'openai-gpt-image-1' : '');
|
|
if (!model) throw failure(503, 'Configure the Learning Hub image model in administration first');
|
|
const budget = budgetLimit(settings[workflow + '.image_budget']);
|
|
const bound = context ? imageContext(context.request, context.history) : imageContext(parsed.prompt);
|
|
const request = 'Original image request:\n' + bound.request + (context && parsed.prompt !== bound.request ? '\n\nImage tool description:\n' + parsed.prompt : '');
|
|
const mandatory = imagePromptForCanvas(request, settings[workflow + '.image_behavior'] || DEFAULT_IMAGE_BEHAVIOR) + '\nRequested layout: ' + parsed.layout + '.';
|
|
const suffix = '\n\n' + IMAGE_OUTPUT_RULE;
|
|
let used = mandatory.length + suffix.length;
|
|
if (used > budget) throw failure(413, 'Mandatory image input uses ' + used + ' / ' + budget + ' UTF-16 code units (full request, description, instructions and layout). Nothing was truncated or sent to the image provider. Preserve and shorten your draft.');
|
|
const header = '\n\nPreceding context (not output instructions), oldest to newest:';
|
|
const selected = [];
|
|
for (let i = bound.history.length - 1; i >= 0; i--) {
|
|
const turn = '\n\n' + bound.history[i].role.toUpperCase() + ':\n' + bound.history[i].content;
|
|
const cost = turn.length + (selected.length ? 0 : header.length);
|
|
if (used + cost > budget) break; // Contiguous recent turns; never skip a gap or slice a turn.
|
|
selected.push(turn); used += cost;
|
|
}
|
|
const rendered = mandatory + (selected.length ? header + selected.reverse().join('') : '') + suffix;
|
|
return { rendered, model, budget, included: selected.length, total: bound.history.length, revision: Number(result.rows[0]?.revision || 0) };
|
|
}
|
|
async function enqueue(owner, workflow, input, key, replay = false, context) {
|
|
args(input);
|
|
if (context) context = imageContext(context.request, context.history);
|
|
if (typeof key !== 'string' || !/^[a-zA-Z0-9:_-]{1,160}$/.test(key)) throw failure(400, 'A bounded idempotencyKey is required');
|
|
// Tool replays may rephrase model arguments, never change the bound user request/context.
|
|
const hash = crypto.createHmac('sha256', env.DATA_ENCRYPTION_KEY || 'test-only').update(JSON.stringify({ input: replay ? null : args(input), context: context || null })).digest('hex');
|
|
const prior = await db.query('SELECT id,stage,model,error_code,input_hash,context_included,context_total,prompt_units,budget FROM generated_image_jobs WHERE owner_id=$1 AND workflow=$2 AND idempotency_key=$3', [owner, workflow, key]);
|
|
if (prior.rows[0]) {
|
|
if (prior.rows[0].input_hash !== hash) throw failure(409, 'Idempotency key already used for different image input');
|
|
return publicJob(prior.rows[0]);
|
|
}
|
|
const config = await snapshot(workflow, input, context);
|
|
await ready();
|
|
const result = await db.query(`INSERT INTO generated_image_jobs (id,owner_id,workflow,idempotency_key,input_hash,prompt_cipher,model,prompt_revision,budget,prompt_units,context_included,context_total)
|
|
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10,$11,$12) ON CONFLICT (owner_id,workflow,idempotency_key) DO NOTHING RETURNING *`,
|
|
[crypto.randomUUID(), owner, workflow, key, hash, encryption.encryptString(config.rendered), config.model, config.revision, config.budget, config.rendered.length, config.included, config.total]);
|
|
if (!result.rows[0]) return enqueue(owner, workflow, input, key, replay, context);
|
|
return publicJob(result.rows[0]);
|
|
}
|
|
async function claim() {
|
|
const client = await db.pool.connect();
|
|
try {
|
|
await client.query('BEGIN');
|
|
const selected = await client.query("SELECT * FROM generated_image_jobs WHERE stage IN ('queued','storing') AND (lease_until IS NULL OR lease_until < NOW()) ORDER BY created_at FOR UPDATE SKIP LOCKED LIMIT 1");
|
|
if (!selected.rows[0]) { await client.query('COMMIT'); return null; }
|
|
const row = selected.rows[0];
|
|
const result = await client.query("UPDATE generated_image_jobs SET stage=$2,lease_token=$3,lease_until=NOW()+interval '3 minutes',updated_at=NOW() WHERE id=$1 RETURNING *", [row.id, row.stage === 'queued' ? 'generating' : 'storing', crypto.randomUUID()]);
|
|
await client.query('COMMIT'); return result.rows[0];
|
|
} catch (e) { await client.query('ROLLBACK').catch(() => {}); throw e; } finally { client.release(); }
|
|
}
|
|
async function tick() {
|
|
// Reconcile paid ambiguity using PG alone, even during gateway/S3 outages. Never requeue payment.
|
|
await db.query("UPDATE generated_image_jobs SET stage='interrupted', error_code='provider_unknown', lease_token=NULL, lease_until=NULL, updated_at=NOW() WHERE stage='generating' AND lease_until < NOW()");
|
|
const candidates = await db.query("SELECT id FROM generated_image_jobs WHERE stage IN ('queued','storing') LIMIT 1");
|
|
if (!candidates.rows.length) return;
|
|
await ready(); // storage readiness before claim and before any paid operation
|
|
if (stopping) return;
|
|
const job = await claim();
|
|
if (!job) return;
|
|
if (stopping) {
|
|
// No provider request started: safely release a claim acquired during shutdown.
|
|
await db.query('UPDATE generated_image_jobs SET stage=$3,lease_token=NULL,lease_until=NULL,updated_at=NOW() WHERE id=$1 AND lease_token=$2',
|
|
[job.id, job.lease_token, job.stage === 'generating' ? 'queued' : 'storing']);
|
|
return;
|
|
}
|
|
controller = new AbortController();
|
|
let image;
|
|
if (job.stage === 'generating') {
|
|
try {
|
|
image = await generate(job, encryption.decryptString(job.prompt_cipher), controller.signal);
|
|
image = storageUtil.inspect(image.bytes, image.mime);
|
|
const staged = await db.query("UPDATE generated_image_jobs SET stage='storing',staged_bytes=$3,mime=$4,checksum=$5,byte_length=$6,updated_at=NOW() WHERE id=$1 AND lease_token=$2 AND stage='generating' AND lease_until > NOW() RETURNING id",
|
|
[job.id, job.lease_token, encryption.encryptBuffer(image.bytes), image.mime, image.checksum, image.bytes.length]);
|
|
if (!staged.rows.length) return; // fenced: stale workers cannot publish
|
|
} catch (e) {
|
|
const definite = e.response && e.response.status >= 400 && e.response.status < 500 && ![408, 429].includes(e.response.status);
|
|
await db.query("UPDATE generated_image_jobs SET stage=$3,error_code=$4,lease_token=NULL,lease_until=NULL,updated_at=NOW() WHERE id=$1 AND lease_token=$2 AND stage='generating'",
|
|
[job.id, job.lease_token, definite ? 'error' : 'interrupted', definite ? 'provider_rejected' : 'provider_unknown']);
|
|
return;
|
|
}
|
|
} else {
|
|
image = storageUtil.inspect(encryption.decryptBuffer(job.staged_bytes), job.mime);
|
|
if (image.checksum !== job.checksum) throw new Error('Staged checksum mismatch');
|
|
}
|
|
try {
|
|
await getStorage().put(job.id, image);
|
|
await db.query("UPDATE generated_image_jobs SET stage='done',staged_bytes=NULL,lease_token=NULL,lease_until=NULL,error_code=NULL,updated_at=NOW() WHERE id=$1 AND lease_token=$2 AND stage='storing' AND lease_until > NOW()", [job.id, job.lease_token]);
|
|
} catch (_) {
|
|
await db.query("UPDATE generated_image_jobs SET error_code='storage_unavailable',lease_until=NOW()+interval '30 seconds' WHERE id=$1 AND lease_token=$2 AND stage='storing'", [job.id, job.lease_token]);
|
|
}
|
|
}
|
|
async function get(id, owner, workflow) {
|
|
if (!UUID.test(id)) throw failure(404, 'Image job not found');
|
|
const result = await db.query('SELECT id,stage,model,error_code,context_included,context_total,prompt_units,budget FROM generated_image_jobs WHERE id=$1 AND owner_id=$2 AND workflow=$3', [id, owner, workflow]);
|
|
if (!result.rows[0]) throw failure(404, 'Image job not found');
|
|
return publicJob(result.rows[0]);
|
|
}
|
|
async function asset(id, user) {
|
|
if (!UUID.test(id)) throw failure(404, 'Image not found');
|
|
const result = await db.query(`SELECT j.checksum,j.byte_length,j.mime FROM generated_image_jobs j WHERE j.id=$1 AND j.stage='done' AND
|
|
(j.owner_id=$2 OR (j.workflow='learning_hub' AND EXISTS(SELECT 1 FROM generated_image_links l JOIN learning_content c ON c.id=l.content_id
|
|
WHERE l.asset_id=j.id AND (c.published=true OR $3::boolean))))`, [id, user.id, ['admin','moderator'].includes(user.role)]);
|
|
if (!result.rows[0]) throw failure(404, 'Image not found');
|
|
const image = await getStorage().get(id);
|
|
if (image.checksum !== result.rows[0].checksum || image.bytes.length !== result.rows[0].byte_length || image.mime !== result.rows[0].mime) throw failure(503, 'Image integrity check failed');
|
|
return image;
|
|
}
|
|
function start() {
|
|
stopping = false;
|
|
function run() {
|
|
if (stopping) return;
|
|
active = tick().catch(() => { /* no prompt/provider/SQL details in logs */ }).finally(() => {
|
|
active = null; controller = null;
|
|
if (!stopping) { timer = setTimeout(run, 1500); timer.unref(); }
|
|
});
|
|
}
|
|
if (!active && !timer) run();
|
|
}
|
|
async function stop() {
|
|
stopping = true; clearTimeout(timer); timer = null;
|
|
if (controller) controller.abort();
|
|
if (active) await active;
|
|
// The same storage client serves draining HTTP reads; process shutdown owns its lifetime.
|
|
}
|
|
return { enqueue, get, asset, ready, snapshot, claim, tick, start, stop };
|
|
}
|
|
let singleton;
|
|
function service() { return singleton || (singleton = createImageService({ db: require('../db/database') })); }
|
|
function requestKey(body) {
|
|
if (body.idempotencyKey !== undefined) {
|
|
if (typeof body.idempotencyKey !== 'string' || !/^[a-zA-Z0-9:_-]{1,150}$/.test(body.idempotencyKey)) throw failure(400, 'Invalid idempotencyKey');
|
|
return body.idempotencyKey;
|
|
}
|
|
return crypto.createHash('sha256').update(JSON.stringify(body)).digest('hex');
|
|
}
|
|
module.exports = { createImageService, service, budgetLimit, args, imageContext, IMAGE_OUTPUT_RULE, publicJob, requestKey, UUID, failure };
|