forked from lunanting/cangku
275 lines
8.5 KiB
JavaScript
275 lines
8.5 KiB
JavaScript
'use strict';
|
||
|
||
const axios = require('axios');
|
||
const planningDatabase = require('./planningDatabase');
|
||
|
||
const ARK_BASE_URL = process.env.ARK_BASE_URL || 'https://ark.cn-beijing.volces.com/api/v3';
|
||
const EMBEDDING_MODEL = process.env.EMBEDDING_MODEL || 'doubao-embedding-vision-250615';
|
||
const DEFAULT_BATCH_SIZE = Math.max(parseInt(process.env.TAOBAO_PRODUCT_EMBEDDING_BATCH_SIZE || '100', 10), 1);
|
||
const DEFAULT_CONCURRENCY = Math.max(parseInt(process.env.TAOBAO_PRODUCT_EMBEDDING_CONCURRENCY || '2', 10), 1);
|
||
const DEFAULT_REQUEST_TIMEOUT_MS = Math.max(parseInt(process.env.TAOBAO_PRODUCT_EMBEDDING_REQUEST_TIMEOUT_MS || '120000', 10), 10000);
|
||
const DEFAULT_ORDER = String(process.env.TAOBAO_PRODUCT_EMBEDDING_ORDER || 'DESC').toUpperCase() === 'ASC' ? 'ASC' : 'DESC';
|
||
const DEFAULT_MIN_ROW_ID = Math.max(parseInt(process.env.TAOBAO_PRODUCT_EMBEDDING_MIN_ROW_ID || '0', 10), 0);
|
||
const RETRY_FAILED_ROWS = ['1', 'true', 'TRUE', 'yes', 'YES'].includes(String(process.env.TAOBAO_PRODUCT_EMBEDDING_RETRY_FAILED || ''));
|
||
|
||
function getArkApiKey() {
|
||
const apiKey = String(process.env.ARK_API_KEY || '').trim();
|
||
if (!apiKey) {
|
||
throw new Error('未配置 ARK_API_KEY，无法执行 taobao_product_records 向量任务');
|
||
}
|
||
return apiKey;
|
||
}
|
||
|
||
function safeJsonParse(value, fallback = null) {
|
||
if (value === null || value === undefined || value === '') {
|
||
return fallback;
|
||
}
|
||
if (typeof value === 'object') {
|
||
return value;
|
||
}
|
||
try {
|
||
return JSON.parse(value);
|
||
} catch (error) {
|
||
return fallback;
|
||
}
|
||
}
|
||
|
||
function normalizeImageUrl(url) {
|
||
const text = String(url || '').trim();
|
||
if (!text) {
|
||
return null;
|
||
}
|
||
if (text.startsWith('//')) {
|
||
return `https:${text}`;
|
||
}
|
||
return text;
|
||
}
|
||
|
||
function resolveImageUrl(mainImageUrl, rawDataJson) {
|
||
const normalizedMainImageUrl = normalizeImageUrl(mainImageUrl);
|
||
if (normalizedMainImageUrl) {
|
||
return normalizedMainImageUrl;
|
||
}
|
||
|
||
const rawData = safeJsonParse(rawDataJson, {});
|
||
if (!rawData || typeof rawData !== 'object') {
|
||
return null;
|
||
}
|
||
|
||
return normalizeImageUrl(
|
||
rawData.pic_url
|
||
|| rawData.picUrl
|
||
|| rawData.pic
|
||
|| rawData.mainItemInfo?.pic
|
||
|| rawData.mainItemInfo?.pic_url
|
||
);
|
||
}
|
||
|
||
async function fetchEmbedding(imageUrl, requestTimeout) {
|
||
const response = await axios.post(
|
||
`${ARK_BASE_URL}/embeddings/multimodal`,
|
||
{
|
||
model: EMBEDDING_MODEL,
|
||
input: [
|
||
{ type: 'text', text: '只关注商品主体外观进行向量化，尽量忽略背景、模特、文字、水印、边框和详情页装饰。相同商品在不同背景下应尽量接近。' },
|
||
{ type: 'image_url', image_url: { url: imageUrl } },
|
||
],
|
||
},
|
||
{
|
||
headers: {
|
||
Authorization: `Bearer ${getArkApiKey()}`,
|
||
'Content-Type': 'application/json',
|
||
},
|
||
timeout: requestTimeout,
|
||
}
|
||
);
|
||
|
||
const embeddingRow = Array.isArray(response.data?.data) ? response.data.data[0] : response.data?.data;
|
||
const vector = Array.isArray(embeddingRow?.embedding) ? embeddingRow.embedding : [];
|
||
return {
|
||
vector,
|
||
model: response.data?.model || EMBEDDING_MODEL,
|
||
};
|
||
}
|
||
|
||
async function getPendingRows(limit, options = {}) {
|
||
const minRowId = Math.max(parseInt(options.minRowId ?? DEFAULT_MIN_ROW_ID, 10) || 0, 0);
|
||
const order = String(options.order || DEFAULT_ORDER).toUpperCase() === 'ASC' ? 'ASC' : 'DESC';
|
||
return planningDatabase.query(
|
||
`
|
||
SELECT id, item_id, main_image_url, raw_data_json
|
||
FROM taobao_product_records
|
||
WHERE id > ?
|
||
AND (embedding_vector_json IS NULL OR TRIM(embedding_vector_json) = '')
|
||
AND (? = 1 OR embedding_model IS NULL OR embedding_model NOT LIKE 'failed:%')
|
||
AND (
|
||
(main_image_url IS NOT NULL AND TRIM(main_image_url) <> '')
|
||
OR (raw_data_json IS NOT NULL AND TRIM(raw_data_json) <> '')
|
||
)
|
||
ORDER BY id ${order}
|
||
LIMIT ?
|
||
`,
|
||
[ minRowId, RETRY_FAILED_ROWS ? 1 : 0, limit ]
|
||
);
|
||
}
|
||
|
||
async function updateEmbeddingSuccess(id, embedding) {
|
||
await planningDatabase.execute(
|
||
`
|
||
UPDATE taobao_product_records
|
||
SET embedding_provider = ?,
|
||
embedding_model = ?,
|
||
embedding_vector_json = ?,
|
||
updated_at = CURRENT_TIMESTAMP
|
||
WHERE id = ?
|
||
`,
|
||
[ 'volcengine-ark', embedding.model, JSON.stringify(embedding.vector), id ]
|
||
);
|
||
}
|
||
|
||
async function updateEmbeddingFailure(id, reason) {
|
||
await planningDatabase.execute(
|
||
`
|
||
UPDATE taobao_product_records
|
||
SET embedding_provider = ?,
|
||
embedding_model = ?,
|
||
updated_at = CURRENT_TIMESTAMP
|
||
WHERE id = ?
|
||
`,
|
||
[ 'volcengine-ark', `failed:${String(reason || '').slice(0, 80)}`, id ]
|
||
);
|
||
}
|
||
|
||
async function processOneRow(row, options) {
|
||
const imageUrl = resolveImageUrl(row.main_image_url, row.raw_data_json);
|
||
if (!imageUrl) {
|
||
await updateEmbeddingFailure(row.id, 'missing_image_url');
|
||
return { success: false, skipped: true, id: row.id, itemId: row.item_id, reason: 'missing_image_url' };
|
||
}
|
||
|
||
try {
|
||
const embedding = await fetchEmbedding(imageUrl, options.requestTimeout);
|
||
if (!Array.isArray(embedding.vector) || embedding.vector.length === 0) {
|
||
await updateEmbeddingFailure(row.id, 'empty_embedding');
|
||
return { success: false, skipped: true, id: row.id, itemId: row.item_id, reason: 'empty_embedding' };
|
||
}
|
||
|
||
await updateEmbeddingSuccess(row.id, embedding);
|
||
return {
|
||
success: true,
|
||
id: row.id,
|
||
itemId: row.item_id,
|
||
dimension: embedding.vector.length,
|
||
model: embedding.model,
|
||
};
|
||
} catch (error) {
|
||
await updateEmbeddingFailure(row.id, error.message || 'request_failed');
|
||
return {
|
||
success: false,
|
||
id: row.id,
|
||
itemId: row.item_id,
|
||
reason: error.message || 'request_failed',
|
||
};
|
||
}
|
||
}
|
||
|
||
async function runWithConcurrency(rows, concurrency, worker) {
|
||
const results = new Array(rows.length);
|
||
let cursor = 0;
|
||
|
||
async function runner() {
|
||
while (true) {
|
||
const currentIndex = cursor;
|
||
cursor += 1;
|
||
if (currentIndex >= rows.length) {
|
||
return;
|
||
}
|
||
results[currentIndex] = await worker(rows[currentIndex], currentIndex);
|
||
}
|
||
}
|
||
|
||
await Promise.all(Array.from({ length: Math.min(concurrency, rows.length) }, () => runner()));
|
||
return results;
|
||
}
|
||
|
||
async function syncTaobaoProductEmbeddings(options = {}) {
|
||
const batchSize = Math.max(parseInt(options.batchSize || DEFAULT_BATCH_SIZE, 10), 1);
|
||
const concurrency = Math.max(parseInt(options.concurrency || DEFAULT_CONCURRENCY, 10), 1);
|
||
const requestTimeout = Math.max(parseInt(options.requestTimeout || DEFAULT_REQUEST_TIMEOUT_MS, 10), 10000);
|
||
const maxBatches = Math.max(parseInt(options.maxBatches || process.env.TAOBAO_PRODUCT_EMBEDDING_MAX_BATCHES || '0', 10), 0);
|
||
const minRowId = Math.max(parseInt(options.minRowId ?? DEFAULT_MIN_ROW_ID, 10) || 0, 0);
|
||
const order = String(options.order || DEFAULT_ORDER).toUpperCase() === 'ASC' ? 'ASC' : 'DESC';
|
||
const stopWhenEmpty = options.stopWhenEmpty !== false;
|
||
|
||
getArkApiKey();
|
||
|
||
const summary = {
|
||
success: true,
|
||
batchSize,
|
||
concurrency,
|
||
requestTimeout,
|
||
minRowId,
|
||
order,
|
||
batches: 0,
|
||
fetched: 0,
|
||
successCount: 0,
|
||
failCount: 0,
|
||
skippedCount: 0,
|
||
successMinId: null,
|
||
successMaxId: null,
|
||
};
|
||
|
||
while (true) {
|
||
if (maxBatches > 0 && summary.batches >= maxBatches) {
|
||
break;
|
||
}
|
||
|
||
const rows = await getPendingRows(batchSize, { minRowId, order });
|
||
if (!rows.length) {
|
||
if (stopWhenEmpty) {
|
||
break;
|
||
}
|
||
return summary;
|
||
}
|
||
|
||
summary.batches += 1;
|
||
summary.fetched += rows.length;
|
||
console.log(`[TaobaoEmbeddingSync] start batch=${summary.batches}, size=${rows.length}`);
|
||
|
||
const results = await runWithConcurrency(rows, concurrency, (row) => processOneRow(row, { requestTimeout }));
|
||
for (const result of results) {
|
||
if (!result) {
|
||
continue;
|
||
}
|
||
if (result.success) {
|
||
summary.successCount += 1;
|
||
const numericId = Number(result.id || 0);
|
||
if (numericId > 0) {
|
||
summary.successMinId = summary.successMinId === null ? numericId : Math.min(summary.successMinId, numericId);
|
||
summary.successMaxId = summary.successMaxId === null ? numericId : Math.max(summary.successMaxId, numericId);
|
||
}
|
||
} else if (result.skipped) {
|
||
summary.skippedCount += 1;
|
||
} else {
|
||
summary.failCount += 1;
|
||
}
|
||
}
|
||
|
||
console.log(
|
||
`[TaobaoEmbeddingSync] finish batch=${summary.batches}, success=${summary.successCount}, fail=${summary.failCount}, skipped=${summary.skippedCount}`
|
||
);
|
||
|
||
if (rows.length < batchSize) {
|
||
break;
|
||
}
|
||
}
|
||
|
||
return summary;
|
||
}
|
||
|
||
module.exports = {
|
||
syncTaobaoProductEmbeddings,
|
||
resolveImageUrl,
|
||
fetchEmbedding,
|
||
};
|