144 lines
4.8 KiB
TypeScript
144 lines
4.8 KiB
TypeScript
import { Worker, Job } from 'bullmq';
|
|
import { getRedisConnection } from '../module/redis.ts';
|
|
import { Prompt, pbService, ai } from '../index.ts';
|
|
import type { ImageCollection } from '../services/pb.service.ts';
|
|
import { updateItemStatus } from '../routes/image-update.ts';
|
|
import { Queue } from 'bullmq';
|
|
import { notify } from '@/module/logger.ts';
|
|
import { addImageGenerateJob } from './image-creator.job.ts';
|
|
|
|
export const PERFECT_PROMPT_JOB = 'perfect-prompt';
|
|
|
|
// 状态常量
|
|
export const PerfectPromptStatus = {
|
|
PENDING: '提示词优化中' as const,
|
|
PLANNING: '计划中' as const,
|
|
FAILED: '失败' as const,
|
|
};
|
|
|
|
// 最大重试次数
|
|
const MAX_RETRIES = 3;
|
|
|
|
export interface PerfectPromptJobData {
|
|
itemId: string;
|
|
prompt: string;
|
|
collectionName?: string;
|
|
data?: Record<string, any>;
|
|
}
|
|
|
|
// 优化提示词的模板
|
|
const DEFAULT_PERFECT_PROMPT = `请你将以下提示词进行完善，使其更加详细和具体，适合用于生成高质量的像素艺术图像。要求如下：
|
|
1. 只返回完善后的提示词，不要包含任何多余的内容或解释。
|
|
2. 确保提示词专注于像素艺术风格，包括但不限于像素化角色、场景和物体的描述。
|
|
3. 使用具体的细节来增强提示词的表现力，例如颜色、构图、光影效果等。
|
|
4. 避免使用与像素艺术无关的术语或描述。
|
|
5. 保持提示词的简洁性，避免过于冗长，但要确保信息量充足。
|
|
6. 如果需要颜色，需要整个图像的颜色更少的描述，而不是复杂的颜色细节, 背景默认纯蓝色。
|
|
7. 使用中文进行描述。
|
|
`;
|
|
|
|
/**
|
|
* 单独添加优化提示词任务
|
|
*/
|
|
export async function addPerfectPromptJob(item: ImageCollection): Promise<void> {
|
|
const connection = getRedisConnection();
|
|
const queue = new Queue(PERFECT_PROMPT_JOB, { connection });
|
|
|
|
const jobData: PerfectPromptJobData = {
|
|
itemId: item.id,
|
|
prompt: item.description || item.summary || item.title || '',
|
|
collectionName: pbService.collectionName,
|
|
};
|
|
|
|
await queue.add(PERFECT_PROMPT_JOB, jobData, {
|
|
attempts: MAX_RETRIES,
|
|
backoff: {
|
|
type: 'exponential',
|
|
delay: 2000,
|
|
},
|
|
removeOnComplete: 100,
|
|
removeOnFail: 100,
|
|
});
|
|
|
|
await updateItemStatus(item.id, PerfectPromptStatus.PENDING);
|
|
await queue.close();
|
|
}
|
|
|
|
/**
|
|
* 运行优化提示词 worker
|
|
*/
|
|
export async function runPerfectPromptWorker(): Promise<void> {
|
|
const connection = getRedisConnection();
|
|
// 获取环境变量中的 API key
|
|
const worker = new Worker(
|
|
PERFECT_PROMPT_JOB,
|
|
async (job: Job<PerfectPromptJobData>) => {
|
|
const { itemId, prompt } = job.data;
|
|
const attemptsMade = job.attemptsMade;
|
|
console.log(`[PerfectPrompt] Processing item: ${itemId}, attempt: ${attemptsMade + 1}/${MAX_RETRIES}`);
|
|
try {
|
|
if (!prompt) {
|
|
throw new Error('Prompt is empty');
|
|
}
|
|
const promptTool = new Prompt({ perfectPrompt: DEFAULT_PERFECT_PROMPT });
|
|
await ai.chat([
|
|
{
|
|
role: 'user',
|
|
content: promptTool.perfect(prompt),
|
|
},
|
|
]);
|
|
const perfectText = promptTool.clearPerfectTags(ai.responseText);
|
|
|
|
if (!perfectText) {
|
|
throw new Error('Generated perfect prompt is empty');
|
|
}
|
|
|
|
console.log(`[PerfectPrompt] Perfect prompt generated for item: ${itemId}`);
|
|
|
|
// 更新状态为已完成，并保存优化后的提示词
|
|
await updateItemStatus(itemId, PerfectPromptStatus.PLANNING, {
|
|
description: perfectText,
|
|
});
|
|
// 任务完成，把任务抛给下一个图片生成队列
|
|
const item = await pbService.collection.getOne(itemId);
|
|
if (item) {
|
|
addImageGenerateJob(item)
|
|
}
|
|
|
|
return { success: true, perfectPrompt: perfectText };
|
|
} catch (error: any) {
|
|
console.error(`[PerfectPrompt] Error: ${error.message}`);
|
|
|
|
// 重试次数用尽，标记为失败
|
|
if (job.attemptsMade >= MAX_RETRIES - 1) {
|
|
await updateItemStatus(itemId, PerfectPromptStatus.FAILED);
|
|
}
|
|
|
|
throw error;
|
|
}
|
|
},
|
|
{
|
|
connection,
|
|
concurrency: 2,
|
|
lockDuration: 60000 * 5, // 锁持续时间 5分钟
|
|
stalledInterval: 30000, // 每30秒检查一次 stalled
|
|
}
|
|
);
|
|
|
|
worker.on('completed', (job) => {
|
|
console.log(`[PerfectPrompt] Job completed: ${job.id}`);
|
|
});
|
|
|
|
worker.on('failed', (job, err) => {
|
|
console.error(`[PerfectPrompt] Job failed: ${job?.id}, error: ${err.message}`);
|
|
notify.notify(`[PerfectPrompt] \nJob failed: ${job?.id}, error: ${err.message}\n Job data: ${JSON.stringify(job?.data)}`);
|
|
if (job && job.attemptsMade >= MAX_RETRIES - 1) {
|
|
worker.close();
|
|
notify.notify(`[PerfectPrompt] Worker stopped after reaching max retries for item ${job.data.itemId}`);
|
|
}
|
|
});
|
|
|
|
console.log('[PerfectPrompt] Worker started');
|
|
}
|
|
|