autocomplete-model/shared/prompt.ts

/**
 * Prompt construction and output post-processing shared by training export,
 * evaluation and the completion daemon, so all three build identical prompts.
 *
 * No dependencies: callers pass in their tokenizer's encode function and the
 * special token ids.
 */

export type FimOrder = 'psm' | 'spm'
export type Encode = (text: string) => number[]

export interface FimTokenIds { prefix: number; suffix: number; middle: number }
export interface HeaderTokenIds { repoName: number; fileSep: number }

/**
 * Prompt token layout (everything before the middle).
 *
 * psm: <pre> prefix <suf> suffix <mid>           Qwen2.5-Coder's native format
 * spm: <pre> <suf> suffix <mid> prefix           Bavarian et al. 2022 SPM variant
 *
 * In spm the prefix runs straight into the middle and ends the prompt, so as
 * the user types, the prompt only grows at the end and the KV cache stays valid.
 */
export function fimPrompt(order: FimOrder, prefix: number[], suffix: number[], ids: FimTokenIds): number[] {
  return order === 'spm'
    ? [ids.prefix, ids.suffix, ...suffix, ids.middle, ...prefix]
    : [ids.prefix, ...prefix, ids.suffix, ...suffix, ids.middle]
}

/**
 * Qwen2.5-Coder repo-level header placed before the FIM prompt:
 *   <repo_name>repo\n [<file_sep>dir/__symbols__.d.ts\n{context}\n] [<file_sep>dir/__errors__.txt\n{errors}\n] <file_sep>path\n
 */
export function repoHeader(repo: string, filePath: string, context: string | undefined, encode: Encode, ids: HeaderTokenIds, errors?: string): number[] {
  const out = [ids.repoName, ...encode(`${repo}\n`)]
  const dir = filePath.includes('/') ? filePath.slice(0, filePath.lastIndexOf('/')) : '.'
  if (context) out.push(ids.fileSep, ...encode(`${dir}/__symbols__.d.ts\n${context}\n`))
  // TS errors near the cursor (shared/diagnostics.ts), for models trained with them (Stage D+).
  if (errors) out.push(ids.fileSep, ...encode(`${dir}/__errors__.txt\n${errors}\n`))
  out.push(ids.fileSep, ...encode(`${filePath}\n`))
  return out
}

/** Last ≤ maxTokens tokens of text, tokenising only a window that starts at a line boundary (training-time truncation). */
export function tailTokens(text: string, maxTokens: number, encode: Encode): number[] {
  let window = text.slice(-maxTokens * 8)
  if (window.length < text.length) window = window.slice(window.indexOf('\n') + 1)
  return encode(window).slice(-maxTokens)
}

/** First ≤ maxTokens tokens of text, window cut at a line boundary. Deterministic, so stable while the suffix is unchanged. */
export function headTokens(text: string, maxTokens: number, encode: Encode): number[] {
  let window = text.slice(0, maxTokens * 8)
  if (window.length < text.length) window = window.slice(0, window.lastIndexOf('\n') + 1)
  return encode(window).slice(0, maxTokens)
}

/**
 * Cache-friendly prefix truncation for live completion.
 *
 * Instead of "the last N tokens" (whose start moves on every keystroke and
 * invalidates the whole KV cache), the prefix starts at a line whose index is
 * a multiple of `chunkLines`, chosen as the earliest such line that fits. The
 * start therefore only moves when the prefix outgrows the budget, about once
 * every `chunkLines` lines of typing.
 */
export function chunkedPrefixTokens(prefix: string, maxTokens: number, encode: Encode, chunkLines = 32, lineOffset = 0): { tokens: number[]; startLine: number } {
  // lineOffset: absolute line number of prefix's first line, when the client sent a truncated document,
  // so chunk boundaries stay on absolute multiples of chunkLines.
  const lineStarts = [0]
  for (let i = prefix.indexOf('\n'); i !== -1; i = prefix.indexOf('\n', i + 1)) lineStarts.push(i + 1)
  const firstChunkLine = (chunkLines - (lineOffset % chunkLines)) % chunkLines
  const candidates: number[] = []
  for (let l = firstChunkLine; l < lineStarts.length; l += chunkLines) candidates.push(l)
  // Estimate first (≈3.2 chars/token for code) to skip hopeless candidates, then verify with the real tokenizer.
  let c = 0
  while (c < candidates.length - 1 && (prefix.length - lineStarts[candidates[c]]) / 3.2 > maxTokens * 0.95) c++
  for (; c < candidates.length; c++) {
    const tokens = encode(prefix.slice(lineStarts[candidates[c]]))
    if (tokens.length <= maxTokens) return { tokens, startLine: lineOffset + candidates[c] }
  }
  // Even the last chunk is over budget (e.g. minified code): fall back to a plain tail.
  return { tokens: tailTokens(prefix, maxTokens, encode), startLine: lineOffset + lineStarts.length - 1 }
}

/**
 * Cleans a raw completion the way an editor client would: drop what the model
 * re-generates of the existing suffix. Never looks at the gold middle.
 *
 * Lines are compared with their indentation: a method's closing `  }` must not
 * be mistaken for the class's `}` that follows in the suffix (that bug removed
 * legitimately predicted closing braces).
 */
export function trimSuffixOverlap(raw: string, suffix: string): string {
  let out = raw
  const sufLine = suffix.split('\n')[0]
  const sufFirstNonBlank = suffix.split('\n').find(l => l.trim())?.trimEnd()
  const lines = out.split('\n')
  if (sufFirstNonBlank) {
    const i = lines.findIndex((l, idx) => idx > 0 && l.trimEnd() === sufFirstNonBlank)
    if (i > 0) out = lines.slice(0, i).join('\n').replace(/\n[ \t]*$/, '')
  }
  if (!out.includes('\n') && sufLine.trim() && out.trimEnd().endsWith(sufLine.trimEnd())) {
    out = out.trimEnd().slice(0, -sufLine.trimEnd().length)
  }
  return out
}

/**
 * Single-line mode when there is real code after the cursor on the same line
 * (a completion there must not add new lines). Closing brackets, quotes and
 * punctuation alone don't count: `foo(▮)` may still complete multi-line args.
 */
export function wantsSingleLine(suffix: string): boolean {
  const restOfLine = suffix.split('\n')[0]
  return /[^\s)\]}'"`;,]/.test(restOfLine)
}

/**
 * Token healing for completions that start mid-identifier (`this.connections.del▮`).
 *
 * Suffix-first prompts end exactly at the cursor, so a half-typed word leaves an unnatural final token and the
 * model tends to treat the word as finished (`(relayUrl)` instead of `ete(relayUrl)`). Instead, cut the partial
 * word off the prompt and force the output to start with it (llama.cpp grammar); the caller strips it again.
 * Bonus: while typing `d`, `de`, `del` the healed prompt is identical, so the KV cache is fully reused.
 */
export function healPrefix(prefix: string): { prefix: string; forced: string } | null {
  const m = /(^|[^\w$])([\w$]{1,32})$/.exec(prefix)
  if (!m) return null
  return { prefix: prefix.slice(0, prefix.length - m[2].length), forced: m[2] }
}

/** GBNF grammar: output must start with `forced`, then anything. */
export function forcedPrefixGrammar(forced: string): string {
  return `root ::= "${forced.replace(/\\/g, '\\\\').replace(/"/g, '\\"')}" rest\nrest ::= [^\\x00]*`
}