pi-mega-compact 0.8.22 → 0.8.23

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
package/LICENSE CHANGED
@@ -1,6 +1,6 @@
1
- BSD 2-Clause License
1
+ BSD 3-Clause License
2
2
 
3
- Copyright (c) 2026 TheArchitectit
3
+ Copyright (c) 2026, TheArchitectit
4
4
 
5
5
  Redistribution and use in source and binary forms, with or without
6
6
  modification, are permitted provided that the following conditions are met:
@@ -12,6 +12,10 @@ modification, are permitted provided that the following conditions are met:
12
12
  this list of conditions and the following disclaimer in the documentation
13
13
  and/or other materials provided with the distribution.
14
14
 
15
+ 3. Neither the name of the copyright holder nor the names of its
16
+ contributors may be used to endorse or promote products derived from
17
+ this software without specific prior written permission.
18
+
15
19
  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
16
20
  AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
17
21
  IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
package/README.md CHANGED
@@ -88,7 +88,7 @@ Testing guide: [`TESTER_GUIDE.md`](TESTER_GUIDE.md)
88
88
 
89
89
  ## License
90
90
 
91
- BSD-2-Clause
91
+ BSD 3-Clause
92
92
 
93
93
  ## ☕ Support
94
94
 
@@ -0,0 +1,172 @@
1
+ /**
2
+ * multilevel.ts — Multi-level RAPTOR retrieval engine (S42A).
3
+ *
4
+ * Upgrades the RAPTOR recall path from flat (leaf-only) to multi-level
5
+ * retrieval across the entire hierarchical tree. Searches ALL levels with
6
+ * configurable level weights, supports leaf expansion for cluster hits,
7
+ * and deduplicates overlapping results.
8
+ *
9
+ * PREVENT-PI-004: pure in-process math (cosine, BFS, extractive). No network.
10
+ * PREVENT-PI-001: produces SearchHit[] that feed into recallAndInline() —
11
+ * affects which checkpoints are recalled, not how messages are dropped.
12
+ */
13
+ import { cosineSimilarity } from "../../embedder.js";
14
+ import { mmrRerank } from "../mmr.js";
15
+ import { leafDescendants } from "./retrieval.js";
16
+ const DEFAULT_LEVEL_WEIGHTS = [1.0, 0.9, 0.8, 0.7, 0.5];
17
+ // ── S42A-2: Level-weighted scoring ─────────────────────────────────────────
18
+ /**
19
+ * Score all RAPTOR tree nodes by cosine similarity to the query, then apply
20
+ * level-specific weights. Returns hits sorted by weighted score descending.
21
+ *
22
+ * Level weights: leaves (level 0) get weight 1.0, level 1 gets 0.9, etc.
23
+ * This ensures detailed leaves score highest while still surfacing higher-level
24
+ * summaries when they're highly relevant.
25
+ */
26
+ export function scoreTreeLevels(query, tree, opts) {
27
+ const { embedder } = opts;
28
+ const weights = opts.levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
29
+ const qv = embedder.embed(query);
30
+ const hits = [];
31
+ // 1. Score all internal (summary) nodes.
32
+ for (const node of tree.nodes.values()) {
33
+ const rawScore = cosineSimilarity(qv, node.embedding);
34
+ const levelWeight = weights[Math.min(node.level, weights.length - 1)];
35
+ hits.push({
36
+ nodeId: node.id,
37
+ level: node.level,
38
+ score: rawScore * levelWeight,
39
+ rawScore,
40
+ isLeaf: false,
41
+ leafIds: node.children,
42
+ summary: node.summary,
43
+ embedding: node.embedding,
44
+ });
45
+ }
46
+ // 2. Score leaf nodes. Leaf ids are not in tree.nodes — they are children
47
+ // referenced by internal nodes. Each leaf's embedding is the level-0
48
+ // parent node that wraps it (same approach as stagedExpansion:95–102).
49
+ const seenLeaves = new Set();
50
+ for (const node of tree.nodes.values()) {
51
+ for (const leafId of node.children) {
52
+ if (seenLeaves.has(leafId) || tree.nodes.has(leafId))
53
+ continue;
54
+ seenLeaves.add(leafId);
55
+ const rawScore = cosineSimilarity(qv, node.embedding);
56
+ const leafWeight = weights[0];
57
+ hits.push({
58
+ nodeId: leafId,
59
+ level: 0,
60
+ score: rawScore * leafWeight,
61
+ rawScore,
62
+ isLeaf: true,
63
+ leafIds: [leafId],
64
+ summary: "", // leaves have no summary — they are raw checkpoint ids
65
+ embedding: node.embedding,
66
+ });
67
+ }
68
+ }
69
+ hits.sort((a, b) => b.score - a.score);
70
+ return hits;
71
+ }
72
+ // ── S42A-3: Leaf expansion ─────────────────────────────────────────────────
73
+ /**
74
+ * Given a set of cluster-level hits, expand each one to include its leaf
75
+ * descendants. Deduplicates: if a leaf is already present as a direct hit,
76
+ * it is not duplicated. Returns the merged set (original hits + expanded leaves).
77
+ */
78
+ export function expandLeafDescendants(hits, tree, maxPerCluster, _embedder, queryVector, levelWeights) {
79
+ const weights = levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
80
+ const existingIds = new Set(hits.map((h) => h.nodeId));
81
+ const expanded = [];
82
+ for (const hit of hits) {
83
+ if (hit.isLeaf) {
84
+ expanded.push(hit);
85
+ continue;
86
+ }
87
+ // Get all leaf descendants for this cluster node.
88
+ const node = tree.nodes.get(hit.nodeId);
89
+ if (!node) {
90
+ expanded.push(hit);
91
+ continue;
92
+ }
93
+ const rawLeafIds = leafDescendants(node, tree);
94
+ // Sort by cosine similarity to query, cap at maxPerCluster.
95
+ const leafHits = rawLeafIds
96
+ .map((lid) => {
97
+ // Leaf embedding = its nearest internal parent's embedding.
98
+ const parent = [...tree.nodes.values()].find((n) => n.children.includes(lid));
99
+ const sim = parent
100
+ ? cosineSimilarity(queryVector, parent.embedding)
101
+ : 0;
102
+ return { lid, sim, parent };
103
+ })
104
+ .sort((a, b) => b.sim - a.sim)
105
+ .slice(0, maxPerCluster)
106
+ .filter((l) => !existingIds.has(l.lid))
107
+ .map((l) => {
108
+ existingIds.add(l.lid);
109
+ const rawScore = l.sim;
110
+ return {
111
+ nodeId: l.lid,
112
+ level: 0,
113
+ score: rawScore * weights[0],
114
+ rawScore,
115
+ isLeaf: true,
116
+ leafIds: [l.lid],
117
+ summary: "",
118
+ embedding: l.parent?.embedding ?? hit.embedding,
119
+ };
120
+ });
121
+ expanded.push(hit, ...leafHits);
122
+ }
123
+ return expanded;
124
+ }
125
+ // ── S42A-4: Result dedup ───────────────────────────────────────────────────
126
+ /**
127
+ * Deduplicate hits: if both a cluster node and its leaf children appear in
128
+ * results, remove the cluster hit (leaves provide more specific context).
129
+ * If no leaves are in the set, keep the cluster hit (it provides the abstract view).
130
+ */
131
+ export function deduplicateMultilevelHits(hits) {
132
+ const leafIds = new Set(hits.filter((h) => h.isLeaf).map((h) => h.nodeId));
133
+ return hits.filter((h) => {
134
+ if (h.isLeaf)
135
+ return true;
136
+ // Cluster hit: keep only if none of its leaf children are present.
137
+ return !h.leafIds.some((lid) => leafIds.has(lid));
138
+ });
139
+ }
140
+ // ── S42A-5: Top-level pipeline ─────────────────────────────────────────────
141
+ /**
142
+ * Full multi-level retrieval pipeline: score → expand → dedup → MMR → top-K.
143
+ * Drop-in replacement for `stagedExpansion()` in the RAPTOR recall path.
144
+ */
145
+ export function multilevelRetrieval(query, tree, opts) {
146
+ if (!tree.rootId)
147
+ return [];
148
+ const { embedder } = opts;
149
+ const weights = opts.levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
150
+ const leafExp = opts.leafExpansion !== false; // default true
151
+ const maxLeafExp = opts.maxLeafExpansion ?? 10;
152
+ const k = opts.k ?? 5;
153
+ const lambda = opts.mmrLambda ?? 0.5;
154
+ const qv = embedder.embed(query);
155
+ // 1. Score all nodes with level weights.
156
+ const scored = scoreTreeLevels(query, tree, { embedder, levelWeights: weights });
157
+ // 2. Top-N candidates for MMR diversity window.
158
+ const topN = scored.slice(0, k * 3);
159
+ // 3. Leaf expansion (optional).
160
+ const expanded = leafExp
161
+ ? expandLeafDescendants(topN, tree, maxLeafExp, embedder, qv, weights)
162
+ : topN;
163
+ // 4. Dedup: remove cluster hits when leaf children are present.
164
+ const deduped = deduplicateMultilevelHits(expanded);
165
+ // 5. MMR rerank to k.
166
+ const mmrItems = deduped.map((h) => ({
167
+ item: h,
168
+ vector: h.embedding,
169
+ relevance: h.score,
170
+ }));
171
+ return mmrRerank(mmrItems, k, lambda);
172
+ }
@@ -0,0 +1,203 @@
1
+ /**
2
+ * multilevel.test.ts — hermetic unit tests for S42A multi-level RAPTOR retrieval.
3
+ *
4
+ * Tests the scoreTreeLevels → expandLeafDescendants → deduplicateMultilevelHits
5
+ * → multilevelRetrieval pipeline. No network, no live store — uses TrigramEmbedder
6
+ * and synthetic RaptorTrees built from makeLeaves().
7
+ */
8
+ import { test } from "node:test";
9
+ import assert from "node:assert/strict";
10
+ import { TrigramEmbedder } from "../../embedder.js";
11
+ import { buildRaptorTree } from "./tree.js";
12
+ import { scoreTreeLevels, expandLeafDescendants, deduplicateMultilevelHits, multilevelRetrieval, } from "./multilevel.js";
13
+ function msg(text) {
14
+ return { role: "user", text };
15
+ }
16
+ /** Build N distinct leaves with deterministic content. */
17
+ function makeLeaves(n, embedder = new TrigramEmbedder()) {
18
+ const leaves = [];
19
+ for (let i = 0; i < n; i++) {
20
+ const text = `topic ${i % 7}: the module ${i} validated the session token and refreshed the cache for region ${i}`;
21
+ leaves.push({
22
+ id: `leaf_${i}`,
23
+ messages: [msg(text)],
24
+ sourceText: text,
25
+ embedding: embedder.embed(text),
26
+ });
27
+ }
28
+ return leaves;
29
+ }
30
+ // ── test: scoreTreeLevels returns nodes at all levels ────────────────────────
31
+ test("scoreTreeLevels returns results at multiple tree levels", () => {
32
+ const embedder = new TrigramEmbedder();
33
+ const leaves = makeLeaves(50);
34
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
35
+ const hits = scoreTreeLevels("the auth module validates the session token", tree, {
36
+ embedder,
37
+ });
38
+ assert.ok(hits.length > 0, "should return hits");
39
+ // Should include both leaf (level 0) and cluster (level ≥ 1) hits.
40
+ const levels = new Set(hits.map((h) => h.level));
41
+ assert.ok(levels.has(0), "should have leaf-level hits");
42
+ if (tree.levels > 1) {
43
+ const hasCluster = [...levels].some((l) => l >= 1);
44
+ assert.ok(hasCluster, "should have cluster-level hits for multi-level tree");
45
+ }
46
+ // All hits should have valid scores.
47
+ for (const h of hits) {
48
+ assert.ok(h.score >= 0 && h.score <= 1, `score ${h.score} out of range`);
49
+ assert.ok(h.rawScore >= 0 && h.rawScore <= 1, `rawScore ${h.rawScore} out of range`);
50
+ assert.ok(h.score <= h.rawScore, "weighted score <= raw score (level weights ≤ 1)");
51
+ }
52
+ // Hits should be sorted by score descending.
53
+ for (let i = 1; i < hits.length; i++) {
54
+ assert.ok(hits[i - 1].score >= hits[i].score, `hits not sorted: hit[${i - 1}].score=${hits[i - 1].score} < hit[${i}].score=${hits[i].score}`);
55
+ }
56
+ });
57
+ // ── test: level weights affect scoring ───────────────────────────────────────
58
+ test("level weights shift scores: higher weight for level → higher weighted score", () => {
59
+ const embedder = new TrigramEmbedder();
60
+ const leaves = makeLeaves(50);
61
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
62
+ const query = "the auth module validates the session token";
63
+ // Uniform weights: all levels scored equally.
64
+ const uniform = scoreTreeLevels(query, tree, {
65
+ embedder,
66
+ levelWeights: [1.0, 1.0, 1.0, 1.0, 1.0],
67
+ });
68
+ // Penalized weights: higher levels penalized.
69
+ const penalized = scoreTreeLevels(query, tree, {
70
+ embedder,
71
+ levelWeights: [1.0, 0.1, 0.1, 0.1, 0.1],
72
+ });
73
+ // With penalized weights, cluster-level hits should score lower.
74
+ const clusterUniform = uniform.filter((h) => h.level >= 1);
75
+ const clusterPenalized = penalized.filter((h) => h.level >= 1);
76
+ if (clusterUniform.length > 0 && clusterPenalized.length > 0) {
77
+ const avgUniform = clusterUniform.reduce((s, h) => s + h.score, 0) / clusterUniform.length;
78
+ const avgPenalized = clusterPenalized.reduce((s, h) => s + h.score, 0) / clusterPenalized.length;
79
+ assert.ok(avgPenalized < avgUniform, `penalized cluster avg (${avgPenalized.toFixed(3)}) should be < uniform (${avgUniform.toFixed(3)})`);
80
+ }
81
+ });
82
+ // ── test: expandLeafDescendants adds leaf hits ──────────────────────────────
83
+ test("expandLeafDescendants adds leaf descendants for cluster hits", () => {
84
+ const embedder = new TrigramEmbedder();
85
+ const leaves = makeLeaves(50);
86
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
87
+ const query = "the auth module validates the session token";
88
+ const scored = scoreTreeLevels(query, tree, { embedder });
89
+ const qv = embedder.embed(query);
90
+ // Take only cluster hits (level ≥ 1).
91
+ const clusterHits = scored.filter((h) => !h.isLeaf).slice(0, 3);
92
+ assert.ok(clusterHits.length > 0, "should have cluster hits");
93
+ const expanded = expandLeafDescendants(clusterHits, tree, 5, // maxPerCluster
94
+ embedder, qv);
95
+ // Expanded set should include leaf hits.
96
+ const leafHits = expanded.filter((h) => h.isLeaf);
97
+ assert.ok(leafHits.length > 0, "should have expanded leaf hits");
98
+ assert.ok(expanded.length > clusterHits.length, `expanded (${expanded.length}) should be > cluster hits (${clusterHits.length})`);
99
+ // No duplicate ids.
100
+ const ids = new Set(expanded.map((h) => h.nodeId));
101
+ assert.equal(ids.size, expanded.length, "no duplicate node ids");
102
+ });
103
+ // ── test: deduplicateMultilevelHits removes cluster when leaves present ─────
104
+ test("deduplicateMultilevelHits removes cluster hits when leaf children are present", () => {
105
+ const embedder = new TrigramEmbedder();
106
+ const leaves = makeLeaves(50);
107
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
108
+ const query = "the auth module validates the session token";
109
+ const scored = scoreTreeLevels(query, tree, { embedder });
110
+ // Take a cluster hit and its leaf children.
111
+ const cluster = scored.find((h) => !h.isLeaf);
112
+ assert.ok(cluster, "should have a cluster hit");
113
+ const leafChildren = cluster.leafIds.slice(0, 2).map((lid) => ({
114
+ nodeId: lid,
115
+ level: 0,
116
+ score: 0.5,
117
+ rawScore: 0.5,
118
+ isLeaf: true,
119
+ leafIds: [lid],
120
+ summary: "",
121
+ embedding: cluster.embedding,
122
+ }));
123
+ const mixed = [cluster, ...leafChildren];
124
+ const deduped = deduplicateMultilevelHits(mixed);
125
+ // Cluster should be removed because its leaf children are present.
126
+ assert.ok(!deduped.find((h) => h.nodeId === cluster.nodeId), "cluster hit should be removed when leaf children are present");
127
+ assert.equal(deduped.length, leafChildren.length, "only leaf hits remain");
128
+ });
129
+ test("deduplicateMultilevelHits keeps cluster hits when no leaf children present", () => {
130
+ const clusterHit = {
131
+ nodeId: "cluster_1",
132
+ level: 1,
133
+ score: 0.8,
134
+ rawScore: 0.8,
135
+ isLeaf: false,
136
+ leafIds: ["leaf_1", "leaf_2"],
137
+ summary: "summarized content",
138
+ embedding: [1, 0, 0],
139
+ };
140
+ const deduped = deduplicateMultilevelHits([clusterHit]);
141
+ assert.equal(deduped.length, 1, "cluster should be kept when no leaf children present");
142
+ assert.equal(deduped[0].nodeId, "cluster_1");
143
+ });
144
+ // ── test: multilevelRetrieval returns cluster + leaf mix ────────────────────
145
+ test("multilevelRetrieval returns a mix of cluster and leaf hits", () => {
146
+ const embedder = new TrigramEmbedder();
147
+ const leaves = makeLeaves(50);
148
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
149
+ const hits = multilevelRetrieval("the auth module validates the session token", tree, {
150
+ embedder,
151
+ k: 5,
152
+ leafExpansion: true,
153
+ maxLeafExpansion: 3,
154
+ });
155
+ assert.ok(hits.length > 0, "should return hits");
156
+ assert.ok(hits.length <= 5, "should respect k=5");
157
+ // Should include leaf hits.
158
+ const leafHits = hits.filter((h) => h.isLeaf);
159
+ assert.ok(leafHits.length > 0, "should include leaf hits");
160
+ // All hits should have valid properties.
161
+ for (const h of hits) {
162
+ assert.ok(h.nodeId, "hit should have nodeId");
163
+ assert.ok(typeof h.level === "number", "hit should have numeric level");
164
+ assert.ok(h.score >= 0, "hit score should be non-negative");
165
+ assert.ok(Array.isArray(h.leafIds), "hit should have leafIds array");
166
+ }
167
+ });
168
+ test("multilevelRetrieval respects k parameter", () => {
169
+ const embedder = new TrigramEmbedder();
170
+ const leaves = makeLeaves(100);
171
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 8 });
172
+ for (const k of [1, 3, 5, 10]) {
173
+ const hits = multilevelRetrieval("the auth module validates the session token", tree, {
174
+ embedder,
175
+ k,
176
+ });
177
+ assert.ok(hits.length <= k, `k=${k}: got ${hits.length} hits, should be ≤ ${k}`);
178
+ }
179
+ });
180
+ test("multilevelRetrieval returns empty for empty tree", () => {
181
+ const embedder = new TrigramEmbedder();
182
+ const emptyTree = { nodes: new Map(), rootId: null, levels: 0, timedOut: false };
183
+ const hits = multilevelRetrieval("any query", emptyTree, { embedder });
184
+ assert.deepEqual(hits, []);
185
+ });
186
+ test("multilevelRetrieval with leafExpansion=false skips leaf expansion", () => {
187
+ const embedder = new TrigramEmbedder();
188
+ const leaves = makeLeaves(50);
189
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
190
+ const hits = multilevelRetrieval("the auth module validates the session token", tree, {
191
+ embedder,
192
+ k: 5,
193
+ leafExpansion: false,
194
+ });
195
+ assert.ok(hits.length > 0, "should still return hits");
196
+ // With leaf expansion off and tree having multiple levels, we may get
197
+ // cluster hits that are NOT expanded. The mix depends on tree structure.
198
+ // Just verify we got valid results.
199
+ for (const h of hits) {
200
+ assert.ok(h.nodeId, "hit should have nodeId");
201
+ assert.ok(h.score >= 0, "hit score should be non-negative");
202
+ }
203
+ });
@@ -18,7 +18,7 @@ function isLeafId(id, tree) {
18
18
  return !tree.nodes.has(id);
19
19
  }
20
20
  /** All leaf (raw) ids reachable beneath a node via BFS. */
21
- function leafDescendants(node, tree) {
21
+ export function leafDescendants(node, tree) {
22
22
  const out = [];
23
23
  const queue = [node];
24
24
  while (queue.length) {
package/package.json CHANGED
@@ -1,9 +1,9 @@
1
1
  {
2
2
  "name": "pi-mega-compact",
3
- "version": "0.8.22",
3
+ "version": "0.8.23",
4
4
  "description": "Layered, local, vector-backed context compressor for pi — supersede/collapse/cluster compaction with deduped inline recall.",
5
5
  "type": "module",
6
- "license": "BSD-2-Clause",
6
+ "license": "BSD-3-Clause",
7
7
  "author": "TheArchitectit",
8
8
  "homepage": "https://github.com/TheArchitectit/pi-mega-compact",
9
9
  "repository": {
@@ -0,0 +1,278 @@
1
+ /**
2
+ * multilevel.test.ts — hermetic unit tests for S42A multi-level RAPTOR retrieval.
3
+ *
4
+ * Tests the scoreTreeLevels → expandLeafDescendants → deduplicateMultilevelHits
5
+ * → multilevelRetrieval pipeline. No network, no live store — uses TrigramEmbedder
6
+ * and synthetic RaptorTrees built from makeLeaves().
7
+ */
8
+
9
+ import { test } from "node:test";
10
+ import assert from "node:assert/strict";
11
+ import { TrigramEmbedder } from "../../embedder.js";
12
+ import { buildRaptorTree, type Leaf } from "./tree.js";
13
+ import {
14
+ scoreTreeLevels,
15
+ expandLeafDescendants,
16
+ deduplicateMultilevelHits,
17
+ multilevelRetrieval,
18
+ type MultilevelHit,
19
+ } from "./multilevel.js";
20
+ import type { EngineMessage } from "../../types.js";
21
+
22
+ function msg(text: string): EngineMessage {
23
+ return { role: "user", text };
24
+ }
25
+
26
+ /** Build N distinct leaves with deterministic content. */
27
+ function makeLeaves(n: number, embedder = new TrigramEmbedder()): Leaf[] {
28
+ const leaves: Leaf[] = [];
29
+ for (let i = 0; i < n; i++) {
30
+ const text = `topic ${i % 7}: the module ${i} validated the session token and refreshed the cache for region ${i}`;
31
+ leaves.push({
32
+ id: `leaf_${i}`,
33
+ messages: [msg(text)],
34
+ sourceText: text,
35
+ embedding: embedder.embed(text),
36
+ });
37
+ }
38
+ return leaves;
39
+ }
40
+
41
+ // ── test: scoreTreeLevels returns nodes at all levels ────────────────────────
42
+
43
+ test("scoreTreeLevels returns results at multiple tree levels", () => {
44
+ const embedder = new TrigramEmbedder();
45
+ const leaves = makeLeaves(50);
46
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
47
+
48
+ const hits = scoreTreeLevels("the auth module validates the session token", tree, {
49
+ embedder,
50
+ });
51
+
52
+ assert.ok(hits.length > 0, "should return hits");
53
+
54
+ // Should include both leaf (level 0) and cluster (level ≥ 1) hits.
55
+ const levels = new Set(hits.map((h) => h.level));
56
+ assert.ok(levels.has(0), "should have leaf-level hits");
57
+ if (tree.levels > 1) {
58
+ const hasCluster = [...levels].some((l) => l >= 1);
59
+ assert.ok(hasCluster, "should have cluster-level hits for multi-level tree");
60
+ }
61
+
62
+ // All hits should have valid scores.
63
+ for (const h of hits) {
64
+ assert.ok(h.score >= 0 && h.score <= 1, `score ${h.score} out of range`);
65
+ assert.ok(h.rawScore >= 0 && h.rawScore <= 1, `rawScore ${h.rawScore} out of range`);
66
+ assert.ok(h.score <= h.rawScore, "weighted score <= raw score (level weights ≤ 1)");
67
+ }
68
+
69
+ // Hits should be sorted by score descending.
70
+ for (let i = 1; i < hits.length; i++) {
71
+ assert.ok(
72
+ hits[i - 1].score >= hits[i].score,
73
+ `hits not sorted: hit[${i - 1}].score=${hits[i - 1].score} < hit[${i}].score=${hits[i].score}`,
74
+ );
75
+ }
76
+ });
77
+
78
+ // ── test: level weights affect scoring ───────────────────────────────────────
79
+
80
+ test("level weights shift scores: higher weight for level → higher weighted score", () => {
81
+ const embedder = new TrigramEmbedder();
82
+ const leaves = makeLeaves(50);
83
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
84
+
85
+ const query = "the auth module validates the session token";
86
+
87
+ // Uniform weights: all levels scored equally.
88
+ const uniform = scoreTreeLevels(query, tree, {
89
+ embedder,
90
+ levelWeights: [1.0, 1.0, 1.0, 1.0, 1.0],
91
+ });
92
+
93
+ // Penalized weights: higher levels penalized.
94
+ const penalized = scoreTreeLevels(query, tree, {
95
+ embedder,
96
+ levelWeights: [1.0, 0.1, 0.1, 0.1, 0.1],
97
+ });
98
+
99
+ // With penalized weights, cluster-level hits should score lower.
100
+ const clusterUniform = uniform.filter((h) => h.level >= 1);
101
+ const clusterPenalized = penalized.filter((h) => h.level >= 1);
102
+
103
+ if (clusterUniform.length > 0 && clusterPenalized.length > 0) {
104
+ const avgUniform =
105
+ clusterUniform.reduce((s, h) => s + h.score, 0) / clusterUniform.length;
106
+ const avgPenalized =
107
+ clusterPenalized.reduce((s, h) => s + h.score, 0) / clusterPenalized.length;
108
+ assert.ok(
109
+ avgPenalized < avgUniform,
110
+ `penalized cluster avg (${avgPenalized.toFixed(3)}) should be < uniform (${avgUniform.toFixed(3)})`,
111
+ );
112
+ }
113
+ });
114
+
115
+ // ── test: expandLeafDescendants adds leaf hits ──────────────────────────────
116
+
117
+ test("expandLeafDescendants adds leaf descendants for cluster hits", () => {
118
+ const embedder = new TrigramEmbedder();
119
+ const leaves = makeLeaves(50);
120
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
121
+
122
+ const query = "the auth module validates the session token";
123
+ const scored = scoreTreeLevels(query, tree, { embedder });
124
+ const qv = embedder.embed(query);
125
+
126
+ // Take only cluster hits (level ≥ 1).
127
+ const clusterHits = scored.filter((h) => !h.isLeaf).slice(0, 3);
128
+ assert.ok(clusterHits.length > 0, "should have cluster hits");
129
+
130
+ const expanded = expandLeafDescendants(
131
+ clusterHits,
132
+ tree,
133
+ 5, // maxPerCluster
134
+ embedder,
135
+ qv,
136
+ );
137
+
138
+ // Expanded set should include leaf hits.
139
+ const leafHits = expanded.filter((h) => h.isLeaf);
140
+ assert.ok(leafHits.length > 0, "should have expanded leaf hits");
141
+ assert.ok(
142
+ expanded.length > clusterHits.length,
143
+ `expanded (${expanded.length}) should be > cluster hits (${clusterHits.length})`,
144
+ );
145
+
146
+ // No duplicate ids.
147
+ const ids = new Set(expanded.map((h) => h.nodeId));
148
+ assert.equal(ids.size, expanded.length, "no duplicate node ids");
149
+ });
150
+
151
+ // ── test: deduplicateMultilevelHits removes cluster when leaves present ─────
152
+
153
+ test("deduplicateMultilevelHits removes cluster hits when leaf children are present", () => {
154
+ const embedder = new TrigramEmbedder();
155
+ const leaves = makeLeaves(50);
156
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
157
+
158
+ const query = "the auth module validates the session token";
159
+ const scored = scoreTreeLevels(query, tree, { embedder });
160
+
161
+ // Take a cluster hit and its leaf children.
162
+ const cluster = scored.find((h) => !h.isLeaf);
163
+ assert.ok(cluster, "should have a cluster hit");
164
+
165
+ const leafChildren: MultilevelHit[] = cluster!.leafIds.slice(0, 2).map((lid) => ({
166
+ nodeId: lid,
167
+ level: 0,
168
+ score: 0.5,
169
+ rawScore: 0.5,
170
+ isLeaf: true,
171
+ leafIds: [lid],
172
+ summary: "",
173
+ embedding: cluster!.embedding,
174
+ }));
175
+
176
+ const mixed = [cluster!, ...leafChildren];
177
+ const deduped = deduplicateMultilevelHits(mixed);
178
+
179
+ // Cluster should be removed because its leaf children are present.
180
+ assert.ok(
181
+ !deduped.find((h) => h.nodeId === cluster!.nodeId),
182
+ "cluster hit should be removed when leaf children are present",
183
+ );
184
+ assert.equal(deduped.length, leafChildren.length, "only leaf hits remain");
185
+ });
186
+
187
+ test("deduplicateMultilevelHits keeps cluster hits when no leaf children present", () => {
188
+ const clusterHit: MultilevelHit = {
189
+ nodeId: "cluster_1",
190
+ level: 1,
191
+ score: 0.8,
192
+ rawScore: 0.8,
193
+ isLeaf: false,
194
+ leafIds: ["leaf_1", "leaf_2"],
195
+ summary: "summarized content",
196
+ embedding: [1, 0, 0],
197
+ };
198
+
199
+ const deduped = deduplicateMultilevelHits([clusterHit]);
200
+ assert.equal(deduped.length, 1, "cluster should be kept when no leaf children present");
201
+ assert.equal(deduped[0].nodeId, "cluster_1");
202
+ });
203
+
204
+ // ── test: multilevelRetrieval returns cluster + leaf mix ────────────────────
205
+
206
+ test("multilevelRetrieval returns a mix of cluster and leaf hits", () => {
207
+ const embedder = new TrigramEmbedder();
208
+ const leaves = makeLeaves(50);
209
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
210
+
211
+ const hits = multilevelRetrieval("the auth module validates the session token", tree, {
212
+ embedder,
213
+ k: 5,
214
+ leafExpansion: true,
215
+ maxLeafExpansion: 3,
216
+ });
217
+
218
+ assert.ok(hits.length > 0, "should return hits");
219
+ assert.ok(hits.length <= 5, "should respect k=5");
220
+
221
+ // Should include leaf hits.
222
+ const leafHits = hits.filter((h) => h.isLeaf);
223
+ assert.ok(leafHits.length > 0, "should include leaf hits");
224
+
225
+ // All hits should have valid properties.
226
+ for (const h of hits) {
227
+ assert.ok(h.nodeId, "hit should have nodeId");
228
+ assert.ok(typeof h.level === "number", "hit should have numeric level");
229
+ assert.ok(h.score >= 0, "hit score should be non-negative");
230
+ assert.ok(Array.isArray(h.leafIds), "hit should have leafIds array");
231
+ }
232
+ });
233
+
234
+ test("multilevelRetrieval respects k parameter", () => {
235
+ const embedder = new TrigramEmbedder();
236
+ const leaves = makeLeaves(100);
237
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 8 });
238
+
239
+ for (const k of [1, 3, 5, 10]) {
240
+ const hits = multilevelRetrieval("the auth module validates the session token", tree, {
241
+ embedder,
242
+ k,
243
+ });
244
+ assert.ok(
245
+ hits.length <= k,
246
+ `k=${k}: got ${hits.length} hits, should be ≤ ${k}`,
247
+ );
248
+ }
249
+ });
250
+
251
+ test("multilevelRetrieval returns empty for empty tree", () => {
252
+ const embedder = new TrigramEmbedder();
253
+ const emptyTree = { nodes: new Map(), rootId: null, levels: 0, timedOut: false };
254
+ const hits = multilevelRetrieval("any query", emptyTree, { embedder });
255
+ assert.deepEqual(hits, []);
256
+ });
257
+
258
+ test("multilevelRetrieval with leafExpansion=false skips leaf expansion", () => {
259
+ const embedder = new TrigramEmbedder();
260
+ const leaves = makeLeaves(50);
261
+ const tree = buildRaptorTree(leaves, { embedder, clustersPerLevel: 4 });
262
+
263
+ const hits = multilevelRetrieval("the auth module validates the session token", tree, {
264
+ embedder,
265
+ k: 5,
266
+ leafExpansion: false,
267
+ });
268
+
269
+ assert.ok(hits.length > 0, "should still return hits");
270
+
271
+ // With leaf expansion off and tree having multiple levels, we may get
272
+ // cluster hits that are NOT expanded. The mix depends on tree structure.
273
+ // Just verify we got valid results.
274
+ for (const h of hits) {
275
+ assert.ok(h.nodeId, "hit should have nodeId");
276
+ assert.ok(h.score >= 0, "hit score should be non-negative");
277
+ }
278
+ });
@@ -0,0 +1,246 @@
1
+ /**
2
+ * multilevel.ts — Multi-level RAPTOR retrieval engine (S42A).
3
+ *
4
+ * Upgrades the RAPTOR recall path from flat (leaf-only) to multi-level
5
+ * retrieval across the entire hierarchical tree. Searches ALL levels with
6
+ * configurable level weights, supports leaf expansion for cluster hits,
7
+ * and deduplicates overlapping results.
8
+ *
9
+ * PREVENT-PI-004: pure in-process math (cosine, BFS, extractive). No network.
10
+ * PREVENT-PI-001: produces SearchHit[] that feed into recallAndInline() —
11
+ * affects which checkpoints are recalled, not how messages are dropped.
12
+ */
13
+
14
+ import type { Embedder, Vector } from "../../embedder.js";
15
+ import { cosineSimilarity } from "../../embedder.js";
16
+ import { mmrRerank } from "../mmr.js";
17
+ import type { RaptorTree } from "./tree.js";
18
+ import { leafDescendants } from "./retrieval.js";
19
+
20
+ // ── S42A-1: Types ──────────────────────────────────────────────────────────
21
+
22
+ export interface MultilevelRetrieveOptions {
23
+ embedder: Embedder;
24
+ /** Weight per tree level (index 0 = leaves, index 1 = level 1, etc.).
25
+ * Default: [1.0, 0.9, 0.8, 0.7, 0.5]. Capped at tree depth.
26
+ * UNCALIBRATED — requires real-data calibration before stable. */
27
+ levelWeights?: number[];
28
+ /** When true, expand cluster hits to include leaf descendants. Default: true. */
29
+ leafExpansion?: boolean;
30
+ /** Max leaf descendants to fetch per cluster hit. Default: 10. */
31
+ maxLeafExpansion?: number;
32
+ /** Final number of results to return. Default: 5. */
33
+ k?: number;
34
+ /** MMR diversity weight. Default: 0.5. */
35
+ mmrLambda?: number;
36
+ }
37
+
38
+ export interface MultilevelHit {
39
+ nodeId: string;
40
+ level: number;
41
+ /** Weighted score after level weighting. */
42
+ score: number;
43
+ /** Raw cosine similarity before level weighting. */
44
+ rawScore: number;
45
+ isLeaf: boolean;
46
+ /** Leaf ids covered by this node (for leaf expansion). */
47
+ leafIds: string[];
48
+ summary: string;
49
+ embedding: Vector;
50
+ }
51
+
52
+ const DEFAULT_LEVEL_WEIGHTS = [1.0, 0.9, 0.8, 0.7, 0.5];
53
+
54
+ // ── S42A-2: Level-weighted scoring ─────────────────────────────────────────
55
+
56
+ /**
57
+ * Score all RAPTOR tree nodes by cosine similarity to the query, then apply
58
+ * level-specific weights. Returns hits sorted by weighted score descending.
59
+ *
60
+ * Level weights: leaves (level 0) get weight 1.0, level 1 gets 0.9, etc.
61
+ * This ensures detailed leaves score highest while still surfacing higher-level
62
+ * summaries when they're highly relevant.
63
+ */
64
+ export function scoreTreeLevels(
65
+ query: string,
66
+ tree: RaptorTree,
67
+ opts: Pick<MultilevelRetrieveOptions, "embedder" | "levelWeights">,
68
+ ): MultilevelHit[] {
69
+ const { embedder } = opts;
70
+ const weights = opts.levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
71
+ const qv = embedder.embed(query);
72
+ const hits: MultilevelHit[] = [];
73
+
74
+ // 1. Score all internal (summary) nodes.
75
+ for (const node of tree.nodes.values()) {
76
+ const rawScore = cosineSimilarity(qv, node.embedding);
77
+ const levelWeight = weights[Math.min(node.level, weights.length - 1)];
78
+ hits.push({
79
+ nodeId: node.id,
80
+ level: node.level,
81
+ score: rawScore * levelWeight,
82
+ rawScore,
83
+ isLeaf: false,
84
+ leafIds: node.children,
85
+ summary: node.summary,
86
+ embedding: node.embedding,
87
+ });
88
+ }
89
+
90
+ // 2. Score leaf nodes. Leaf ids are not in tree.nodes — they are children
91
+ // referenced by internal nodes. Each leaf's embedding is the level-0
92
+ // parent node that wraps it (same approach as stagedExpansion:95–102).
93
+ const seenLeaves = new Set<string>();
94
+ for (const node of tree.nodes.values()) {
95
+ for (const leafId of node.children) {
96
+ if (seenLeaves.has(leafId) || tree.nodes.has(leafId)) continue;
97
+ seenLeaves.add(leafId);
98
+ const rawScore = cosineSimilarity(qv, node.embedding);
99
+ const leafWeight = weights[0];
100
+ hits.push({
101
+ nodeId: leafId,
102
+ level: 0,
103
+ score: rawScore * leafWeight,
104
+ rawScore,
105
+ isLeaf: true,
106
+ leafIds: [leafId],
107
+ summary: "", // leaves have no summary — they are raw checkpoint ids
108
+ embedding: node.embedding,
109
+ });
110
+ }
111
+ }
112
+
113
+ hits.sort((a, b) => b.score - a.score);
114
+ return hits;
115
+ }
116
+
117
+ // ── S42A-3: Leaf expansion ─────────────────────────────────────────────────
118
+
119
+ /**
120
+ * Given a set of cluster-level hits, expand each one to include its leaf
121
+ * descendants. Deduplicates: if a leaf is already present as a direct hit,
122
+ * it is not duplicated. Returns the merged set (original hits + expanded leaves).
123
+ */
124
+ export function expandLeafDescendants(
125
+ hits: MultilevelHit[],
126
+ tree: RaptorTree,
127
+ maxPerCluster: number,
128
+ _embedder: Embedder,
129
+ queryVector: Vector,
130
+ levelWeights?: number[],
131
+ ): MultilevelHit[] {
132
+ const weights = levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
133
+ const existingIds = new Set(hits.map((h) => h.nodeId));
134
+ const expanded: MultilevelHit[] = [];
135
+
136
+ for (const hit of hits) {
137
+ if (hit.isLeaf) {
138
+ expanded.push(hit);
139
+ continue;
140
+ }
141
+
142
+ // Get all leaf descendants for this cluster node.
143
+ const node = tree.nodes.get(hit.nodeId);
144
+ if (!node) {
145
+ expanded.push(hit);
146
+ continue;
147
+ }
148
+
149
+ const rawLeafIds = leafDescendants(node, tree);
150
+
151
+ // Sort by cosine similarity to query, cap at maxPerCluster.
152
+ const leafHits: MultilevelHit[] = rawLeafIds
153
+ .map((lid) => {
154
+ // Leaf embedding = its nearest internal parent's embedding.
155
+ const parent = [...tree.nodes.values()].find((n) =>
156
+ n.children.includes(lid),
157
+ );
158
+ const sim = parent
159
+ ? cosineSimilarity(queryVector, parent.embedding)
160
+ : 0;
161
+ return { lid, sim, parent };
162
+ })
163
+ .sort((a, b) => b.sim - a.sim)
164
+ .slice(0, maxPerCluster)
165
+ .filter((l) => !existingIds.has(l.lid))
166
+ .map((l) => {
167
+ existingIds.add(l.lid);
168
+ const rawScore = l.sim;
169
+ return {
170
+ nodeId: l.lid,
171
+ level: 0,
172
+ score: rawScore * weights[0],
173
+ rawScore,
174
+ isLeaf: true,
175
+ leafIds: [l.lid],
176
+ summary: "",
177
+ embedding: l.parent?.embedding ?? hit.embedding,
178
+ } as MultilevelHit;
179
+ });
180
+
181
+ expanded.push(hit, ...leafHits);
182
+ }
183
+
184
+ return expanded;
185
+ }
186
+
187
+ // ── S42A-4: Result dedup ───────────────────────────────────────────────────
188
+
189
+ /**
190
+ * Deduplicate hits: if both a cluster node and its leaf children appear in
191
+ * results, remove the cluster hit (leaves provide more specific context).
192
+ * If no leaves are in the set, keep the cluster hit (it provides the abstract view).
193
+ */
194
+ export function deduplicateMultilevelHits(hits: MultilevelHit[]): MultilevelHit[] {
195
+ const leafIds = new Set(hits.filter((h) => h.isLeaf).map((h) => h.nodeId));
196
+ return hits.filter((h) => {
197
+ if (h.isLeaf) return true;
198
+ // Cluster hit: keep only if none of its leaf children are present.
199
+ return !h.leafIds.some((lid) => leafIds.has(lid));
200
+ });
201
+ }
202
+
203
+ // ── S42A-5: Top-level pipeline ─────────────────────────────────────────────
204
+
205
+ /**
206
+ * Full multi-level retrieval pipeline: score → expand → dedup → MMR → top-K.
207
+ * Drop-in replacement for `stagedExpansion()` in the RAPTOR recall path.
208
+ */
209
+ export function multilevelRetrieval(
210
+ query: string,
211
+ tree: RaptorTree,
212
+ opts: MultilevelRetrieveOptions,
213
+ ): MultilevelHit[] {
214
+ if (!tree.rootId) return [];
215
+
216
+ const { embedder } = opts;
217
+ const weights = opts.levelWeights ?? DEFAULT_LEVEL_WEIGHTS;
218
+ const leafExp = opts.leafExpansion !== false; // default true
219
+ const maxLeafExp = opts.maxLeafExpansion ?? 10;
220
+ const k = opts.k ?? 5;
221
+ const lambda = opts.mmrLambda ?? 0.5;
222
+
223
+ const qv = embedder.embed(query);
224
+
225
+ // 1. Score all nodes with level weights.
226
+ const scored = scoreTreeLevels(query, tree, { embedder, levelWeights: weights });
227
+
228
+ // 2. Top-N candidates for MMR diversity window.
229
+ const topN = scored.slice(0, k * 3);
230
+
231
+ // 3. Leaf expansion (optional).
232
+ const expanded = leafExp
233
+ ? expandLeafDescendants(topN, tree, maxLeafExp, embedder, qv, weights)
234
+ : topN;
235
+
236
+ // 4. Dedup: remove cluster hits when leaf children are present.
237
+ const deduped = deduplicateMultilevelHits(expanded);
238
+
239
+ // 5. MMR rerank to k.
240
+ const mmrItems = deduped.map((h) => ({
241
+ item: h,
242
+ vector: h.embedding as Vector,
243
+ relevance: h.score,
244
+ }));
245
+ return mmrRerank(mmrItems, k, lambda);
246
+ }
@@ -33,7 +33,7 @@ function isLeafId(id: string, tree: RaptorTree): boolean {
33
33
  }
34
34
 
35
35
  /** All leaf (raw) ids reachable beneath a node via BFS. */
36
- function leafDescendants(node: RaptorNode, tree: RaptorTree): string[] {
36
+ export function leafDescendants(node: RaptorNode, tree: RaptorTree): string[] {
37
37
  const out: string[] = [];
38
38
  const queue = [node];
39
39
  while (queue.length) {