Add NVIDIA NIM provider

This commit is contained in:
Mario Zechner
2026-06-02 15:09:36 +02:00
parent 13898f048f
commit 6014801229
20 changed files with 700 additions and 66 deletions

View File

@@ -15,14 +15,14 @@ import { createAgentSession } from "../src/core/sdk.ts";
import { SessionManager } from "../src/core/session-manager.ts";
import { SettingsManager } from "../src/core/settings-manager.ts";
describe("createAgentSession OpenRouter attribution headers", () => {
describe("createAgentSession provider attribution headers", () => {
let tempDir: string;
let cwd: string;
let agentDir: string;
let originalTelemetryEnv: string | undefined;
beforeEach(() => {
tempDir = join(tmpdir(), `pi-sdk-openrouter-test-${Date.now()}-${Math.random().toString(36).slice(2)}`);
tempDir = join(tmpdir(), `pi-sdk-attribution-test-${Date.now()}-${Math.random().toString(36).slice(2)}`);
cwd = join(tempDir, "project");
agentDir = join(tempDir, "agent");
mkdirSync(cwd, { recursive: true });
@@ -42,9 +42,9 @@ describe("createAgentSession OpenRouter attribution headers", () => {
}
});
function createModel(provider: string, baseUrl: string): Model<Api> {
function createModel(provider: string, baseUrl: string, id = `${provider}-test-model`): Model<Api> {
return {
id: `${provider}-test-model`,
id,
name: `${provider} Test Model`,
api: "openai-completions",
provider,
@@ -172,6 +172,14 @@ describe("createAgentSession OpenRouter attribution headers", () => {
expect(headers?.["X-OpenRouter-Categories"]).toBe("cli-agent");
});
it("preserves legacy OpenRouter base URL substring attribution matching", async () => {
const headers = await captureHeaders(createModel("custom-openrouter", "not-a-url-openrouter.ai"));
expect(headers?.["HTTP-Referer"]).toBe("https://pi.dev");
expect(headers?.["X-OpenRouter-Title"]).toBe("pi");
expect(headers?.["X-OpenRouter-Categories"]).toBe("cli-agent");
});
it("lets provider and request headers override the defaults", async () => {
const headers = await captureHeaders(createModel("openrouter", "https://openrouter.ai/api/v1"), {
providerHeaders: {
@@ -188,6 +196,56 @@ describe("createAgentSession OpenRouter attribution headers", () => {
expect(headers?.["X-OpenRouter-Categories"]).toBe("provider-category");
});
it("adds default attribution headers for direct NVIDIA NIM endpoints", async () => {
const headers = await captureHeaders(createModel("custom-nim", "https://integrate.api.nvidia.com/v1"));
expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBe("Pi");
});
it("adds default attribution headers for the NVIDIA provider", async () => {
const headers = await captureHeaders(createModel("nvidia", "https://example.test/v1"));
expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBe("Pi");
});
it("does not add NVIDIA NIM attribution headers when telemetry is disabled", async () => {
const headers = await captureHeaders(createModel("nvidia", "https://integrate.api.nvidia.com/v1"), {
telemetryEnabled: false,
});
expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBeUndefined();
});
it("lets provider and request headers override NVIDIA NIM defaults", async () => {
const headers = await captureHeaders(createModel("nvidia", "https://integrate.api.nvidia.com/v1"), {
providerHeaders: {
"X-BILLING-INVOKE-ORIGIN": "Provider",
},
requestHeaders: {
"X-BILLING-INVOKE-ORIGIN": "Request",
},
});
expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBe("Request");
});
it("does not add NVIDIA NIM attribution headers for NVIDIA models routed through OpenRouter", async () => {
const headers = await captureHeaders(
createModel("openrouter", "https://openrouter.ai/api/v1", "nvidia/nemotron-3-super-120b-a12b"),
);
expect(headers?.["HTTP-Referer"]).toBe("https://pi.dev");
expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBeUndefined();
});
it("does not add NVIDIA NIM attribution headers for NVIDIA models routed through Vercel AI Gateway", async () => {
const headers = await captureHeaders(
createModel("vercel-ai-gateway", "https://ai-gateway.vercel.sh/v1", "nvidia/nemotron-3-super-120b-a12b"),
);
expect(headers?.["X-BILLING-INVOKE-ORIGIN"]).toBeUndefined();
});
it("adds OpenCode session headers", async () => {
const headers = await captureHeaders(createModel("opencode", "https://opencode.ai/zen/v1"), {
sessionId: "opencode-session",