shumengya/sproutclaw
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add NVIDIA NIM provider

Browse Source
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
Download Patch File Download Diff File Expand all files Collapse all files

1
packages/coding-agent/CHANGELOG.md
View File

@@ -4,6 +4,7 @@
### Added
- Added NVIDIA NIM provider selection, setup documentation, and direct NIM request attribution headers.
- Added `ctx.mode` to extension contexts so extensions can distinguish TUI, RPC, JSON, and print mode.
- Added `ctx.getSystemPromptOptions()` for extension commands to inspect the current base system prompt inputs.

5
packages/coding-agent/README.md
View File

@@ -113,6 +113,7 @@ For each built-in provider, pi maintains a list of tool-capable models, updated
- OpenAI
- Azure OpenAI
- DeepSeek
- NVIDIA NIM
- Google Gemini
- Google Vertex
- Amazon Bedrock
@@ -290,7 +291,7 @@ See [docs/settings.md](docs/settings.md) for all options.
Pi has two separate startup features:
- **Update check:** fetches `https://pi.dev/api/latest-version` to check whether a newer Pi version exists. Disable it with `PI_SKIP_VERSION_CHECK=1`. Disabling update checks only turns off this check.
- **Install/update telemetry:** after first install or a changelog-detected update, sends an anonymous version ping to `https://pi.dev/api/report-install`. Opt out by setting `enableInstallTelemetry` to `false` in `settings.json`, or by setting `PI_TELEMETRY=0`. This does not disable update checks; Pi may still contact `pi.dev` for the latest version unless update checks are disabled or offline mode is enabled.
- **Install/update telemetry:** after first install or a changelog-detected update, sends an anonymous version ping to `https://pi.dev/api/report-install`. This setting also controls optional provider attribution headers for OpenRouter, Cloudflare, and direct NVIDIA NIM requests. Opt out by setting `enableInstallTelemetry` to `false` in `settings.json`, or by setting `PI_TELEMETRY=0`. This does not disable update checks; Pi may still contact `pi.dev` for the latest version unless update checks are disabled or offline mode is enabled.
Use `--offline` or `PI_OFFLINE=1` to disable all startup network operations described here, including update checks, package update checks, and install/update telemetry.
@@ -641,7 +642,7 @@ pi --thinking high "Solve this complex problem"
| `PI_PACKAGE_DIR` | Override package directory (useful for Nix/Guix where store paths tokenize poorly) |
| `PI_OFFLINE` | Disable startup network operations, including update checks, package update checks, and install/update telemetry |
| `PI_SKIP_VERSION_CHECK` | Skip the Pi version update check at startup. This prevents the `pi.dev` latest-version request |
| `PI_TELEMETRY` | Override install/update telemetry. Use `1`/`true`/`yes` to enable or `0`/`false`/`no` to disable. This does not disable update checks |
| `PI_TELEMETRY` | Override install/update telemetry and provider attribution headers. Use `1`/`true`/`yes` to enable or `0`/`false`/`no` to disable. This does not disable update checks |
| `PI_CACHE_RETENTION` | Set to `long` for extended prompt cache (Anthropic: 1h, OpenAI: 24h) |
| `VISUAL`, `EDITOR` | External editor for Ctrl+G |

2
packages/coding-agent/docs/providers.md
View File

@@ -52,6 +52,7 @@ pi
| Azure OpenAI Responses | `AZURE_OPENAI_API_KEY` | `azure-openai-responses` |
| OpenAI | `OPENAI_API_KEY` | `openai` |
| DeepSeek | `DEEPSEEK_API_KEY` | `deepseek` |
| NVIDIA NIM | `NVIDIA_API_KEY` | `nvidia` |
| Google Gemini | `GEMINI_API_KEY` | `google` |
| Mistral | `MISTRAL_API_KEY` | `mistral` |
| Groq | `GROQ_API_KEY` | `groq` |
@@ -86,6 +87,7 @@ Store credentials in `~/.pi/agent/auth.json`:
"anthropic": { "type": "api_key", "key": "sk-ant-..." },
"openai": { "type": "api_key", "key": "sk-..." },
"deepseek": { "type": "api_key", "key": "sk-..." },
"nvidia": { "type": "api_key", "key": "nvapi-..." },
"google": { "type": "api_key", "key": "..." },
"opencode": { "type": "api_key", "key": "..." },
"opencode-go": { "type": "api_key", "key": "..." },

2
packages/coding-agent/docs/usage.md
View File

@@ -275,7 +275,7 @@ pi --exclude-tools ask_question
| `PI_PACKAGE_DIR` | Override package directory, useful for Nix/Guix store paths |
| `PI_OFFLINE` | Disable startup network operations, including update checks, package update checks, and install/update telemetry |
| `PI_SKIP_VERSION_CHECK` | Skip the Pi version update check at startup. This prevents the `pi.dev` latest-version request |
| `PI_TELEMETRY` | Override install/update telemetry: `1`/`true`/`yes` or `0`/`false`/`no`. This does not disable update checks |
| `PI_TELEMETRY` | Override install/update telemetry and provider attribution headers: `1`/`true`/`yes` or `0`/`false`/`no`. This does not disable update checks |
| `PI_CACHE_RETENTION` | Set to `long` for extended prompt cache where supported |
| `VISUAL`, `EDITOR` | External editor for Ctrl+G |

1
packages/coding-agent/src/cli/args.ts
View File

@@ -335,6 +335,7 @@ ${chalk.bold("Environment Variables:")}
AZURE_OPENAI_API_VERSION - Azure OpenAI API version (default: v1)
AZURE_OPENAI_DEPLOYMENT_NAME_MAP - Azure OpenAI model=deployment map (comma-separated)
DEEPSEEK_API_KEY - DeepSeek API key
NVIDIA_API_KEY - NVIDIA NIM API key
GEMINI_API_KEY - Google Gemini API key
GROQ_API_KEY - Groq API key
CEREBRAS_API_KEY - Cerebras API key

1
packages/coding-agent/src/core/agent-session.ts
View File

@@ -2787,6 +2787,7 @@ export class AgentSession {
customInstructions,
replaceInstructions,
reserveTokens: branchSummarySettings.reserveTokens,
streamFn: this.agent.streamFn,
});
if (result.aborted) {
return { cancelled: true, aborted: true };

31
packages/coding-agent/src/core/compaction/branch-summarization.ts
View File

@@ -5,8 +5,8 @@
* a summary of the branch being left so context isn't lost.
*/
import type { AgentMessage } from "@earendil-works/pi-agent-core";
import type { Model } from "@earendil-works/pi-ai";
import type { AgentMessage, StreamFn } from "@earendil-works/pi-agent-core";
import type { Model, SimpleStreamOptions } from "@earendil-works/pi-ai";
import { completeSimple } from "@earendil-works/pi-ai";
import {
convertToLlm,
@@ -77,6 +77,8 @@ export interface GenerateBranchSummaryOptions {
replaceInstructions?: boolean;
/** Tokens reserved for prompt + LLM response (default 16384) */
reserveTokens?: number;
/** Optional session stream function. Used to preserve SDK request behavior without mutating agent state. */
streamFn?: StreamFn;
}
// ============================================================================
@@ -284,7 +286,16 @@ export async function generateBranchSummary(
entries: SessionEntry[],
options: GenerateBranchSummaryOptions,
): Promise<BranchSummaryResult> {
const { model, apiKey, headers, signal, customInstructions, replaceInstructions, reserveTokens = 16384 } = options;
const {
model,
apiKey,
headers,
signal,
customInstructions,
replaceInstructions,
reserveTokens = 16384,
streamFn,
} = options;
// Token budget = context window minus reserved space for prompt + response
const contextWindow = model.contextWindow || 128000;
@@ -320,12 +331,14 @@ export async function generateBranchSummary(
},
];
// Call LLM for summarization
const response = await completeSimple(
model,
{ systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages },
{ apiKey, headers, signal, maxTokens: 2048 },
);
// Call LLM for summarization. Prefer the session stream function so SDK
// request behavior (timeouts, retries, attribution headers) stays consistent
// without running through agent state/events.
const context = { systemPrompt: SUMMARIZATION_SYSTEM_PROMPT, messages: summarizationMessages };
const requestOptions: SimpleStreamOptions = { apiKey, headers, signal, maxTokens: 2048 };
const response = streamFn
? await (await streamFn(model, context, requestOptions)).result()
: await completeSimple(model, context, requestOptions);
// Check if aborted or errored
if (response.stopReason === "aborted") {

1
packages/coding-agent/src/core/model-resolver.ts
View File

@@ -17,6 +17,7 @@ export const defaultModelPerProvider: Record<KnownProvider, string> = {
openai: "gpt-5.4",
"azure-openai-responses": "gpt-5.4",
"openai-codex": "gpt-5.5",
nvidia: "nvidia/nemotron-3-super-120b-a12b",
deepseek: "deepseek-v4-pro",
google: "gemini-3.1-pro-preview",
"google-vertex": "gemini-3.1-pro-preview",

97
packages/coding-agent/src/core/provider-attribution.ts Normal file
View File

@@ -0,0 +1,97 @@
import type { Api, Model } from "@earendil-works/pi-ai";
import type { SettingsManager } from "./settings-manager.ts";
import { isInstallTelemetryEnabled } from "./telemetry.ts";
const OPENROUTER_HOST = "openrouter.ai";
const NVIDIA_NIM_HOST = "integrate.api.nvidia.com";
const CLOUDFLARE_API_HOST = "api.cloudflare.com";
const CLOUDFLARE_AI_GATEWAY_HOST = "gateway.ai.cloudflare.com";
const OPENCODE_HOST = "opencode.ai";
function matchesHost(baseUrl: string, expectedHost: string): boolean {
try {
return new URL(baseUrl).hostname === expectedHost;
} catch {
return false;
}
}
function isOpenRouterModel(model: Model<Api>): boolean {
return model.provider === "openrouter" || model.baseUrl.includes(OPENROUTER_HOST);
}
function isNvidiaNimModel(model: Model<Api>): boolean {
return model.provider === "nvidia" || matchesHost(model.baseUrl, NVIDIA_NIM_HOST);
}
function isCloudflareModel(model: Model<Api>): boolean {
return (
model.provider === "cloudflare-workers-ai" ||
model.provider === "cloudflare-ai-gateway" ||
matchesHost(model.baseUrl, CLOUDFLARE_API_HOST) ||
matchesHost(model.baseUrl, CLOUDFLARE_AI_GATEWAY_HOST)
);
}
function getDefaultAttributionHeaders(
model: Model<Api>,
settingsManager: SettingsManager,
): Record<string, string> | undefined {
if (!isInstallTelemetryEnabled(settingsManager)) {
return undefined;
}
if (isOpenRouterModel(model)) {
return {
"HTTP-Referer": "https://pi.dev",
"X-OpenRouter-Title": "pi",
"X-OpenRouter-Categories": "cli-agent",
};
}
if (isNvidiaNimModel(model)) {
return {
"X-BILLING-INVOKE-ORIGIN": "Pi",
};
}
if (isCloudflareModel(model)) {
return {
"User-Agent": "pi-coding-agent",
};
}
return undefined;
}
function getSessionHeaders(model: Model<Api>, sessionId: string | undefined): Record<string, string> | undefined {
if (!sessionId) return undefined;
if (
model.provider !== "opencode" &&
model.provider !== "opencode-go" &&
!matchesHost(model.baseUrl, OPENCODE_HOST)
) {
return undefined;
}
return { "x-opencode-session": sessionId, "x-opencode-client": "pi" };
}
export function mergeProviderAttributionHeaders(
model: Model<Api>,
settingsManager: SettingsManager,
sessionId: string | undefined,
...headerSources: Array<Record<string, string> | undefined>
): Record<string, string> | undefined {
const merged = {
...getSessionHeaders(model, sessionId),
...getDefaultAttributionHeaders(model, settingsManager),
};
for (const headers of headerSources) {
if (headers) {
Object.assign(merged, headers);
}
}
return Object.keys(merged).length > 0 ? merged : undefined;
}

1
packages/coding-agent/src/core/provider-display-names.ts
View File

@@ -17,6 +17,7 @@ export const BUILT_IN_PROVIDER_DISPLAY_NAMES: Record<string, string> = {
"minimax-cn": "MiniMax (China)",
moonshotai: "Moonshot AI",
"moonshotai-cn": "Moonshot AI (China)",
nvidia: "NVIDIA NIM",
opencode: "OpenCode Zen",
"opencode-go": "OpenCode Go",
openai: "OpenAI",

52
packages/coding-agent/src/core/sdk.ts
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) => {

66
packages/coding-agent/test/sdk-openrouter-attribution.test.ts
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",