dsh-slopo/lib/probe.js

/**
 * Pre-flight checks for the embedding backend.
 *
 * Slopo reports a bad embedding configuration as two words on stdout —
 * `Error: Embedding failed.` — with the real reason buried in its agent log.
 * That is a slow way to find out that a model string has the wrong provider
 * prefix or that the configured dimensions do not match what the model returns.
 *
 * These checks run *before* a real embed so the answer arrives in one tool call
 * instead of after a failed index.
 */

/**
 * LiteLLM provider prefixes we are confident are prefixes.
 *
 * This allow-list exists because the alternative — stripping the first path
 * segment unconditionally — mangles a model that legitimately contains a slash.
 * `unclemusclez/jina-embeddings-v2-base-code:latest` is a real Ollama model name
 * with no provider prefix at all, and stripping it would produce nonsense. An
 * unknown first segment is therefore left alone, which fails safe: the probe
 * reports what it sent, so a wrong guess is visible rather than silent.
 */
const KNOWN_PROVIDERS = new Set([
  'openai',
  'openai_like',
  'ollama',
  'ollama_chat',
  'hosted_vllm',
  'azure',
  'azure_ai',
  'cohere',
  'voyage',
  'jina_ai',
  'mistral',
  'gemini',
  'vertex_ai',
  'bedrock',
  'huggingface',
  'together_ai',
  'deepinfra',
  'openrouter',
  'xinference',
  'nvidia_nim',
  'sagemaker',
  'watsonx',
  'databricks',
  'custom',
])

/**
 * The model name an OpenAI-compatible endpoint expects, without LiteLLM's
 * provider prefix.
 *
 * @param {string} model
 * @returns {string}
 */
export function bareModelName(model) {
  const trimmed = String(model).trim()
  const slash = trimmed.indexOf('/')
  if (slash <= 0) return trimmed
  const prefix = trimmed.slice(0, slash)
  return KNOWN_PROVIDERS.has(prefix) ? trimmed.slice(slash + 1) : trimmed
}

/**
 * The `/embeddings` URL for a configured base, tolerating a base that already
 * carries the version segment.
 *
 * @param {string} apiBase
 * @returns {string}
 */
export function embeddingUrl(apiBase) {
  const base = String(apiBase).trim().replace(/\/+$/, '')
  if (base === '') return ''
  return /\/v\d+$/.test(base) ? `${base}/embeddings` : `${base}/v1/embeddings`
}

/**
 * Make one embedding request and compare the returned width with the
 * configured one.
 *
 * Never throws: every failure is a reported outcome, because this runs inside a
 * status check whose whole job is to describe what is wrong.
 *
 * @param {{
 *   apiBase: string,
 *   model: string,
 *   dimensions: number,
 *   apiKey?: string,
 *   timeoutMs?: number,
 *   fetchImpl?: typeof fetch,
 * }} input
 * @returns {Promise<Record<string, unknown>>}
 */
export async function probeEmbeddingEndpoint({
  apiBase,
  model,
  dimensions,
  apiKey = '',
  timeoutMs = 8_000,
  fetchImpl = globalThis.fetch,
}) {
  if (typeof apiBase !== 'string' || apiBase.trim() === '') {
    return { checked: false, skipped: 'no embeddingApiBase configured, so there is nothing to probe' }
  }
  if (typeof fetchImpl !== 'function') {
    return { checked: false, skipped: 'this runtime has no fetch' }
  }
  const url = embeddingUrl(apiBase)
  const sentModel = bareModelName(model)
  const controller = new AbortController()
  const timer = setTimeout(() => controller.abort(), timeoutMs)
  try {
    const headers = { 'content-type': 'application/json' }
    if (apiKey !== '') headers.authorization = `Bearer ${apiKey}`
    const response = await fetchImpl(url, {
      method: 'POST',
      headers,
      body: JSON.stringify({ model: sentModel, input: ['slopo connectivity probe'] }),
      signal: controller.signal,
    })
    const text = await response.text()
    if (!response.ok) {
      return {
        checked: true,
        reachable: true,
        ok: false,
        url,
        sentModel,
        status: response.status,
        error: text.slice(0, 400),
      }
    }
    let parsed
    try {
      parsed = JSON.parse(text)
    } catch {
      return { checked: true, reachable: true, ok: false, url, sentModel, error: 'the endpoint returned non-JSON' }
    }
    const vector = parsed?.data?.[0]?.embedding
    if (!Array.isArray(vector)) {
      return { checked: true, reachable: true, ok: false, url, sentModel, error: 'the response carried no embedding' }
    }
    const actual = vector.length
    const expected = Number(dimensions)
    return {
      checked: true,
      reachable: true,
      ok: actual === expected,
      url,
      sentModel,
      dimensions: actual,
      expected,
      error:
        actual === expected
          ? undefined
          : `the model returns ${actual} dimensions but embeddingDimensions is ${expected}`,
    }
  } catch (error) {
    const aborted = error?.name === 'AbortError'
    return {
      checked: true,
      reachable: false,
      ok: false,
      url,
      sentModel,
      error: aborted ? `no response within ${timeoutMs}ms` : String(error?.message ?? error),
    }
  } finally {
    clearTimeout(timer)
  }
}