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

@@ -11,11 +11,11 @@ import type { ExtensionRunner, LoadExtensionsResult, SessionStartEvent, ToolDefi
import { convertToLlm } from "./messages.ts";
import { ModelRegistry } from "./model-registry.ts";
import { findInitialModel } from "./model-resolver.ts";
import { mergeProviderAttributionHeaders } from "./provider-attribution.ts";
import type { ResourceLoader } from "./resource-loader.ts";
import { DefaultResourceLoader } from "./resource-loader.ts";
import { getDefaultSessionDir, SessionManager } from "./session-manager.ts";
import { SettingsManager } from "./settings-manager.ts";
import { isInstallTelemetryEnabled } from "./telemetry.ts";
import { time } from "./timings.ts";
import {
createBashTool,
@@ -128,44 +128,6 @@ function getDefaultAgentDir(): string {
return getAgentDir();
}
function getAttributionHeaders(
model: Model<any>,
settingsManager: SettingsManager,
sessionId?: string,
): Record<string, string> | undefined {
if (
sessionId &&
(model.provider === "opencode" || model.provider === "opencode-go" || model.baseUrl.includes("opencode.ai"))
) {
return { "x-opencode-session": sessionId, "x-opencode-client": "pi" };
}
if (!isInstallTelemetryEnabled(settingsManager)) {
return undefined;
}
if (model.provider === "openrouter" || model.baseUrl.includes("openrouter.ai")) {
return {
"HTTP-Referer": "https://pi.dev",
"X-OpenRouter-Title": "pi",
"X-OpenRouter-Categories": "cli-agent",
};
}
if (
model.provider === "cloudflare-workers-ai" ||
model.provider === "cloudflare-ai-gateway" ||
model.baseUrl.includes("api.cloudflare.com") ||
model.baseUrl.includes("gateway.ai.cloudflare.com")
) {
return {
"User-Agent": "pi-coding-agent",
};
}
return undefined;
}
/**
* Create an AgentSession with the specified options.
*
@@ -349,7 +311,6 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
const timeoutMs = options?.timeoutMs ?? providerRetrySettings.timeoutMs ?? effectiveTimeoutMs;
const websocketConnectTimeoutMs =
options?.websocketConnectTimeoutMs ?? settingsManager.getWebSocketConnectTimeoutMs();
const attributionHeaders = getAttributionHeaders(model, settingsManager, options?.sessionId);
return streamSimple(model, context, {
...options,
apiKey: auth.apiKey,
@@ -357,10 +318,13 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {}
websocketConnectTimeoutMs,
maxRetries: options?.maxRetries ?? providerRetrySettings.maxRetries,
maxRetryDelayMs: options?.maxRetryDelayMs ?? providerRetrySettings.maxRetryDelayMs,
headers:
attributionHeaders || auth.headers || options?.headers
? { ...attributionHeaders, ...auth.headers, ...options?.headers }
: undefined,
headers: mergeProviderAttributionHeaders(
model,
settingsManager,
options?.sessionId,
auth.headers,
options?.headers,
),
});
},
onPayload: async (payload, _model) => {