dsh-slopo/test/probe.test.js

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

import { bareModelName, embeddingUrl, probeEmbeddingEndpoint } from '../lib/probe.js'

/** A fetch stand-in returning one canned response. */
function fakeFetch({ ok = true, status = 200, body = {}, throws, onCall } = {}) {
  return async (url, options) => {
    if (onCall !== undefined) onCall(url, options)
    if (throws !== undefined) throw throws
    return {
      ok,
      status,
      async text() {
        return typeof body === 'string' ? body : JSON.stringify(body)
      },
    }
  }
}

test('bareModelName strips a known provider prefix', () => {
  assert.equal(bareModelName('ollama/unclemusclez/jina-embeddings-v2-base-code:latest'), 'unclemusclez/jina-embeddings-v2-base-code:latest')
  assert.equal(bareModelName('openai/text-embedding-3-small'), 'text-embedding-3-small')
  assert.equal(bareModelName('hosted_vllm/meta-llama/Llama-3-8B'), 'meta-llama/Llama-3-8B')
})

test('bareModelName leaves a model that merely contains a slash alone', () => {
  // An unknown first segment is far more likely to be part of the model name
  // than a provider prefix, and mangling it would make the probe lie.
  assert.equal(bareModelName('unclemusclez/jina-embeddings-v2-base-code:latest'), 'unclemusclez/jina-embeddings-v2-base-code:latest')
  assert.equal(bareModelName('text-embedding-3-small'), 'text-embedding-3-small')
})

test('embeddingUrl tolerates a base that already carries the version segment', () => {
  assert.equal(embeddingUrl('http://host:11434'), 'http://host:11434/v1/embeddings')
  assert.equal(embeddingUrl('http://host:11434/'), 'http://host:11434/v1/embeddings')
  assert.equal(embeddingUrl('http://host:11434/v1'), 'http://host:11434/v1/embeddings')
  assert.equal(embeddingUrl('  http://host:11434/v2/  '), 'http://host:11434/v2/embeddings')
  assert.equal(embeddingUrl(''), '')
})

test('probeEmbeddingEndpoint reports a matching dimension as ok', async () => {
  let seen
  const result = await probeEmbeddingEndpoint({
    apiBase: 'http://host:11434',
    model: 'ollama/unclemusclez/jina:latest',
    dimensions: 768,
    fetchImpl: fakeFetch({
      body: { data: [{ embedding: new Array(768).fill(0.1) }] },
      onCall: (url, options) => {
        seen = { url, options }
      },
    }),
  })
  assert.equal(result.ok, true)
  assert.equal(result.dimensions, 768)
  assert.equal(result.expected, 768)
  assert.equal(result.url, 'http://host:11434/v1/embeddings')
  // The prefix is LiteLLM's; the endpoint must receive the bare model name.
  assert.equal(JSON.parse(seen.options.body).model, 'unclemusclez/jina:latest')
})

test('probeEmbeddingEndpoint catches a dimension mismatch, which is what slopo freezes', async () => {
  const result = await probeEmbeddingEndpoint({
    apiBase: 'http://host:11434',
    model: 'ollama/m:latest',
    dimensions: 1024,
    fetchImpl: fakeFetch({ body: { data: [{ embedding: new Array(768).fill(0) }] } }),
  })
  assert.equal(result.ok, false)
  assert.equal(result.dimensions, 768)
  assert.equal(result.expected, 1024)
  assert.match(result.error, /returns 768 dimensions but embeddingDimensions is 1024/)
})

test('probeEmbeddingEndpoint reports the width when dimensions is not set yet', async () => {
  // This is the case a new user is in, and the reason the probe runs before
  // every required setting exists: nobody can look up how wide an arbitrary
  // model is, so `dimensions` in the result IS the answer.
  const result = await probeEmbeddingEndpoint({
    apiBase: 'http://host:11434',
    model: 'ollama/m:latest',
    dimensions: 0,
    fetchImpl: fakeFetch({ body: { data: [{ embedding: new Array(768).fill(0) }] } }),
  })
  assert.equal(result.ok, false, 'unset is not a match')
  assert.equal(result.dimensions, 768, 'the observed width is reported regardless')
  assert.equal(result.expected, 0)
  assert.match(result.error, /returns 768 dimensions and embeddingDimensions is not set/)
})

test('probeEmbeddingEndpoint reports an HTTP error instead of throwing', async () => {
  const result = await probeEmbeddingEndpoint({
    apiBase: 'http://host:11434',
    model: 'ollama/m:latest',
    dimensions: 768,
    fetchImpl: fakeFetch({ ok: false, status: 404, body: 'model not found' }),
  })
  assert.equal(result.ok, false)
  assert.equal(result.status, 404)
  assert.match(result.error, /model not found/)
})

test('probeEmbeddingEndpoint reports an unreachable endpoint instead of throwing', async () => {
  const result = await probeEmbeddingEndpoint({
    apiBase: 'http://host:11434',
    model: 'ollama/m:latest',
    dimensions: 768,
    fetchImpl: fakeFetch({ throws: new Error('connect ECONNREFUSED') }),
  })
  assert.equal(result.ok, false)
  assert.equal(result.reachable, false)
  assert.match(result.error, /ECONNREFUSED/)
})

test('probeEmbeddingEndpoint reports a timeout as a timeout', async () => {
  const abort = new Error('aborted')
  abort.name = 'AbortError'
  const result = await probeEmbeddingEndpoint({
    apiBase: 'http://host:11434',
    model: 'm',
    dimensions: 768,
    timeoutMs: 25,
    fetchImpl: fakeFetch({ throws: abort }),
  })
  assert.equal(result.ok, false)
  assert.match(result.error, /no response within 25ms/)
})

test('probeEmbeddingEndpoint skips cleanly with nothing to probe', async () => {
  const result = await probeEmbeddingEndpoint({ apiBase: '', model: 'm', dimensions: 768 })
  assert.equal(result.checked, false)
  assert.match(result.skipped, /no embeddingApiBase/)
})

test('probeEmbeddingEndpoint sends the key as a bearer token only when there is one', async () => {
  const seen = []
  const call = (apiKey) =>
    probeEmbeddingEndpoint({
      apiBase: 'http://host',
      model: 'm',
      dimensions: 1,
      apiKey,
      fetchImpl: fakeFetch({
        body: { data: [{ embedding: [1] }] },
        onCall: (_url, options) => seen.push(options.headers),
      }),
    })
  await call('secret-key')
  await call('')
  assert.equal(seen[0].authorization, 'Bearer secret-key')
  assert.equal(Object.hasOwn(seen[1], 'authorization'), false)
})