init

lib/embed.js

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