Files
LittleWhiteBox/modules/story-summary/vector/pipeline/state-integration.js

308 lines
9.4 KiB
JavaScript
Raw Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
// ============================================================================
// state-integration.js - L0 记忆锚点管理
// 支持增量提取、清空、取消
// ============================================================================
import { getContext } from '../../../../../../../extensions.js';
import { xbLog } from '../../../../core/debug-core.js';
import {
saveStateAtoms,
saveStateVectors,
deleteStateAtomsFromFloor,
deleteStateVectorsFromFloor,
getStateAtoms,
clearStateAtoms,
clearStateVectors,
getL0FloorStatus,
setL0FloorStatus,
clearL0Index,
deleteL0IndexFromFloor,
} from '../storage/state-store.js';
import { embed } from '../llm/siliconflow.js';
import { extractAtomsForRound, cancelBatchExtraction } from '../llm/atom-extraction.js';
import { getVectorConfig } from '../../data/config.js';
import { getEngineFingerprint } from '../utils/embedder.js';
import { filterText } from '../utils/text-filter.js';
const MODULE_ID = 'state-integration';
let initialized = false;
export function cancelL0Extraction() {
cancelBatchExtraction();
}
// ============================================================================
// 初始化
// ============================================================================
export function initStateIntegration() {
if (initialized) return;
initialized = true;
globalThis.LWB_StateRollbackHook = handleStateRollback;
xbLog.info(MODULE_ID, 'L0 状态层集成已初始化');
}
// ============================================================================
// 统计
// ============================================================================
export async function getAnchorStats() {
const { chat } = getContext();
if (!chat?.length) {
return { extracted: 0, total: 0, pending: 0, empty: 0, fail: 0 };
}
const aiFloors = [];
for (let i = 0; i < chat.length; i++) {
if (!chat[i]?.is_user) aiFloors.push(i);
}
let ok = 0;
let empty = 0;
let fail = 0;
for (const f of aiFloors) {
const s = getL0FloorStatus(f);
if (!s) continue;
if (s.status === 'ok') ok++;
else if (s.status === 'empty') empty++;
else if (s.status === 'fail') fail++;
}
const total = aiFloors.length;
const completed = ok + empty;
const pending = Math.max(0, total - completed);
return { extracted: completed, total, pending, empty, fail };
}
// ============================================================================
// 增量提取
// ============================================================================
function buildL0InputText(userMessage, aiMessage) {
const parts = [];
const userName = userMessage?.name || '用户';
const aiName = aiMessage?.name || '角色';
if (userMessage?.mes?.trim()) {
parts.push(`【用户:${userName}\n${filterText(userMessage.mes).trim()}`);
}
if (aiMessage?.mes?.trim()) {
parts.push(`【角色:${aiName}\n${filterText(aiMessage.mes).trim()}`);
}
return parts.join('\n\n---\n\n').trim();
}
export async function incrementalExtractAtoms(chatId, chat, onProgress) {
if (!chatId || !chat?.length) return { built: 0 };
const vectorCfg = getVectorConfig();
if (!vectorCfg?.enabled) return { built: 0 };
const pendingPairs = [];
for (let i = 0; i < chat.length; i++) {
const msg = chat[i];
if (!msg || msg.is_user) continue;
const st = getL0FloorStatus(i);
if (st?.status === 'ok' || st?.status === 'empty') {
continue;
}
const userMsg = (i > 0 && chat[i - 1]?.is_user) ? chat[i - 1] : null;
const inputText = buildL0InputText(userMsg, msg);
if (!inputText) {
setL0FloorStatus(i, { status: 'empty', reason: 'filtered_empty', atoms: 0 });
continue;
}
pendingPairs.push({ userMsg, aiMsg: msg, aiFloor: i });
}
if (!pendingPairs.length) {
onProgress?.(0, 0, '已全部提取');
return { built: 0 };
}
xbLog.info(MODULE_ID, `增量 L0 提取pending=${pendingPairs.length}`);
let completed = 0;
const total = pendingPairs.length;
let builtAtoms = 0;
for (const pair of pendingPairs) {
const floor = pair.aiFloor;
const prev = getL0FloorStatus(floor);
try {
const atoms = await extractAtomsForRound(pair.userMsg, pair.aiMsg, floor, { timeout: 20000 });
if (atoms == null) {
throw new Error('llm_failed');
}
if (!atoms.length) {
setL0FloorStatus(floor, { status: 'empty', reason: 'llm_empty', atoms: 0 });
} else {
atoms.forEach(a => a.chatId = chatId);
saveStateAtoms(atoms);
await vectorizeAtoms(chatId, atoms);
setL0FloorStatus(floor, { status: 'ok', atoms: atoms.length });
builtAtoms += atoms.length;
}
} catch (e) {
setL0FloorStatus(floor, {
status: 'fail',
attempts: (prev?.attempts || 0) + 1,
reason: String(e?.message || e).replace(/\s+/g, ' ').slice(0, 120),
});
} finally {
completed++;
onProgress?.(`L0: ${completed}/${total}`, completed, total);
}
}
xbLog.info(MODULE_ID, `增量 L0 完成atoms=${builtAtoms}, floors=${pendingPairs.length}`);
return { built: builtAtoms };
}
async function vectorizeAtoms(chatId, atoms) {
if (!atoms?.length) return;
const vectorCfg = getVectorConfig();
if (!vectorCfg?.enabled) return;
const texts = atoms.map(a => a.semantic);
const fingerprint = getEngineFingerprint(vectorCfg);
try {
const vectors = await embed(texts, { timeout: 30000 });
const items = atoms.map((a, i) => ({
atomId: a.atomId,
floor: a.floor,
vector: vectors[i],
}));
await saveStateVectors(chatId, items, fingerprint);
xbLog.info(MODULE_ID, `L0 向量化完成: ${items.length}`);
} catch (e) {
xbLog.error(MODULE_ID, 'L0 向量化失败', e);
}
}
// ============================================================================
// 清空
// ============================================================================
export async function clearAllAtomsAndVectors(chatId) {
clearStateAtoms();
clearL0Index();
if (chatId) {
await clearStateVectors(chatId);
}
xbLog.info(MODULE_ID, '已清空所有记忆锚点');
}
// ============================================================================
// 实时增量AI 消息后触发)- 保留原有逻辑
// ============================================================================
let extractionQueue = [];
let isProcessing = false;
export async function extractAndStoreAtomsForRound(aiFloor, aiMessage, userMessage) {
const { chatId } = getContext();
if (!chatId) return;
const vectorCfg = getVectorConfig();
if (!vectorCfg?.enabled) return;
extractionQueue.push({ aiFloor, aiMessage, userMessage, chatId });
processQueue();
}
async function processQueue() {
if (isProcessing || extractionQueue.length === 0) return;
isProcessing = true;
while (extractionQueue.length > 0) {
const { aiFloor, aiMessage, userMessage, chatId } = extractionQueue.shift();
try {
const atoms = await extractAtomsForRound(userMessage, aiMessage, aiFloor, { timeout: 12000 });
if (!atoms?.length) {
xbLog.info(MODULE_ID, `floor ${aiFloor}: 无有效 atoms`);
continue;
}
atoms.forEach(a => a.chatId = chatId);
saveStateAtoms(atoms);
await vectorizeAtoms(chatId, atoms);
xbLog.info(MODULE_ID, `floor ${aiFloor}: ${atoms.length} atoms 已存储`);
} catch (e) {
xbLog.error(MODULE_ID, `floor ${aiFloor} 处理失败`, e);
}
}
isProcessing = false;
}
// ============================================================================
// 回滚钩子
// ============================================================================
async function handleStateRollback(floor) {
xbLog.info(MODULE_ID, `收到回滚请求: floor >= ${floor}`);
const { chatId } = getContext();
deleteStateAtomsFromFloor(floor);
deleteL0IndexFromFloor(floor);
if (chatId) {
await deleteStateVectorsFromFloor(chatId, floor);
}
}
// ============================================================================
// 兼容旧接口
// ============================================================================
export async function batchExtractAndStoreAtoms(chatId, chat, onProgress) {
if (!chatId || !chat?.length) return { built: 0 };
const vectorCfg = getVectorConfig();
if (!vectorCfg?.enabled) return { built: 0 };
xbLog.info(MODULE_ID, `开始批量 L0 提取: ${chat.length} 条消息`);
clearStateAtoms();
clearL0Index();
await clearStateVectors(chatId);
return await incrementalExtractAtoms(chatId, chat, onProgress);
}
export async function rebuildStateVectors(chatId, vectorCfg) {
if (!chatId || !vectorCfg?.enabled) return { built: 0 };
const atoms = getStateAtoms();
if (!atoms.length) return { built: 0 };
xbLog.info(MODULE_ID, `重建 L0 向量: ${atoms.length} 条 atom`);
await clearStateVectors(chatId);
await vectorizeAtoms(chatId, atoms);
return { built: atoms.length };
}