dsh-recall/lib/rank.js

/**
 * Ranking retrieved chunks.
 *
 * **Pure vector search is not good enough for code, and this is the whole reason
 * the module exists.** An embedding of `getUserById` sits close to every other
 * lookup function in the repository, so a query naming it returns ten plausible
 * lookups and buries the actual definition. A keyword match has the opposite
 * failure: it finds the definition instantly and misses the query that describes
 * behaviour without naming it.
 *
 * The two are combined by **reciprocal rank fusion** rather than by adding
 * normalised scores. Cosine similarity and a term-frequency score live on
 * different scales, and the usual fix — min-max normalising each into [0,1] —
 * is actively harmful here: when every candidate is unrelated, the similarities
 * are tightly clustered, and stretching that cluster across the full range turns
 * noise into a confident-looking ranking. Rank fusion is scale-free, so it
 * cannot make that mistake, and a chunk found by only one of the two methods
 * still surfaces.
 */

/** RRF damping. 60 is the value from the original TREC work; larger flattens. */
const RRF_K = 60

/** Terms shorter than this carry no signal in code and match everywhere. */
const MIN_TERM_LENGTH = 2

/** Words that appear in almost every source file and therefore rank nothing. */
const STOPWORDS = new Set([
  'the', 'and', 'for', 'with', 'that', 'this', 'from', 'into', 'when', 'then',
  'const', 'let', 'var', 'function', 'return', 'import', 'export', 'class',
  'def', 'self', 'new', 'use', 'are', 'was', 'were', 'not', 'but', 'you', 'how',
  'does', 'what', 'where', 'which', 'who', 'why', 'can', 'should', 'would',
])

/** Lowercased haystacks, keyed by the matrix object they describe. */
const haystacks = new WeakMap()

/**
 * Split an identifier into its words.
 *
 * `getUserById` must match the query terms `get`, `user`, `by` and `id`, because
 * a caller writing a search phrase naturally writes them separately. Without
 * this, the single most useful lexical signal in a codebase — a name spelled in
 * camelCase — is invisible.
 *
 * @param {string} identifier
 * @returns {string[]} lowercased words
 */
export function splitIdentifier(identifier) {
  return identifier
    .replace(/([a-z0-9])([A-Z])/g, '$1 $2')
    .replace(/([A-Z]+)([A-Z][a-z])/g, '$1 $2')
    .split(/[^A-Za-z0-9]+/)
    .filter((word) => word !== '')
}

/**
 * The search terms of a query.
 *
 * Both the raw tokens and the sub-words of each identifier are kept, so a query
 * of `getUserById` matches a chunk that spells it `get_user_by_id` as well as
 * one that spells it identically.
 *
 * @param {string} query
 * @returns {string[]} de-duplicated terms, longest first
 */
export function queryTerms(query) {
  const found = new Set()
  for (const raw of String(query).split(/[^A-Za-z0-9_$]+/)) {
    if (raw === '') continue
    const lower = raw.toLowerCase()
    if (lower.length >= MIN_TERM_LENGTH && !STOPWORDS.has(lower)) found.add(lower)
    if (raw.length >= MIN_TERM_LENGTH) {
      for (const word of splitIdentifier(raw)) {
        const part = word.toLowerCase()
        if (part.length >= MIN_TERM_LENGTH && !STOPWORDS.has(part)) found.add(part)
      }
    }
  }
  return [...found].sort((left, right) => right.length - left.length)
}

/** The lowercased text a chunk is matched against, built once per matrix. */
function haystackFor(matrix, row) {
  let list = haystacks.get(matrix)
  if (list === undefined) {
    list = new Array(matrix.rows.length)
    haystacks.set(matrix, list)
  }
  if (list[row.id] === undefined) {
    const symbol = row.symbol === undefined ? '' : ` ${row.symbol}`
    list[row.id] = `${row.path}${symbol} ${row.text}`.toLowerCase()
  }
  return list[row.id]
}

/**
 * Score every chunk against the query terms.
 *
 * Matches in the path or the symbol are weighted above matches in the body,
 * because a question about behaviour usually names the thing, and the name is
 * what the path and symbol carry.
 *
 * @param {object} matrix
 * @param {string[]} terms
 * @returns {{lexical: Float64Array, matched: Uint8Array}} raw scores and a hit flag
 */
export function lexicalScores(matrix, terms) {
  const count = matrix.rows.length
  const lexical = new Float64Array(count)
  const matched = new Uint8Array(count)
  if (terms.length === 0) return { lexical, matched }

  // Document frequency per term, for the inverse-document-frequency weight.
  const frequency = new Map()
  for (const term of terms) {
    let seen = 0
    for (let at = 0; at < count; at += 1) {
      if (haystackFor(matrix, matrix.rows[at]).includes(term)) seen += 1
    }
    frequency.set(term, seen)
  }

  for (let at = 0; at < count; at += 1) {
    const row = matrix.rows[at]
    const hay = haystackFor(matrix, row)
    const header = `${row.path} ${row.symbol ?? ''}`.toLowerCase()
    let score = 0
    for (const term of terms) {
      const df = frequency.get(term) ?? 0
      if (df === 0) continue
      // Standard IDF, damped so a term present in every file still contributes.
      const idf = Math.log(1 + count / (1 + df))
      let occurrences = 0
      let from = hay.indexOf(term)
      while (from !== -1) {
        occurrences += 1
        from = hay.indexOf(term, from + term.length)
      }
      if (occurrences === 0) continue
      const weight = header.includes(term) ? 3 : 1
      score += idf * (1 + Math.log(occurrences)) * weight
    }
    if (score > 0) {
      lexical[at] = score
      matched[at] = 1
    }
  }
  return { lexical, matched }
}

/** Indices ordered by score, descending, ties broken by index for determinism. */
function orderBy(score) {
  const indices = Array.from({ length: score.length }, (_, at) => at)
  indices.sort((left, right) => score[right] - score[left] || left - right)
  return indices
}

/**
 * The excerpt a result shows.
 *
 * Head-truncated rather than tail-truncated: a function's signature, its
 * parameters and its first statements are what identify it, while the tail of a
 * long function is usually a closing brace and a return.
 *
 * @param {string} text chunk text
 * @param {number} maxLines cap
 * @returns {{excerpt: string, truncated: boolean}}
 */
export function excerptOf(text, maxLines) {
  const lines = text.split('\n')
  if (lines.length <= maxLines) return { excerpt: text, truncated: false }
  return { excerpt: lines.slice(0, maxLines).join('\n'), truncated: true }
}

/**
 * The fused order of every chunk for one query, before selection.
 *
 * Split out from ranking so several queries can be ranked against the same rows
 * and merged: query expansion produces a handful of phrasings, each gets an
 * order, and the orders are fused. Selecting results is a separate step because
 * it is the only part that cares about limits and budgets.
 *
 * @param {object} input
 * @param {object} input.matrix vector matrix and rows
 * @param {Float32Array} input.queryVector unit query vector
 * @param {string} input.query raw query text
 * @param {number} [input.vectorWeight] fusion weight for the vector ranking
 * @param {number} [input.lexicalWeight] fusion weight for the keyword ranking
 * @returns {{order: number[], similarity: Float64Array, vectorBest: number, lexicalBest: number}}
 */
export function rankOrder(input) {
  const { matrix, queryVector, query, vectorWeight = 0.7, lexicalWeight = 0.3 } = input
  const count = matrix.rows.length
  if (count === 0) {
    return { order: [], similarity: new Float64Array(0), vectorBest: 0, lexicalBest: 0 }
  }

  const dimensions = matrix.dimensions
  const similarity = new Float64Array(count)
  for (let at = 0; at < count; at += 1) {
    let sum = 0
    const offset = at * dimensions
    for (let dim = 0; dim < dimensions; dim += 1) sum += matrix.vectors[offset + dim] * queryVector[dim]
    similarity[at] = sum
  }

  const terms = queryTerms(query)
  const { lexical } = lexicalScores(matrix, terms)

  const vectorOrder = orderBy(similarity)
  const lexicalOrder = orderBy(lexical)
  const vectorRank = new Int32Array(count).fill(count)
  const lexicalRank = new Int32Array(count).fill(count)
  for (let rank = 0; rank < count; rank += 1) vectorRank[vectorOrder[rank]] = rank
  for (let rank = 0; rank < count; rank += 1) lexicalRank[lexicalOrder[rank]] = rank

  const fused = new Float64Array(count)
  for (let at = 0; at < count; at += 1) {
    fused[at] = vectorWeight / (RRF_K + vectorRank[at]) + lexicalWeight / (RRF_K + lexicalRank[at])
  }

  return {
    order: orderBy(fused),
    similarity,
    vectorBest: Number(similarity[vectorOrder[0]]?.toFixed(4) ?? 0),
    lexicalBest: Number(lexical[lexicalOrder[0]]?.toFixed(4) ?? 0),
  }
}

/**
 * Merge several rankings of the same rows into one order, by rank.
 *
 * Used for query expansion: each alternative phrasing ranks every chunk, and a
 * chunk that several phrasings agree on should outrank one that a single
 * phrasing happened to like. Rank fusion again rather than adding scores,
 * because the rankings come from different query vectors and their similarities
 * are not on a comparable scale.
 *
 * A chunk missing from a ranking is simply absent from that ranking's
 * contribution rather than penalised, so a phrasing that finds nothing does not
 * drag anything down.
 *
 * @param {readonly (readonly number[])[]} rankings row indices, each in rank order
 * @returns {number[]} row indices, best first
 */
export function fuseRankings(rankings) {
  const present = rankings.filter((ranking) => Array.isArray(ranking) && ranking.length > 0)
  if (present.length === 0) return []
  if (present.length === 1) return [...present[0]]

  const scores = new Map()
  for (const ranking of present) {
    for (let rank = 0; rank < ranking.length; rank += 1) {
      const row = ranking[rank]
      scores.set(row, (scores.get(row) ?? 0) + 1 / (RRF_K + rank))
    }
  }
  // Ties break by row index, so the same input always gives the same output.
  return [...scores.keys()].sort((left, right) => scores.get(right) - scores.get(left) || left - right)
}

/**
 * Turn a fused order into the results a caller sees.
 *
 * @param {number[]} order row indices, best first
 * @param {object} input
 * @param {object} input.matrix
 * @param {Float64Array} input.similarity per-row similarity, for reporting
 * @param {number} [input.limit]
 * @param {number} [input.perFile]
 * @param {number} [input.excerptLines]
 * @param {number} [input.totalLines]
 * @returns {object[]} results
 */
export function selectFrom(order, input) {
  const { matrix, similarity, limit = 8, perFile = 2, excerptLines = 80, totalLines = 400 } = input
  const perFileUsed = new Map()
  const results = []
  // A shared line budget across the whole result set. The per-excerpt cap alone
  // does not bound the response — `limit` does — and a chunk is already a
  // semantic unit, so cutting one in half is how a search turns into a
  // follow-up read that costs more than the lines it saved. Allocation is by
  // rank, so the best answer gets the room.
  let remaining = totalLines
  for (const at of order) {
    if (results.length >= limit) break
    // Always return at least one result, however small the budget.
    if (remaining <= 0 && results.length > 0) break
    const row = matrix.rows[at]
    if (row === undefined) continue
    const used = perFileUsed.get(row.path) ?? 0
    if (used >= perFile) continue
    perFileUsed.set(row.path, used + 1)
    const allowance = Math.max(1, Math.min(excerptLines, remaining))
    const { excerpt, truncated } = excerptOf(row.text, allowance)
    remaining -= excerpt.split('\n').length
    results.push({
      path: row.path,
      startLine: row.startLine,
      endLine: row.endLine,
      ...(row.symbol === undefined ? {} : { symbol: row.symbol }),
      similarity: Number(similarity[at].toFixed(4)),
      excerpt,
      truncated,
    })
  }
  return results
}

/**
 * Rank chunks for a query.
 *
 * @param {object} input
 * @param {object} input.matrix vector matrix and rows
 * @param {Float32Array} input.queryVector unit query vector
 * @param {string} input.query raw query text
 * @param {number} [input.limit] results to return
 * @param {number} [input.perFile] most chunks one file may contribute
 * @param {number} [input.excerptLines] lines per excerpt
 * @param {number} [input.totalLines] total line budget
 * @param {number} [input.vectorWeight] fusion weight for the vector ranking
 * @param {number} [input.lexicalWeight] fusion weight for the keyword ranking
 * @returns {{results: object[], vectorBest: number, lexicalBest: number}}
 */
export function rankChunks(input) {
  const ranked = rankOrder(input)
  const results = selectFrom(ranked.order, {
    matrix: input.matrix,
    similarity: ranked.similarity,
    limit: input.limit,
    perFile: input.perFile,
    excerptLines: input.excerptLines,
    totalLines: input.totalLines,
  })
  return { results, vectorBest: ranked.vectorBest, lexicalBest: ranked.lexicalBest }
}