// ═══════════════════════════════════════════════════════════════════════════ // Reranker - 硅基 bge-reranker-v2-m3 // 对候选文档进行精排,过滤与 query 不相关的内容 // ═══════════════════════════════════════════════════════════════════════════ import { xbLog } from '../../../../core/debug-core.js'; import { getVectorConfig } from '../../data/config.js'; const MODULE_ID = 'reranker'; const DEFAULT_RERANK_URL = 'https://api.siliconflow.cn/v1'; const RERANK_MODEL = 'BAAI/bge-reranker-v2-m3'; const DEFAULT_TIMEOUT = 15000; const MAX_DOCUMENTS = 100; // API 限制 const RERANK_BATCH_SIZE = 20; const RERANK_MAX_CONCURRENCY = 5; let rerankKeyIndex = 0; function getRerankApiConfig() { const cfg = getVectorConfig() || {}; return cfg.rerankApi || { provider: 'siliconflow', url: DEFAULT_RERANK_URL, key: '', model: RERANK_MODEL, }; } function getNextRerankKey(rawKey) { const keys = String(rawKey || '') .split(/[,;|\n]+/) .map(k => k.trim()) .filter(Boolean); if (!keys.length) return ''; if (keys.length === 1) return keys[0]; const idx = rerankKeyIndex % keys.length; rerankKeyIndex = (rerankKeyIndex + 1) % keys.length; return keys[idx]; } /** * 对文档列表进行 Rerank 精排 * * @param {string} query - 查询文本 * @param {Array} documents - 文档文本列表 * @param {object} options - 选项 * @param {number} options.topN - 返回前 N 个结果,默认 40 * @param {number} options.timeout - 超时时间,默认 15000ms * @param {AbortSignal} options.signal - 取消信号 * @returns {Promise>} 排序后的结果 */ export async function rerank(query, documents, options = {}) { const { topN = 40, timeout = DEFAULT_TIMEOUT, signal } = options; if (!query?.trim()) { xbLog.warn(MODULE_ID, 'query 为空,跳过 rerank'); return { results: documents.map((_, i) => ({ index: i, relevance_score: 0 })), failed: true }; } if (!documents?.length) { return { results: [], failed: false }; } const apiCfg = getRerankApiConfig(); const key = getNextRerankKey(apiCfg.key); if (!key) { xbLog.warn(MODULE_ID, '未配置 API Key,跳过 rerank'); return { results: documents.map((_, i) => ({ index: i, relevance_score: 0 })), failed: true }; } // 截断超长文档列表 const truncatedDocs = documents.slice(0, MAX_DOCUMENTS); if (documents.length > MAX_DOCUMENTS) { xbLog.warn(MODULE_ID, `文档数 ${documents.length} 超过限制 ${MAX_DOCUMENTS},已截断`); } // 过滤空文档,记录原始索引 const validDocs = []; const indexMap = []; // validDocs index → original index for (let i = 0; i < truncatedDocs.length; i++) { const text = String(truncatedDocs[i] || '').trim(); if (text) { validDocs.push(text); indexMap.push(i); } } if (!validDocs.length) { xbLog.warn(MODULE_ID, '无有效文档,跳过 rerank'); return { results: [], failed: false }; } const controller = new AbortController(); const timeoutId = setTimeout(() => controller.abort(), timeout); try { const T0 = performance.now(); const baseUrl = String(apiCfg.url || DEFAULT_RERANK_URL).replace(/\/+$/, ''); const response = await fetch(`${baseUrl}/rerank`, { method: 'POST', headers: { 'Authorization': `Bearer ${key}`, 'Content-Type': 'application/json', }, body: JSON.stringify({ model: String(apiCfg.model || RERANK_MODEL), // Zero-darkbox: do not silently truncate query. query, documents: validDocs, top_n: Math.min(topN, validDocs.length), return_documents: false, }), signal: signal || controller.signal, }); clearTimeout(timeoutId); if (!response.ok) { const errorText = await response.text().catch(() => ''); throw new Error(`Rerank API ${response.status}: ${errorText.slice(0, 200)}`); } const data = await response.json(); const results = data.results || []; // 映射回原始索引 const mapped = results.map(r => ({ index: indexMap[r.index], relevance_score: r.relevance_score ?? 0, })); const elapsed = Math.round(performance.now() - T0); xbLog.info(MODULE_ID, `Rerank 完成: ${validDocs.length} docs → ${results.length} selected (${elapsed}ms)`); return { results: mapped, failed: false }; } catch (e) { clearTimeout(timeoutId); if (e?.name === 'AbortError') { xbLog.warn(MODULE_ID, 'Rerank 超时或取消'); } else { xbLog.error(MODULE_ID, 'Rerank 失败', e); } // 降级:返回原顺序,分数均匀分布 return { results: documents.slice(0, topN).map((_, i) => ({ index: i, relevance_score: 0, })), failed: true, }; } } /** * 对 chunk 对象列表进行 Rerank * * @param {string} query - 查询文本 * @param {Array} chunks - chunk 对象列表,需要有 text 字段 * @param {object} options - 选项 * @returns {Promise>} 排序后的 chunk 列表,带 _rerankScore 字段 */ export async function rerankChunks(query, chunks, options = {}) { const { topN = 40, minScore = 0.1 } = options; if (!chunks?.length) return []; const texts = chunks.map(c => c.text || c.semantic || ''); // ─── 单批:直接调用 ─── if (texts.length <= RERANK_BATCH_SIZE) { const { results, failed } = await rerank(query, texts, { topN: Math.min(topN, texts.length), timeout: options.timeout, signal: options.signal, }); if (failed) { return chunks.map(c => ({ ...c, _rerankScore: 0, _rerankFailed: true })); } return results .filter(r => r.relevance_score >= minScore) .sort((a, b) => b.relevance_score - a.relevance_score) .slice(0, topN) .map(r => ({ ...chunks[r.index], _rerankScore: r.relevance_score, })); } // ─── 多批:拆分 → 并发 → 合并 ─── const batches = []; for (let i = 0; i < texts.length; i += RERANK_BATCH_SIZE) { batches.push({ texts: texts.slice(i, i + RERANK_BATCH_SIZE), offset: i, }); } const concurrency = Math.min(batches.length, RERANK_MAX_CONCURRENCY); xbLog.info(MODULE_ID, `并发 Rerank: ${batches.length} 批 × ≤${RERANK_BATCH_SIZE} docs, concurrency=${concurrency}`); const batchResults = new Array(batches.length); let failedBatches = 0; const runBatch = async (batchIdx) => { const batch = batches[batchIdx]; const { results, failed } = await rerank(query, batch.texts, { topN: batch.texts.length, timeout: options.timeout, signal: options.signal, }); if (failed) { failedBatches++; // 单批降级:保留原始顺序,score=0 batchResults[batchIdx] = batch.texts.map((_, i) => ({ globalIndex: batch.offset + i, relevance_score: 0, _batchFailed: true, })); } else { batchResults[batchIdx] = results.map(r => ({ globalIndex: batch.offset + r.index, relevance_score: r.relevance_score, })); } }; // 并发池 let nextIdx = 0; const worker = async () => { while (nextIdx < batches.length) { const idx = nextIdx++; await runBatch(idx); } }; await Promise.all(Array.from({ length: concurrency }, () => worker())); // 全部失败 → 整体降级 if (failedBatches === batches.length) { xbLog.warn(MODULE_ID, `全部 ${batches.length} 批 rerank 失败,整体降级`); return chunks.slice(0, topN).map(c => ({ ...c, _rerankScore: 0, _rerankFailed: true, })); } // 合并所有批次结果 const merged = batchResults.flat(); const selected = merged .filter(r => r._batchFailed || r.relevance_score >= minScore) .sort((a, b) => b.relevance_score - a.relevance_score) .slice(0, topN) .map(r => ({ ...chunks[r.globalIndex], _rerankScore: r.relevance_score, ...(r._batchFailed ? { _rerankFailed: true } : {}), })); xbLog.info(MODULE_ID, `Rerank 合并: ${merged.length} candidates, ${failedBatches}/${batches.length} 批失败, 选中 ${selected.length}` ); return selected; } /** * 测试 Rerank 服务连接 */ export async function testRerankService(apiConfig = {}) { const next = { provider: String(apiConfig.provider || 'siliconflow').trim(), url: String(apiConfig.url || DEFAULT_RERANK_URL).trim(), key: String(apiConfig.key || '').trim(), model: String(apiConfig.model || RERANK_MODEL).trim(), }; if (!next.key) { throw new Error('请配置 Rerank API Key'); } const key = getNextRerankKey(next.key); const controller = new AbortController(); const timeoutId = setTimeout(() => controller.abort(), 15000); try { const baseUrl = String(next.url || DEFAULT_RERANK_URL).replace(/\/+$/, ''); const response = await fetch(`${baseUrl}/rerank`, { method: 'POST', headers: { 'Authorization': `Bearer ${key}`, 'Content-Type': 'application/json', }, body: JSON.stringify({ model: next.model, query: '测试查询', documents: ['测试文档1', '测试文档2'], top_n: 2, return_documents: false, }), signal: controller.signal, }); clearTimeout(timeoutId); if (!response.ok) { const errorText = await response.text().catch(() => ''); throw new Error(`Rerank API ${response.status}: ${errorText.slice(0, 200)}`); } const data = await response.json(); const results = Array.isArray(data.results) ? data.results : []; return { success: true, message: `连接成功:返回 ${results.length} 个结果`, }; } catch (e) { throw new Error(`连接失败: ${e.message}`); } finally { clearTimeout(timeoutId); } }