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

// ============================================================================
// 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 };
}