dsh-recall/test/embed.test.js

import { strict as assert } from 'node:assert'
import { test } from 'node:test'

import { EmbeddingError, clampInput, createEmbedder, dot, normalize } from '../lib/embed.js'
import { makeStubFetch } from './helpers.js'

function embedder(fetchImpl, overrides = {}) {
  return createEmbedder({
    apiBase: 'http://stub',
    model: 'test-model',
    dimensions: 768,
    fetchImpl,
    concurrency: 2,
    ...overrides,
  })
}

test('normalize produces a unit vector and leaves zero alone', () => {
  const unit = normalize(Float32Array.from([3, 4]))
  assert.equal(Math.round(Math.hypot(unit[0], unit[1]) * 1000) / 1000, 1)
  assert.deepEqual(Array.from(normalize(Float32Array.from([0, 0]))), [0, 0])
})

test('dot is the cosine similarity of unit vectors', () => {
  assert.equal(dot(Float32Array.from([1, 0]), Float32Array.from([1, 0])), 1)
  assert.equal(dot(Float32Array.from([1, 0]), Float32Array.from([0, 1])), 0)
})

test('clampInput truncates to the character budget', () => {
  assert.equal(clampInput('abc').length, 3)
  assert.equal(clampInput('x'.repeat(50_000)).length, 24_000)
})

test('embed uses the batch endpoint and preserves order', async () => {
  const fetchImpl = makeStubFetch()
  const result = await embedder(fetchImpl).embed(['alpha alpha', 'beta beta'])
  assert.equal(result.length, 2)
  assert.equal(fetchImpl.calls.filter((call) => call.url.endsWith('/api/embed')).length, 1)
  // The two texts share no tokens, so their vectors must be orthogonal-ish.
  assert.ok(dot(result[0], result[1]) < 0.05)
})

test('embed returns unit vectors', async () => {
  const [vector] = await embedder(makeStubFetch()).embed(['some real code here'])
  let sum = 0
  for (const value of vector) sum += value * value
  assert.ok(Math.abs(sum - 1) < 1e-5)
})

test('an endpoint without the batch API falls back to single requests', async () => {
  const fetchImpl = makeStubFetch({ unsupportedBatch: true })
  const subject = embedder(fetchImpl)
  const result = await subject.embed(['one two', 'three four'])
  assert.equal(result.length, 2)
  assert.equal(subject.batchEnabled, false)
  assert.equal(fetchImpl.calls.filter((call) => call.url.endsWith('/api/embeddings')).length, 2)
})

test('a dimension mismatch is fatal and explains itself', async () => {
  const fetchImpl = makeStubFetch({ width: 128 })
  await assert.rejects(
    () => embedder(fetchImpl).embed(['x']),
    (error) => {
      assert.ok(error instanceof EmbeddingError)
      assert.equal(error.code, 'DIMENSION_MISMATCH')
      assert.match(error.message, /128-dimensional .* 768 are configured/)
      return true
    },
  )
})

test('a persistent server error surfaces as UNREACHABLE after retrying', async () => {
  let attempts = 0
  const fetchImpl = async () => {
    attempts += 1
    return { ok: false, status: 500, text: async () => 'boom', json: async () => ({}) }
  }
  await assert.rejects(
    () => embedder(fetchImpl).embed(['x']),
    (error) => {
      assert.equal(error.code, 'UNREACHABLE')
      return true
    },
  )
  assert.equal(attempts, 3, 'the request should be attempted MAX_ATTEMPTS times')
})

test('a transient failure is retried and then succeeds', async () => {
  const good = makeStubFetch()
  let calls = 0
  const fetchImpl = async (url, init) => {
    calls += 1
    if (calls === 1) throw new Error('socket hang up')
    return good(url, init)
  }
  const [vector] = await embedder(fetchImpl).embed(['retry me'])
  assert.equal(vector.length, 768)
  assert.equal(calls, 2)
})

test('a client error is not retried', async () => {
  let attempts = 0
  const fetchImpl = async () => {
    attempts += 1
    return { ok: false, status: 400, text: async () => 'bad model', json: async () => ({}) }
  }
  await assert.rejects(
    () => embedder(fetchImpl).embed(['x']),
    (error) => {
      assert.equal(error.code, 'HTTP_ERROR')
      return true
    },
  )
  assert.equal(attempts, 1, 'a 400 is a configuration fault, not a transient one')
})

test('embedding nothing makes no request', async () => {
  const fetchImpl = makeStubFetch()
  const result = await embedder(fetchImpl).embed([])
  assert.deepEqual(result, [])
  assert.equal(fetchImpl.calls.length, 0)
})

test('cancellation stops the work', async () => {
  const fetchImpl = makeStubFetch()
  const controller = new AbortController()
  controller.abort()
  await assert.rejects(() => embedder(fetchImpl).embed(['x'], controller.signal))
})

test('keepAlive and maxAttempts are sent only when configured', async () => {
  const withKeepAlive = makeStubFetch()
  await embedder(withKeepAlive, { keepAlive: '10m' }).embed(['x'])
  assert.equal(withKeepAlive.calls[0].body.keep_alive, '10m')

  const without = makeStubFetch()
  await embedder(without).embed(['x'])
  assert.equal('keep_alive' in without.calls[0].body, false, 'an unset keep-alive must not be sent')
})

test('the keep-alive also reaches the single-text endpoint', async () => {
  const fetchImpl = makeStubFetch({ unsupportedBatch: true })
  await embedder(fetchImpl, { keepAlive: '2m' }).embed(['x'])
  const legacy = fetchImpl.calls.find((call) => call.url.endsWith('/api/embeddings'))
  assert.equal(legacy.body.keep_alive, '2m')
})

test('maxAttempts of one gives a probe a single chance', async () => {
  let attempts = 0
  const fetchImpl = async () => {
    attempts += 1
    return { ok: false, status: 500, text: async () => 'boom', json: async () => ({}) }
  }
  await assert.rejects(() => embedder(fetchImpl, { maxAttempts: 1 }).embed(['x']))
  assert.equal(attempts, 1)
})

test('a socket that accepts and never answers fails on the timeout, not forever', async () => {
  // The real failure mode observed in development: connection opens, no bytes.
  const fetchImpl = (url, init) =>
    new Promise((_resolve, reject) => {
      init.signal.addEventListener('abort', () => reject(init.signal.reason), { once: true })
    })
  const started = Date.now()
  await assert.rejects(
    () => embedder(fetchImpl, { timeoutMs: 300, maxAttempts: 1 }).embed(['x']),
    (error) => {
      assert.equal(error.code, 'UNREACHABLE')
      return true
    },
  )
  assert.ok(Date.now() - started < 5_000, 'the timeout must bound the wait')
})