From 3cc34c92eec2540ef36d2513eff8d4e4bf73bad9 Mon Sep 17 00:00:00 2001 From: Time Attakc <89218912+time-attack@users.noreply.github.com> Date: Tue, 21 Jul 2026 00:17:26 -0700 Subject: [PATCH] feat(ai): add NVIDIA NIM provider recipe (#2965) (#3022) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds NVIDIA NIM / API Catalog as a first-class OpenAI-compatible AI recipe: - chat via nvidia/nemotron-3-super-120b-a12b (conservative capability claims: no tools, no subagent loop until proven) - hosted embedding models incl. nvidia/llama-nemotron-embed-1b-v2 with Matryoshka-style dimension overrides (1024/1280/1536/2048) and fixed natural dims for the other catalog models - asymmetric input_type mapping (document -> passage, query -> query) via a gateway compat fetch shim, since the generic openai-compatible recipe cannot infer that provider-specific requirement - base URL https://integrate.api.nvidia.com/v1 verified live (OpenAI-shaped /v1/models, all five recipe model ids present in the catalog) Changed from the original PR: dropped the recipe's custom resolveAuth — it duplicated defaultResolveAuth's Authorization-Bearer behavior exactly and violated the IRON RULE that only Azure overrides resolveAuth (test/ai/recipes-existing-regression.test.ts). NVIDIA_API_KEY now flows through defaultResolveAuth via auth_env.required, and the recipe test pins resolveAuth === undefined + the default Bearer resolution + the missing-key AIConfigError. Also scrubbed a private downstream-agent name from ported comments per the repo privacy rule. Takeover of #2965 by @ravehorn. Co-authored-by: Garry Tan Co-authored-by: SAGE Codex Co-authored-by: Claude Fable 5 --- src/core/ai/dims.ts | 43 ++++++++++++++++ src/core/ai/gateway.ts | 26 ++++++++++ src/core/ai/recipes/index.ts | 2 + src/core/ai/recipes/nvidia.ts | 71 ++++++++++++++++++++++++++ src/core/embedding-dim-check.ts | 25 ++++++++- test/ai/recipe-nvidia.test.ts | 89 +++++++++++++++++++++++++++++++++ 6 files changed, 254 insertions(+), 2 deletions(-) create mode 100644 src/core/ai/recipes/nvidia.ts create mode 100644 test/ai/recipe-nvidia.test.ts diff --git a/src/core/ai/dims.ts b/src/core/ai/dims.ts index 4f170e8fc..7961516a5 100644 --- a/src/core/ai/dims.ts +++ b/src/core/ai/dims.ts @@ -90,6 +90,38 @@ export function isValidOpenAITextEmbedding3Dim(modelId: string, dims: number): b return Number.isInteger(dims) && dims >= 1 && dims <= max; } +// NVIDIA NIM hosted embedding models use asymmetric input_type values. Most +// emit fixed natural dimensions, but llama-nemotron-embed-1b-v2 accepts +// Matryoshka-style dimension overrides (e.g. matching an existing 1280d +// brain column without re-embedding through another provider). +const NVIDIA_EMBEDDING_DIMS: Record = { + 'nvidia/nv-embedqa-e5-v5': 1024, + 'nvidia/llama-nemotron-embed-1b-v2': 2048, + 'nvidia/nv-embed-v1': 4096, + 'nvidia/nv-embedcode-7b-v1': 4096, +}; + +const NVIDIA_EMBEDDING_DIM_OPTIONS: Record = { + 'nvidia/llama-nemotron-embed-1b-v2': [1024, 1280, 1536, 2048], +}; + +export function isNvidiaEmbeddingModel(modelId: string): boolean { + return modelId in NVIDIA_EMBEDDING_DIMS; +} + +export function nvidiaEmbeddingDim(modelId: string): number | undefined { + return NVIDIA_EMBEDDING_DIMS[modelId]; +} + +export function nvidiaEmbeddingDimOptions(modelId: string): number[] | undefined { + return NVIDIA_EMBEDDING_DIM_OPTIONS[modelId]; +} + +export function supportsNvidiaEmbeddingDimension(modelId: string, dims: number): boolean { + const options = nvidiaEmbeddingDimOptions(modelId); + return !!options && options.includes(dims); +} + /** * Build the providerOptions blob for embedMany() that pins output dimensions. * @@ -194,6 +226,17 @@ export function dimsProviderOptions( }, }; } + // NVIDIA NIM hosted embeddings are OpenAI-compatible but require + // asymmetric input_type. Use passage for indexing/document-side vectors + // and query for search-side vectors. Only llama-nemotron-embed-1b-v2 + // supports a dimensions override; fixed-dim models reject it. + if (isNvidiaEmbeddingModel(modelId)) { + const opts: Record = { + input_type: inputType === 'query' ? 'query' : 'passage', + }; + if (supportsNvidiaEmbeddingDimension(modelId, dims)) opts.dimensions = dims; + return { openaiCompatible: opts }; + } // OpenAI text-embedding-3 family on the openai-compatible adapter // (Azure OpenAI hosts these via its OpenAI-compatible /embeddings // endpoint). The provider defaults to the model's native size (3072 diff --git a/src/core/ai/gateway.ts b/src/core/ai/gateway.ts index 08ee43cea..3752a7919 100644 --- a/src/core/ai/gateway.ts +++ b/src/core/ai/gateway.ts @@ -1059,6 +1059,30 @@ const voyageCompatFetch = (async (input: RequestInfo | URL, init?: RequestInit) * float[] (not base64), so the Layer 2 cap compares against the JSON * payload size of each embedding rather than a base64 string length. */ +/** + * NVIDIA NIM compatibility shim. NVIDIA uses the OpenAI embeddings wire + * shape but requires asymmetric input_type values: query for retrieval and + * passage for indexed documents. The generic gateway store carries + * query/document across the AI SDK boundary; map document to passage here. + */ +const nvidiaCompatFetch = (async (input: RequestInfo | URL, init?: RequestInit) => { + let baseInit: RequestInit = init ?? {}; + if (baseInit.body && typeof baseInit.body === 'string') { + try { + const parsed = JSON.parse(baseInit.body); + if (parsed && typeof parsed === 'object' && parsed.input_type === undefined) { + parsed.input_type = __embedInputTypeStore.getStore() === 'query' ? 'query' : 'passage'; + const headers = new Headers(baseInit.headers ?? {}); + headers.delete('content-length'); + baseInit = { ...baseInit, body: JSON.stringify(parsed), headers }; + } + } catch { + // Preserve the provider response when the SDK body is unexpectedly non-JSON. + } + } + return fetch(input as any, baseInit); +}) as unknown as typeof fetch; + const zeroEntropyCompatFetch = (async (input: RequestInfo | URL, init?: RequestInit) => { // OUTBOUND: normalize URL, rewrite path /embeddings → /models/embed, then // rewrite body. fetch accepts RequestInfo (string | Request) | URL; we @@ -1319,6 +1343,8 @@ function instantiateEmbedding(recipe: Recipe, modelId: string, cfg: AIGatewayCon ? voyageCompatFetch : recipe.id === 'zeroentropyai' ? zeroEntropyCompatFetch + : recipe.id === 'nvidia' + ? nvidiaCompatFetch : openAICompatAsymmetricFetch); const client = createOpenAICompatible({ name: recipe.id, diff --git a/src/core/ai/recipes/index.ts b/src/core/ai/recipes/index.ts index 2e6954089..eb751ec61 100644 --- a/src/core/ai/recipes/index.ts +++ b/src/core/ai/recipes/index.ts @@ -25,6 +25,7 @@ import { zeroentropyai } from './zeroentropyai.ts'; import { llamaServerReranker } from './llama-server-reranker.ts'; import { moonshot } from './moonshot.ts'; import { mistral } from './mistral.ts'; +import { nvidia } from './nvidia.ts'; const ALL: Recipe[] = [ openai, @@ -46,6 +47,7 @@ const ALL: Recipe[] = [ zeroentropyai, moonshot, mistral, + nvidia, ]; /** Map from `provider:id` key to recipe. */ diff --git a/src/core/ai/recipes/nvidia.ts b/src/core/ai/recipes/nvidia.ts new file mode 100644 index 000000000..1a8916a70 --- /dev/null +++ b/src/core/ai/recipes/nvidia.ts @@ -0,0 +1,71 @@ +import type { Recipe } from '../types.ts'; + +/** + * NVIDIA NIM / API Catalog exposes OpenAI-compatible /v1/chat/completions + * and /v1/embeddings APIs. + * + * Retrieval models use asymmetric encoding. The gateway maps gbrain's + * document/query distinction to NVIDIA's wire values: + * document -> input_type: passage + * query -> input_type: query + * + * The model ids below intentionally keep NVIDIA's full catalog ids because + * the hosted endpoint expects values like `nvidia/nv-embedqa-e5-v5` in the + * request body. Short aliases are provided for CLI ergonomics. + */ +export const nvidia: Recipe = { + id: 'nvidia', + name: 'NVIDIA NIM', + tier: 'openai-compat', + implementation: 'openai-compatible', + base_url_default: 'https://integrate.api.nvidia.com/v1', + auth_env: { + required: ['NVIDIA_API_KEY'], + setup_url: 'https://build.nvidia.com', + }, + aliases: { + 'nv-embedqa-e5-v5': 'nvidia/nv-embedqa-e5-v5', + 'llama-nemotron-embed-1b-v2': 'nvidia/llama-nemotron-embed-1b-v2', + 'nemotron-3-super': 'nvidia/nemotron-3-super-120b-a12b', + 'nemotron-3-super-120b-a12b': 'nvidia/nemotron-3-super-120b-a12b', + 'nv-embed-v1': 'nvidia/nv-embed-v1', + 'nv-embedcode-7b-v1': 'nvidia/nv-embedcode-7b-v1', + }, + // No resolveAuth override: NVIDIA is plain `Authorization: Bearer `, + // which defaultResolveAuth derives from auth_env.required. IRON RULE + // (test/ai/recipes-existing-regression.test.ts): only Azure overrides + // resolveAuth. + touchpoints: { + chat: { + models: [ + 'nvidia/nemotron-3-super-120b-a12b', + ], + supports_tools: false, + supports_subagent_loop: false, + // Do not treat Nemotron as a Minions subagent driver until tool-calling + // and replay stability are proven through a separate adapter test. + max_context_tokens: 128000, + price_last_verified: '2026-05-24', + }, + embedding: { + models: [ + 'nvidia/nv-embedqa-e5-v5', + 'nvidia/llama-nemotron-embed-1b-v2', + 'nvidia/nv-embed-v1', + 'nvidia/nv-embedcode-7b-v1', + ], + // Default to the lightest tested hosted model. Larger NVIDIA models are + // supported via explicit embedding_dimensions (2048 or 4096). + default_dims: 1024, + dims_options: [1024, 2048, 4096], + // Conservative split; hosted NVIDIA embedding endpoints require + // input_type and may reject large payloads before tokenizing. + max_batch_tokens: 8192, + chars_per_token: 4, + safety_factor: 0.75, + cost_per_1m_tokens_usd: undefined, + price_last_verified: '2026-05-24', + }, + }, + setup_hint: 'Get an API key at https://build.nvidia.com, then `export NVIDIA_API_KEY=...`.', +}; diff --git a/src/core/embedding-dim-check.ts b/src/core/embedding-dim-check.ts index b8d022a8f..e863f748e 100644 --- a/src/core/embedding-dim-check.ts +++ b/src/core/embedding-dim-check.ts @@ -29,6 +29,9 @@ import { isOpenAITextEmbedding3Model, isValidOpenAITextEmbedding3Dim, maxOpenAITextEmbedding3Dim, + nvidiaEmbeddingDim, + nvidiaEmbeddingDimOptions, + supportsNvidiaEmbeddingDimension, } from './ai/dims.ts'; /** @@ -366,7 +369,9 @@ function validateDimAgainstTouchpoint( dimsOptions: number[] | undefined, requestedDims: number | undefined, ): ResolveSchemaDimResult { - const dim = requestedDims ?? defaultDims; + const nvidiaNaturalDims = recipe.id === 'nvidia' ? nvidiaEmbeddingDim(modelId) : undefined; + const effectiveDefaultDims = nvidiaNaturalDims ?? defaultDims; + const dim = requestedDims ?? effectiveDefaultDims; if (!Number.isInteger(dim) || dim <= 0) { return { @@ -396,7 +401,7 @@ function validateDimAgainstTouchpoint( dim, model: `${recipe.id}:${modelId}`, provider: recipe.id, - recipeDefault: defaultDims, + recipeDefault: effectiveDefaultDims, }; } @@ -411,6 +416,22 @@ function isCustomDimValidForProvider( requestedDims: number, dimsOptions: number[] | undefined, ): CustomDimCheck { + // NVIDIA models are mixed: some fixed-dim, one Matryoshka-style. Handle + // them before generic recipe dims_options so llama-nemotron can use 1280d. + if (recipe.id === 'nvidia') { + const naturalDims = nvidiaEmbeddingDim(modelId); + if (naturalDims !== undefined && requestedDims === naturalDims) return { valid: true, error: '' }; + if (supportsNvidiaEmbeddingDimension(modelId, requestedDims)) return { valid: true, error: '' }; + const options = nvidiaEmbeddingDimOptions(modelId); + return { + valid: false, + error: + `NVIDIA model "${modelId}" does not support dimensions ${requestedDims}. ` + + `Natural dimensions: ${naturalDims ?? 'unknown'}. ` + + (options ? `Supported overrides: ${options.join(', ')}.` : 'No dimension overrides are supported for this NVIDIA model.'), + }; + } + // Tier 1: recipe-declared dims_options. if (dimsOptions && dimsOptions.length > 0) { if (dimsOptions.includes(requestedDims)) return { valid: true, error: '' }; diff --git a/test/ai/recipe-nvidia.test.ts b/test/ai/recipe-nvidia.test.ts new file mode 100644 index 000000000..0dc295ce5 --- /dev/null +++ b/test/ai/recipe-nvidia.test.ts @@ -0,0 +1,89 @@ +import { describe, expect, test } from 'bun:test'; +import { + dimsProviderOptions, + nvidiaEmbeddingDimOptions, + supportsNvidiaEmbeddingDimension, +} from '../../src/core/ai/dims.ts'; +import { getRecipe, RECIPES } from '../../src/core/ai/recipes/index.ts'; +import { nvidia } from '../../src/core/ai/recipes/nvidia.ts'; +import { defaultResolveAuth } from '../../src/core/ai/gateway.ts'; +import { AIConfigError } from '../../src/core/ai/errors.ts'; + +describe('recipe: nvidia', () => { + test('registered with OpenAI-compatible NIM endpoint', () => { + expect(RECIPES.has('nvidia')).toBe(true); + expect(getRecipe('nvidia')).toBe(nvidia); + expect(nvidia.id).toBe('nvidia'); + expect(nvidia.tier).toBe('openai-compat'); + expect(nvidia.implementation).toBe('openai-compatible'); + expect(nvidia.base_url_default).toBe('https://integrate.api.nvidia.com/v1'); + }); + + test('auth flows through defaultResolveAuth — NVIDIA_API_KEY as bearer token', () => { + // IRON RULE: only Azure overrides resolveAuth. NVIDIA is plain + // Authorization Bearer, so the recipe must NOT declare its own resolver; + // defaultResolveAuth derives the header from auth_env.required. + expect(nvidia.resolveAuth).toBeUndefined(); + expect(nvidia.auth_env?.required).toEqual(['NVIDIA_API_KEY']); + expect(defaultResolveAuth(nvidia, { NVIDIA_API_KEY: 'fake-nvidia' }, 'embedding')).toEqual({ + headerName: 'Authorization', + token: 'Bearer fake-nvidia', + }); + expect(() => defaultResolveAuth(nvidia, {}, 'chat')).toThrow(AIConfigError); + }); + + test('chat touchpoint declares Nemotron 3 Super without subagent-loop claims', () => { + const chat = nvidia.touchpoints.chat!; + expect(chat.models).toContain('nvidia/nemotron-3-super-120b-a12b'); + expect(chat.supports_tools).toBe(false); + expect(chat.supports_subagent_loop).toBe(false); + expect(chat.max_context_tokens).toBe(128000); + }); + + test('embedding touchpoint declares tested NVIDIA models and natural dimensions', () => { + const e = nvidia.touchpoints.embedding!; + expect(e.models).toContain('nvidia/nv-embedqa-e5-v5'); + expect(e.models).toContain('nvidia/llama-nemotron-embed-1b-v2'); + expect(e.models).toContain('nvidia/nv-embed-v1'); + expect(e.models).toContain('nvidia/nv-embedcode-7b-v1'); + expect(e.default_dims).toBe(1024); + expect(e.dims_options).toEqual([1024, 2048, 4096]); + expect(e.max_batch_tokens).toBeGreaterThan(0); + }); + + test('aliases allow short model names while preserving NVIDIA catalog ids', () => { + expect(nvidia.aliases?.['nv-embedqa-e5-v5']).toBe('nvidia/nv-embedqa-e5-v5'); + expect(nvidia.aliases?.['llama-nemotron-embed-1b-v2']).toBe('nvidia/llama-nemotron-embed-1b-v2'); + expect(nvidia.aliases?.['nemotron-3-super']).toBe('nvidia/nemotron-3-super-120b-a12b'); + expect(nvidia.aliases?.['nemotron-3-super-120b-a12b']).toBe('nvidia/nemotron-3-super-120b-a12b'); + }); + + test('dimsProviderOptions emits passage input_type by default for NVIDIA embeddings', () => { + expect(dimsProviderOptions('openai-compatible', 'nvidia/nv-embedqa-e5-v5', 1024)).toEqual({ + openaiCompatible: { input_type: 'passage' }, + }); + }); + + test('dimsProviderOptions maps query/document inputType for NVIDIA embeddings', () => { + expect(dimsProviderOptions('openai-compatible', 'nvidia/nv-embedqa-e5-v5', 1024, 'query')).toEqual({ + openaiCompatible: { input_type: 'query' }, + }); + expect(dimsProviderOptions('openai-compatible', 'nvidia/nv-embedqa-e5-v5', 1024, 'document')).toEqual({ + openaiCompatible: { input_type: 'passage' }, + }); + }); + + test('llama-nemotron supports a 1280d Matryoshka dimension override', () => { + expect(nvidiaEmbeddingDimOptions('nvidia/llama-nemotron-embed-1b-v2')).toContain(1280); + expect(supportsNvidiaEmbeddingDimension('nvidia/llama-nemotron-embed-1b-v2', 1280)).toBe(true); + expect(dimsProviderOptions('openai-compatible', 'nvidia/llama-nemotron-embed-1b-v2', 1280, 'query')).toEqual({ + openaiCompatible: { input_type: 'query', dimensions: 1280 }, + }); + }); + + test('fixed-dim NVIDIA models omit dimensions because they reject overrides', () => { + const opts = dimsProviderOptions('openai-compatible', 'nvidia/nv-embedqa-e5-v5', 1024, 'query'); + expect(opts).toEqual({ openaiCompatible: { input_type: 'query' } }); + expect(JSON.stringify(opts)).not.toContain('dimensions'); + }); +});