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

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");