init

lib/rank.js

@@ -0,0 +1,337 @@
+/**
+ * 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 }
+}