// ═══════════════════════════════════════════════════════════════════════════ // Reranker - 硅基 bge-reranker-v2-m3 // 对候选文档进行精排,过滤与 query 不相关的内容 // ═══════════════════════════════════════════════════════════════════════════ import { xbLog } from '../../../../core/debug-core.js'; import { getApiKey } from './siliconflow.js'; const MODULE_ID = 'reranker'; const RERANK_URL = 'https://api.siliconflow.cn/v1/rerank'; const RERANK_MODEL = 'BAAI/bge-reranker-v2-m3'; const DEFAULT_TIMEOUT = 15000; const MAX_DOCUMENTS = 100; // API 限制 /** * 对文档列表进行 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 key = getApiKey(); 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 response = await fetch(RERANK_URL, { method: 'POST', headers: { 'Authorization': `Bearer ${key}`, 'Content-Type': 'application/json', }, body: JSON.stringify({ 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 []; if (chunks.length <= topN) { const texts = chunks.map(c => c.text || c.semantic || ''); const { results, failed } = await rerank(query, texts, { topN: chunks.length, ...options }); if (failed) { return chunks.map(c => ({ ...c, _rerankScore: 0, _rerankFailed: true })); } const scoreMap = new Map(results.map(r => [r.index, r.relevance_score])); return chunks.map((c, i) => ({ ...c, _rerankScore: scoreMap.get(i) ?? 0, })).sort((a, b) => b._rerankScore - a._rerankScore); } const texts = chunks.map(c => c.text || c.semantic || ''); const { results, failed } = await rerank(query, texts, { topN, ...options }); if (failed) { return chunks.slice(0, topN).map(c => ({ ...c, _rerankScore: 0, _rerankFailed: true, })); } return results .filter(r => r.relevance_score >= minScore) .sort((a, b) => b.relevance_score - a.relevance_score) .map(r => ({ ...chunks[r.index], _rerankScore: r.relevance_score, })); } /** * 测试 Rerank 服务连接 */ export async function testRerankService() { const key = getApiKey(); if (!key) { throw new Error('请配置硅基 API Key'); } try { const { results } = await rerank('测试查询', ['测试文档1', '测试文档2'], { topN: 2 }); return { success: true, message: `连接成功,返回 ${results.length} 个结果`, }; } catch (e) { throw new Error(`连接失败: ${e.message}`); } }