chore: update story summary and lint fixes
This commit is contained in:
184
modules/story-summary/vector/llm/reranker.js
Normal file
184
modules/story-summary/vector/llm/reranker.js
Normal file
@@ -0,0 +1,184 @@
|
||||
// ═══════════════════════════════════════════════════════════════════════════
|
||||
// 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<string>} documents - 文档文本列表
|
||||
* @param {object} options - 选项
|
||||
* @param {number} options.topN - 返回前 N 个结果,默认 40
|
||||
* @param {number} options.timeout - 超时时间,默认 15000ms
|
||||
* @param {AbortSignal} options.signal - 取消信号
|
||||
* @returns {Promise<Array<{index: number, relevance_score: number}>>} 排序后的结果
|
||||
*/
|
||||
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 documents.map((_, i) => ({ index: i, relevance_score: 0.5 }));
|
||||
}
|
||||
|
||||
if (!documents?.length) {
|
||||
return [];
|
||||
}
|
||||
|
||||
const key = getApiKey();
|
||||
if (!key) {
|
||||
xbLog.warn(MODULE_ID, '未配置 API Key,跳过 rerank');
|
||||
return documents.map((_, i) => ({ index: i, relevance_score: 0.5 }));
|
||||
}
|
||||
|
||||
// 截断超长文档列表
|
||||
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 [];
|
||||
}
|
||||
|
||||
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,
|
||||
query: query.slice(0, 1000), // 限制 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 mapped;
|
||||
|
||||
} catch (e) {
|
||||
clearTimeout(timeoutId);
|
||||
|
||||
if (e?.name === 'AbortError') {
|
||||
xbLog.warn(MODULE_ID, 'Rerank 超时或取消');
|
||||
} else {
|
||||
xbLog.error(MODULE_ID, 'Rerank 失败', e);
|
||||
}
|
||||
|
||||
// 降级:返回原顺序,分数均匀分布
|
||||
return documents.slice(0, topN).map((_, i) => ({
|
||||
index: i,
|
||||
relevance_score: 1 - (i / documents.length) * 0.5,
|
||||
}));
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* 对 chunk 对象列表进行 Rerank
|
||||
*
|
||||
* @param {string} query - 查询文本
|
||||
* @param {Array<object>} chunks - chunk 对象列表,需要有 text 字段
|
||||
* @param {object} options - 选项
|
||||
* @returns {Promise<Array<object>>} 排序后的 chunk 列表,带 _rerankScore 字段
|
||||
*/
|
||||
export async function rerankChunks(query, chunks, options = {}) {
|
||||
const { topN = 40, minScore = 0.1 } = options;
|
||||
|
||||
if (!chunks?.length) return [];
|
||||
if (chunks.length <= topN) {
|
||||
// 数量不超限,仍然 rerank 以获取分数,但不过滤
|
||||
const texts = chunks.map(c => c.text || c.semantic || '');
|
||||
const results = await rerank(query, texts, { topN: chunks.length, ...options });
|
||||
|
||||
const scoreMap = new Map(results.map(r => [r.index, r.relevance_score]));
|
||||
return chunks.map((c, i) => ({
|
||||
...c,
|
||||
_rerankScore: scoreMap.get(i) ?? 0.5,
|
||||
})).sort((a, b) => b._rerankScore - a._rerankScore);
|
||||
}
|
||||
|
||||
const texts = chunks.map(c => c.text || c.semantic || '');
|
||||
const results = await rerank(query, texts, { topN, ...options });
|
||||
|
||||
// 过滤低分 + 排序
|
||||
const selected = 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,
|
||||
}));
|
||||
|
||||
return selected;
|
||||
}
|
||||
|
||||
/**
|
||||
* 测试 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}`);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user