Blame for node/reader.ts

// The Canopy reader on one screen, in Node: the exported ONNX model (scripts/export_reader.py) plus its
// tokenizer, encoded exactly as canopy/classifier/model.py:encode() and decoded as canopy/reader/evaluate.py.
// A reference for Daintree's integration; src/parity.ts checks it against the Python outputs.
import { readFileSync } from "node:fs";
import { basename, dirname, join, resolve } from "node:path";
import * as ort from "onnxruntime-node";
import { AutoTokenizer, env } from "@huggingface/transformers";

export interface ReaderMeta {
  max_len: number;
  header_tokens: number;
  line_tokens: number;
  question_candidate_rows: number;
  status_candidate_rows: number;
  inputs: string[];
  outputs: string[];
  categories: string[];
  classes: Record<string, string[]>;
  thresholds: { summary_first: number; summary_more: number; option: number; gate: number; gate_categories: string[] };
}

export interface Reading {
  category: string;
  category_probs: number[];
  needs_attention: number; // gated and rescaled: >= 0.5 means the card needs the user
  needs_attention_raw: number;
  blocked_on_user: number;
  notable: number;
  question_line: number; // line indices into the prepared screen, -1 for none
  status_line: number; // the headline
  action_line: number;
  task_line: number;
  summary_lines: number[]; // most important first
  option_lines: number[]; // screen order
  key_lines: number[];
  tests: string;
  changes: string;
  risk: string;
  instruction: string;
  progress: string;
  failure_repeats: string;
  score: string;
}

const cp = (s: string) => [...s].length; // Python len(): code points, not UTF-16 units
// Python's str.strip() whitespace (str.isspace), which differs from JS \s: it includes \x1c-\x1f and \x85 and
// excludes \ufeff. Training used Python, so blank-line skipping and candidate rows must use the same set.
const PY_WS = "\t\n\v\f\r\x1c\x1d\x1e\x1f \x85\xa0\u1680\u2000-\u200a\u2028\u2029\u202f\u205f\u3000";
const STRIP = new RegExp(`^[${PY_WS}]+|[${PY_WS}]+$`, "gu");
const strip = (s: string) => s.replace(STRIP, "");
const sigmoid = (x: number) => 1 / (1 + Math.exp(-x));

function softmax(xs: number[]): number[] {
  const m = Math.max(...xs);
  const e = xs.map((x) => Math.exp(x - m));
  const s = e.reduce((a, b) => a + b, 0);
  return e.map((x) => x / s);
}

export function lineCandidates(lines: string[], rows: number): number[] {
  const start = Math.max(0, lines.length - rows);
  const out: number[] = [];
  for (let i = start; i < lines.length; i++) if (cp(strip(lines[i])) > 3) out.push(i);
  return out;
}

export class CanopyReader {
  meta: ReaderMeta;
  session: ort.InferenceSession;
  tok: any;
  cls: number;
  sep: number;
  pad: number;

  private constructor(meta: ReaderMeta, session: ort.InferenceSession, tok: any) {
    this.meta = meta;
    this.session = session;
    this.tok = tok;
    this.cls = tok.model.tokens_to_ids.get(tok.cls_token ?? "[CLS]");
    this.sep = tok.model.tokens_to_ids.get(tok.sep_token ?? "[SEP]");
    this.pad = tok.model.tokens_to_ids.get(tok.pad_token ?? "[PAD]");
  }

  /** `dir` holds reader.json, the tokenizer files and model.<variant>.onnx. */
  static async load(dir: string, variant = "int8", threads = 0): Promise<CanopyReader> {
    const meta = JSON.parse(readFileSync(join(dir, "reader.json"), "utf8")) as ReaderMeta;
    env.allowRemoteModels = false;
    env.localModelPath = dirname(resolve(dir)) + "/";
    const tok = await AutoTokenizer.from_pretrained(basename(resolve(dir)), { local_files_only: true } as any);
    const opts: ort.InferenceSession.SessionOptions = { graphOptimizationLevel: "all" };
    if (threads) opts.intraOpNumThreads = threads;
    const session = await ort.InferenceSession.create(join(dir, `model.${variant}.onnx`), opts);
    return new CanopyReader(meta, session, tok);
  }

  private ids(text: string): number[] {
    return this.tok.encode(text, { add_special_tokens: false });
  }

  /** As encode(): [CLS] state-json [SEP] then "[i] line\n" pieces filled from the bottom up, [SEP]. */
  /** `state` must hold integer numbers only (as Daintree sends): JSON.stringify writes 1.0 as 1, where the
   * Python training encoder wrote 1.0, so a float in the state would tokenize differently. */
  encode(state: object, lines: string[]) {
    const m = this.meta;
    const header = this.ids(JSON.stringify(state)).slice(0, m.header_tokens);
    let budget = m.max_len - header.length - 3;
    const pieces: [number, number[]][] = [];
    for (let i = lines.length - 1; i >= 0; i--) {
      if (!strip(lines[i])) continue;
      const ids = this.ids(`[${i}] ${lines[i]}\n`).slice(0, m.line_tokens);
      if (ids.length > budget) break;
      budget -= ids.length;
      pieces.push([i, ids]);
    }
    pieces.reverse();
    const ids = [this.cls, ...header, this.sep];
    const lineIndex: number[] = [];
    const linePos: number[] = [];
    for (const [i, p] of pieces) {
      lineIndex.push(i);
      linePos.push(ids.length);
      ids.push(...p);
    }
    ids.push(this.sep);
    const q = new Set(lineCandidates(lines, m.question_candidate_rows));
    const s = new Set(lineCandidates(lines, m.status_candidate_rows));
    return { ids, lineIndex, linePos, qMask: lineIndex.map((i) => q.has(i)), sMask: lineIndex.map((i) => s.has(i)) };
  }

  async read(state: object, lines: string[]): Promise<Reading> {
    const e = this.encode(state, lines);
    const n = e.ids.length;
    const k = Math.max(1, e.lineIndex.length);
    const pad = <T>(xs: T[], fill: T) => [...xs, ...Array(k - xs.length).fill(fill)];
    const feeds: Record<string, ort.Tensor> = {
      input_ids: new ort.Tensor("int64", BigInt64Array.from(e.ids.map(BigInt)), [1, n]),
      attention_mask: new ort.Tensor("int64", BigInt64Array.from(e.ids.map(() => 1n)), [1, n]),
      line_pos: new ort.Tensor("int64", BigInt64Array.from(pad(e.linePos, 0).map(BigInt)), [1, k]),
      line_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.lineIndex.map(() => 1), 0)), [1, k]),
      q_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.qMask.map(Number), 0)), [1, k]),
      s_mask: new ort.Tensor("bool", Uint8Array.from(pad(e.sMask.map(Number), 0)), [1, k]),
    };
    const out = await this.session.run(feeds);
    const arr = (name: string) => Array.from(out[name].data as Float32Array, Number);
    return this.decode(arr, e.lineIndex, e.qMask);
  }

  decode(arr: (name: string) => number[], idx: number[], qMask: boolean[]): Reading {
    const m = idx.length;
    const t = this.meta.thresholds;
    const pick = (name: string) => {
      const scores = arr(name);
      const cand = [...scores.slice(0, m), scores[scores.length - 1]];
      let j = 0;
      for (let x = 1; x < cand.length; x++) if (cand[x] > cand[j]) j = x;
      return j === m ? -1 : idx[j];
    };
    const logit = arr("summary_lines").slice(0, m);
    const pSum = logit.map(sigmoid);
    const ranked = [...logit.keys()].sort((a, b) => logit[b] - logit[a] || a - b); // by logit: no saturation ties
    const summary: number[] = [];
    for (const j of ranked.slice(0, 3)) {
      if (pSum[j] >= (summary.length ? t.summary_more : t.summary_first)) summary.push(idx[j]);
      else break;
    }
    const pOpt = arr("option_lines").slice(0, m).map(sigmoid);
    const key = arr("key_lines").slice(0, m);
    const probs = softmax(arr("category"));
    const raw = sigmoid(arr("needs_attention")[0]);
    const needs = t.gate_categories.reduce((a, c) => a + probs[this.meta.categories.indexOf(c)], 0);
    const g = Math.max(raw, needs);
    const argmax = (xs: number[]) => xs.reduce((b, x, i) => (x > xs[b] ? i : b), 0);
    const cls = (name: string) => this.meta.classes[name][argmax(arr(name))];
    return {
      category: this.meta.categories[argmax(probs)],
      category_probs: probs,
      needs_attention: g < t.gate ? (0.5 * g) / t.gate : 0.5 + (0.5 * (g - t.gate)) / (1 - t.gate),
      needs_attention_raw: raw,
      blocked_on_user: sigmoid(arr("blocked_on_user")[0]),
      notable: sigmoid(arr("notable")[0]),
      question_line: pick("question_line"),
      status_line: pick("status_line"),
      action_line: pick("action_line"),
      task_line: pick("task_line"),
      summary_lines: summary,
      option_lines: idx.filter((_, j) => qMask[j] && pOpt[j] >= t.option),
      key_lines: idx.filter((_, j) => key[j] > 0),
      tests: cls("tests"),
      changes: cls("changes"),
      risk: cls("risk"),
      instruction: cls("instruction"),
      progress: cls("progress"),
      failure_repeats: cls("failure_repeats"),
      score: cls("score"),
    };
  }
}

// --- The card as Daintree shows it: line picks turned into clean text ---------------------------------

// Bullets, tree and box glyphs, spinners, prompt markers: stripped only when whitespace (or the end) follows,
// so `**bold**` and `>file` survive.
const LEAD = /^(?:[⏺•●✦⎿└├│┃║■✻✶✢✳✽✺·*❯›>▸▶◆◇](?:\s+|$))+/u;
const BOX_TAIL = /\s*[│┃║]\s*$/u;
const TIMER = /\s*\((?:\d[^)]*)?(?:esc to \w+|thinking|thought for|ctrl\+\w to \w+)[^)]*\)\s*$/iu;
const SENTENCE_END = /[.!?:;)"'`]$/u;

/** One screen line as card text: glyphs, borders and timers off, words kept. null when nothing is left. */
export function displayText(line: string | undefined): string | null {
  if (line == null) return null;
  const text = strip(line).replace(LEAD, "").replace(BOX_TAIL, "").replace(TIMER, "").trim();
  return cp(text) >= 2 ? text : null;
}

export interface Card {
  category: string;
  needsAttention: boolean;
  score: number;
  headline: string | null;
  summary: string[]; // up to 3, most important first; a sentence wrapped over lines is one entry
  question: string | null;
  options: string[];
  action: string | null;
  task: string | null;
  tests: string;
  changes: string;
  progress: number | null;
}

/** Does screen line `j` carry on the sentence of the line above it? Yes when it was picked with it, or
 * when it has no bullet of its own and starts in lowercase or indented (an agent's own line break). */
function continues(lines: string[], j: number, picks: number[]): boolean {
  const next = lines[j];
  if (next == null || !strip(next)) return false;
  if (LEAD.test(strip(next))) return false; // a new bullet is a new item, even when it was picked too
  if (picks.includes(j)) return true;
  return /^\s{2,}\S/u.test(next) || /^\p{Ll}/u.test(strip(next));
}

/** Line `i` as text, with up to two following lines joined when the sentence carries on over them. */
function sentenceAt(lines: string[], i: number, picks: number[], used: Set<number>): string | null {
  let text = displayText(lines[i]);
  used.add(i);
  for (let j = i + 1; text && j <= i + 2 && !SENTENCE_END.test(text) && !used.has(j) && continues(lines, j, picks); j++) {
    used.add(j);
    const more = displayText(lines[j]);
    if (more) text += " " + more;
  }
  return text;
}

/** Summary picks as text, most important first, each a whole sentence where it wraps. */
function summaryText(lines: string[], picks: number[], used: Set<number>): string[] {
  const out: string[] = [];
  for (const i of picks) {
    if (used.has(i)) continue;
    const text = sentenceAt(lines, i, picks, used);
    if (text) out.push(text);
  }
  return out;
}

export function card(r: Reading, lines: string[]): Card {
  const asks = r.category === "approval" || r.category === "question";
  const line = (i: number) => (i >= 0 ? displayText(lines[i]) : null);
  const used = new Set<number>();
  const headline = r.status_line >= 0 ? sentenceAt(lines, r.status_line, r.summary_lines, used) : null;
  return {
    category: r.category,
    needsAttention: r.needs_attention >= 0.5,
    score: Number(r.score),
    headline,
    summary: summaryText(lines, r.summary_lines, used),
    question: asks ? line(r.question_line) : null,
    options: asks ? r.option_lines.map((i) => displayText(lines[i])).filter((x): x is string => !!x) : [],
    action: r.category === "approval" ? line(r.action_line) : null,
    task: line(r.task_line),
    tests: r.tests,
    changes: r.changes,
    progress: r.progress === "none" ? null : Number(r.progress),
  };
}