125 lines
5.1 KiB
JavaScript
125 lines
5.1 KiB
JavaScript
/**
|
|
* 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
|
|
},
|
|
}
|
|
}
|