@memory-river/core 0.2.0

package/dist/retrieval/cross-encoder-gate.js

@@ -0,0 +1,239 @@
+import * as fs from "node:fs";
+import * as os from "node:os";
+import * as path from "node:path";
+import { performance } from "node:perf_hooks";
+const MMARCO_MMINILM_RERANKER_MODEL = {
+    modelId: "cross-encoder/mmarco-mMiniLMv2-L12-H384-v1",
+    revision: "main",
+    onnxCandidates: [
+        "onnx/model_quint8_avx2.onnx",
+        "onnx/model_quantized.onnx",
+        "onnx/model.onnx",
+    ],
+};
+const DEFAULT_TOP_K = 5;
+const DEFAULT_ZH_LOGIT = -7.0;
+const DEFAULT_EN_LOGIT = 3.47;
+const DEFAULT_THREADS = 4;
+const DEFAULT_BATCH_SIZE = 8;
+const DEFAULT_MAX_LENGTH = 512;
+const DEFAULT_PASSAGE_TOKEN_LIMIT = 320;
+const CJK_RE = /[\u3400-\u4dbf\u4e00-\u9fff\uf900-\ufaff\u3040-\u30ff\uac00-\ud7af]/u;
+let scorerPromise = null;
+let warningEmitted = false;
+let scorerOverride = null;
+export function isCragCrossEncoderGateEnabled(env = process.env) {
+    // Enabled by default (fix for the dead cosine gate); disable with MR_CRAG_CROSS_ENCODER=0.
+    return env.MR_CRAG_CROSS_ENCODER !== "0" && env.ENABLE_CRAG_GATE !== "0";
+}
+export function containsCjk(text) {
+    return CJK_RE.test(text);
+}
+export function cragGateThresholdForText(query, passage, env = process.env) {
+    const raw = containsCjk(`${query}\n${passage}`)
+        ? env.MR_CRAG_GATE_ZH_LOGIT
+        : env.MR_CRAG_GATE_EN_LOGIT;
+    const parsed = raw === undefined ? Number.NaN : Number(raw);
+    if (Number.isFinite(parsed))
+        return parsed;
+    return containsCjk(`${query}\n${passage}`) ? DEFAULT_ZH_LOGIT : DEFAULT_EN_LOGIT;
+}
+export function cragGateTopK(env = process.env) {
+    const parsed = Number(env.MR_CRAG_GATE_TOPK);
+    return Number.isInteger(parsed) && parsed > 0 ? parsed : DEFAULT_TOP_K;
+}
+function candidateScore(candidate) {
+    return Number(candidate.finalScore ?? candidate.fusedScore ?? 0);
+}
+export async function scoreCandidates(query, candidates, options) {
+    const gateOptions = options ?? {};
+    const env = gateOptions.env ?? process.env;
+    if (!isCragCrossEncoderGateEnabled(env) || candidates.length === 0)
+        return null;
+    const topK = gateOptions.topK ?? cragGateTopK(env);
+    const shortlist = candidates
+        .map((candidate, index) => ({ candidate, index, score: candidateScore(candidate) }))
+        .sort((a, b) => b.score - a.score || a.index - b.index)
+        .slice(0, Math.min(topK, candidates.length));
+    if (shortlist.length === 0)
+        return null;
+    const scorer = gateOptions.scorer ?? await getCrossEncoderScorer(gateOptions);
+    if (!scorer)
+        return null;
+    const pairs = shortlist.map(({ candidate }) => ({
+        query,
+        passage: candidate.entry.text,
+    }));
+    const started = performance.now();
+    const { logits } = await scorer.scorePairs(pairs);
+    const timingMs = performance.now() - started;
+    return {
+        scored: shortlist.map((shortlisted, index) => ({
+            candidate: shortlisted.candidate,
+            logit: Number(logits[index]),
+        })),
+        timingMs,
+    };
+}
+export async function applyCragCrossEncoderGate(query, candidates, options = {}) {
+    const env = options.env ?? process.env;
+    const result = await scoreCandidates(query, candidates, options);
+    if (result === null)
+        return candidates;
+    const rejected = new Set();
+    for (const { candidate, logit } of result.scored) {
+        const threshold = cragGateThresholdForText(query, candidate.entry.text, env);
+        if (Number.isFinite(logit) && logit < threshold) {
+            rejected.add(candidate);
+        }
+    }
+    return candidates.filter(candidate => !rejected.has(candidate));
+}
+async function getCrossEncoderScorer(options) {
+    if (scorerOverride)
+        return scorerOverride;
+    if (!scorerPromise) {
+        scorerPromise = loadCrossEncoderScorer(options).catch((err) => {
+            const logger = options.logger ?? console;
+            if (!warningEmitted) {
+                warningEmitted = true;
+                logger.warn("[CRAG] cross-encoder gate unavailable; passing candidates through:", err?.message ?? err);
+            }
+            return null;
+        });
+    }
+    return scorerPromise;
+}
+function defaultCacheDir(env = process.env) {
+    return env.TRANSFORMERS_CACHE
+        ?? env.HF_HOME
+        ?? path.join(os.homedir(), ".cache", "huggingface");
+}
+function modelDirCandidates(spec, options) {
+    const cacheDir = options.cacheDir ?? defaultCacheDir(options.env);
+    return [...new Set([
+            ...(options.modelDir ? [options.modelDir] : []),
+            path.join(cacheDir, spec.modelId),
+            path.join(cacheDir, ...spec.modelId.split("/")),
+        ])];
+}
+function findOnnx(modelDir, candidates) {
+    for (const relative of candidates) {
+        const candidate = path.join(modelDir, relative);
+        if (fs.existsSync(candidate))
+            return candidate;
+    }
+    return null;
+}
+function resolveModel(options) {
+    for (const modelDir of modelDirCandidates(MMARCO_MMINILM_RERANKER_MODEL, options)) {
+        const onnxPath = findOnnx(modelDir, MMARCO_MMINILM_RERANKER_MODEL.onnxCandidates);
+        if (onnxPath)
+            return { spec: MMARCO_MMINILM_RERANKER_MODEL, modelDir, onnxPath };
+    }
+    throw new Error(`${MMARCO_MMINILM_RERANKER_MODEL.modelId} ONNX not found under ` +
+        modelDirCandidates(MMARCO_MMINILM_RERANKER_MODEL, options).join(", "));
+}
+function innerTokenIds(tokenizer, text) {
+    const encoded = tokenizer.encode(text);
+    if (encoded.length >= 2 && encoded[0] === 0 && encoded[encoded.length - 1] === 2) {
+        return encoded.slice(1, -1);
+    }
+    return encoded;
+}
+function buildPairIds(input) {
+    const queryTokens = innerTokenIds(input.tokenizer, input.query);
+    const passageTokens = innerTokenIds(input.tokenizer, input.passage);
+    const maxPayload = Math.max(0, input.maxLength - 4);
+    const queryBudget = Math.min(queryTokens.length, maxPayload);
+    const query = queryTokens.slice(0, queryBudget);
+    const remaining = Math.max(0, maxPayload - query.length);
+    const passageBudget = Math.min(input.passageTokenLimit, remaining);
+    const passage = passageTokens.slice(0, passageBudget);
+    return [0, ...query, 2, 2, ...passage, 2];
+}
+function makeTensor(ort, values, dims) {
+    return new ort.Tensor("int64", BigInt64Array.from(values.map(value => BigInt(value))), [...dims]);
+}
+async function loadCrossEncoderScorer(options) {
+    const resolved = resolveModel(options);
+    const started = performance.now();
+    const [{ AutoTokenizer, env }, ortNamespace] = await Promise.all([
+        import("@xenova/transformers"),
+        import("onnxruntime-node"),
+    ]);
+    const ort = ortNamespace.default ?? ortNamespace;
+    env.allowRemoteModels = false;
+    env.localModelPath = options.cacheDir ?? defaultCacheDir(options.env);
+    env.cacheDir = options.cacheDir ?? defaultCacheDir(options.env);
+    const tokenizer = await AutoTokenizer.from_pretrained(resolved.spec.modelId, {
+        revision: resolved.spec.revision,
+    });
+    const session = await ort.InferenceSession.create(resolved.onnxPath, {
+        executionProviders: ["cpu"],
+        intraOpNumThreads: DEFAULT_THREADS,
+        interOpNumThreads: 1,
+        graphOptimizationLevel: "all",
+    });
+    const scorer = new OnnxCrossEncoderScorer(ort, session, tokenizer);
+    await scorer.scorePairs([{ query: "warmup query", passage: "warmup passage" }]);
+    options.logger?.log?.(`[CRAG] cross-encoder gate loaded ${resolved.onnxPath} in ${Math.round(performance.now() - started)}ms`);
+    return scorer;
+}
+class OnnxCrossEncoderScorer {
+    ort;
+    session;
+    tokenizer;
+    constructor(ort, session, tokenizer) {
+        this.ort = ort;
+        this.session = session;
+        this.tokenizer = tokenizer;
+    }
+    async scorePairs(pairs) {
+        const logits = [];
+        for (let offset = 0; offset < pairs.length; offset += DEFAULT_BATCH_SIZE) {
+            const batch = pairs.slice(offset, offset + DEFAULT_BATCH_SIZE);
+            logits.push(...await this.runBatch(batch));
+        }
+        return { logits };
+    }
+    async runBatch(pairs) {
+        if (pairs.length === 0)
+            return [];
+        const encoded = pairs.map(pair => buildPairIds({
+            tokenizer: this.tokenizer,
+            query: pair.query,
+            passage: pair.passage,
+            maxLength: DEFAULT_MAX_LENGTH,
+            passageTokenLimit: DEFAULT_PASSAGE_TOKEN_LIMIT,
+        }));
+        const seqLength = Math.max(...encoded.map(item => item.length));
+        const inputIds = [];
+        const attentionMask = [];
+        for (const ids of encoded) {
+            inputIds.push(...ids);
+            attentionMask.push(...ids.map(id => id === 1 ? 0 : 1));
+            for (let index = ids.length; index < seqLength; index++) {
+                inputIds.push(1);
+                attentionMask.push(0);
+            }
+        }
+        const outputs = await this.session.run({
+            input_ids: makeTensor(this.ort, inputIds, [pairs.length, seqLength]),
+            attention_mask: makeTensor(this.ort, attentionMask, [pairs.length, seqLength]),
+        });
+        if (!outputs.logits)
+            throw new Error("ONNX session did not return logits");
+        return Array.from(outputs.logits.data).map(Number);
+    }
+}
+export function __setCragCrossEncoderScorerForTests(scorer) {
+    scorerOverride = scorer;
+    scorerPromise = null;
+    warningEmitted = false;
+}
+export function __resetCragCrossEncoderForTests() {
+    scorerOverride = null;
+    scorerPromise = null;
+    warningEmitted = false;
+}

package/dist/retrieval/retriever-v4.d.ts

@@ -0,0 +1,78 @@
+/**
+ * Retriever - Hybrid Search (正確實現 RRF Fusion) + CRAG Evaluator
+ * memory-lance-v4
+ * * 核心功能：
+ * - 混合檢索 (向量 + BM25) 與 RRF Fusion
+ * - V4 記憶權重計分 (相似度 + 重要度 + 健康度)
+ * - 本地毫秒級 CRAG 評估器 (MiniLM 餘弦相似度)
+ * - 赫布理論 (Hebbian Learning) 自動強化記憶
+ */
+import { MemoryStore, type MemorySearchResult } from "../store/store-v4.js";
+import { Embedder } from "../providers/embedder-v5.js";
+import { HooksEngine, type HookTriggerResult } from "../cognition/hooks-engine.js";
+import type { PluginConfig, MemoryEntry } from "../types.js";
+export interface HybridSearchResponse {
+    results: MemorySearchResult[];
+    hookOriginIds: string[];
+    hookOriginKeywords: Record<string, string>;
+    queryHash: string;
+}
+type CausalChainNode = {
+    entry: MemoryEntry;
+    hopFromSeed: number;
+    origin: 'parent' | 'child' | 'seed';
+};
+export declare class Retriever {
+    private store;
+    private embedder;
+    private readonly rerankerCacheDir;
+    private vectorWeight;
+    private bm25Weight;
+    private candidatePoolMultiplier;
+    private hooksEngine;
+    private boostHealthQueue;
+    private boostHealthQueueRunning;
+    private recallMetadataQueue;
+    private recallMetadataQueueRunning;
+    constructor(store: MemoryStore, embedder: Embedder, config: PluginConfig["retrieval"], rerankerCacheDir: string, hooksEngine?: HooksEngine);
+    hybridSearch(query: string, limit?: number): Promise<HybridSearchResponse>;
+    hybridSearchWithoutBoost(query: string, limit?: number): Promise<HybridSearchResponse>;
+    private hybridSearchInternal;
+    private recordHookEffectiveness;
+    private enqueueRecallMetadata;
+    private drainRecallMetadataQueue;
+    private enqueueBoostHealth;
+    private drainBoostHealthQueue;
+    private recordQueueDrop;
+    private recordBackgroundWriteFailure;
+    private extractFragmentId;
+    vectorOnly(query: string, limit?: number): Promise<MemorySearchResult[]>;
+    ftsOnly(query: string, limit?: number): Promise<MemorySearchResult[]>;
+    /**
+     * 追蹤因果鏈：根據 parentId 向上找原因、向下找結果
+     * @param memoryId 起點記憶 ID
+     * @param depth 深度（預設 2：爺爺→爸爸→我→孩子→孫子）
+     * @returns 因果鏈上的所有 MemoryEntry
+     */
+    getCausalChain(memoryId: string, depth?: number): Promise<CausalChainNode[]>;
+    setHooksEngine(engine: HooksEngine): void;
+    getStore(): MemoryStore;
+    updateMemoryRecord(id: string, updates: {
+        text?: string;
+        importance?: number;
+    }): Promise<void>;
+    triggerHooks(query: string): Promise<HookTriggerResult>;
+    cragEvaluate(query: string, results: MemorySearchResult[], enableHebbian?: boolean, hookInjectedIds?: Set<string>): Promise<MemorySearchResult[]>;
+    /**
+     * Entity Synergy Merge 實作
+     *
+     * 對所有 partial 記憶：
+     *  1. 兩兩抽取實體，計算 entity overlap
+     *  2. overlap 超過門檻的記憶對 → 合併文字
+     *  3. 合併後用 MiniLM 重新評估 relevancy
+     *  4. 通過的升級進結果池
+     */
+    private trySynergyMerge;
+    private filterHooks;
+}
+export {};