peardata/ui/qvac/context.js

/**
 * Context window budgeting for small local QVAC models.
 * Rough token estimate + history compaction before completion.
 */

/**
 * Conservative token estimate (chars / 3.2) for English + JSON tool payloads.
 * @param {unknown} text
 */
export function estimateTokens(text) {
  const s = typeof text === 'string' ? text : JSON.stringify(text ?? '')
  return Math.max(1, Math.ceil(s.length / 3.2))
}

/**
 * @param {Array<{ role?: string, content?: string, name?: string }>} messages
 */
export function estimateMessagesTokens(messages) {
  let n = 0
  for (const m of messages || []) {
    n += 4 // role framing
    n += estimateTokens(m.content || '')
    if (m.name) n += estimateTokens(m.name)
  }
  return n
}

/**
 * @param {any[]} tools
 */
export function estimateToolsTokens(tools) {
  if (!tools?.length) return 0
  return estimateTokens(JSON.stringify(tools)) + 32
}

/**
 * Compact chat history to fit a token budget.
 * Always keeps: first system message(s), last user message, recent turns.
 * Truncates large tool payloads; drops oldest middle messages.
 *
 * @param {Array<{ role: string, content: string, name?: string }>} messages
 * @param {{
 *   maxTokens?: number,
 *   keepRecentUserTurns?: number,
 *   maxToolChars?: number,
 *   maxMsgChars?: number,
 * }} [opts]
 */
export function compactMessages(messages, opts = {}) {
  const maxTokens = Math.max(512, Number(opts.maxTokens) || 2800)
  const keepRecentUserTurns = Math.max(1, Number(opts.keepRecentUserTurns) || 3)
  const maxToolChars = Math.max(400, Number(opts.maxToolChars) || 2500)
  const maxMsgChars = Math.max(800, Number(opts.maxMsgChars) || 4000)

  const src = (messages || []).map((m) => ({
    role: m.role,
    content: String(m.content ?? ''),
    ...(m.name ? { name: m.name } : {}),
  }))

  if (!src.length) return src

  // Split leading system messages
  /** @type {typeof src} */
  const systems = []
  let i = 0
  while (i < src.length && src[i].role === 'system') {
    systems.push(truncateMsg(src[i], maxMsgChars * 2))
    i++
  }
  const rest = src.slice(i)

  // Find last user message index in rest
  let lastUser = -1
  for (let j = rest.length - 1; j >= 0; j--) {
    if (rest[j].role === 'user') {
      lastUser = j
      break
    }
  }

  // Truncate tool / long assistant bodies first
  const trimmed = rest.map((m) => {
    if (m.role === 'tool') return truncateMsg(m, maxToolChars)
    if (m.role === 'assistant') return truncateMsg(m, maxMsgChars)
    return truncateMsg(m, maxMsgChars)
  })

  // Keep recent window ending at last message, including last N user turns
  let start = 0
  if (lastUser >= 0) {
    let users = 0
    start = lastUser
    for (let j = lastUser; j >= 0; j--) {
      if (trimmed[j].role === 'user') {
        users++
        start = j
        if (users >= keepRecentUserTurns) break
      }
    }
  }
  let window = trimmed.slice(start)

  // Drop from front until under budget (keep systems + window)
  const pack = () => [...systems, ...window]
  while (window.length > 2 && estimateMessagesTokens(pack()) > maxTokens) {
    // Prefer dropping oldest non-user if possible
    if (window[0]?.role !== 'user' || window.length > 4) {
      window = window.slice(1)
    } else {
      window = window.slice(1)
    }
  }

  // Still over budget: hard-trim contents
  if (estimateMessagesTokens(pack()) > maxTokens) {
    window = window.map((m) =>
      truncateMsg(m, m.role === 'tool' ? 600 : m.role === 'system' ? 2000 : 1200)
    )
  }

  // Still over: keep only systems + last user + trailing assistant/tool chain
  if (estimateMessagesTokens(pack()) > maxTokens) {
    const last = window[window.length - 1]
    const lastU = [...window].reverse().find((m) => m.role === 'user')
    window = [lastU, last].filter(Boolean)
    // dedupe if same ref
    if (window.length === 2 && window[0] === window[1]) window = [window[0]]
  }

  // Add a short note if we dropped history
  const dropped = rest.length - window.length
  if (dropped > 0 && systems[0]) {
    systems[0] = {
      ...systems[0],
      content:
        systems[0].content +
        `\n\n[Context compacted: ${dropped} earlier messages omitted to fit the model window.]`,
    }
  }

  return pack()
}

/**
 * @param {{ role: string, content: string, name?: string }} m
 * @param {number} maxChars
 */
function truncateMsg(m, maxChars) {
  const c = String(m.content || '')
  if (c.length <= maxChars) return m
  return {
    ...m,
    content: c.slice(0, maxChars - 20) + '\n…[truncated]',
  }
}

/**
 * Budget for prompt given model context size.
 * Reserves space for generation + tool schemas.
 * @param {number} ctxSize
 * @param {number} toolsTokens
 */
export function promptBudget(ctxSize, toolsTokens = 0) {
  const ctx = Math.max(2048, Number(ctxSize) || 4096)
  // Leave room for model output + thinking
  const reserveOut = Math.min(1024, Math.floor(ctx * 0.25))
  const reserveTools = Math.min(toolsTokens, Math.floor(ctx * 0.2))
  const budget = ctx - reserveOut - reserveTools - 64
  return Math.max(800, budget)
}

/**
 * Detect context-overflow style errors from QVAC / llama.cpp.
 * @param {string} msg
 */
export function isContextOverflowError(msg) {
  const s = String(msg || '').toLowerCase()
  return (
    s.includes('context window') ||
    s.includes('context length') ||
    s.includes('exceeds the model') ||
    s.includes('prompt is too long') ||
    s.includes('n_keep') ||
    s.includes('too many tokens')
  )
}