331 lines
10 KiB
JavaScript
331 lines
10 KiB
JavaScript
import PQueue from 'p-queue';
|
||
import sharp from 'sharp';
|
||
import { generateText } from 'ai';
|
||
import { createAnthropic } from '@ai-sdk/anthropic';
|
||
import { createOpenAI } from '@ai-sdk/openai';
|
||
import { createGoogleGenerativeAI } from '@ai-sdk/google';
|
||
import { createOpenAICompatible } from '@ai-sdk/openai-compatible';
|
||
import { Job } from '../db/models.js';
|
||
import { broadcast } from '../ws/broadcast.js';
|
||
import { findModel, DEFAULT_MODEL_ID, normalizeModelId } from '../models.js';
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Provider instances
|
||
// ---------------------------------------------------------------------------
|
||
const anthropic = createAnthropic({
|
||
apiKey: process.env.ANTHROPIC_API_KEY,
|
||
});
|
||
|
||
const openai = createOpenAI({
|
||
apiKey: process.env.OPENAI_API_KEY,
|
||
});
|
||
|
||
const google = createGoogleGenerativeAI({
|
||
apiKey: process.env.GOOGLE_API_KEY,
|
||
});
|
||
|
||
// Ollama Cloud exposes an OpenAI-compatible /v1 endpoint.
|
||
// Using @ai-sdk/openai-compatible avoids the local-Ollama schema validation
|
||
// in ollama-ai-provider which requires fields (eval_duration etc.) that
|
||
// Ollama Cloud doesn't return.
|
||
const ollamaCloud = createOpenAICompatible({
|
||
name: 'ollama-cloud',
|
||
baseURL: 'https://ollama.com/v1',
|
||
headers: {
|
||
Authorization: `Bearer ${process.env.OLLAMA_API_KEY ?? ''}`,
|
||
},
|
||
});
|
||
|
||
const PROVIDERS = {
|
||
anthropic: (id) => anthropic(id),
|
||
openai: (id) => openai(id),
|
||
google: (id) => google(id),
|
||
ollama: (id) => ollamaCloud(id),
|
||
};
|
||
|
||
function resolveModel(modelId) {
|
||
const normalized = normalizeModelId(modelId);
|
||
const meta = findModel(normalized) ?? findModel(DEFAULT_MODEL_ID);
|
||
return PROVIDERS[meta.provider](meta.id);
|
||
}
|
||
|
||
function resolveModelMeta(modelId) {
|
||
const normalized = normalizeModelId(modelId);
|
||
const meta = findModel(normalized) ?? findModel(DEFAULT_MODEL_ID);
|
||
return { normalized, meta };
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// p-queue: in-process queue, no external server needed
|
||
// ---------------------------------------------------------------------------
|
||
export const queue = new PQueue({
|
||
concurrency: parseInt(process.env.JOB_CONCURRENCY ?? '3', 10),
|
||
});
|
||
|
||
queue.on('add', () => broadcastQueueStats());
|
||
queue.on('active', () => broadcastQueueStats());
|
||
queue.on('next', () => broadcastQueueStats());
|
||
queue.on('idle', () => broadcastQueueStats());
|
||
|
||
function broadcastQueueStats() {
|
||
broadcast({ type: 'queue_stats', pending: queue.size, running: queue.pending });
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Helpers
|
||
// ---------------------------------------------------------------------------
|
||
|
||
/** Strip image data before sending job info to clients. */
|
||
function sanitizeJob(job) {
|
||
const json = job.toJSON();
|
||
delete json.imageDataUrl;
|
||
return json;
|
||
}
|
||
|
||
async function setStatus(job, status, extra = {}) {
|
||
await job.update({ status, ...extra });
|
||
broadcast({ type: 'job_update', job: sanitizeJob(job) });
|
||
}
|
||
|
||
/** Padding factor applied to each side of the detected bounding box. */
|
||
const BBOX_PADDING = parseFloat(process.env.BBOX_PADDING ?? '0.15');
|
||
|
||
/**
|
||
* Extract the raw image buffer + mime type from a data-URL string.
|
||
*/
|
||
function parseDataUrl(dataUrl) {
|
||
const match = dataUrl.match(/^data:([^;]+);base64,(.+)$/s);
|
||
if (!match) throw new Error('Invalid data URL');
|
||
return { mimeType: match[1], buffer: Buffer.from(match[2], 'base64') };
|
||
}
|
||
|
||
/**
|
||
* Build a data-URL from a buffer + mime type.
|
||
*/
|
||
function toDataUrl(buffer, mimeType) {
|
||
return `data:${mimeType};base64,${buffer.toString('base64')}`;
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Pass 1 — Ask the LLM to locate the subject described in the prompt
|
||
// ---------------------------------------------------------------------------
|
||
const BBOX_SYSTEM_PROMPT = `You are a vision assistant that locates objects in images.
|
||
The user will give you an image and a description of what they are interested in.
|
||
Your job is to identify the main subject described and return the bounding box
|
||
coordinates for that subject as a proportion of the image dimensions (values 0-1).
|
||
|
||
You MUST respond with ONLY a JSON object in this exact format, no other text:
|
||
{
|
||
"subject": "<short description of what you identified>",
|
||
"bbox": {
|
||
"x": <left edge 0-1>,
|
||
"y": <top edge 0-1>,
|
||
"width": <box width 0-1>,
|
||
"height": <box height 0-1>
|
||
}
|
||
}
|
||
|
||
If the prompt doesn't clearly reference a specific subject in the image, or if the
|
||
subject fills most of the image already, return the full image:
|
||
{ "subject": "full image", "bbox": { "x": 0, "y": 0, "width": 1, "height": 1 } }`;
|
||
|
||
async function detectSubjectBbox(model, imageDataUrl, prompt) {
|
||
const { text } = await generateText({
|
||
model,
|
||
system: BBOX_SYSTEM_PROMPT,
|
||
messages: [
|
||
{
|
||
role: 'user',
|
||
content: [
|
||
{
|
||
type: 'text',
|
||
text: `Locate the main subject referenced in this prompt and return its bounding box.\n\nPrompt: "${prompt}"`,
|
||
},
|
||
{ type: 'image', image: imageDataUrl },
|
||
],
|
||
},
|
||
],
|
||
maxTokens: 256,
|
||
});
|
||
|
||
// Parse JSON from the response — tolerate markdown fences
|
||
const cleaned = text.replace(/```json\s*/g, '').replace(/```/g, '').trim();
|
||
const parsed = JSON.parse(cleaned);
|
||
|
||
// Validate and clamp values
|
||
const bbox = parsed.bbox;
|
||
if (
|
||
typeof bbox?.x !== 'number' ||
|
||
typeof bbox?.y !== 'number' ||
|
||
typeof bbox?.width !== 'number' ||
|
||
typeof bbox?.height !== 'number'
|
||
) {
|
||
throw new Error('LLM returned invalid bbox structure');
|
||
}
|
||
|
||
const clamp = (v) => Math.max(0, Math.min(1, v));
|
||
return {
|
||
subject: parsed.subject ?? 'unknown',
|
||
bbox: {
|
||
x: clamp(bbox.x),
|
||
y: clamp(bbox.y),
|
||
width: clamp(bbox.width),
|
||
height: clamp(bbox.height),
|
||
},
|
||
};
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Crop helper — takes a data-URL, bounding box (0-1), and padding factor
|
||
// ---------------------------------------------------------------------------
|
||
async function cropImage(imageDataUrl, bbox, padding = BBOX_PADDING) {
|
||
const { mimeType, buffer } = parseDataUrl(imageDataUrl);
|
||
const metadata = await sharp(buffer).metadata();
|
||
const imgW = metadata.width;
|
||
const imgH = metadata.height;
|
||
|
||
// Convert proportional bbox to pixel values
|
||
let left = bbox.x * imgW;
|
||
let top = bbox.y * imgH;
|
||
let width = bbox.width * imgW;
|
||
let height = bbox.height * imgH;
|
||
|
||
// Add padding (proportional to the box dimensions)
|
||
const padX = width * padding;
|
||
const padY = height * padding;
|
||
left = Math.max(0, left - padX);
|
||
top = Math.max(0, top - padY);
|
||
width = Math.min(imgW - left, width + padX * 2);
|
||
height = Math.min(imgH - top, height + padY * 2);
|
||
|
||
// Round to integers for sharp
|
||
left = Math.round(left);
|
||
top = Math.round(top);
|
||
width = Math.round(width);
|
||
height = Math.round(height);
|
||
|
||
// Guard against degenerate crops
|
||
if (width < 1 || height < 1) {
|
||
return { croppedDataUrl: imageDataUrl, cropPixels: null };
|
||
}
|
||
|
||
const croppedBuf = await sharp(buffer)
|
||
.extract({ left, top, width, height })
|
||
.toBuffer();
|
||
|
||
return {
|
||
croppedDataUrl: toDataUrl(croppedBuf, mimeType),
|
||
cropPixels: { left, top, width, height, imgW, imgH },
|
||
};
|
||
}
|
||
|
||
/**
|
||
* Scrub image data from the job row so we don't persist user photos.
|
||
* Called in both the success and error paths.
|
||
*/
|
||
async function clearImageData(job) {
|
||
await job.update({ imageDataUrl: null });
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Main runner — two-pass: detect bbox → crop → analyse cropped image
|
||
// ---------------------------------------------------------------------------
|
||
async function runJob(jobId) {
|
||
const job = await Job.findByPk(jobId);
|
||
if (!job) return;
|
||
|
||
await setStatus(job, 'running');
|
||
|
||
try {
|
||
resolveModelMeta(job.model);
|
||
const model = resolveModel(job.model);
|
||
|
||
// ---- Pass 1: detect subject bounding box ----
|
||
let bboxResult;
|
||
try {
|
||
bboxResult = await detectSubjectBbox(model, job.imageDataUrl, job.prompt);
|
||
} catch (bboxErr) {
|
||
// If bbox detection fails, fall back to the full image
|
||
console.warn(`Bbox detection failed, using full image: ${bboxErr.message}`);
|
||
bboxResult = {
|
||
subject: 'full image (fallback)',
|
||
bbox: { x: 0, y: 0, width: 1, height: 1 },
|
||
};
|
||
}
|
||
|
||
// Persist the detected bbox for debugging / UI display
|
||
await job.update({
|
||
detectedSubject: bboxResult.subject,
|
||
bboxX: bboxResult.bbox.x,
|
||
bboxY: bboxResult.bbox.y,
|
||
bboxWidth: bboxResult.bbox.width,
|
||
bboxHeight: bboxResult.bbox.height,
|
||
});
|
||
broadcast({ type: 'job_update', job: sanitizeJob(job) });
|
||
|
||
// ---- Crop the image around the detected subject ----
|
||
const isFullImage =
|
||
bboxResult.bbox.x === 0 &&
|
||
bboxResult.bbox.y === 0 &&
|
||
bboxResult.bbox.width === 1 &&
|
||
bboxResult.bbox.height === 1;
|
||
|
||
let imageForAnalysis = job.imageDataUrl;
|
||
|
||
if (!isFullImage) {
|
||
const { croppedDataUrl, cropPixels } = await cropImage(
|
||
job.imageDataUrl,
|
||
bboxResult.bbox,
|
||
);
|
||
imageForAnalysis = croppedDataUrl;
|
||
|
||
if (cropPixels) {
|
||
console.log(
|
||
`Cropped to ${cropPixels.width}×${cropPixels.height} ` +
|
||
`from ${cropPixels.imgW}×${cropPixels.imgH} ` +
|
||
`(subject: "${bboxResult.subject}")`,
|
||
);
|
||
}
|
||
}
|
||
|
||
// ---- Pass 2: run the original prompt against the (cropped) image ----
|
||
const { text, usage } = await generateText({
|
||
model,
|
||
messages: [
|
||
{
|
||
role: 'user',
|
||
content: [
|
||
{ type: 'text', text: job.prompt },
|
||
{ type: 'image', image: imageForAnalysis },
|
||
],
|
||
},
|
||
],
|
||
maxTokens: 1024,
|
||
});
|
||
|
||
// Scrub image before marking done — don't persist photos
|
||
await clearImageData(job);
|
||
|
||
await setStatus(job, 'done', {
|
||
result: text,
|
||
inputTokens: usage?.promptTokens ?? null,
|
||
outputTokens: usage?.completionTokens ?? null,
|
||
});
|
||
} catch (err) {
|
||
// Scrub image even on failure
|
||
await clearImageData(job).catch(() => {});
|
||
|
||
const detail = err?.message
|
||
? `${err.name ?? 'Error'}: ${err.message}${err.cause ? `\nCause: ${JSON.stringify(err.cause, null, 2)}` : ''}`
|
||
: String(err);
|
||
await setStatus(job, 'error', { errorMessage: detail });
|
||
}
|
||
}
|
||
|
||
// ---------------------------------------------------------------------------
|
||
// Public API
|
||
// ---------------------------------------------------------------------------
|
||
export function enqueueJob(jobId) {
|
||
queue.add(() => runJob(jobId));
|
||
}
|