269 lines
9.8 KiB
JavaScript
269 lines
9.8 KiB
JavaScript
/**
|
|
* The embedding client.
|
|
*
|
|
* Speaks the Ollama HTTP API, which also covers any OpenAI-compatible gateway
|
|
* that implements `/api/embed`, and is deliberately the *only* outbound call
|
|
* this plugin makes. There is no vendor SDK and no configuration surface that
|
|
* could redirect code anywhere except the endpoint the operator named.
|
|
*
|
|
* Two properties the rest of the plugin depends on:
|
|
*
|
|
* **Every vector is unit length.** Cosine similarity is then a plain dot
|
|
* product, which is what makes ranking over tens of thousands of chunks a few
|
|
* milliseconds of arithmetic rather than a normalisation pass per comparison.
|
|
*
|
|
* **A dimension mismatch is fatal and loud.** A vector of the wrong width
|
|
* cannot be compared with the others, and silently storing one would corrupt
|
|
* the index in a way that looks like bad ranking rather than a broken config —
|
|
* the hardest possible failure to diagnose from the outside.
|
|
*/
|
|
|
|
/** Failure carrying a stable code, so the tool can explain rather than dump a stack. */
|
|
export class EmbeddingError extends Error {
|
|
/**
|
|
* @param {string} message human-readable cause
|
|
* @param {string} code stable identifier
|
|
*/
|
|
constructor(message, code) {
|
|
super(message)
|
|
this.name = 'EmbeddingError'
|
|
this.code = code
|
|
}
|
|
}
|
|
|
|
/** Characters sent per chunk. Comfortably inside an 8k-token code model's window. */
|
|
const MAX_INPUT_CHARS = 24_000
|
|
|
|
/** Attempts per request, including the first. */
|
|
const MAX_ATTEMPTS = 3
|
|
|
|
/** Base backoff between attempts. */
|
|
const RETRY_BASE_MS = 250
|
|
|
|
/** Truncate one input to the character budget. */
|
|
export function clampInput(text) {
|
|
return text.length <= MAX_INPUT_CHARS ? text : text.slice(0, MAX_INPUT_CHARS)
|
|
}
|
|
|
|
/** Scale a vector to unit length. A zero vector is returned unchanged. */
|
|
export function normalize(vector) {
|
|
let sum = 0
|
|
for (let at = 0; at < vector.length; at += 1) sum += vector[at] * vector[at]
|
|
if (sum === 0) return vector
|
|
const scale = 1 / Math.sqrt(sum)
|
|
for (let at = 0; at < vector.length; at += 1) vector[at] *= scale
|
|
return vector
|
|
}
|
|
|
|
/** Dot product, which is cosine similarity for unit vectors. */
|
|
export function dot(left, right) {
|
|
let sum = 0
|
|
const length = Math.min(left.length, right.length)
|
|
for (let at = 0; at < length; at += 1) sum += left[at] * right[at]
|
|
return sum
|
|
}
|
|
|
|
/** Sleep that rejects promptly when the caller's signal fires. */
|
|
function delay(ms, signal) {
|
|
return new Promise((resolve, reject) => {
|
|
const timer = setTimeout(() => {
|
|
signal?.removeEventListener('abort', onAbort)
|
|
resolve()
|
|
}, ms)
|
|
const onAbort = () => {
|
|
clearTimeout(timer)
|
|
reject(signal.reason ?? new Error('aborted'))
|
|
}
|
|
if (signal?.aborted) onAbort()
|
|
else signal?.addEventListener('abort', onAbort, { once: true })
|
|
})
|
|
}
|
|
|
|
/**
|
|
* Build an embedder bound to one endpoint.
|
|
*
|
|
* @param {object} options
|
|
* @param {string} options.apiBase base URL, without a trailing slash
|
|
* @param {string} options.model model id as the endpoint knows it
|
|
* @param {number} options.dimensions expected vector width
|
|
* @param {string} [options.apiKey] bearer token, when the endpoint wants one
|
|
* @param {number} [options.batchSize] inputs per request on the batch endpoint
|
|
* @param {number} [options.concurrency] requests in flight
|
|
* @param {number} [options.timeoutMs] per-request timeout
|
|
* @param {typeof fetch} [options.fetchImpl] injection point for tests
|
|
*/
|
|
export function createEmbedder(options) {
|
|
const {
|
|
apiBase,
|
|
model,
|
|
dimensions,
|
|
apiKey = '',
|
|
batchSize = 16,
|
|
concurrency = 2,
|
|
timeoutMs = 30_000,
|
|
maxAttempts = MAX_ATTEMPTS,
|
|
keepAlive = '',
|
|
fetchImpl = fetch,
|
|
} = options
|
|
|
|
const base = apiBase.replace(/\/+$/, '')
|
|
const headers = { 'content-type': 'application/json' }
|
|
if (apiKey !== '') headers.authorization = `Bearer ${apiKey}`
|
|
|
|
/**
|
|
* Add Ollama's `keep_alive` to a request body, when one is configured.
|
|
*
|
|
* Two reasons. An index run issues hundreds of requests, and without a
|
|
* keep-alive the model can be evicted between them and reloaded from disk
|
|
* repeatedly. And on a host already holding another model resident, a load
|
|
* that has to wait for memory is exactly the situation where a bounded,
|
|
* explicit residency beats Ollama's own default.
|
|
*/
|
|
const withKeepAlive = (body) => (keepAlive === '' ? body : { ...body, keep_alive: keepAlive })
|
|
|
|
// Discovered on first use: the batch endpoint is preferred, but an older
|
|
// Ollama build answers it with a 404 and then every request pays that
|
|
// discovery again. Remembering the answer keeps the fallback cheap.
|
|
let batchEndpointWorks = true
|
|
|
|
/** One HTTP request, with retries on transport faults and 5xx. */
|
|
async function request(path, body, signal) {
|
|
let lastError
|
|
for (let attempt = 0; attempt < maxAttempts; attempt += 1) {
|
|
if (signal?.aborted) throw signal.reason ?? new Error('aborted')
|
|
const timeout = AbortSignal.timeout(timeoutMs)
|
|
const combined = signal === undefined ? timeout : AbortSignal.any([signal, timeout])
|
|
try {
|
|
const response = await fetchImpl(`${base}${path}`, {
|
|
method: 'POST',
|
|
headers,
|
|
body: JSON.stringify(body),
|
|
signal: combined,
|
|
})
|
|
if (response.status === 404 || response.status === 405) {
|
|
const error = new EmbeddingError(`${path} is not available on ${base}`, 'ENDPOINT_UNSUPPORTED')
|
|
error.status = response.status
|
|
throw error
|
|
}
|
|
if (!response.ok) {
|
|
const text = await response.text().catch(() => '')
|
|
const error = new EmbeddingError(
|
|
`${path} responded ${response.status}${text === '' ? '' : `: ${text.slice(0, 300)}`}`,
|
|
'HTTP_ERROR',
|
|
)
|
|
error.status = response.status
|
|
throw error
|
|
}
|
|
return await response.json()
|
|
} catch (error) {
|
|
// A 4xx that is not a retryable status is a configuration fault, not a
|
|
// transient one, and retrying it just delays the explanation.
|
|
if (error instanceof EmbeddingError && error.status !== undefined && error.status < 500) throw error
|
|
if (signal?.aborted) throw signal.reason ?? error
|
|
lastError = error
|
|
if (attempt + 1 < MAX_ATTEMPTS) await delay(RETRY_BASE_MS * 2 ** attempt, signal)
|
|
}
|
|
}
|
|
throw new EmbeddingError(
|
|
`could not reach ${base}: ${lastError?.message ?? lastError}`,
|
|
'UNREACHABLE',
|
|
)
|
|
}
|
|
|
|
/** Check a returned vector's width against the configured one. */
|
|
function check(vector, index) {
|
|
if (!Array.isArray(vector)) {
|
|
throw new EmbeddingError(`response item ${index} was not an array`, 'BAD_RESPONSE')
|
|
}
|
|
if (dimensions > 0 && vector.length !== dimensions) {
|
|
throw new EmbeddingError(
|
|
`model "${model}" returned ${vector.length}-dimensional vectors but ${dimensions} are configured`,
|
|
'DIMENSION_MISMATCH',
|
|
)
|
|
}
|
|
return normalize(Float32Array.from(vector))
|
|
}
|
|
|
|
/** Embed one batch through whichever endpoint answers. */
|
|
async function embedBatch(texts, signal) {
|
|
const inputs = texts.map(clampInput)
|
|
if (inputs.length === 0) return []
|
|
if (batchEndpointWorks) {
|
|
try {
|
|
const payload = await request('/api/embed', withKeepAlive({ model, input: inputs }), signal)
|
|
if (!Array.isArray(payload?.embeddings) || payload.embeddings.length !== inputs.length) {
|
|
throw new EmbeddingError(
|
|
`/api/embed returned ${payload?.embeddings?.length ?? 0} vectors for ${inputs.length} inputs`,
|
|
'BAD_RESPONSE',
|
|
)
|
|
}
|
|
return payload.embeddings.map((vector, index) => check(vector, index))
|
|
} catch (error) {
|
|
// Only a missing endpoint falls back. A dimension mismatch or a bad
|
|
// response is a real fault and must not be papered over by a second
|
|
// endpoint that would fail the same way more slowly.
|
|
if (!(error instanceof EmbeddingError) || error.code !== 'ENDPOINT_UNSUPPORTED') throw error
|
|
batchEndpointWorks = false
|
|
}
|
|
}
|
|
const out = []
|
|
for (const input of inputs) {
|
|
const payload = await request('/api/embeddings', withKeepAlive({ model, prompt: input }), signal)
|
|
out.push(check(payload?.embedding, out.length))
|
|
}
|
|
return out
|
|
}
|
|
|
|
return {
|
|
model,
|
|
dimensions,
|
|
base,
|
|
/** Whether the batch endpoint has been ruled out. */
|
|
get batchEnabled() {
|
|
return batchEndpointWorks
|
|
},
|
|
/**
|
|
* Embed many texts, preserving order.
|
|
*
|
|
* @param {string[]} texts inputs
|
|
* @param {AbortSignal} [signal] caller cancellation
|
|
* @param {(done: number, total: number) => void} [onProgress]
|
|
* @returns {Promise<Float32Array[]>} unit vectors in input order
|
|
*/
|
|
async embed(texts, signal, onProgress) {
|
|
const out = new Array(texts.length)
|
|
const size = Math.max(1, Math.min(batchSize, texts.length))
|
|
const batches = []
|
|
for (let at = 0; at < texts.length; at += size) {
|
|
batches.push({ start: at, items: texts.slice(at, at + size) })
|
|
}
|
|
let done = 0
|
|
let cursor = 0
|
|
const workers = new Array(Math.max(1, Math.min(concurrency, batches.length))).fill(0).map(async () => {
|
|
while (cursor < batches.length) {
|
|
if (signal?.aborted) throw signal.reason ?? new Error('aborted')
|
|
const batch = batches[cursor]
|
|
cursor += 1
|
|
const vectors = await embedBatch(batch.items, signal)
|
|
for (let at = 0; at < vectors.length; at += 1) out[batch.start + at] = vectors[at]
|
|
done += vectors.length
|
|
onProgress?.(done, texts.length)
|
|
}
|
|
})
|
|
await Promise.all(workers)
|
|
return out
|
|
},
|
|
/**
|
|
* Embed one query string.
|
|
*
|
|
* @param {string} text query
|
|
* @param {AbortSignal} [signal] caller cancellation
|
|
* @returns {Promise<Float32Array>} unit vector
|
|
*/
|
|
async embedOne(text, signal) {
|
|
const [vector] = await embedBatch([text], signal)
|
|
return vector
|
|
},
|
|
}
|
|
}
|