Files
diegosouzapw--omniroute/tests/unit/embeddings-nvidia-input-type.test.ts
2026-07-13 13:39:12 +08:00

109 lines
3.5 KiB
TypeScript

import test from "node:test";
import assert from "node:assert/strict";
import { mkdtempSync } from "node:fs";
import { tmpdir } from "node:os";
import { join } from "node:path";
process.env.DATA_DIR = mkdtempSync(join(tmpdir(), "omniroute-embeddings-nvidia-"));
const { handleEmbedding } = await import("../../open-sse/handlers/embeddings.ts");
// Issue #1378: NVIDIA NIM asymmetric embedding models (e.g. nvidia/nv-embedqa-e5-v5)
// require an `input_type` parameter ("query" | "passage"); without it the upstream
// returns 400 "'input_type' parameter is required". OmniRoute must inject the
// registered model-level default when the client omits input_type, and must respect
// a client-supplied input_type when present.
function captureFetch(captured: { body?: Record<string, unknown> }) {
return async (url: unknown, options: { headers?: unknown; body?: unknown } = {}) => {
captured.body = JSON.parse(String(options.body || "{}"));
return new Response(
JSON.stringify({
data: [{ object: "embedding", embedding: [0.1, 0.2], index: 0 }],
usage: { prompt_tokens: 4, total_tokens: 4 },
}),
{ status: 200, headers: { "content-type": "application/json" } }
);
};
}
test("handleEmbedding injects NVIDIA asymmetric default input_type when client omits it", async () => {
const originalFetch = globalThis.fetch;
const captured: { body?: Record<string, unknown> } = {};
globalThis.fetch = captureFetch(captured) as typeof fetch;
try {
const result = await handleEmbedding({
body: {
model: "nvidia/nvidia/nv-embedqa-e5-v5",
input: "What is the capital of France?",
},
credentials: { apiKey: "nvidia-key" },
log: null,
});
assert.equal(result.success, true);
// The model-level default input_type must be forwarded to the upstream body.
assert.equal(
captured.body?.input_type,
"query",
"expected NVIDIA asymmetric model default input_type to be injected"
);
} finally {
globalThis.fetch = originalFetch;
}
});
test("handleEmbedding respects a client-supplied input_type (does not overwrite)", async () => {
const originalFetch = globalThis.fetch;
const captured: { body?: Record<string, unknown> } = {};
globalThis.fetch = captureFetch(captured) as typeof fetch;
try {
const result = await handleEmbedding({
body: {
model: "nvidia/nvidia/nv-embedqa-e5-v5",
input: "Paris is the capital of France.",
input_type: "passage",
},
credentials: { apiKey: "nvidia-key" },
log: null,
});
assert.equal(result.success, true);
assert.equal(
captured.body?.input_type,
"passage",
"expected client-supplied input_type to be respected, not overwritten by the default"
);
} finally {
globalThis.fetch = originalFetch;
}
});
test("handleEmbedding does not inject input_type for symmetric models without a default", async () => {
const originalFetch = globalThis.fetch;
const captured: { body?: Record<string, unknown> } = {};
globalThis.fetch = captureFetch(captured) as typeof fetch;
try {
const result = await handleEmbedding({
body: {
model: "openai/text-embedding-3-small",
input: "hello world",
},
credentials: { apiKey: "openai-key" },
log: null,
});
assert.equal(result.success, true);
assert.equal(
"input_type" in (captured.body || {}),
false,
"symmetric models without a default must not receive an injected input_type"
);
} finally {
globalThis.fetch = originalFetch;
}
});