@hicaru/pi-rlm 0.3.19 → 0.3.21

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.
Files changed (60) hide show
  1. package/README.md +8 -5
  2. package/package.json +1 -1
  3. package/src/bridge/handlers/completion.ts +3 -0
  4. package/src/bridge/handlers/emitting.ts +0 -4
  5. package/src/bridge/handlers/rlm-query.ts +3 -3
  6. package/src/bridge/handlers/task-registry.ts +46 -19
  7. package/src/bridge/handlers/types.ts +6 -3
  8. package/src/bridge/model.ts +4 -0
  9. package/src/commands/rlm.ts +14 -7
  10. package/src/config/defaults.ts +41 -13
  11. package/src/config/settings.ts +7 -3
  12. package/src/config/skillstate.ts +236 -44
  13. package/src/context/merge.ts +10 -3
  14. package/src/context/namespace.ts +6 -2
  15. package/src/context/refresh.ts +32 -11
  16. package/src/core/answer.ts +15 -0
  17. package/src/core/budget.ts +39 -17
  18. package/src/core/compaction.ts +85 -9
  19. package/src/core/engine.ts +117 -32
  20. package/src/core/iteration.ts +4 -0
  21. package/src/core/limits.ts +10 -14
  22. package/src/core/root-context.ts +83 -19
  23. package/src/core/root-digest.ts +48 -11
  24. package/src/core/root-state.ts +39 -12
  25. package/src/core/run-state.ts +86 -14
  26. package/src/core/session-archive.ts +174 -0
  27. package/src/core/types.ts +13 -2
  28. package/src/index.ts +142 -12
  29. package/src/mode/rlm-mode.ts +2 -2
  30. package/src/prompts/glossary.ts +36 -5
  31. package/src/prompts/native.ts +8 -2
  32. package/src/prompts/user.ts +6 -4
  33. package/src/sandbox/py/__pycache__/retrieval.cpython-314.pyc +0 -0
  34. package/src/sandbox/py/__pycache__/scaffold.cpython-314.pyc +0 -0
  35. package/src/sandbox/py/retrieval.py +202 -36
  36. package/src/sandbox/py/scaffold.py +20 -5
  37. package/src/sandbox/py/worker.py +1 -1
  38. package/src/sandbox/sandbox-manager.ts +19 -0
  39. package/src/sandbox/sandbox.ts +13 -1
  40. package/src/text/parsing.ts +133 -2
  41. package/src/text/tokens.ts +39 -4
  42. package/src/tool/repl-details.ts +2 -2
  43. package/src/tool/repl-render.ts +38 -2
  44. package/src/tool/repl-tool.ts +37 -23
  45. package/src/tool/rlm-aggregator.ts +1 -1
  46. package/src/tool/rlm-details.ts +1 -2
  47. package/src/tool/rlm-events.ts +3 -6
  48. package/src/tool/rlm-tool.ts +1 -1
  49. package/src/tool/subcall-render.ts +7 -4
  50. package/src/tool/subcall-store.ts +5 -18
  51. package/src/ui/config-panel.ts +4 -19
  52. package/src/ui/intro.ts +1 -1
  53. package/src/ui/panel/run-registry.ts +2 -2
  54. package/src/ui/python-highlight.ts +49 -0
  55. package/src/ui/stage-cards.ts +192 -0
  56. package/src/ui/tree/tree-model.ts +69 -19
  57. package/src/ui/tree/tree-rows.ts +2 -1
  58. package/src/util/abort.ts +34 -0
  59. package/src/util/bm25.ts +170 -21
  60. package/src/util/errors.ts +1 -1
@@ -0,0 +1,192 @@
1
+ /**
2
+ * stage-cards — [rlm.stage] transcript cards for orchestrator stage transitions.
3
+ *
4
+ * Same visual language as zebra-catch's finding cards: a `**◆ …**` header with
5
+ * status + stats, a body paragraph, `− tag (source)` bullets — sent via
6
+ * `pi.sendMessage({ customType, display: true })` so pi draws the labeled box.
7
+ * The renderer honors `options.expanded`, which pi's CustomMessageComponent
8
+ * re-invokes on the user's ctrl+o (app.tools.expand) — collapsed shows the
9
+ * header + expand hint, expanded shows the full card. Digest bodies embed the
10
+ * root-digest text VERBATIM (glossary wording — never re-worded here).
11
+ *
12
+ * Pure builders + one renderer; emission points live in the rlmExtension
13
+ * closure (src/index.ts) — no session state at module load.
14
+ */
15
+
16
+ import { Box, Markdown, Spacer, Text } from "@earendil-works/pi-tui";
17
+ import type { MessageRenderer, Theme } from "@earendil-works/pi-coding-agent";
18
+ import type { SkillNoteInput } from "../config/skillstate.ts";
19
+ import { previewText } from "../text/preview.ts";
20
+ import { expandHint } from "../tool/subcall-render.ts";
21
+ import { formatTokens } from "./theme.ts";
22
+ import { markdownTheme } from "./theme-adapter.ts";
23
+
24
+ /** The transcript tag pi renders for these cards — `[rlm.stage]`. */
25
+ export const STAGE_CUSTOM_TYPE = "rlm.stage";
26
+
27
+ /** Distill bullets shown before the "+N more" ellipsis. */
28
+ const DISTILL_BULLET_CAP = 8;
29
+ /** Note text cap inside a distill bullet — the full note lives in skill.state itself. */
30
+ const DISTILL_TEXT_CHARS = 100;
31
+
32
+ /** One orchestrator stage transition. Discriminated — the renderer never guesses. */
33
+ export type StageCardDetails =
34
+ | {
35
+ readonly kind: "digest";
36
+ /** 1-based compaction index (rootDigests counter at emit time). */
37
+ readonly index: number;
38
+ readonly turnsFolded: number;
39
+ /** Host-consumed token estimate for the displaced span (V1 soak probe). */
40
+ readonly tokensBefore: number;
41
+ /** Our own estimate over the same span — divergence here is the probe's signal. */
42
+ readonly tokensBeforeRecomputed: number;
43
+ /** The root-digest summary text, embedded verbatim when expanded. */
44
+ readonly summary: string;
45
+ }
46
+ | {
47
+ readonly kind: "degrade";
48
+ readonly reason: string;
49
+ readonly idleTurns: number;
50
+ readonly idleMax: number;
51
+ }
52
+ | {
53
+ readonly kind: "recover";
54
+ readonly fencesAccepted: number;
55
+ readonly fencesTotal: number;
56
+ }
57
+ | {
58
+ readonly kind: "distill";
59
+ /** The notes this session contributed (pre-id/hits inputs, as handed to SkillStore.merge). */
60
+ readonly merged: readonly SkillNoteInput[];
61
+ readonly total: number;
62
+ readonly byTag: Readonly<Record<string, number>>;
63
+ };
64
+
65
+ function isRecord(value: unknown): value is Record<string, unknown> {
66
+ return typeof value === "object" && value !== null;
67
+ }
68
+
69
+ /** Runtime guard for details replayed from old/corrupt session files — fail-soft, never throws. */
70
+ export function isStageCardDetails(value: unknown): value is StageCardDetails {
71
+ if (!isRecord(value)) return false;
72
+ switch (value.kind) {
73
+ case "digest":
74
+ return (
75
+ typeof value.index === "number" &&
76
+ typeof value.turnsFolded === "number" &&
77
+ typeof value.tokensBefore === "number" &&
78
+ typeof value.tokensBeforeRecomputed === "number" &&
79
+ typeof value.summary === "string"
80
+ );
81
+ case "degrade":
82
+ return (
83
+ typeof value.reason === "string" &&
84
+ typeof value.idleTurns === "number" &&
85
+ typeof value.idleMax === "number"
86
+ );
87
+ case "recover":
88
+ return typeof value.fencesAccepted === "number" && typeof value.fencesTotal === "number";
89
+ case "distill": {
90
+ if (typeof value.total !== "number" || !Array.isArray(value.merged) || !isRecord(value.byTag)) return false;
91
+ if (!Object.values(value.byTag).every((v) => typeof v === "number")) return false;
92
+ return value.merged.every((note) => isRecord(note) && typeof note.text === "string");
93
+ }
94
+ default:
95
+ return false;
96
+ }
97
+ }
98
+
99
+ /** "gotcha 4 · recipe 3" — tags with a nonzero count, insertion order. */
100
+ function byTagPart(byTag: Readonly<Record<string, number>>): string {
101
+ const parts: string[] = [];
102
+ for (const [tag, count] of Object.entries(byTag)) {
103
+ if (count > 0) parts.push(`${tag} ${String(count)}`);
104
+ }
105
+ return parts.length === 0 ? "" : ` (${parts.join(" · ")})`;
106
+ }
107
+
108
+ /** The `**◆ …**` headline — the only thing visible while collapsed. */
109
+ export function stageCardHeaderLine(details: StageCardDetails): string {
110
+ switch (details.kind) {
111
+ case "digest":
112
+ return (
113
+ `**◆ digest #${String(details.index)}** folded ${String(details.turnsFolded)} turns · ` +
114
+ `${formatTokens(details.tokensBefore)} tok (recomputed ${formatTokens(details.tokensBeforeRecomputed)})`
115
+ );
116
+ case "degrade":
117
+ return (
118
+ `**⚠ Σ degraded** idle ${String(details.idleTurns)}/${String(details.idleMax)} fence-eligible turns · ` +
119
+ previewText(details.reason, 80)
120
+ );
121
+ case "recover":
122
+ return `**◆ Σ recovered** fences accepted ${String(details.fencesAccepted)}/${String(details.fencesTotal)}`;
123
+ case "distill":
124
+ return (
125
+ `**◆ skill.state** +${String(details.merged.length)} notes distilled · ` +
126
+ `${String(details.total)} total${byTagPart(details.byTag)}`
127
+ );
128
+ }
129
+ }
130
+
131
+ /** Body under the header — blank for recover (the headline says it all). */
132
+ function stageCardBody(details: StageCardDetails): string {
133
+ switch (details.kind) {
134
+ case "digest":
135
+ return details.summary;
136
+ case "degrade":
137
+ return "− splices paused; tool outcomes still feed Σ (observe floor)";
138
+ case "distill": {
139
+ const shown = details.merged.slice(0, DISTILL_BULLET_CAP);
140
+ const lines = new Array<string>(shown.length);
141
+ for (let i = 0; i < shown.length; i++) {
142
+ const note = shown[i];
143
+ if (note === undefined) continue;
144
+ lines[i] = `− ${note.tags?.[0] ?? "note"}: ${previewText(note.text, DISTILL_TEXT_CHARS)}`;
145
+ }
146
+ const rest = details.merged.length - shown.length;
147
+ return rest > 0 ? [...lines, `− +${String(rest)} more`].join("\n") : lines.join("\n");
148
+ }
149
+ case "recover":
150
+ return "";
151
+ }
152
+ }
153
+
154
+ /** Full card content — what sendMessage stores and what expanded rendering shows. */
155
+ export function stageCardMarkdown(details: StageCardDetails): string {
156
+ const body = stageCardBody(details);
157
+ return body === "" ? stageCardHeaderLine(details) : `${stageCardHeaderLine(details)}\n\n${body}`;
158
+ }
159
+
160
+ /** Markdown block styled through the injected theme's adapter. */
161
+ function themedMarkdown(text: string, theme: Theme): Markdown {
162
+ return new Markdown(text, 0, 0, markdownTheme(theme), {
163
+ color: (t) => theme.fg("customMessageText", t),
164
+ });
165
+ }
166
+
167
+ /**
168
+ * The [rlm.stage] message renderer — mirrors pi's default custom-message box
169
+ * (labeled Box + Markdown) but honors `expanded`: collapsed shows just the
170
+ * headline + expand hint, expanded the full card. Returns undefined on
171
+ * stale/corrupt details so pi falls back to default rendering (fail-soft,
172
+ * host contract).
173
+ */
174
+ export const renderStageCard: MessageRenderer = (message, options, theme) => {
175
+ try {
176
+ if (!isStageCardDetails(message.details)) return undefined;
177
+ const box = new Box(1, 1, (t) => theme.bg("customMessageBg", t));
178
+ box.addChild(new Text(theme.fg("customMessageLabel", theme.bold(`[${STAGE_CUSTOM_TYPE}]`)), 0, 0));
179
+ box.addChild(new Spacer(1));
180
+ box.addChild(themedMarkdown(
181
+ options.expanded ? stageCardMarkdown(message.details) : stageCardHeaderLine(message.details),
182
+ theme,
183
+ ));
184
+ if (!options.expanded) {
185
+ box.addChild(new Spacer(1));
186
+ box.addChild(new Text(expandHint(theme), 0, 0));
187
+ }
188
+ return box;
189
+ } catch {
190
+ return undefined; // default rendering shows the stored markdown instead
191
+ }
192
+ };
@@ -5,18 +5,21 @@
5
5
  * the result and only rebuilds when the underlying store reports a change.
6
6
  *
7
7
  * Nothing is ever hidden: every sub-call renders as its own row (parity with
8
- * pi, which shows each concurrent tool call individually) — except runs of
9
- * CONSECUTIVE IDENTICAL sibling leaves (same label+model+status), which
10
- * collapse into one expandable "label ×N" group row so a 20-item llm_batch is
11
- * one line, not 20 and a wholesale batch failure is one `✗ label ×N` line.
12
- * Errors group exactly like successes; distinct failures keep their own rows
13
- * and reasons, and singletons render as plain rows. Interleaved siblings (✗
14
- * with different keys between) stay in encounter order — position is never
15
- * rewritten. Collapsed subtrees are skipped at the user's explicit request
16
- * (chevron flips). Token rows are own-spend only a row never blends models.
8
+ * pi, which shows each concurrent tool call individually) — except identical
9
+ * sibling llm leaves (same label+model+status), which consolidate into ONE
10
+ * expandable "label ×N" group row placed at the first member's position: a
11
+ * 16-item batch failure is a single `✗ label ×N · reason` line however many
12
+ * other rows interleave the run. rlm nodes and nodes with children never
13
+ * group; grouped llm leaves move up to the group head, every other row keeps
14
+ * its encounter order. Errors group exactly like successes; diverging reasons
15
+ * collapse to "N failure reasons" (per-item reasons stay in the detail
16
+ * modal). Singletons render as plain rows. Collapsed subtrees are skipped at
17
+ * the user's explicit request (chevron flips). Token rows are own-spend only
18
+ * — a row never blends models.
17
19
  */
18
20
 
19
21
  import type { RlmSubcall, RlmRunStatus, SubcallPhase, SubcallStatus } from "../../tool/rlm-details.ts";
22
+ import { ERROR_PREFIX, isErrorText } from "../../util/errors.ts";
20
23
 
21
24
  /** Immutable per-run view the model consumes (built by RunRegistry from a live store). */
22
25
  export interface RunSnapshot {
@@ -77,6 +80,8 @@ export interface GroupRow {
77
80
  readonly icon: SubcallStatus | "queued";
78
81
  readonly expandable: boolean;
79
82
  readonly expanded: boolean;
83
+ /** First-line failure reason shared by every member — error groups only, else "N failure reasons". */
84
+ readonly reason?: string;
80
85
  }
81
86
 
82
87
  /** Internal build-time entry: a real node or an accumulating group. */
@@ -94,21 +99,65 @@ const groupable = (sc: RlmSubcall, byParent: ReadonlyMap<string | undefined, Rlm
94
99
 
95
100
  const groupKey = (sc: RlmSubcall): string => `${sc.label}|${sc.model ?? ""}|${sc.status}`;
96
101
 
97
- /** Merge consecutive identical sibling leaves into group entries; keep order. */
102
+ /** Longest reason shown inline in a group header full text lives in the modal. */
103
+ const REASON_MAX_CHARS = 48;
104
+
105
+ /** First line of an error text, "Error: " prefix stripped — undefined when there is nothing usable. */
106
+ function errorLineOf(sc: RlmSubcall): string | undefined {
107
+ const raw = sc.detail ?? sc.resultPreview;
108
+ if (raw === undefined || raw === "") return undefined;
109
+ const body = isErrorText(raw) ? raw.slice(ERROR_PREFIX.length + 1) : raw;
110
+ const nl = body.indexOf("\n");
111
+ const line = (nl === -1 ? body : body.slice(0, nl)).trim();
112
+ return line === "" ? undefined : line;
113
+ }
114
+
115
+ /** The reason all members share, or "N failure reasons" when their errors diverge. */
116
+ function groupReason(members: readonly RlmSubcall[]): string | undefined {
117
+ let first: string | undefined;
118
+ const distinct = new Set<string>();
119
+ for (const m of members) {
120
+ const line = errorLineOf(m);
121
+ if (line === undefined) continue;
122
+ if (first === undefined) first = line;
123
+ distinct.add(line);
124
+ }
125
+ if (first === undefined || distinct.size === 0) return undefined;
126
+ const reason = distinct.size === 1 ? first : `${String(distinct.size)} failure reasons`;
127
+ return reason.length > REASON_MAX_CHARS ? `${reason.slice(0, REASON_MAX_CHARS - 1)}…` : reason;
128
+ }
129
+
130
+ /**
131
+ * Consolidate every identical sibling leaf (same label+model+status) into ONE
132
+ * group entry at its first member's position — a 16-item batch failure stays
133
+ * a single `✗ label ×16` line however many other rows interleave the run.
134
+ * Non-grouping siblings keep their encounter order; members keep start order
135
+ * for the expanded view.
136
+ */
98
137
  function partition(children: readonly RlmSubcall[], byParent: ReadonlyMap<string | undefined, RlmSubcall[]>): readonly Entry[] {
138
+ const groups = new Map<string, Extract<Entry, { type: "group" }>>();
99
139
  const out: Entry[] = [];
100
140
  for (const sc of children) {
101
- if (groupable(sc, byParent)) {
102
- const key = groupKey(sc);
103
- const last = out[out.length - 1];
104
- if (last !== undefined && last.type === "group" && last.key === key) {
105
- last.members.push(sc);
106
- continue;
107
- }
108
- out.push({ type: "group", key, label: sc.label, model: sc.model, status: sc.status, members: [sc] });
109
- } else {
141
+ if (!groupable(sc, byParent)) {
110
142
  out.push({ type: "node", sc });
143
+ continue;
144
+ }
145
+ const key = groupKey(sc);
146
+ const existing = groups.get(key);
147
+ if (existing !== undefined) {
148
+ existing.members.push(sc);
149
+ continue;
111
150
  }
151
+ const group: Extract<Entry, { type: "group" }> = {
152
+ type: "group",
153
+ key,
154
+ label: sc.label,
155
+ model: sc.model,
156
+ status: sc.status,
157
+ members: [sc],
158
+ };
159
+ groups.set(key, group);
160
+ out.push(group);
112
161
  }
113
162
  return out;
114
163
  }
@@ -218,6 +267,7 @@ export function buildRows(
218
267
  icon: iconOf(entry.status, entry.members.some((m) => m.phase === "queued") ? "queued" : undefined),
219
268
  expandable: true,
220
269
  expanded,
270
+ reason: entry.status === "error" ? groupReason(entry.members) : undefined,
221
271
  });
222
272
  if (!expanded) return;
223
273
  for (let i = 0; i < entry.members.length; i++) {
@@ -55,7 +55,8 @@ function formatGroup(row: GroupRow, selected: boolean, width: number, theme: The
55
55
  const chevron = row.expanded ? GLYPHS.expanded : GLYPHS.collapsed;
56
56
  const cursor = selected ? theme.fg("accent", "❯") : " ";
57
57
  const icon = row.icon === "done" ? theme.fg("success", GLYPHS.done) : row.icon === "error" ? theme.fg("error", GLYPHS.error) : theme.fg("warning", spinnerFrame());
58
- const left = `${cursor} ${row.prefix}${chevron} ${icon} ${row.label} ×${row.count}`;
58
+ const reason = row.reason === undefined ? "" : theme.fg("warning", ` · ${row.reason}`);
59
+ const left = `${cursor} ${row.prefix}${chevron} ${icon} ${row.label} ×${row.count}${reason}`;
59
60
  return assembleLine(left, row.tokens, row.tokensIn, row.tokensOut, row.model, selected, width, theme);
60
61
  }
61
62
 
@@ -0,0 +1,34 @@
1
+ /**
2
+ * abort — combine abort signals without leaking listeners.
3
+ *
4
+ * pi hands every tool execution a per-call signal (esc) while the session adds
5
+ * longer-lived controllers (/rlm-stop rotation) — in-flight work must react to both.
6
+ * The longer-lived signal outlives any single call, so its relay listener MUST detach
7
+ * when the awaited work settles, or every call accumulates one listener forever.
8
+ */
9
+
10
+ export interface RacedSignal {
11
+ /** Fires when either input aborts; undefined when no input signal was given. */
12
+ readonly signal: AbortSignal | undefined;
13
+ /** Detach relay listeners — call when the awaited work settles. */
14
+ dispose(): void;
15
+ }
16
+
17
+ export function raceAbort(a: AbortSignal | undefined, b: AbortSignal | undefined): RacedSignal {
18
+ if (a === undefined || b === undefined) return { signal: a ?? b, dispose: () => {} };
19
+ const combined = new AbortController();
20
+ const relay = (src: AbortSignal): (() => void) => () => combined.abort(src.reason);
21
+ const onA = relay(a);
22
+ const onB = relay(b);
23
+ if (a.aborted) combined.abort(a.reason);
24
+ if (b.aborted) combined.abort(b.reason);
25
+ a.addEventListener("abort", onA);
26
+ b.addEventListener("abort", onB);
27
+ return {
28
+ signal: combined.signal,
29
+ dispose: () => {
30
+ a.removeEventListener("abort", onA);
31
+ b.removeEventListener("abort", onB);
32
+ },
33
+ };
34
+ }
package/src/util/bm25.ts CHANGED
@@ -3,8 +3,10 @@
3
3
  *
4
4
  * BM25 DUALITY, documented not accidental: this file and `sandbox/py/retrieval.py`
5
5
  * (`_Bm25Index`) implement the SAME scoring for two runtimes — identical constants, identical
6
- * tokenizer, identical idf/norm formulas — so a note ranked here lands in the same order the
7
- * sandbox-side `search` would put it. Keep the two files in lockstep (AGENTS.md convention).
6
+ * tokenizer + stemmer, identical idf/norm formulas, identical PRF/phrase-bonus pipeline — so a
7
+ * note ranked here lands in the same order the sandbox-side `search` would put it. Keep the two
8
+ * files in lockstep (AGENTS.md convention). Python-only extras (window overlap, glob
9
+ * pre-filter, adjacent-window merge) are window mechanics, not scoring, and stop at that file.
8
10
  */
9
11
 
10
12
  const BM25_K1 = 1.2; // mirrors retrieval.py:_BM25_K1
@@ -14,18 +16,69 @@ const BM25_B = 0.75; // mirrors retrieval.py:_BM25_B
14
16
  const TOKEN_SPLIT = /[^0-9A-Za-z]+/;
15
17
  const CAMEL_SPLIT = /(?<=[a-z0-9])(?=[A-Z])/;
16
18
 
19
+ // BM25 V2 constants (twin: retrieval.py — identical values there).
20
+ const PRF_FEEDBACK_DOCS = 3; // top first-pass docs harvested for expansion terms
21
+ const PRF_EXPANSION_TERMS = 8; // max terms added by pseudo-relevance feedback
22
+ const PRF_EXPANSION_WEIGHT = 0.4; // Rocchio beta: expansion terms contribute at this weight
23
+ const PRF_MIN_DOCS = 12; // below this the corpus is too small to harvest from
24
+ const PHRASE_BONUS_WEIGHT = 0.25; // per adjacent query-bigram occurrence in a doc
25
+ const RERANK_POOL_MULT = 3; // phrase-bonus pool = top (k * mult) docs, capped
26
+ const RERANK_POOL_CAP = 60;
27
+
28
+ /**
29
+ * Light deterministic suffix stripper (TWIN: retrieval.py `_stem` — identical rules).
30
+ * A matching aid, not linguistics: different surface forms converge (files/file → fil,
31
+ * running/run → run, studies/study → studi); identical forms always map to themselves.
32
+ */
33
+ function stem(t: string): string {
34
+ if (t.length <= 3) return t;
35
+ let r: string;
36
+ if (t.endsWith("ies")) {
37
+ r = t.slice(0, -3) + "i"; // studies → studi
38
+ } else if (t.endsWith("sses")) {
39
+ r = t.slice(0, -2); // classes → class
40
+ } else if (t.endsWith("es")) {
41
+ const stem2 = t.slice(0, -2);
42
+ r = /(x|ch|sh)$/.test(stem2) ? stem2 : t.slice(0, -1); // boxes → box, files → file
43
+ } else if (t.endsWith("s") && !/(ss|us|is)$/.test(t)) {
44
+ r = t.length > 4 ? t.slice(0, -1) : t; // cats → cat; keeps "was"/"its"
45
+ } else {
46
+ r = t;
47
+ }
48
+ if (r.endsWith("ing") && r.length >= 6) {
49
+ let base = r.slice(0, -3); // running → runn
50
+ if (base.length >= 4) {
51
+ if (base.length >= 2 && base[base.length - 1] === base[base.length - 2]) {
52
+ base = base.slice(0, -1); // runn → run
53
+ }
54
+ r = base; // "string" stays whole
55
+ }
56
+ } else if (r.endsWith("ed") && r.length >= 5) {
57
+ let base = r.slice(0, -2); // mapped → mapp
58
+ if (base.length >= 4) {
59
+ if (base.length >= 2 && base[base.length - 1] === base[base.length - 2]) {
60
+ base = base.slice(0, -1); // mapp → map
61
+ }
62
+ r = base;
63
+ }
64
+ }
65
+ if (r.endsWith("y") && r.length > 3) r = r.slice(0, -1) + "i"; // study → studi (meets studies)
66
+ if (r.endsWith("e") && r.length > 3) r = r.slice(0, -1); // file → fil (meets files)
67
+ return r.length >= 3 ? r : t;
68
+ }
69
+
17
70
  export function bm25Tokenize(text: string): readonly string[] {
18
71
  const out: string[] = [];
19
72
  for (const raw of text.split(TOKEN_SPLIT)) {
20
73
  if (raw === "") continue;
21
74
  const lowered = raw.toLowerCase();
22
- out.push(lowered);
75
+ out.push(stem(lowered));
23
76
  if (raw.length > 3) {
24
77
  const parts = raw.split(CAMEL_SPLIT);
25
78
  if (parts.length > 1) {
26
79
  for (const part of parts) {
27
80
  const piece = part.toLowerCase();
28
- if (piece !== "" && piece !== lowered) out.push(piece);
81
+ if (piece !== "" && piece !== lowered) out.push(stem(piece));
29
82
  }
30
83
  }
31
84
  }
@@ -43,14 +96,42 @@ export interface Bm25Hit<T> {
43
96
  readonly score: number;
44
97
  }
45
98
 
99
+ /** Ranking knobs (frozen option bags; defaults enable the full V2 pipeline). */
100
+ export interface Bm25RankOptions {
101
+ /** Pseudo-relevance-feedback query expansion (auto-off below PRF_MIN_DOCS docs). */
102
+ readonly prf?: boolean;
103
+ /** Adjacent-bigram phrase bonus over a re-rank pool. */
104
+ readonly phraseBonus?: boolean;
105
+ }
106
+ export const BM25_RANK_DEFAULTS: Readonly<Bm25RankOptions> = Object.freeze({
107
+ prf: true,
108
+ phraseBonus: true,
109
+ });
110
+
111
+ /**
112
+ * Positive-scoring docs as (idx, score) pairs, best first. The Python twin's score dict only
113
+ * ever holds docs with ≥1 matching term; a dense TS array would otherwise let zero-score docs
114
+ * into PRF feedback and the phrase pool — a divergence the parity suite catches.
115
+ */
116
+ function positivePairs(scores: readonly number[]): readonly (readonly [number, number])[] {
117
+ const pairs: (readonly [number, number])[] = [];
118
+ for (let i = 0; i < scores.length; i++) {
119
+ if (scores[i] > 0) pairs.push([i, scores[i]]);
120
+ }
121
+ pairs.sort((a, b) => b[1] - a[1] || a[0] - b[0]);
122
+ return pairs;
123
+ }
124
+
46
125
  /**
47
126
  * Rank entries against a query, best first, top-k, score > 0 only.
48
- * Pre-allocated score/index arrays; no growth in the scoring loops.
127
+ * Pipeline mirrors retrieval.py:search weighted scoring PRF expansion → phrase bonus
128
+ * over a re-rank pool → top-k. Pre-allocated arrays; no growth in the scoring loops.
49
129
  */
50
130
  export function bm25Rank<T>(
51
131
  query: string,
52
132
  entries: readonly Bm25Entry<T>[],
53
133
  k: number,
134
+ options: Bm25RankOptions = BM25_RANK_DEFAULTS,
54
135
  ): readonly Bm25Hit<T>[] {
55
136
  const n = entries.length;
56
137
  if (n === 0 || k <= 0) return [];
@@ -70,28 +151,96 @@ export function bm25Rank<T>(
70
151
  let totalLen = 0;
71
152
  for (let i = 0; i < n; i++) totalLen += docLens[i];
72
153
  const avgLen = totalLen > 0 ? totalLen / n : 1.0;
154
+ const idfOf = (term: string): number => {
155
+ const df = postings.get(term)?.length ?? 0;
156
+ return Math.log(1.0 + (n - df + 0.5) / (df + 0.5));
157
+ };
158
+ const score = (weights: ReadonlyMap<string, number>): number[] => {
159
+ const scores = new Array<number>(n).fill(0);
160
+ for (const [term, weight] of weights) {
161
+ const posting = postings.get(term);
162
+ if (posting === undefined) continue;
163
+ const idf = idfOf(term);
164
+ for (const [idx, tf] of posting) {
165
+ const norm = BM25_K1 * (1.0 - BM25_B + BM25_B * (docLens[idx] / avgLen));
166
+ scores[idx] += weight * idf * (tf * (BM25_K1 + 1.0)) / (tf + norm);
167
+ }
168
+ }
169
+ return scores;
170
+ };
73
171
 
74
- const scores = new Array<number>(n).fill(0);
75
- const seen = new Set<string>();
172
+ const weights = new Map<string, number>();
76
173
  for (const term of bm25Tokenize(query)) {
77
- if (seen.has(term)) continue; // Python twin scores set(terms)
78
- seen.add(term);
79
- const posting = postings.get(term);
80
- if (posting === undefined) continue;
81
- const df = posting.length;
82
- const idf = Math.log(1.0 + (n - df + 0.5) / (df + 0.5));
83
- for (const [idx, tf] of posting) {
84
- const norm = BM25_K1 * (1.0 - BM25_B + BM25_B * (docLens[idx] / avgLen));
85
- scores[idx] += (idf * (tf * (BM25_K1 + 1.0))) / (tf + norm);
174
+ if (!weights.has(term)) weights.set(term, 1.0); // Python twin: first occurrence wins
175
+ }
176
+ let scores = score(weights);
177
+ const basePositive = positivePairs(scores);
178
+ if (options.prf !== false && n >= PRF_MIN_DOCS && basePositive.length > 0) {
179
+ const feedback = basePositive.slice(0, PRF_FEEDBACK_DOCS).map(([idx]) => idx);
180
+ const tf = new Map<string, number>();
181
+ for (const idx of feedback) {
182
+ for (const term of bm25Tokenize(entries[idx].text)) {
183
+ tf.set(term, (tf.get(term) ?? 0) + 1);
184
+ }
185
+ }
186
+ // Exclude query terms BEFORE slicing — the twin (retrieval.py:_expansion_terms) excludes
187
+ // first, then takes the top-8; slicing first would let query terms eat expansion slots.
188
+ const ranked: readonly (readonly [number, string])[] = Array.from(tf, ([term, f]) => [f * idfOf(term), term] as const)
189
+ .filter(([, term]) => !weights.has(term))
190
+ .sort((a, b) => b[0] - a[0] || (a[1] < b[1] ? -1 : a[1] > b[1] ? 1 : 0))
191
+ .slice(0, PRF_EXPANSION_TERMS);
192
+ let added = false;
193
+ for (const [, term] of ranked) {
194
+ if (!weights.has(term)) {
195
+ weights.set(term, PRF_EXPANSION_WEIGHT);
196
+ added = true;
197
+ }
86
198
  }
199
+ if (added) scores = score(weights);
200
+ }
201
+ if (scores.every((s) => s <= 0)) return [];
202
+
203
+ // Pool of (idx, score) pairs — the phrase bonus MUTATES the score (twin parity), so the
204
+ // reported value must come from the pool pairs, not the pre-bonus score array.
205
+ const positive = positivePairs(scores);
206
+ const poolN = Math.min(positive.length, RERANK_POOL_CAP, Math.max(k * RERANK_POOL_MULT, PRF_FEEDBACK_DOCS));
207
+ let pool = positive.slice(0, poolN);
208
+ const seq = bm25Tokenize(query); // raw sequence — repeated terms make self-bigrams, like Python
209
+ const bigrams: (readonly [string, string])[] = [];
210
+ if (options.phraseBonus !== false) {
211
+ for (let i = 0; i < seq.length - 1; i++) {
212
+ if (seq[i] !== seq[i + 1] && !bigrams.some(([a, b]) => a === seq[i] && b === seq[i + 1])) {
213
+ bigrams.push([seq[i], seq[i + 1]]);
214
+ }
215
+ }
216
+ }
217
+ if (bigrams.length > 0) {
218
+ const boosted = pool.map(([idx, s]) => {
219
+ const toks = bm25Tokenize(entries[idx].text);
220
+ const pos = new Map<string, number[]>();
221
+ for (let i = 0; i < toks.length; i++) {
222
+ const list = pos.get(toks[i]);
223
+ if (list !== undefined) list.push(i);
224
+ else pos.set(toks[i], [i]);
225
+ }
226
+ let bonus = 0;
227
+ for (const [a, b] of bigrams) {
228
+ const pa = pos.get(a);
229
+ const pb = pos.get(b);
230
+ if (pa !== undefined && pb !== undefined && pa.some((x) => pb.some((y) => y === x + 1))) {
231
+ bonus += PHRASE_BONUS_WEIGHT * Math.max(idfOf(a), idfOf(b));
232
+ }
233
+ }
234
+ return [idx, s + bonus] as const;
235
+ });
236
+ boosted.sort((x, y) => y[1] - x[1] || x[0] - y[0]);
237
+ pool = boosted;
87
238
  }
88
239
 
89
- const order = Array.from({ length: n }, (_, i) => i);
90
- order.sort((a, b) => scores[b] - scores[a] || a - b); // deterministic tie-break
91
240
  const out: Bm25Hit<T>[] = [];
92
- for (let i = 0; i < n && out.length < k; i++) {
93
- const idx = order[i];
94
- if (scores[idx] > 0) out.push({ item: entries[idx].item, score: scores[idx] });
241
+ for (const [idx, s] of pool) {
242
+ if (out.length >= k) break;
243
+ if (s > 0) out.push({ item: entries[idx].item, score: s });
95
244
  }
96
245
  return out;
97
246
  }
@@ -12,7 +12,7 @@ export function err<T = never, E = string>(error: E): Result<T, E> {
12
12
  return { ok: false, error };
13
13
  }
14
14
 
15
- const ERROR_PREFIX = "Error:";
15
+ export const ERROR_PREFIX = "Error:";
16
16
 
17
17
  export function formatError(message: string): string {
18
18
  return `${ERROR_PREFIX} ${message}`;