autocomplete-model/shared/nextedit.ts

/**
 * Next-edit prediction prompt format, shared by the training pipeline and the daemon.
 *
 *   <|repo_name|>{repo}\n
 *   [<|file_sep|>{dir}/__errors__.txt\n{errors}\n]          (shared/diagnostics.ts, optional)
 *   [<|file_sep|>__recent_edits__.diff\n{diff}\n]           (most recent edit(s), optional)
 *   <|file_sep|>{path}\n
 *   {file before region}<REGION_START>{region, with <CURSOR> at the cursor}<REGION_END>{file after region}<OUTPUT_START>
 *
 * Target: the rewritten region, then <OUTPUT_END>; or <NO_EDIT><OUTPUT_END> when nothing should change.
 * Markers reuse Qwen2.5's vision special tokens, which never occur in code.
 */
import type { Encode } from './prompt.ts'

export const NEXT_EDIT_TOKENS = {
  regionStart: '<|quad_start|>',
  regionEnd: '<|quad_end|>',
  cursor: '<|box_start|>',
  outputStart: '<|object_ref_start|>',
  outputEnd: '<|object_ref_end|>',
  noEdit: '<|vision_pad|>',
} as const

export type NextEditTokenIds = Record<keyof typeof NEXT_EDIT_TOKENS | 'repoName' | 'fileSep', number>

export interface NextEditInput {
  repo: string
  path: string
  /** Current file text (after the user's recent edit). */
  file: string
  /** Editable region: [startLine, endLine) in `file`, 0-based. */
  regionStartLine: number
  regionEndLine: number
  /** Cursor offset in `file` (must lie inside the region). */
  cursor: number
  recentEdits?: string
  errors?: string
}

/** Character offsets of the start of each line, plus file length. */
export function lineOffsets(text: string): number[] {
  const out = [0]
  for (let i = text.indexOf('\n'); i !== -1; i = text.indexOf('\n', i + 1)) out.push(i + 1)
  return out
}

/** Region text for [startLine, endLine): whole lines, including the final newline when present. */
export function regionText(text: string, startLine: number, endLine: number): string {
  const offs = lineOffsets(text)
  return text.slice(offs[startLine], endLine < offs.length ? offs[endLine] : text.length)
}

/**
 * Compact line diff of one edit for the "recent edits" context:
 *   @@ path:line
 *   -old
 *   +new
 */
export function lineDiff(path: string, before: string, after: string, maxLines = 12): string {
  const a = before.split('\n'), b = after.split('\n')
  let p = 0
  while (p < a.length && p < b.length && a[p] === b[p]) p++
  let s = 0
  while (s < a.length - p && s < b.length - p && a[a.length - 1 - s] === b[b.length - 1 - s]) s++
  const removed = a.slice(p, a.length - s).slice(0, maxLines)
  const added = b.slice(p, b.length - s).slice(0, maxLines)
  if (!removed.length && !added.length) return ''
  return [`@@ ${path}:${p + 1}`, ...removed.map(l => `-${l}`), ...added.map(l => `+${l}`)].join('\n')
}

/** Changed line range of `after` relative to `before`, in `after`'s line numbers: [start, end). */
export function changedLines(before: string, after: string): [number, number] | null {
  const a = before.split('\n'), b = after.split('\n')
  let p = 0
  while (p < a.length && p < b.length && a[p] === b[p]) p++
  if (p === a.length && p === b.length) return null
  let s = 0
  while (s < a.length - p && s < b.length - p && a[a.length - 1 - s] === b[b.length - 1 - s]) s++
  return [p, Math.max(p + 1, b.length - s)]
}

/**
 * Picks the editable region: covers the cursor line and every line that must change (both in `file`
 * coordinates), plus a margin; null if that would exceed maxLines.
 */
export function chooseRegion(file: string, cursorLine: number, mustCover: [number, number][], opts: { margin?: number; maxLines?: number } = {}): [number, number] | null {
  const margin = opts.margin ?? 3
  const maxLines = opts.maxLines ?? 40
  const nLines = file.split('\n').length
  let start = cursorLine, end = cursorLine + 1
  for (const [a, b] of mustCover) { start = Math.min(start, a); end = Math.max(end, b) }
  if (end - start > maxLines) return null
  start = Math.max(0, start - margin)
  end = Math.min(nLines, end + margin)
  return [start, end]
}

export function nextEditTokenIds(tokenToId: (t: string) => number | undefined): NextEditTokenIds {
  const id = (t: string) => {
    const v = tokenToId(t)
    if (v === undefined) throw new Error(`tokenizer is missing ${t}`)
    return v
  }
  return {
    regionStart: id(NEXT_EDIT_TOKENS.regionStart), regionEnd: id(NEXT_EDIT_TOKENS.regionEnd), cursor: id(NEXT_EDIT_TOKENS.cursor),
    outputStart: id(NEXT_EDIT_TOKENS.outputStart), outputEnd: id(NEXT_EDIT_TOKENS.outputEnd), noEdit: id(NEXT_EDIT_TOKENS.noEdit),
    repoName: id('<|repo_name|>'), fileSep: id('<|file_sep|>'),
  }
}

/** Prompt tokens up to and including OUTPUT_START. Context before/after the region is truncated at line boundaries. */
export function buildNextEditPrompt(input: NextEditInput, encode: Encode, ids: NextEditTokenIds, opts: { beforeTokens?: number; afterTokens?: number } = {}): number[] {
  const beforeTokens = opts.beforeTokens ?? 1024
  const afterTokens = opts.afterTokens ?? 384
  const offs = lineOffsets(input.file)
  const rStart = offs[input.regionStartLine]
  const rEnd = input.regionEndLine < offs.length ? offs[input.regionEndLine] : input.file.length
  if (input.cursor < rStart || input.cursor > rEnd) throw new Error('cursor outside region')

  const out: number[] = [ids.repoName, ...encode(`${input.repo}\n`)]
  const dir = input.path.includes('/') ? input.path.slice(0, input.path.lastIndexOf('/')) : '.'
  if (input.errors) out.push(ids.fileSep, ...encode(`${dir}/__errors__.txt\n${input.errors}\n`))
  if (input.recentEdits) out.push(ids.fileSep, ...encode(`__recent_edits__.diff\n${input.recentEdits}\n`))
  out.push(ids.fileSep, ...encode(`${input.path}\n`))

  // Before the region: last lines that fit; after: first lines that fit.
  // (Window to ~6 chars/token first, then drop whole lines in proportion to the excess: a few encodes, not one per line.)
  let before = input.file.slice(0, rStart)
  if (before.length > beforeTokens * 6) before = before.slice(before.indexOf('\n', before.length - beforeTokens * 6) + 1)
  let beforeIds = encode(before)
  while (beforeIds.length > beforeTokens && before.includes('\n')) {
    const excess = beforeIds.length - beforeTokens  // undershoot (×0.8), then single lines: same cut as line-by-line
    const drop = excess > 32 ? Math.floor(0.8 * before.length * excess / beforeIds.length) : 1
    const nl = before.indexOf('\n', drop - 1)
    before = nl < 0 ? '' : before.slice(nl + 1)
    beforeIds = encode(before)
  }
  let after = input.file.slice(rEnd)
  let afterIds = encode(after.slice(0, afterTokens * 8))
  if (afterIds.length > afterTokens) {
    after = after.slice(0, afterTokens * 8)
    while (afterIds.length > afterTokens && after.includes('\n')) {
      after = after.slice(0, after.lastIndexOf('\n', after.length - 2) + 1)
      afterIds = encode(after)
    }
  }
  out.push(...beforeIds, ids.regionStart,
    ...encode(input.file.slice(rStart, input.cursor)), ids.cursor, ...encode(input.file.slice(input.cursor, rEnd)),
    ids.regionEnd, ...afterIds, ids.outputStart)
  return out
}

/** Target tokens: rewritten region (or NO_EDIT), then OUTPUT_END. The caller appends <|endoftext|>. */
export function nextEditTarget(newRegion: string | null, encode: Encode, ids: NextEditTokenIds): number[] {
  return newRegion === null ? [ids.noEdit, ids.outputEnd] : [...encode(newRegion), ids.outputEnd]
}

/**
 * Model output → rewritten region, or null for "no edit". llama-server leaves special tokens out of `content`, so the
 * NO_EDIT token usually arrives as an empty string; never treat empty output as "replace the region with nothing".
 */
export function parseNextEditOutput(raw: string, oldText: string): string | null {
  if (raw.startsWith(NEXT_EDIT_TOKENS.noEdit) || raw.trim() === '' || raw === oldText) return null
  return raw
}