init

lib/rerank.js

@@ -0,0 +1,124 @@
+/**
+ * Re-ranking candidates on the local model.
+ *
+ * Cosine similarity ranks a chunk against a query in isolation. A small model
+ * reading the query *and* the candidate together does better — it can see that a
+ * chunk is about the right topic but the wrong side of it, which is exactly the
+ * mistake embeddings make.
+ *
+ * ## `undefined` is not zero
+ *
+ * The re-ranker answers with a score per candidate, or with **`undefined`** when
+ * it has no opinion — slow, wedged, unparseable, or scoring a different number
+ * of candidates than it was given. The caller must treat those differently:
+ * replacing a good embedding score with a zero because the model timed out would
+ * turn a working search into an empty one. No opinion means keep the ranking you
+ * had.
+ *
+ * This is why the module returns a bare array rather than an object with a
+ * status: there are exactly two outcomes, and one of them is "carry on".
+ */
+
+/** Attempts. Re-ranking is an improvement, not a requirement. */
+const MAX_ATTEMPTS = 2
+
+/** Characters sent per candidate. The head of a chunk identifies it. */
+const MAX_CANDIDATE_CHARS = 700
+
+/** Candidates scored in one call. Longer prompts get slower and vaguer. */
+const MAX_CANDIDATES = 12
+
+const PROMPT = `You score how well each code excerpt answers a developer's search query.
+
+Give each excerpt a whole number from 0 to 10:
+- 10: this is the code the query is asking about.
+- 5: same area, but not what was asked.
+- 0: unrelated.
+
+Answer with JSON only, no prose: {"scores": [n, n, ...]} with exactly one score per excerpt, in order.
+
+Query: {{QUERY}}
+
+{{CANDIDATES}}`
+
+/** Pull the first JSON array of numbers out of a model response. */
+export function parseScores(body) {
+  const text = String(body ?? '')
+  // Prefer the documented shape, then fall back to any bare array: models drift
+  // between the two and both are unambiguous.
+  const wrapped = /\{[^{}]*"scores"\s*:\s*\[([^\]]*)\]/i.exec(text)
+  const bare = wrapped === null ? /\[([^\]]*)\]/.exec(text) : null
+  const inner = wrapped?.[1] ?? bare?.[1]
+  if (inner === undefined) return undefined
+  const numbers = inner
+    .split(',')
+    .map((part) => Number.parseFloat(part.trim()))
+    .filter((value) => Number.isFinite(value))
+  return numbers.length === 0 ? undefined : numbers
+}
+
+/**
+ * Build a re-ranker 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 call
+ * @param {typeof fetch} [options.fetchImpl] injection point for tests
+ */
+export function createReranker(options) {
+  const { apiBase, model, timeoutMs = 15_000, fetchImpl = fetch } = options
+  const base = String(apiBase).replace(/\/+$/, '')
+
+  return {
+    model,
+    /**
+     * Score candidates against a query.
+     *
+     * @param {string} query the search query
+     * @param {readonly string[]} candidates excerpt texts, in rank order
+     * @param {AbortSignal} [signal] caller cancellation
+     * @returns {Promise<number[] | undefined>} one score in 0..1 per candidate, in
+     *   order, or `undefined` when the model offered no usable opinion
+     */
+    async score(query, candidates, signal) {
+      const texts = candidates.slice(0, MAX_CANDIDATES)
+      if (texts.length === 0 || String(query ?? '').trim() === '') return undefined
+      const rendered = texts
+        .map((text, at) => `Excerpt ${at + 1}:\n${String(text).slice(0, MAX_CANDIDATE_CHARS)}`)
+        .join('\n\n')
+      const prompt = PROMPT.replace('{{QUERY}}', String(query).trim()).replace('{{CANDIDATES}}', rendered)
+
+      for (let attempt = 0; attempt < MAX_ATTEMPTS; attempt += 1) {
+        if (signal?.aborted) return undefined
+        try {
+          // See the expander: a non-finite timeout must not become a silent
+          // "no opinion" that looks exactly like a slow model.
+          const deadline = AbortSignal.timeout(Number.isFinite(timeoutMs) && timeoutMs > 0 ? timeoutMs : 15_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,
+              stream: false,
+              options: { num_predict: 80, temperature: 0 },
+            }),
+            signal: combined,
+          })
+          if (!response.ok) throw new Error(`${base} responded ${response.status}`)
+          const payload = await response.json()
+          const scores = parseScores(payload?.response)
+          // A count mismatch is not a partial answer — it is a model that lost
+          // track of the list, and applying it would score the wrong candidates.
+          if (scores === undefined || scores.length !== texts.length) throw new Error('unusable score count')
+          return scores.map((value) => Math.max(0, Math.min(10, value)) / 10)
+        } catch {
+          if (signal?.aborted) return undefined
+        }
+      }
+      return undefined
+    },
+  }
+}