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>path\n
 */
export function repoHeader(repo: string, filePath: string, context: string | undefined, encode: Encode, ids: HeaderTokenIds): number[] {
  const out = [ids.repoName, ...encode(`${repo}\n`)]
  if (context) {
    const dir = filePath.includes('/') ? filePath.slice(0, filePath.lastIndexOf('/')) : '.'
    out.push(ids.fileSep, ...encode(`${dir}/__symbols__.d.ts\n${context}\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)
}