103 lines
3.8 KiB
JavaScript
103 lines
3.8 KiB
JavaScript
|
// ═══════════════════════════════════════════════════════════════════════════
|
|||
|
|
// query-expansion.js - 完整输入，不截断
|
|||
|
|
// ═══════════════════════════════════════════════════════════════════════════
|
|||
|
|
|
|||
|
|
import { callLLM, parseJson } from './llm-service.js';
|
|||
|
|
import { xbLog } from '../../../../core/debug-core.js';
|
|||
|
|
import { filterText } from '../utils/text-filter.js';
|
|||
|
|
|
|||
|
|
const MODULE_ID = 'query-expansion';
|
|||
|
|
const SESSION_ID = 'xb6';
|
|||
|
|
|
|||
|
|
const SYSTEM_PROMPT = `你是检索词生成器。根据最近对话，输出用于检索历史剧情的关键词。
|
|||
|
|
|
|||
|
|
只输出JSON：
|
|||
|
|
{"e":["显式人物/地名"],"i":["隐含人物/情绪/话题"],"q":["检索短句"]}
|
|||
|
|
|
|||
|
|
规则：
|
|||
|
|
- e: 对话中明确提到的人名/地名，1-4个
|
|||
|
|
- i: 推断出的相关人物/情绪/话题，1-5个
|
|||
|
|
- q: 用于向量检索的短句，2-3个，每个15字内
|
|||
|
|
- 关注：正在讨论什么、涉及谁、情绪氛围`;
|
|||
|
|
|
|||
|
|
/**
|
|||
|
|
* Query Expansion
|
|||
|
|
* @param {Array} messages - 完整消息数组（最后2-3轮）
|
|||
|
|
*/
|
|||
|
|
export async function expandQuery(messages, options = {}) {
|
|||
|
|
const { timeout = 6000 } = options;
|
|||
|
|
|
|||
|
|
if (!messages?.length) {
|
|||
|
|
return { entities: [], implicit: [], queries: [] };
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 完整格式化，不截断
|
|||
|
|
const input = messages.map(m => {
|
|||
|
|
const speaker = m.is_user ? '用户' : (m.name || '角色');
|
|||
|
|
const text = filterText(m.mes || '').trim();
|
|||
|
|
return `【${speaker}】\n${text}`;
|
|||
|
|
}).join('\n\n');
|
|||
|
|
|
|||
|
|
const T0 = performance.now();
|
|||
|
|
|
|||
|
|
try {
|
|||
|
|
const response = await callLLM([
|
|||
|
|
{ role: 'system', content: SYSTEM_PROMPT },
|
|||
|
|
{ role: 'user', content: input },
|
|||
|
|
], {
|
|||
|
|
temperature: 0.15,
|
|||
|
|
max_tokens: 250,
|
|||
|
|
timeout,
|
|||
|
|
sessionId: SESSION_ID,
|
|||
|
|
});
|
|||
|
|
|
|||
|
|
const parsed = parseJson(response);
|
|||
|
|
if (!parsed) {
|
|||
|
|
xbLog.warn(MODULE_ID, 'JSON解析失败', response?.slice(0, 200));
|
|||
|
|
return { entities: [], implicit: [], queries: [] };
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
const result = {
|
|||
|
|
entities: Array.isArray(parsed.e) ? parsed.e.slice(0, 5) : [],
|
|||
|
|
implicit: Array.isArray(parsed.i) ? parsed.i.slice(0, 6) : [],
|
|||
|
|
queries: Array.isArray(parsed.q) ? parsed.q.slice(0, 4) : [],
|
|||
|
|
};
|
|||
|
|
|
|||
|
|
xbLog.info(MODULE_ID, `完成 (${Math.round(performance.now() - T0)}ms) e=${result.entities.length} i=${result.implicit.length} q=${result.queries.length}`);
|
|||
|
|
return result;
|
|||
|
|
|
|||
|
|
} catch (e) {
|
|||
|
|
xbLog.error(MODULE_ID, '调用失败', e);
|
|||
|
|
return { entities: [], implicit: [], queries: [] };
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// 缓存
|
|||
|
|
const cache = new Map();
|
|||
|
|
const CACHE_TTL = 300000;
|
|||
|
|
|
|||
|
|
function hashMessages(messages) {
|
|||
|
|
const text = messages.slice(-2).map(m => (m.mes || '').slice(0, 100)).join('|');
|
|||
|
|
let h = 0;
|
|||
|
|
for (let i = 0; i < text.length; i++) h = ((h << 5) - h + text.charCodeAt(i)) | 0;
|
|||
|
|
return h.toString(36);
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
export async function expandQueryCached(messages, options = {}) {
|
|||
|
|
const key = hashMessages(messages);
|
|||
|
|
const cached = cache.get(key);
|
|||
|
|
if (cached && Date.now() - cached.time < CACHE_TTL) return cached.result;
|
|||
|
|
|
|||
|
|
const result = await expandQuery(messages, options);
|
|||
|
|
if (result.entities.length || result.queries.length) {
|
|||
|
|
if (cache.size > 50) cache.delete(cache.keys().next().value);
|
|||
|
|
cache.set(key, { result, time: Date.now() });
|
|||
|
|
}
|
|||
|
|
return result;
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
export function buildSearchText(expansion) {
|
|||
|
|
return [...(expansion.entities || []), ...(expansion.implicit || []), ...(expansion.queries || [])]
|
|||
|
|
.filter(Boolean).join(' ');
|
|||
|
|
}
|