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(' ');
|
||
}
|