dsh-recall/lib/expand.js

/**
 * Query expansion, on the local model.
 *
 * A developer's first query is often a description rather than a name — "how
 * does a room decide it is occupied" — and that phrasing shares few tokens with
 * the code that answers it. Asking a small local model for a handful of
 * alternative phrasings, embedding each, and fusing the rankings costs one local
 * generation and a few local embeddings, and is free in money.
 *
 * ## Every failure is a fallback, never an error
 *
 * This runs on the critical path of a search, against a model on somebody else's
 * box. It can be slow, wedged, or absent. None of those may fail the search, so
 * the contract is narrow: **the returned list always contains the original
 * query**, and anything that goes wrong simply returns it alone. The caller can
 * treat expansion as a bonus that is either there or not.
 *
 * The internal timeout is deliberately separate from the caller's signal: the
 * caller cancelling means stop everything, while the model being slow just means
 * search with what we have.
 */

/** Attempts per request. Expansion is a bonus; one more try is enough. */
const MAX_ATTEMPTS = 2

/** Characters of query allowed through. Longer is a paste, not a question. */
const MAX_QUERY_CHARS = 500

/** A variant this long is a paragraph, not a search phrase. */
const MAX_VARIANT_CHARS = 200

/** Variants shorter than this match everything and rank nothing. */
const MIN_VARIANT_CHARS = 3

/**
 * Openers that mark a line as the model talking rather than a phrasing.
 *
 * A small model asked for a list sometimes answers with a sentence about the
 * request — "I cannot help with that." A line like that is not a weak variant,
 * it is not a variant at all: embedding it would search the codebase for the
 * model's apology and rank whatever matched.
 *
 * `i\b` and the refusal check below are the load-bearing ones. A question-shaped
 * variant ("how does a room decide it is occupied") is entirely legitimate, so
 * the test is first-person and refusal vocabulary, not sentence punctuation.
 */
const COMMENTARY = /^(?:i\b|i'm\b|i am\b|sorry\b|unfortunately\b|note\b|explanation\b|here\b|sure\b|of course\b|certainly\b|these\b|those\b|the following\b|as an?\b|phrasings?\b|query\b)/i

/** Refusal vocabulary, wherever it sits in the line. */
const REFUSAL = /\b(?:cannot help|can't help|don't have|do not have|not able to|unable to|no way to)\b/i

const PROMPT = `You turn a developer's search query into alternative phrasings that would match source code.

Rules:
- One phrasing per line. No numbering, no bullets, no commentary, no blank lines.
- Use the vocabulary the code itself would use: identifiers, domain nouns, verb phrases.
- Vary the angle: a likely function name, a likely filename fragment, a description of the behaviour.
- Do not repeat the original query.

Query: {{QUERY}}

Phrasings:`

/**
 * Reduce one model line to a usable phrase, or nothing.
 *
 * Models are reliably inclined to number a list, wrap it in quotes, or add a
 * closing remark. None of that is worth an error path, so it is all stripped.
 */
export function cleanVariant(line) {
  let text = String(line ?? '')
    .replace(/^\s*(?:[-*•]|\d+[.)])\s*/, '')
    .trim()
  // A trailing aside ("this should find the handler") is common; drop it.
  text = text.split(/\s+[—–-]{1,2}\s+/)[0].trim()
  if (text.length >= 2 && ((text.startsWith('"') && text.endsWith('"')) || (text.startsWith("'") && text.endsWith("'")))) {
    text = text.slice(1, -1).trim()
  }
  if (text.length < MIN_VARIANT_CHARS || text.length > MAX_VARIANT_CHARS) return undefined
  if (COMMENTARY.test(text) || REFUSAL.test(text)) return undefined
  return text
}

/**
 * Parse a model response into alternative queries, best effort.
 *
 * @param {string} body raw model output
 * @param {string} query the original, excluded from the results
 * @param {number} count most variants to keep
 * @returns {string[]} deduplicated variants, in the order the model gave them
 */
export function parseVariants(body, query, count) {
  const seen = new Set([String(query).trim().toLowerCase()])
  const out = []
  for (const line of String(body ?? '').split('\n')) {
    if (out.length >= count) break
    const variant = cleanVariant(line)
    if (variant === undefined) continue
    const key = variant.toLowerCase()
    if (seen.has(key)) continue
    seen.add(key)
    out.push(variant)
  }
  return out
}

/**
 * Build an expander bound to one local model.
 *
 * @param {object} options
 * @param {string} options.apiBase endpoint, no trailing slash
 * @param {string} options.model model id
 * @param {number} [options.timeoutMs] budget for the whole expansion
 * @param {number} [options.count] most variants to ask for
 * @param {typeof fetch} [options.fetchImpl] injection point for tests
 */
export function createExpander(options) {
  const { apiBase, model, timeoutMs = 8_000, count = 4, fetchImpl = fetch } = options
  const base = String(apiBase).replace(/\/+$/, '')

  /** One request, bounded by its own deadline. */
  async function request(query, budget, signal) {
    // A non-finite budget would make AbortSignal.timeout throw, and the caller's
    // catch would turn that into a silent fallback — a misconfiguration wearing
    // the disguise of a slow model. Clamp instead.
    const deadline = AbortSignal.timeout(Number.isFinite(budget) && budget > 0 ? budget : 8_000)
    const combined = signal === undefined ? deadline : AbortSignal.any([signal, deadline])
    const response = await fetchImpl(`${base}/api/generate`, {
      method: 'POST',
      headers: { 'content-type': 'application/json' },
      body: JSON.stringify({
        model,
        prompt: PROMPT.replace('{{QUERY}}', query),
        stream: false,
        // A handful of short lines; a model allowed to ramble will.
        options: { num_predict: 120, temperature: 0.4 },
      }),
      signal: combined,
    })
    if (!response.ok) throw new Error(`${base} responded ${response.status}`)
    const payload = await response.json()
    return typeof payload?.response === 'string' ? payload.response : ''
  }

  return {
    model,
    /**
     * Alternative phrasings for a query.
     *
     * @param {string} query the original
     * @param {AbortSignal} [signal] caller cancellation
     * @returns {Promise<string[]>} the original first, then whatever else worked
     */
    async expand(query, signal) {
      const original = String(query ?? '').trim()
      if (original === '' || original.length > MAX_QUERY_CHARS) return [original]
      for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt += 1) {
        if (signal?.aborted) return [original]
        try {
          const body = await request(original, timeoutMs, signal)
          return [original, ...parseVariants(body, original, count)]
        } catch {
          // Slow, wedged, or unreachable all mean the same thing here: search
          // with the query we were given rather than failing the call.
          if (signal?.aborted) return [original]
        }
      }
      return [original]
    },
  }
}