253 lines
10 KiB
JavaScript
253 lines
10 KiB
JavaScript
const express = require('express');
|
|
const { v4: uuidv4 } = require('uuid');
|
|
const { generateImage } = require('../lib/generate-image');
|
|
|
|
function createGenerateRoutes(db) {
|
|
const router = express.Router();
|
|
const runningJobs = new Set();
|
|
|
|
async function callAI(prompt, maxTokens = 2000) {
|
|
const aiConfig = {};
|
|
const keys = ['ai_api_key', 'ai_base_url', 'ai_model', 'ai_temperature'];
|
|
for (const key of keys) {
|
|
const row = db.prepare('SELECT value FROM configs WHERE key = ?').get(key);
|
|
aiConfig[key] = row?.value || '';
|
|
}
|
|
if (!aiConfig.ai_api_key) throw new Error('AI API key not configured. Please set it in Settings.');
|
|
|
|
const response = await fetch(`${aiConfig.ai_base_url || 'https://api.deepseek.com'}/chat/completions`, {
|
|
method: 'POST',
|
|
headers: {
|
|
'Content-Type': 'application/json',
|
|
'Authorization': `Bearer ${aiConfig.ai_api_key}`,
|
|
},
|
|
body: JSON.stringify({
|
|
model: aiConfig.ai_model || 'deepseek-chat',
|
|
messages: [{ role: 'user', content: prompt }],
|
|
temperature: parseFloat(aiConfig.ai_temperature || '0.8'),
|
|
max_tokens: maxTokens,
|
|
}),
|
|
});
|
|
|
|
if (!response.ok) {
|
|
const err = await response.text();
|
|
throw new Error(`AI API error: ${response.status} ${err}`);
|
|
}
|
|
|
|
const data = await response.json();
|
|
return data.choices?.[0]?.message?.content || '';
|
|
}
|
|
|
|
async function generateTitle(productTitle) {
|
|
const titleRow = db.prepare('SELECT value FROM configs WHERE key = ?').get('title_prompt_template');
|
|
const fallbackRow = db.prepare('SELECT value FROM configs WHERE key = ?').get('prompt_template');
|
|
const template = titleRow?.value || fallbackRow?.value || '为商品写一个小红书标题：{{title}}';
|
|
const prompt = template.replace(/\{\{title\}\}/g, productTitle);
|
|
const generated = await callAI(prompt, 500);
|
|
const title = generated.replace(/^#+\s*/, '').replace(/[\n\r]/g, '').trim();
|
|
return title || productTitle;
|
|
}
|
|
|
|
async function generateContent(productTitle) {
|
|
const contentRow = db.prepare('SELECT value FROM configs WHERE key = ?').get('content_prompt_template');
|
|
const fallbackRow = db.prepare('SELECT value FROM configs WHERE key = ?').get('prompt_template');
|
|
const template = contentRow?.value || fallbackRow?.value || '为商品写一篇小红书正文：{{title}}';
|
|
const prompt = template.replace(/\{\{title\}\}/g, productTitle);
|
|
const generated = await callAI(prompt, 2000);
|
|
const topicMatches = generated.match(/#[^\s#]+/g) || [];
|
|
const topics = topicMatches.map((topic) => topic.replace('#', '').trim());
|
|
return { content: generated.trim(), topics };
|
|
}
|
|
|
|
async function generateText(productTitle) {
|
|
const title = await generateTitle(productTitle);
|
|
const { content, topics } = await generateContent(productTitle);
|
|
return { title, content, topics };
|
|
}
|
|
|
|
async function createNoteForProduct(product, options = {}) {
|
|
const existing = db.prepare(
|
|
"SELECT id, title, content, topics, image_paths, image_prompt FROM notes WHERE product_id = ? AND status = 'generating' ORDER BY created_at DESC LIMIT 1"
|
|
).get(product.id);
|
|
if (existing) {
|
|
return {
|
|
noteId: existing.id,
|
|
title: existing.title,
|
|
content: existing.content,
|
|
topics: JSON.parse(existing.topics || '[]'),
|
|
imagePaths: JSON.parse(existing.image_paths || '[]'),
|
|
imagePrompt: existing.image_prompt || '',
|
|
};
|
|
}
|
|
|
|
const noteId = uuidv4();
|
|
db.prepare("INSERT INTO notes (id, product_id, shop_id, title, content, status) VALUES (?, ?, ?, ?, '', 'generating')")
|
|
.run(noteId, product.id, product.shop_id, product.title);
|
|
|
|
const { title, content, topics } = await generateText(product.title);
|
|
const imageResult = await generateImage(db, noteId, product.title, options);
|
|
const imagePaths = imageResult.imagePaths || [];
|
|
const imagePrompt = imageResult.imagePrompt || '';
|
|
|
|
db.prepare(
|
|
"UPDATE notes SET title=?, content=?, topics=?, image_paths=?, image_prompt=?, status='generated', updated_at=CURRENT_TIMESTAMP WHERE id=?"
|
|
).run(title, content, JSON.stringify(topics), JSON.stringify(imagePaths), imagePrompt, noteId);
|
|
db.prepare("UPDATE products SET status='generated', updated_at=CURRENT_TIMESTAMP WHERE id=?").run(product.id);
|
|
|
|
return { noteId, title, content, topics, imagePaths, imagePrompt };
|
|
}
|
|
|
|
function updateJob(jobId, fields) {
|
|
const keys = Object.keys(fields);
|
|
if (keys.length === 0) return;
|
|
const setSql = keys.map((key) => `${key} = ?`).join(', ');
|
|
const values = keys.map((key) => fields[key]);
|
|
values.push(jobId);
|
|
db.prepare(`UPDATE generation_jobs SET ${setSql}, updated_at=CURRENT_TIMESTAMP WHERE id = ?`).run(...values);
|
|
}
|
|
|
|
async function runGenerationJob(jobId, productIds, countPerProduct, options = {}) {
|
|
if (runningJobs.has(jobId)) return;
|
|
runningJobs.add(jobId);
|
|
try {
|
|
updateJob(jobId, { status: 'running', message: '开始生成' });
|
|
let completedNotes = 0;
|
|
let failedNotes = 0;
|
|
|
|
for (const productId of productIds) {
|
|
const product = db.prepare('SELECT * FROM products WHERE id = ?').get(productId);
|
|
if (!product) {
|
|
failedNotes += countPerProduct;
|
|
completedNotes += countPerProduct;
|
|
updateJob(jobId, {
|
|
completed_notes: completedNotes,
|
|
failed_notes: failedNotes,
|
|
current_product_id: productId,
|
|
current_product_title: '',
|
|
message: '商品不存在',
|
|
});
|
|
continue;
|
|
}
|
|
|
|
db.prepare("UPDATE products SET status='generating', updated_at=CURRENT_TIMESTAMP WHERE id=?").run(product.id);
|
|
updateJob(jobId, {
|
|
current_product_id: product.id,
|
|
current_product_title: product.title,
|
|
message: `正在生成：${product.title}`,
|
|
});
|
|
|
|
for (let index = 0; index < countPerProduct; index += 1) {
|
|
try {
|
|
await createNoteForProduct(product, options);
|
|
} catch (error) {
|
|
failedNotes += 1;
|
|
} finally {
|
|
completedNotes += 1;
|
|
updateJob(jobId, {
|
|
completed_notes: completedNotes,
|
|
failed_notes: failedNotes,
|
|
message: `进度 ${completedNotes}/${productIds.length * countPerProduct}`,
|
|
});
|
|
}
|
|
}
|
|
|
|
db.prepare("UPDATE products SET status='generated', updated_at=CURRENT_TIMESTAMP WHERE id=?").run(product.id);
|
|
}
|
|
|
|
updateJob(jobId, {
|
|
status: 'completed',
|
|
message: failedNotes > 0 ? `完成，失败 ${failedNotes} 条` : '全部完成',
|
|
});
|
|
} catch (error) {
|
|
updateJob(jobId, { status: 'failed', message: error.message });
|
|
} finally {
|
|
runningJobs.delete(jobId);
|
|
}
|
|
}
|
|
|
|
function enqueueGeneration(jobId, productIds, countPerProduct, options) {
|
|
setImmediate(() => {
|
|
runGenerationJob(jobId, productIds, countPerProduct, options);
|
|
});
|
|
}
|
|
|
|
router.get('/jobs', (req, res) => {
|
|
const { shop_id } = req.query;
|
|
let sql = 'SELECT * FROM generation_jobs WHERE 1=1';
|
|
const params = [];
|
|
if (shop_id) {
|
|
sql += ' AND shop_id = ?';
|
|
params.push(shop_id);
|
|
}
|
|
sql += ' ORDER BY created_at DESC LIMIT 20';
|
|
res.json(db.prepare(sql).all(...params));
|
|
});
|
|
|
|
router.post('/note', (req, res) => {
|
|
const { product_id, shop_id, count_per_product, image_count, image_size, image_prompt_template } = req.body;
|
|
if (!product_id || !shop_id) {
|
|
return res.status(400).json({ error: 'product_id and shop_id required' });
|
|
}
|
|
|
|
const product = db.prepare('SELECT * FROM products WHERE id = ?').get(product_id);
|
|
if (!product) return res.status(404).json({ error: 'Product not found' });
|
|
|
|
const repeatCount = Math.max(1, parseInt(count_per_product || '1', 10) || 1);
|
|
const jobId = uuidv4();
|
|
db.prepare(
|
|
'INSERT INTO generation_jobs (id, shop_id, status, total_products, total_notes, message) VALUES (?, ?, ?, ?, ?, ?)'
|
|
).run(jobId, shop_id, 'pending', 1, repeatCount, '等待开始');
|
|
enqueueGeneration(jobId, [product_id], repeatCount, { image_count, image_size, image_prompt_template });
|
|
res.status(201).json({ job_id: jobId, status: 'pending' });
|
|
});
|
|
|
|
router.post('/image', async (req, res) => {
|
|
const { note_id } = req.body;
|
|
if (!note_id) return res.status(400).json({ error: 'note_id required' });
|
|
|
|
const note = db.prepare('SELECT * FROM notes WHERE id = ?').get(note_id);
|
|
if (!note) return res.status(404).json({ error: 'Note not found' });
|
|
|
|
try {
|
|
const imageResult = await generateImage(db, note_id, note.title);
|
|
const imagePaths = imageResult.imagePaths || [];
|
|
const existingPaths = JSON.parse(note.image_paths || '[]');
|
|
const allPaths = [...existingPaths, ...imagePaths];
|
|
db.prepare("UPDATE notes SET image_paths=?, image_prompt=?, updated_at=CURRENT_TIMESTAMP WHERE id=?")
|
|
.run(JSON.stringify(allPaths), imageResult.imagePrompt || note.image_prompt || '', note_id);
|
|
res.json({ image_paths: imagePaths, total: allPaths.length });
|
|
} catch (error) {
|
|
res.status(500).json({ error: error.message });
|
|
}
|
|
});
|
|
|
|
async function batchGenerate(req, res) {
|
|
const { shop_id, product_ids, count_per_product, image_count, image_size, image_prompt_template } = req.body;
|
|
if (!shop_id) return res.status(400).json({ error: 'shop_id required' });
|
|
|
|
let sql = 'SELECT * FROM products WHERE shop_id = ?';
|
|
const params = [shop_id];
|
|
if (Array.isArray(product_ids) && product_ids.length > 0) {
|
|
sql += ` AND id IN (${product_ids.map(() => '?').join(',')})`;
|
|
params.push(...product_ids);
|
|
}
|
|
const products = db.prepare(sql).all(...params);
|
|
if (products.length === 0) return res.json({ count: 0, message: 'No pending products' });
|
|
|
|
const repeatCount = Math.max(1, parseInt(count_per_product || '1', 10) || 1);
|
|
const jobId = uuidv4();
|
|
db.prepare(
|
|
'INSERT INTO generation_jobs (id, shop_id, status, total_products, total_notes, message) VALUES (?, ?, ?, ?, ?, ?)'
|
|
).run(jobId, shop_id, 'pending', products.length, products.length * repeatCount, '等待开始');
|
|
enqueueGeneration(jobId, products.map((product) => product.id), repeatCount, { image_count, image_size, image_prompt_template });
|
|
res.status(201).json({ job_id: jobId, total: products.length, total_notes: products.length * repeatCount });
|
|
}
|
|
|
|
router.post('/', batchGenerate);
|
|
router.post('/batch', batchGenerate);
|
|
|
|
return router;
|
|
}
|
|
|
|
module.exports = createGenerateRoutes;
|