116 lines
3.7 KiB
JavaScript
116 lines
3.7 KiB
JavaScript
import PQueue from 'p-queue';
|
|
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 } 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 meta = findModel(modelId) ?? findModel(DEFAULT_MODEL_ID);
|
|
return PROVIDERS[meta.provider](meta.id);
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// 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('next', () => broadcastQueueStats());
|
|
queue.on('idle', () => broadcastQueueStats());
|
|
|
|
function broadcastQueueStats() {
|
|
broadcast({ type: 'queue_stats', pending: queue.size, running: queue.pending });
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Helpers
|
|
// ---------------------------------------------------------------------------
|
|
async function setStatus(job, status, extra = {}) {
|
|
await job.update({ status, ...extra });
|
|
broadcast({ type: 'job_update', job: job.toJSON() });
|
|
}
|
|
|
|
// ---------------------------------------------------------------------------
|
|
// Main runner
|
|
// ---------------------------------------------------------------------------
|
|
async function runJob(jobId) {
|
|
const job = await Job.findByPk(jobId);
|
|
if (!job) return;
|
|
|
|
await setStatus(job, 'running');
|
|
|
|
try {
|
|
const { text, usage } = await generateText({
|
|
model: resolveModel(job.model),
|
|
messages: [
|
|
{
|
|
role: 'user',
|
|
content: [
|
|
{ type: 'text', text: job.prompt },
|
|
{ type: 'image', image: job.imageDataUrl },
|
|
],
|
|
},
|
|
],
|
|
maxTokens: 1024,
|
|
});
|
|
|
|
await setStatus(job, 'done', {
|
|
result: text,
|
|
inputTokens: usage?.promptTokens ?? null,
|
|
outputTokens: usage?.completionTokens ?? null,
|
|
});
|
|
} catch (err) {
|
|
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));
|
|
}
|