feat(api/search): neighbor + group-aware budgeting for highlights

Replace the service-side raw-line char budget with client-side budgeting over
the model's scored spans (highlight-budget.ts):
- neighbor policy: each selected line may pull in its ±1 adjacent lines for
  context, capped at 35% of the budget so context can't crowd out answer lines.
- group budget: merge adjacent selected lines into blocks, rank blocks by best
  span score, emit best blocks in page order until the char budget is reached —
  so many tiny unrelated lines don't each eat the budget.

generateHighlightsBatch now returns scored spans (index+score) and no longer
sends max_highlight_chars (budgeting moved client-side). highlights.ts runs
selectHighlightIndices before assembleAnswer.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: mogery

apps/api/src/search/highlight-budget.test.ts

@@ -0,0 +1,93 @@
+import { selectHighlightIndices } from "./highlight-budget";
+
+describe("selectHighlightIndices", () => {
+  it("returns [] when there are no scored spans", () => {
+    expect(selectHighlightIndices([10, 10], [])).toEqual([]);
+  });
+
+  it("ignores out-of-range / invalid span indices", () => {
+    const out = selectHighlightIndices(
+      [10, 10],
+      [
+        { index: 5, score: 0.9 },
+        { index: -1, score: 0.8 },
+      ],
+    );
+    expect(out).toEqual([]);
+  });
+
+  it("expands a selected line with its ±1 neighbors", () => {
+    const out = selectHighlightIndices(
+      [10, 10, 10, 10, 10],
+      [{ index: 2, score: 0.9 }],
+    );
+    expect(out).toEqual([1, 2, 3]);
+  });
+
+  it("does not add neighbors that exceed the neighbor budget", () => {
+    // Neighbors are 100 chars each; neighbor budget = floor(100 * 0.35) = 35.
+    const out = selectHighlightIndices(
+      [100, 100, 10, 100, 100],
+      [{ index: 2, score: 0.9 }],
+      { maxChars: 100, neighborBudgetFraction: 0.35 },
+    );
+    expect(out).toEqual([2]);
+  });
+
+  it("gives neighbor budget to the higher-scoring line first", () => {
+    // Budget allows exactly one 50-char neighbor. index 4 (0.9) should claim it
+    // (its neighbor 3), leaving index 0 (0.4) without one.
+    const out = selectHighlightIndices(
+      [50, 50, 50, 50, 50],
+      [
+        { index: 0, score: 0.4 },
+        { index: 4, score: 0.9 },
+      ],
+      { maxChars: 1000, neighborBudgetFraction: 0.05 },
+    );
+    expect(out).toEqual([0, 3, 4]);
+  });
+
+  it("merges adjacent selections (incl. a bridging neighbor) into one block", () => {
+    const out = selectHighlightIndices(
+      [10, 10, 10, 10],
+      [
+        { index: 0, score: 0.3 },
+        { index: 2, score: 0.9 },
+      ],
+    );
+    expect(out).toEqual([0, 1, 2, 3]);
+  });
+
+  it("emits the highest-scoring block first and drops blocks that exceed budget", () => {
+    // Two non-adjacent blocks of 100 chars; budget fits only one.
+    const out = selectHighlightIndices(
+      [100, 100, 100, 100, 100, 100, 100],
+      [
+        { index: 0, score: 0.5 },
+        { index: 6, score: 0.9 },
+      ],
+      { maxChars: 150, neighborBudgetFraction: 0 },
+    );
+    expect(out).toEqual([6]);
+  });
+
+  it("returns chosen blocks in page order, not score order", () => {
+    const out = selectHighlightIndices(
+      [100, 100, 100, 100, 100, 100, 100],
+      [
+        { index: 0, score: 0.5 },
+        { index: 6, score: 0.9 },
+      ],
+      { maxChars: 500, neighborBudgetFraction: 0 },
+    );
+    expect(out).toEqual([0, 6]);
+  });
+
+  it("always keeps the top block even if it alone exceeds the budget", () => {
+    const out = selectHighlightIndices([1000], [{ index: 0, score: 0.5 }], {
+      maxChars: 800,
+    });
+    expect(out).toEqual([0]);
+  });
+});

apps/api/src/search/highlight-budget.ts

@@ -0,0 +1,96 @@
+import type { ScoredSpan } from "./highlight-model";
+
+// Turn the model's scored spans into the final set of line indices to emit,
+// within a character budget. Two policies, applied in order:
+//
+//  1. Neighbor policy: each selected line may pull in its ±1 adjacent lines for
+//     context, but neighbors collectively get at most NEIGHBOR_BUDGET_FRACTION
+//     of the budget — so context never crowds out actual answer lines.
+//  2. Group budget: merge adjacent selected lines into contiguous blocks, rank
+//     blocks by their best span score, and emit the best blocks (in page order)
+//     until the budget is reached. Budgeting by block rather than raw line keeps
+//     a coherent passage together instead of letting many tiny unrelated lines
+//     each eat into the budget.
+
+// Total character budget for the assembled snippet (approximate — measured on
+// raw line text, not the markdown assembleAnswer adds back).
+const MAX_HIGHLIGHT_CHARS = 800;
+// Neighbor (context) lines may use at most this fraction of the budget.
+const NEIGHBOR_BUDGET_FRACTION = 0.35;
+
+interface Block {
+  indices: number[];
+  chars: number;
+  score: number; // best core-span score in the block
+}
+
+/**
+ * Choose which line indices to emit. `lineLengths[i]` is the character length of
+ * line i (sentences[i].text.length); `scored` are the model's selected spans.
+ * Returns indices in ascending (page) order, ready for assembleAnswer.
+ */
+export function selectHighlightIndices(
+  lineLengths: number[],
+  scored: ScoredSpan[],
+  opts: { maxChars?: number; neighborBudgetFraction?: number } = {},
+): number[] {
+  const maxChars = opts.maxChars ?? MAX_HIGHLIGHT_CHARS;
+  const neighborFraction =
+    opts.neighborBudgetFraction ?? NEIGHBOR_BUDGET_FRACTION;
+  const n = lineLengths.length;
+
+  // Core selected lines (valid, in range). Keep the best score per index.
+  const coreScore = new Map<number, number>();
+  for (const s of scored) {
+    if (!Number.isInteger(s.index) || s.index < 0 || s.index >= n) continue;
+    const prev = coreScore.get(s.index);
+    if (prev === undefined || s.score > prev) coreScore.set(s.index, s.score);
+  }
+  if (coreScore.size === 0) return [];
+
+  // Process core lines best-score first so higher-value lines claim neighbor
+  // budget before lower-value ones.
+  const coreByScore = [...coreScore.entries()].sort((a, b) => b[1] - a[1]);
+
+  // Neighbor expansion (±1 line), capped at a fraction of the total budget.
+  const neighborBudget = Math.floor(maxChars * neighborFraction);
+  const selected = new Set<number>(coreScore.keys());
+  let neighborChars = 0;
+  for (const [idx] of coreByScore) {
+    for (const nb of [idx - 1, idx + 1]) {
+      if (nb < 0 || nb >= n) continue;
+      if (selected.has(nb)) continue;
+      if (neighborChars + lineLengths[nb] > neighborBudget) continue;
+      selected.add(nb);
+      neighborChars += lineLengths[nb];
+    }
+  }
+
+  // Merge adjacent selected indices into contiguous blocks.
+  const sorted = [...selected].sort((a, b) => a - b);
+  const blocks: Block[] = [];
+  for (const idx of sorted) {
+    const last = blocks[blocks.length - 1];
+    const score = coreScore.get(idx) ?? -Infinity;
+    if (last && idx === last.indices[last.indices.length - 1] + 1) {
+      last.indices.push(idx);
+      last.chars += lineLengths[idx];
+      last.score = Math.max(last.score, score);
+    } else {
+      blocks.push({ indices: [idx], chars: lineLengths[idx], score });
+    }
+  }
+
+  // Rank blocks by best score and emit until the budget is reached. Always keep
+  // the top block so we return something even if it alone exceeds the budget.
+  blocks.sort((a, b) => b.score - a.score);
+  const chosen: Block[] = [];
+  let total = 0;
+  for (const b of blocks) {
+    if (chosen.length > 0 && total + b.chars > maxChars) continue;
+    chosen.push(b);
+    total += b.chars;
+  }
+
+  return chosen.flatMap(b => b.indices).sort((a, b) => a - b);
+}

apps/api/src/search/highlight-model.test.ts

@@ -63,13 +63,14 @@ describe("generateHighlightsBatch", () => {
       query: "q1",
       lines: ["a one", "b one"],
     });
-    // Threshold + top_k + budget are always sent so the cutoff matches the trained format.
+    // Threshold + top_k are sent; char budgeting is done client-side (group-aware),
+    // so max_highlight_chars is intentionally NOT sent.
     expect(typeof sent.requests[0].threshold).toBe("number");
     expect(typeof sent.requests[0].top_k).toBe("number");
-    expect(typeof sent.requests[0].max_highlight_chars).toBe("number");
+    expect(sent.requests[0].max_highlight_chars).toBeUndefined();
   });
 
-  it("returns the selected span indices per page, aligned by request order", async () => {
+  it("returns the scored spans per page, aligned by request order", async () => {
     mockFetchOnce({
       results: [
         {
@@ -90,7 +91,13 @@ describe("generateHighlightsBatch", () => {
       { logger },
     );
 
-    expect(out).toEqual([[0, 2], [1]]);
+    expect(out).toEqual([
+      [
+        { index: 0, score: 0.9 },
+        { index: 2, score: 0.5 },
+      ],
+      [{ index: 1, score: 0.7 }],
+    ]);
   });
 
   it("returns an empty array for a page with no selected spans", async () => {
@@ -107,26 +114,31 @@ describe("generateHighlightsBatch", () => {
     expect(out).toEqual([[], []]);
   });
 
-  it("ignores malformed (non-integer / missing) indices", async () => {
+  it("ignores malformed (non-integer index / missing score) spans", async () => {
     mockFetchOnce({
       results: [
         {
           highlights: [
             { index: 1, score: 0.9 },
-            { score: 0.8 }, // missing index
-            { index: 2.5, score: 0.7 }, // non-integer
+            { index: 4 }, // missing score
+            { index: 2.5, score: 0.7 }, // non-integer index
             { index: 3, score: 0.6 },
           ],
         },
       ],
     });
 
     const out = await generateHighlightsBatch(
-      [{ query: "q", lines: ["a", "b", "c", "d"] }],
+      [{ query: "q", lines: ["a", "b", "c", "d", "e"] }],
       { logger },
     );
 
-    expect(out).toEqual([[1, 3]]);
+    expect(out).toEqual([
+      [
+        { index: 1, score: 0.9 },
+        { index: 3, score: 0.6 },
+      ],
+    ]);
   });
 
   it("falls back to nulls (one per item) when the service errors", async () => {

apps/api/src/search/highlight-model.ts

@@ -5,22 +5,29 @@ import { config } from "../config";
 // trained on the line-like spans `parseMarkdownToSentences` produces, so callers
 // pass those spans as explicit `lines`, the service scores each against the
 // query, and returns the selected spans as `highlights[]` with their `index`
-// back into the input `lines` array. Callers map those indices through
-// `assembleAnswer` to rebuild structure (tables/code). We push every page
-// through one `/batch_highlight` call (resident GPU model). Endpoint + token are
-// config-gated (HIGHLIGHT_MODEL_URL / HIGHLIGHT_MODEL_TOKEN); callers must
-// confirm they're set via highlightsEnvReady() before invoking.
+// back into the input `lines` array (and a `score`). Callers run those scored
+// spans through the neighbor/group budgeter, then `assembleAnswer` to rebuild
+// structure (tables/code). We push every page through one `/batch_highlight`
+// call (resident GPU model). Endpoint + token are config-gated
+// (HIGHLIGHT_MODEL_URL / HIGHLIGHT_MODEL_TOKEN); callers must confirm they're
+// set via highlightsEnvReady() before invoking.
 
 // Experimental score cutoff for keeping a span, tuned against the line-span
 // format the model was trained on.
 const HIGHLIGHT_THRESHOLD = 0.08;
 // Cap on how many spans to keep per page after thresholding (highest-scoring).
+// We do char budgeting ourselves (group-aware, in selectHighlightIndices), so we
+// don't send max_highlight_chars — the service would budget by raw lines.
 const HIGHLIGHT_TOP_K = 12;
-// Character budget per page; the service keeps the highest-scoring spans until
-// the budget is reached (measured on raw span text).
-const MAX_HIGHLIGHT_CHARS = 800;
 const REQUEST_TIMEOUT_MS = 30000;
 
+// One scored candidate span: `index` into the page's `lines`, `score` from the
+// model. Consumed by the neighbor/group budgeter (selectHighlightIndices).
+export interface ScoredSpan {
+  index: number;
+  score: number;
+}
+
 interface HighlightItem {
   query: string;
   // Candidate spans for the page, in document order — from
@@ -41,14 +48,18 @@ interface BatchHighlightResponse {
   results?: HighlightResult[];
 }
 
-// Pull the selected span indices out of one page's result. The service already
-// applied the threshold / char budget, so we just collect the kept indices.
-function selectedIndices(result: HighlightResult | undefined): number[] {
+// Pull the scored spans out of one page's result, keeping only well-formed ones
+// (integer index, numeric score). The service already applied threshold + top_k.
+function scoredSpans(result: HighlightResult | undefined): ScoredSpan[] {
   const highlights = result?.highlights ?? [];
-  const out: number[] = [];
+  const out: ScoredSpan[] = [];
   for (const h of highlights) {
-    if (typeof h.index === "number" && Number.isInteger(h.index)) {
-      out.push(h.index);
+    if (
+      typeof h.index === "number" &&
+      Number.isInteger(h.index) &&
+      typeof h.score === "number"
+    ) {
+      out.push({ index: h.index, score: h.score });
     }
   }
   return out;
@@ -57,14 +68,14 @@ function selectedIndices(result: HighlightResult | undefined): number[] {
 /**
  * Score many pages' candidate spans in a single batch call to the
  * query-highlights model. Returns an array aligned to `items`: for each page,
- * the indices (into that page's `lines`) the model selected, or null when the
- * whole batch fails (so callers can keep the provider snippet). A page with no
- * spans clearing the threshold yields an empty array, not null.
+ * the scored spans (index into that page's `lines`, plus score) the model
+ * selected, or null when the whole batch fails (so callers can keep the provider
+ * snippet). A page with no spans clearing the threshold yields an empty array.
  */
 export async function generateHighlightsBatch(
   items: HighlightItem[],
   opts: { logger: Logger },
-): Promise<(number[] | null)[]> {
+): Promise<(ScoredSpan[] | null)[]> {
   if (items.length === 0) return [];
 
   const controller = new AbortController();
@@ -83,7 +94,6 @@ export async function generateHighlightsBatch(
           lines: item.lines,
           threshold: HIGHLIGHT_THRESHOLD,
           top_k: HIGHLIGHT_TOP_K,
-          max_highlight_chars: MAX_HIGHLIGHT_CHARS,
         })),
       }),
       signal: controller.signal,
@@ -98,15 +108,15 @@ export async function generateHighlightsBatch(
 
     const data = (await res.json()) as BatchHighlightResponse;
     const results = data.results ?? [];
-    const indices = items.map((_, i) => selectedIndices(results[i]));
+    const spans = items.map((_, i) => scoredSpans(results[i]));
 
     opts.logger.info("query highlights batch generated", {
       pages: items.length,
-      withHighlights: indices.filter(x => x.length > 0).length,
+      withHighlights: spans.filter(x => x.length > 0).length,
       elapsedMs: Date.now() - start,
     });
 
-    return indices;
+    return spans;
   } catch (error) {
     opts.logger.warn("query highlights batch failed", {
       error: error instanceof Error ? error.message : String(error),