feat: 参数化检索链路与注入语义

README.md

@@ -185,10 +185,12 @@ git clone https://github.com/pjm0616/ST-Bionic-Memory-Ecology.git st-bme
 
 | 设置 | 默认 | 说明 |
 |------|------|------|
-| 召回上下文轮数 | 3 | 让 AI 根据最近几轮来召回 |
-| 向量候选数 | 15 | 向量搜索返回多少候选 |
-| 扩散深度 | 2 | 在图上沿关系扩散几层 |
-| 最终注入数 | 12 | 最终注入多少条记忆 |
+| 提取上下文轮数 | 2 | 按轮计的提取上下文，通常约等于向前补 4 层普通消息 |
+| 向量预筛 Top-K | 20 | 向量预筛阶段最多保留多少个候选 |
+| LLM 精排候选池 | 30 | 进入 LLM 精排阶段前的候选池大小 |
+| LLM 最终选择上限 | 8 | LLM 精排后最多保留多少条记忆 |
+| 图扩散 Top-K | 100 | 图扩散阶段最多保留多少个候选 |
+| 注入深度 | 9999 | 当前走 IN_CHAT@Depth，数值越大越靠前插入 |
 | Token 预算 | 1500 | 注入的最大 token 估算 |
 
 ---
@@ -230,6 +232,8 @@ git clone https://github.com/pjm0616/ST-Bionic-Memory-Ecology.git st-bme
 - 面板外观
 
 桌面端会显示左侧竖向子导航，右侧显示宽版配置表单；移动端则改成顶部横向子页切换。
+检索流水线现在可以分别配置向量预筛、图扩散、混合评分和 LLM 精排。
+注入深度使用 `IN_CHAT@Depth` 语义，默认 `9999` 表示尽量靠前插入，减少对最近几层对话的直接控制感。
 
 ### 图谱可视化
 桌面端右侧大区域显示力导向图谱，节点可拖拽、缩放、点击查看详情。支持 4 套配色主题切换。

index.js

@@ -4,6 +4,7 @@
 import {
   eventSource,
   event_types,
+  extension_prompt_types,
   getRequestHeaders,
   saveSettingsDebounced,
 } from "../../../../script.js";
@@ -75,13 +76,17 @@ const defaultSettings = {
 
   // 召回设置
   recallEnabled: true,
-  recallTopK: 15, // 混合评分 Top-K
+  recallTopK: 20, // 向量预筛 Top-K
   recallMaxNodes: 8, // LLM 召回最大节点数
   recallEnableLLM: true, // 是否启用 LLM 精确召回
+  recallEnableVectorPrefilter: true, // 是否启用向量预筛
+  recallEnableGraphDiffusion: true, // 是否启用图扩散
+  recallDiffusionTopK: 100, // 图扩散阶段保留的候选上限
+  recallLlmCandidatePool: 30, // 传给 LLM 精排的候选池大小
 
   // 注入设置
   injectPosition: "atDepth", // 注入位置
-  injectDepth: 4, // 注入深度（atDepth 模式）
+  injectDepth: 9999, // IN_CHAT@Depth 注入深度，数值越大越靠前
   injectRole: 0, // 0=system, 1=user, 2=assistant
 
   // 混合评分权重
@@ -273,7 +278,12 @@ function clearInjectionState() {
 
   try {
     const context = getContext();
-    context.setExtensionPrompt(MODULE_NAME, "", 1, 0);
+    context.setExtensionPrompt(
+      MODULE_NAME,
+      "",
+      extension_prompt_types.IN_CHAT,
+      0,
+    );
   } catch (error) {
     console.warn("[ST-BME] 清理旧注入失败:", error);
   }
@@ -538,7 +548,12 @@ function updateModuleSettings(patch = {}) {
   ) {
     try {
       const context = getContext();
-      context.setExtensionPrompt(MODULE_NAME, "", 1, 0);
+      context.setExtensionPrompt(
+        MODULE_NAME,
+        "",
+        extension_prompt_types.IN_CHAT,
+        0,
+      );
       lastInjectionContent = "";
       lastRecalledItems = [];
     } catch (error) {
@@ -1194,6 +1209,10 @@ async function runRecall() {
         topK: settings.recallTopK,
         maxRecallNodes: settings.recallMaxNodes,
         enableLLMRecall: settings.recallEnableLLM,
+        enableVectorPrefilter: settings.recallEnableVectorPrefilter,
+        enableGraphDiffusion: settings.recallEnableGraphDiffusion,
+        diffusionTopK: settings.recallDiffusionTopK,
+        llmCandidatePool: settings.recallLlmCandidatePool,
         recallPrompt: settings.recallPrompt || undefined,
         weights: {
           graphWeight: settings.graphWeight,
@@ -1224,8 +1243,8 @@ async function runRecall() {
     context.setExtensionPrompt(
       MODULE_NAME,
       injectionText,
-      1, // extension_prompt_types.IN_PROMPT
-      clampInt(settings.injectDepth, 4, 0, 9999),
+      extension_prompt_types.IN_CHAT, // 当前注入走 IN_CHAT@Depth
+      clampInt(settings.injectDepth, 9999, 0, 9999),
     );
 
     // 保存召回结果和访问强化

panel.html

@@ -512,13 +512,6 @@
                       </span>
                       <input id="bme-setting-recall-enabled" type="checkbox" />
                     </label>
-                    <label class="bme-toggle-item" for="bme-setting-recall-llm">
-                      <span class="bme-toggle-copy">
-                        <span class="bme-toggle-title">启用 LLM 精确召回</span>
-                        <span class="bme-toggle-desc">让候选记忆经过额外的 LLM 精排筛选。</span>
-                      </span>
-                      <input id="bme-setting-recall-llm" type="checkbox" />
-                    </label>
                     <label class="bme-toggle-item" for="bme-setting-evolution-enabled">
                       <span class="bme-toggle-copy">
                         <span class="bme-toggle-title">启用记忆进化</span>
@@ -614,7 +607,7 @@
                   <div class="bme-config-card-head">
                     <div>
                       <div class="bme-config-card-title">提取</div>
-                      <div class="bme-config-card-subtitle">控制自动提取的频率和上下文窗口。</div>
+                      <div class="bme-config-card-subtitle">控制自动提取的频率和按轮计的上下文窗口。默认 2 轮，通常约等于向前补 4 层普通消息。</div>
                     </div>
                   </div>
                   <div class="bme-config-row">
@@ -637,31 +630,57 @@
                   <div class="bme-config-card-head">
                     <div>
                       <div class="bme-config-card-title">召回与注入</div>
-                      <div class="bme-config-card-subtitle">控制候选规模和最终注入层级。</div>
+                      <div class="bme-config-card-subtitle">控制最终注入层级。当前注入走 IN_CHAT@Depth，数值越大越靠前。</div>
                     </div>
                     <div class="bme-config-guard-note">在“功能开关”中启用后生效。</div>
                   </div>
                   <div class="bme-config-row">
-                    <label for="bme-setting-recall-top-k">召回候选上限</label>
-                    <input id="bme-setting-recall-top-k" class="bme-config-input" type="number" min="1" max="100" />
-                  </div>
-                  <div class="bme-config-row">
-                    <label for="bme-setting-inject-depth">注入深度</label>
+                    <label for="bme-setting-inject-depth">注入深度（数值越大越靠前）</label>
                     <input id="bme-setting-inject-depth" class="bme-config-input" type="number" min="0" max="9999" />
                   </div>
                 </div>
 
-                <div class="bme-config-card bme-guarded-card" data-guard-settings="recallEnabled,recallEnableLLM">
+                <div
+                  class="bme-config-card bme-guarded-card bme-stage-card"
+                  data-guard-settings="recallEnabled"
+                  data-stage-toggle-id="bme-setting-recall-vector-prefilter-enabled"
+                >
                   <div class="bme-config-card-head">
                     <div>
-                      <div class="bme-config-card-title">LLM 精确召回</div>
-                      <div class="bme-config-card-subtitle">控制进入 LLM 精排阶段的节点数量。</div>
+                      <div class="bme-config-card-title">向量预筛</div>
+                      <div class="bme-config-card-subtitle">控制是否启用向量候选检索，以及向量阶段输出的 Top-K 宽度。</div>
                     </div>
                     <div class="bme-config-guard-note">在“功能开关”中启用后生效。</div>
                   </div>
-                  <div class="bme-config-row">
-                    <label for="bme-setting-recall-max-nodes">LLM 精确召回上限</label>
-                    <input id="bme-setting-recall-max-nodes" class="bme-config-input" type="number" min="1" max="50" />
+                  <label class="bme-inline-checkbox" for="bme-setting-recall-vector-prefilter-enabled">
+                    <input id="bme-setting-recall-vector-prefilter-enabled" type="checkbox" />
+                    <span>启用向量预筛</span>
+                  </label>
+                  <div class="bme-config-row bme-stage-param">
+                    <label for="bme-setting-recall-top-k">向量预筛 Top-K</label>
+                    <input id="bme-setting-recall-top-k" class="bme-config-input" type="number" min="1" max="100" />
+                  </div>
+                </div>
+
+                <div
+                  class="bme-config-card bme-guarded-card bme-stage-card"
+                  data-guard-settings="recallEnabled"
+                  data-stage-toggle-id="bme-setting-recall-graph-diffusion-enabled"
+                >
+                  <div class="bme-config-card-head">
+                    <div>
+                      <div class="bme-config-card-title">图扩散</div>
+                      <div class="bme-config-card-subtitle">控制是否沿图结构继续扩散候选，以及扩散阶段保留的候选上限。</div>
+                    </div>
+                    <div class="bme-config-guard-note">在“功能开关”中启用后生效。</div>
+                  </div>
+                  <label class="bme-inline-checkbox" for="bme-setting-recall-graph-diffusion-enabled">
+                    <input id="bme-setting-recall-graph-diffusion-enabled" type="checkbox" />
+                    <span>启用图扩散</span>
+                  </label>
+                  <div class="bme-config-row bme-stage-param">
+                    <label for="bme-setting-recall-diffusion-top-k">图扩散 Top-K</label>
+                    <input id="bme-setting-recall-diffusion-top-k" class="bme-config-input" type="number" min="1" max="300" />
                   </div>
                 </div>
 
@@ -669,7 +688,7 @@
                   <div class="bme-config-card-head">
                     <div>
                       <div class="bme-config-card-title">混合评分</div>
-                      <div class="bme-config-card-subtitle">平衡图扩散、向量相似和重要度在召回中的占比。</div>
+                      <div class="bme-config-card-subtitle">评分层始终运行；下面 3 个权重共同决定图扩散、向量相似和重要度在召回中的占比。</div>
                     </div>
                     <div class="bme-config-guard-note">在“功能开关”中启用后生效。</div>
                   </div>
@@ -694,6 +713,32 @@
                   </div>
                 </div>
 
+                <div
+                  class="bme-config-card bme-guarded-card bme-stage-card"
+                  data-guard-settings="recallEnabled"
+                  data-stage-toggle-id="bme-setting-recall-llm"
+                >
+                  <div class="bme-config-card-head">
+                    <div>
+                      <div class="bme-config-card-title">LLM 精确召回</div>
+                      <div class="bme-config-card-subtitle">控制是否启用 LLM 精排，以及传给 LLM 的候选池大小与最终保留上限。</div>
+                    </div>
+                    <div class="bme-config-guard-note">在“功能开关”中启用后生效。</div>
+                  </div>
+                  <label class="bme-inline-checkbox" for="bme-setting-recall-llm">
+                    <input id="bme-setting-recall-llm" type="checkbox" />
+                    <span>启用 LLM 精排</span>
+                  </label>
+                  <div class="bme-config-row bme-stage-param">
+                    <label for="bme-setting-recall-llm-candidate-pool">LLM 精排候选池</label>
+                    <input id="bme-setting-recall-llm-candidate-pool" class="bme-config-input" type="number" min="1" max="100" />
+                  </div>
+                  <div class="bme-config-row bme-stage-param">
+                    <label for="bme-setting-recall-max-nodes">LLM 最终选择上限</label>
+                    <input id="bme-setting-recall-max-nodes" class="bme-config-input" type="number" min="1" max="50" />
+                  </div>
+                </div>
+
                 <div class="bme-config-card bme-guarded-card" data-guard-settings="enableEvolution">
                   <div class="bme-config-card-head">
                     <div>

panel.js

@@ -663,6 +663,14 @@ function _refreshConfigTab() {
     _setCheckboxValue("bme-setting-enabled", settings.enabled ?? false);
     _setCheckboxValue("bme-setting-recall-enabled", settings.recallEnabled ?? true);
     _setCheckboxValue("bme-setting-recall-llm", settings.recallEnableLLM ?? true);
+    _setCheckboxValue(
+        "bme-setting-recall-vector-prefilter-enabled",
+        settings.recallEnableVectorPrefilter ?? true,
+    );
+    _setCheckboxValue(
+        "bme-setting-recall-graph-diffusion-enabled",
+        settings.recallEnableGraphDiffusion ?? true,
+    );
     _setCheckboxValue("bme-setting-evolution-enabled", settings.enableEvolution ?? true);
     _setCheckboxValue(
         "bme-setting-precise-conflict-enabled",
@@ -699,9 +707,17 @@ function _refreshConfigTab() {
         "bme-setting-extract-context-turns",
         settings.extractContextTurns ?? 2,
     );
-    _setInputValue("bme-setting-recall-top-k", settings.recallTopK ?? 15);
+    _setInputValue("bme-setting-recall-top-k", settings.recallTopK ?? 20);
     _setInputValue("bme-setting-recall-max-nodes", settings.recallMaxNodes ?? 8);
-    _setInputValue("bme-setting-inject-depth", settings.injectDepth ?? 4);
+    _setInputValue(
+        "bme-setting-recall-diffusion-top-k",
+        settings.recallDiffusionTopK ?? 100,
+    );
+    _setInputValue(
+        "bme-setting-recall-llm-candidate-pool",
+        settings.recallLlmCandidatePool ?? 30,
+    );
+    _setInputValue("bme-setting-inject-depth", settings.injectDepth ?? 9999);
     _setInputValue("bme-setting-graph-weight", settings.graphWeight ?? 0.6);
     _setInputValue("bme-setting-vector-weight", settings.vectorWeight ?? 0.3);
     _setInputValue(
@@ -780,6 +796,7 @@ function _refreshConfigTab() {
     _setInputValue("bme-setting-reflection-prompt", settings.reflectionPrompt || DEFAULT_PROMPTS.reflection);
 
     _refreshGuardedConfigStates(settings);
+    _refreshStageCardStates(settings);
     _refreshPromptCardStates(settings);
     _highlightThemeChoice(settings.panelTheme || "crimson");
     _syncConfigSectionState();
@@ -803,10 +820,20 @@ function _bindConfigControls() {
     bindCheckbox("bme-setting-recall-enabled", (checked) => {
         _patchSettings({ recallEnabled: checked });
         _refreshGuardedConfigStates();
+        _refreshStageCardStates();
     });
     bindCheckbox("bme-setting-recall-llm", (checked) => {
         _patchSettings({ recallEnableLLM: checked });
         _refreshGuardedConfigStates();
+        _refreshStageCardStates();
+    });
+    bindCheckbox("bme-setting-recall-vector-prefilter-enabled", (checked) => {
+        _patchSettings({ recallEnableVectorPrefilter: checked });
+        _refreshStageCardStates();
+    });
+    bindCheckbox("bme-setting-recall-graph-diffusion-enabled", (checked) => {
+        _patchSettings({ recallEnableGraphDiffusion: checked });
+        _refreshStageCardStates();
     });
     bindCheckbox("bme-setting-evolution-enabled", (checked) => {
         _patchSettings({ enableEvolution: checked });
@@ -849,13 +876,19 @@ function _bindConfigControls() {
     bindNumber("bme-setting-extract-context-turns", 2, 0, 20, (value) =>
         _patchSettings({ extractContextTurns: value }),
     );
-    bindNumber("bme-setting-recall-top-k", 15, 1, 100, (value) =>
+    bindNumber("bme-setting-recall-top-k", 20, 1, 100, (value) =>
         _patchSettings({ recallTopK: value }),
     );
     bindNumber("bme-setting-recall-max-nodes", 8, 1, 50, (value) =>
         _patchSettings({ recallMaxNodes: value }),
     );
-    bindNumber("bme-setting-inject-depth", 4, 0, 9999, (value) =>
+    bindNumber("bme-setting-recall-diffusion-top-k", 100, 1, 300, (value) =>
+        _patchSettings({ recallDiffusionTopK: value }),
+    );
+    bindNumber("bme-setting-recall-llm-candidate-pool", 30, 1, 100, (value) =>
+        _patchSettings({ recallLlmCandidatePool: value }),
+    );
+    bindNumber("bme-setting-inject-depth", 9999, 0, 9999, (value) =>
         _patchSettings({ injectDepth: value }),
     );
     bindFloat("bme-setting-graph-weight", 0.6, 0, 1, (value) =>
@@ -1131,6 +1164,28 @@ function _refreshGuardedConfigStates(settings = _getSettings?.() || {}) {
     });
 }
 
+function _refreshStageCardStates(settings = _getSettings?.() || {}) {
+    if (!panelEl) return;
+    panelEl.querySelectorAll(".bme-stage-card").forEach((card) => {
+        const toggleId = card.dataset.stageToggleId;
+        const toggle = toggleId ? document.getElementById(toggleId) : null;
+        const cardDisabled = card.classList.contains("is-disabled");
+        const stageEnabled =
+            toggleId === "bme-setting-recall-llm"
+                ? settings.recallEnableLLM ?? true
+                : toggle
+                  ? Boolean(toggle.checked)
+                  : true;
+
+        card.classList.toggle("stage-disabled", !cardDisabled && !stageEnabled);
+        card.querySelectorAll(".bme-stage-param").forEach((section) => {
+            section.querySelectorAll("input, select, textarea, button").forEach((element) => {
+                element.disabled = cardDisabled || !stageEnabled;
+            });
+        });
+    });
+}
+
 function _refreshPromptCardStates(settings = _getSettings?.() || {}) {
     if (!panelEl) return;
     panelEl.querySelectorAll(".bme-prompt-card").forEach((card) => {

retriever.js

@@ -13,15 +13,6 @@ import {
 import { callLLMForJSON } from "./llm.js";
 import { findSimilarNodesByText, validateVectorConfig } from "./vector-index.js";
 
-/**
- * 自适应阈值
- */
-const STRATEGY_THRESHOLDS = {
-  SMALL: 20, // < 20 节点：跳过向量，全图 + LLM
-  MEDIUM: 200, // 20-200 节点：向量 + 图扩散 + 评分（不调 LLM）
-  // > 200 节点：三层全开
-};
-
 /**
  * 三层混合检索管线
  *
@@ -42,9 +33,13 @@ export async function retrieve({
   schema,
   options = {},
 }) {
-  const topK = options.topK ?? 15;
+  const topK = options.topK ?? 20;
   const maxRecallNodes = options.maxRecallNodes ?? 8;
   const enableLLMRecall = options.enableLLMRecall ?? true;
+  const enableVectorPrefilter = options.enableVectorPrefilter ?? true;
+  const enableGraphDiffusion = options.enableGraphDiffusion ?? true;
+  const diffusionTopK = options.diffusionTopK ?? 100;
+  const llmCandidatePool = options.llmCandidatePool ?? 30;
   const weights = options.weights ?? {};
 
   // v2 options
@@ -69,6 +64,11 @@ export async function retrieve({
   const nodeCount = activeNodes.length;
   const normalizedTopK = Math.max(1, topK);
   const normalizedMaxRecallNodes = Math.max(1, maxRecallNodes);
+  const normalizedDiffusionTopK = Math.max(1, diffusionTopK);
+  const normalizedLlmCandidatePool = Math.max(
+    normalizedMaxRecallNodes,
+    llmCandidatePool,
+  );
   console.log(
     `[ST-BME] 检索开始: ${nodeCount} 个活跃节点${enableVisibility ? " (认知边界已启用)" : ""}`,
   );
@@ -83,7 +83,7 @@ export async function retrieve({
 
   // ========== 第 1 层：向量预筛 ==========
   if (
-    nodeCount >= STRATEGY_THRESHOLDS.SMALL &&
+    enableVectorPrefilter &&
     validateVectorConfig(embeddingConfig).valid
   ) {
     console.log("[ST-BME] 第1层: 向量预筛");
@@ -97,7 +97,7 @@ export async function retrieve({
   }
 
   // ========== 第 2 层：图扩散 ==========
-  if (nodeCount >= STRATEGY_THRESHOLDS.SMALL) {
+  if (enableGraphDiffusion) {
     console.log("[ST-BME] 第2层: PEDSA 图扩散");
     const entityAnchors = extractEntityAnchors(userMessage, activeNodes);
 
@@ -139,7 +139,7 @@ export async function retrieve({
       diffusionResults = diffuseAndRank(adjacencyMap, uniqueSeeds, {
         maxSteps: 2,
         decayFactor: 0.6,
-        topK: 100,
+        topK: normalizedDiffusionTopK,
       }).filter((item) => {
         const node = getNode(graph, item.nodeId);
         return node && !node.archived;
@@ -167,8 +167,8 @@ export async function retrieve({
     scoreMap.set(d.nodeId, entry);
   }
 
-  // 小图模式：所有节点都参与评分
-  if (nodeCount < STRATEGY_THRESHOLDS.SMALL) {
+  // 两个上游阶段都未产出候选时，退回到全部活跃节点参与评分
+  if (scoreMap.size === 0) {
     for (const node of activeNodes) {
       if (!scoreMap.has(node.id)) {
         scoreMap.set(node.id, { graphScore: 0, vectorScore: 0 });
@@ -198,10 +198,7 @@ export async function retrieve({
   scoredNodes.sort((a, b) => b.finalScore - a.finalScore);
 
   // 决定是否使用 LLM 精确召回
-  useLLM =
-    enableLLMRecall &&
-    (nodeCount < STRATEGY_THRESHOLDS.SMALL || // 小图：直接 LLM
-      nodeCount > STRATEGY_THRESHOLDS.MEDIUM); // 大图：LLM 精确
+  useLLM = enableLLMRecall;
 
   let selectedNodeIds;
 
@@ -209,7 +206,7 @@ export async function retrieve({
     console.log("[ST-BME] LLM 精确召回");
     const candidateNodes = scoredNodes.slice(
       0,
-      Math.min(30, scoredNodes.length),
+      Math.min(normalizedLlmCandidatePool, scoredNodes.length),
     );
     selectedNodeIds = await llmRecall(
       userMessage,

style.css

@@ -1006,6 +1006,32 @@
     opacity: 0.72;
 }
 
+.bme-stage-card.stage-disabled {
+    border-color: rgba(255, 255, 255, 0.08);
+}
+
+.bme-inline-checkbox {
+    display: inline-flex;
+    align-items: center;
+    gap: 10px;
+    margin-bottom: 12px;
+    font-size: 12px;
+    font-weight: 700;
+    color: var(--bme-on-surface);
+    cursor: pointer;
+}
+
+.bme-inline-checkbox input {
+    width: 16px;
+    height: 16px;
+    margin: 0;
+    accent-color: var(--bme-primary);
+}
+
+.bme-stage-card.stage-disabled .bme-stage-param {
+    opacity: 0.6;
+}
+
 .bme-config-placeholder {
     background: var(--bme-surface-low);
     border: 1px dashed var(--bme-border);

tests/default-settings.mjs

@@ -0,0 +1,37 @@
+import assert from "node:assert/strict";
+import fs from "node:fs/promises";
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+import vm from "node:vm";
+
+async function loadDefaultSettings() {
+  const __dirname = path.dirname(fileURLToPath(import.meta.url));
+  const indexPath = path.resolve(__dirname, "../index.js");
+  const source = await fs.readFile(indexPath, "utf8");
+  const settingsMatch = source.match(/const defaultSettings = \{[\s\S]*?^\};/m);
+
+  if (!settingsMatch) {
+    throw new Error("无法从 index.js 提取 defaultSettings");
+  }
+
+  const context = vm.createContext({});
+  const script = new vm.Script(`
+${settingsMatch[0]}
+this.defaultSettings = defaultSettings;
+`);
+  script.runInContext(context);
+  return context.defaultSettings;
+}
+
+const defaultSettings = await loadDefaultSettings();
+
+assert.equal(defaultSettings.extractContextTurns, 2);
+assert.equal(defaultSettings.recallTopK, 20);
+assert.equal(defaultSettings.recallMaxNodes, 8);
+assert.equal(defaultSettings.recallEnableVectorPrefilter, true);
+assert.equal(defaultSettings.recallEnableGraphDiffusion, true);
+assert.equal(defaultSettings.recallDiffusionTopK, 100);
+assert.equal(defaultSettings.recallLlmCandidatePool, 30);
+assert.equal(defaultSettings.injectDepth, 9999);
+
+console.log("default-settings tests passed");

tests/retrieval-config.mjs

@@ -0,0 +1,191 @@
+import assert from "node:assert/strict";
+import fs from "node:fs/promises";
+import path from "node:path";
+import { fileURLToPath } from "node:url";
+import vm from "node:vm";
+
+async function loadRetrieve(stubs) {
+  const __dirname = path.dirname(fileURLToPath(import.meta.url));
+  const retrieverPath = path.resolve(__dirname, "../retriever.js");
+  const source = await fs.readFile(retrieverPath, "utf8");
+  const transformed = `${source
+    .replace(/^import[\s\S]*?from\s+["'][^"']+["'];\r?\n/gm, "")
+    .replace("export async function retrieve", "async function retrieve")}
+this.retrieve = retrieve;
+`;
+
+  const context = vm.createContext({
+    console: { log() {}, error() {}, warn() {} },
+    ...stubs,
+  });
+  new vm.Script(transformed).runInContext(context);
+  return context.retrieve;
+}
+
+function createGraph() {
+  const nodes = [
+    {
+      id: "rule-1",
+      type: "rule",
+      importance: 9,
+      createdTime: 1,
+      archived: false,
+      fields: { title: "规则一" },
+      seqRange: [1, 1],
+    },
+    {
+      id: "rule-2",
+      type: "rule",
+      importance: 7,
+      createdTime: 2,
+      archived: false,
+      fields: { title: "规则二" },
+      seqRange: [2, 2],
+    },
+    {
+      id: "rule-3",
+      type: "rule",
+      importance: 3,
+      createdTime: 3,
+      archived: false,
+      fields: { title: "规则三" },
+      seqRange: [3, 3],
+    },
+  ];
+  return { nodes, edges: [] };
+}
+
+function createGraphHelpers(graph) {
+  return {
+    getActiveNodes(target, type = null) {
+      const source = target?.nodes || graph.nodes;
+      return source.filter(
+        (node) => !node.archived && (!type || node.type === type),
+      );
+    },
+    getNode(target, id) {
+      return (target?.nodes || graph.nodes).find((node) => node.id === id) || null;
+    },
+    getNodeEdges(target, nodeId) {
+      return (target?.edges || graph.edges).filter(
+        (edge) => edge.fromId === nodeId || edge.toId === nodeId,
+      );
+    },
+    buildTemporalAdjacencyMap() {
+      return new Map();
+    },
+  };
+}
+
+const schema = [{ id: "rule", label: "规则", alwaysInject: false }];
+
+const state = {
+  vectorCalls: [],
+  diffusionCalls: [],
+  llmCalls: [],
+  llmCandidateCount: 0,
+};
+
+const graph = createGraph();
+const helpers = createGraphHelpers(graph);
+const retrieve = await loadRetrieve({
+  ...helpers,
+  hybridScore: ({ graphScore = 0, vectorScore = 0, importance = 0 }) =>
+    graphScore + vectorScore + importance,
+  reinforceAccessBatch() {},
+  validateVectorConfig() {
+    return { valid: true };
+  },
+  async findSimilarNodesByText(_graph, _message, _embeddingConfig, topK) {
+    state.vectorCalls.push(topK);
+    return [
+      { nodeId: "rule-1", score: 0.9 },
+      { nodeId: "rule-2", score: 0.8 },
+      { nodeId: "rule-3", score: 0.7 },
+    ];
+  },
+  diffuseAndRank(_adjacencyMap, seeds, options) {
+    state.diffusionCalls.push({ seeds, options });
+    return [
+      { nodeId: "rule-2", energy: 1.2 },
+      { nodeId: "rule-3", energy: 0.9 },
+    ];
+  },
+  async callLLMForJSON({ userPrompt }) {
+    state.llmCalls.push(userPrompt);
+    state.llmCandidateCount = userPrompt
+      .split("\n")
+      .filter((line) => line.trim().startsWith("[")).length;
+    return { selected_ids: ["rule-2", "rule-1"] };
+  },
+});
+
+state.vectorCalls.length = 0;
+state.diffusionCalls.length = 0;
+state.llmCalls.length = 0;
+const noStageResult = await retrieve({
+  graph,
+  userMessage: "只看当前规则",
+  recentMessages: [],
+  embeddingConfig: {},
+  schema,
+  options: {
+    topK: 2,
+    maxRecallNodes: 2,
+    enableVectorPrefilter: false,
+    enableGraphDiffusion: false,
+    enableLLMRecall: false,
+  },
+});
+assert.equal(state.vectorCalls.length, 0);
+assert.equal(state.diffusionCalls.length, 0);
+assert.equal(state.llmCalls.length, 0);
+assert.deepEqual(Array.from(noStageResult.selectedNodeIds), ["rule-1", "rule-2"]);
+
+state.vectorCalls.length = 0;
+state.diffusionCalls.length = 0;
+state.llmCalls.length = 0;
+state.llmCandidateCount = 0;
+const llmPoolResult = await retrieve({
+  graph,
+  userMessage: "请根据规则给出结论",
+  recentMessages: ["用户：现在该怎么做？"],
+  embeddingConfig: {},
+  schema,
+  options: {
+    topK: 4,
+    maxRecallNodes: 2,
+    enableVectorPrefilter: true,
+    enableGraphDiffusion: false,
+    enableLLMRecall: true,
+    llmCandidatePool: 2,
+  },
+});
+assert.deepEqual(state.vectorCalls, [4]);
+assert.equal(state.diffusionCalls.length, 0);
+assert.equal(state.llmCandidateCount, 2);
+assert.deepEqual(Array.from(llmPoolResult.selectedNodeIds), ["rule-2", "rule-1"]);
+
+state.vectorCalls.length = 0;
+state.diffusionCalls.length = 0;
+state.llmCalls.length = 0;
+await retrieve({
+  graph,
+  userMessage: "规则一和规则二有什么关联",
+  recentMessages: [],
+  embeddingConfig: {},
+  schema,
+  options: {
+    topK: 3,
+    maxRecallNodes: 2,
+    enableVectorPrefilter: true,
+    enableGraphDiffusion: true,
+    diffusionTopK: 7,
+    enableLLMRecall: false,
+  },
+});
+assert.deepEqual(state.vectorCalls, [3]);
+assert.equal(state.diffusionCalls.length, 1);
+assert.equal(state.diffusionCalls[0].options.topK, 7);
+
+console.log("retrieval-config tests passed");