Add NVIDIA NIM provider
This commit is contained in:
@@ -2,6 +2,10 @@
|
||||
|
||||
## [Unreleased]
|
||||
|
||||
### Added
|
||||
|
||||
- Added NVIDIA NIM as a built-in OpenAI-compatible provider, exposing public NIM models that support tool use.
|
||||
|
||||
### Fixed
|
||||
|
||||
- Fixed Amazon Bedrock requests to replace blank required user/tool-result text with a placeholder and skip blank replay text blocks ([#4975](https://github.com/earendil-works/pi/issues/4975)).
|
||||
|
||||
@@ -54,6 +54,7 @@ Unified LLM API with automatic model discovery, provider configuration, token an
|
||||
- **Azure OpenAI (Responses)**
|
||||
- **OpenAI Codex** (ChatGPT Plus/Pro subscription, requires OAuth, see below)
|
||||
- **DeepSeek**
|
||||
- **NVIDIA NIM**
|
||||
- **Anthropic**
|
||||
- **Google**
|
||||
- **Vertex AI** (Gemini via Vertex AI)
|
||||
@@ -801,7 +802,7 @@ A **provider** offers models through a specific API. For example:
|
||||
- **Google** models use the `google-generative-ai` API
|
||||
- **OpenAI** models use the `openai-responses` API
|
||||
- **Mistral** models use the `mistral-conversations` API
|
||||
- **xAI, Cerebras, Groq, Together AI, etc.** models use the `openai-completions` API (OpenAI-compatible)
|
||||
- **xAI, Cerebras, Groq, NVIDIA NIM, Together AI, etc.** models use the `openai-completions` API (OpenAI-compatible)
|
||||
|
||||
### Querying Providers and Models
|
||||
|
||||
@@ -923,7 +924,7 @@ const ollamaReasoningModel: Model<'openai-completions'> = {
|
||||
|
||||
### OpenAI Compatibility Settings
|
||||
|
||||
The `openai-completions` API is implemented by many providers with minor differences. By default, the library auto-detects compatibility settings based on `baseUrl` for a small set of known OpenAI-compatible providers (Cerebras, xAI, Chutes, DeepSeek, Together AI, zAi, OpenCode, Cloudflare Workers AI, etc.). For custom proxies or unknown endpoints, you can override these settings via the `compat` field. For `openai-responses` models, the compat field only supports Responses-specific flags.
|
||||
The `openai-completions` API is implemented by many providers with minor differences. By default, the library auto-detects compatibility settings based on `baseUrl` for a small set of known OpenAI-compatible providers (Cerebras, xAI, Chutes, DeepSeek, NVIDIA NIM, Together AI, zAi, OpenCode, Cloudflare Workers AI, etc.). For custom proxies or unknown endpoints, you can override these settings via the `compat` field. For `openai-responses` models, the compat field only supports Responses-specific flags.
|
||||
|
||||
```typescript
|
||||
interface OpenAICompletionsCompat {
|
||||
@@ -1102,6 +1103,7 @@ In Node.js environments, you can set environment variables to avoid passing API
|
||||
| Azure OpenAI | `AZURE_OPENAI_API_KEY` + `AZURE_OPENAI_BASE_URL` (e.g. `https://{resource}.openai.azure.com`) or `AZURE_OPENAI_RESOURCE_NAME`. Supports `*.openai.azure.com` and `*.cognitiveservices.azure.com`; root endpoints auto-normalize to `/openai/v1`. Optional: `AZURE_OPENAI_API_VERSION` (default `v1`), `AZURE_OPENAI_DEPLOYMENT_NAME_MAP`. |
|
||||
| Anthropic | `ANTHROPIC_API_KEY` or `ANTHROPIC_OAUTH_TOKEN` |
|
||||
| DeepSeek | `DEEPSEEK_API_KEY` |
|
||||
| NVIDIA NIM | `NVIDIA_API_KEY` |
|
||||
| Google | `GEMINI_API_KEY` |
|
||||
| Vertex AI | `GOOGLE_CLOUD_API_KEY` or `GOOGLE_CLOUD_PROJECT` (or `GCLOUD_PROJECT`) + `GOOGLE_CLOUD_LOCATION` + ADC |
|
||||
| Mistral | `MISTRAL_API_KEY` |
|
||||
|
||||
@@ -32,12 +32,17 @@ interface ModelsDevModel {
|
||||
};
|
||||
modalities?: {
|
||||
input?: string[];
|
||||
output?: string[];
|
||||
};
|
||||
provider?: {
|
||||
npm?: string;
|
||||
};
|
||||
}
|
||||
|
||||
interface NvidiaNimModelListItem {
|
||||
id: string;
|
||||
}
|
||||
|
||||
interface AiGatewayModel {
|
||||
id: string;
|
||||
name?: string;
|
||||
@@ -117,6 +122,39 @@ const TOGETHER_TOGGLE_REASONING_LEVEL_MAP = {
|
||||
|
||||
const AI_GATEWAY_MODELS_URL = "https://ai-gateway.vercel.sh/v1";
|
||||
const AI_GATEWAY_BASE_URL = "https://ai-gateway.vercel.sh";
|
||||
const NVIDIA_BASE_URL = "https://integrate.api.nvidia.com/v1";
|
||||
const NVIDIA_HEADERS = {
|
||||
"NVCF-POLL-SECONDS": "3600",
|
||||
} as const;
|
||||
const NVIDIA_OPENAI_COMPAT: OpenAICompletionsCompat = {
|
||||
supportsStore: false,
|
||||
supportsDeveloperRole: false,
|
||||
supportsReasoningEffort: false,
|
||||
maxTokensField: "max_tokens",
|
||||
supportsStrictMode: false,
|
||||
supportsLongCacheRetention: false,
|
||||
};
|
||||
const NVIDIA_NIM_UNSUPPORTED_MODELS = new Set([
|
||||
"abacusai/dracarys-llama-3.1-70b-instruct",
|
||||
"bytedance/seed-oss-36b-instruct",
|
||||
"deepseek-ai/deepseek-v4-flash",
|
||||
"deepseek-ai/deepseek-v4-pro",
|
||||
"google/gemma-2-2b-it",
|
||||
"google/gemma-3n-e2b-it",
|
||||
"google/gemma-3n-e4b-it",
|
||||
"google/gemma-4-31b-it",
|
||||
"meta/llama-3.2-1b-instruct",
|
||||
"meta/llama-4-maverick-17b-128e-instruct",
|
||||
"microsoft/phi-4-mini-instruct",
|
||||
"minimaxai/minimax-m2.7",
|
||||
"mistralai/mistral-nemotron",
|
||||
"nvidia/nemotron-mini-4b-instruct",
|
||||
"qwen/qwen3-next-80b-a3b-instruct",
|
||||
"qwen/qwen3.5-122b-a10b",
|
||||
"qwen/qwen3.5-397b-a17b",
|
||||
"sarvamai/sarvam-m",
|
||||
"upstage/solar-10.7b-instruct",
|
||||
]);
|
||||
const ZAI_TOOL_STREAM_UNSUPPORTED_MODELS = new Set(["glm-4.5", "glm-4.5-air", "glm-4.5-flash", "glm-4.5v"]);
|
||||
const EAGER_TOOL_INPUT_STREAMING_UNSUPPORTED_ANTHROPIC_MODELS = new Set([
|
||||
"github-copilot:claude-haiku-4.5",
|
||||
@@ -312,6 +350,30 @@ function getBedrockBaseUrl(modelId: string): string {
|
||||
: "https://bedrock-runtime.us-east-1.amazonaws.com";
|
||||
}
|
||||
|
||||
function normalizeNvidiaModelId(modelId: string): string {
|
||||
return modelId.toLowerCase().replaceAll("_", ".");
|
||||
}
|
||||
|
||||
async function fetchNvidiaNimModelIds(): Promise<Map<string, string>> {
|
||||
try {
|
||||
console.log("Fetching models from NVIDIA NIM API...");
|
||||
const response = await fetch(`${NVIDIA_BASE_URL}/models`);
|
||||
const data = (await response.json()) as { data?: NvidiaNimModelListItem[] };
|
||||
const modelIds = new Map<string, string>();
|
||||
|
||||
for (const model of data.data ?? []) {
|
||||
modelIds.set(model.id, model.id);
|
||||
modelIds.set(normalizeNvidiaModelId(model.id), model.id);
|
||||
}
|
||||
|
||||
console.log(`Fetched ${data.data?.length ?? 0} model IDs from NVIDIA NIM`);
|
||||
return modelIds;
|
||||
} catch (error) {
|
||||
console.error("Failed to fetch NVIDIA NIM models:", error);
|
||||
return new Map();
|
||||
}
|
||||
}
|
||||
|
||||
async function fetchOpenRouterModels(): Promise<Model<any>[]> {
|
||||
try {
|
||||
console.log("Fetching models from OpenRouter API...");
|
||||
@@ -435,6 +497,7 @@ async function loadModelsDevData(): Promise<Model<any>[]> {
|
||||
const data = await response.json();
|
||||
|
||||
const models: Model<any>[] = [];
|
||||
const nvidiaNimModelIds = data.nvidia?.models ? await fetchNvidiaNimModelIds() : new Map<string, string>();
|
||||
|
||||
// Process Amazon Bedrock models
|
||||
if (data["amazon-bedrock"]?.models) {
|
||||
@@ -836,6 +899,40 @@ async function loadModelsDevData(): Promise<Model<any>[]> {
|
||||
}
|
||||
}
|
||||
|
||||
// Process NVIDIA NIM models
|
||||
if (data.nvidia?.models) {
|
||||
for (const [modelId, model] of Object.entries(data.nvidia.models)) {
|
||||
const m = model as ModelsDevModel;
|
||||
if (m.tool_call !== true) continue;
|
||||
if (!m.modalities?.input?.includes("text")) continue;
|
||||
if (!m.modalities?.output?.includes("text")) continue;
|
||||
|
||||
const liveModelId = nvidiaNimModelIds.get(modelId) ?? nvidiaNimModelIds.get(normalizeNvidiaModelId(modelId));
|
||||
if (!liveModelId) continue;
|
||||
if (NVIDIA_NIM_UNSUPPORTED_MODELS.has(liveModelId)) continue;
|
||||
|
||||
models.push({
|
||||
id: liveModelId,
|
||||
name: m.name || liveModelId,
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: NVIDIA_BASE_URL,
|
||||
headers: { ...NVIDIA_HEADERS },
|
||||
reasoning: m.reasoning === true,
|
||||
input: m.modalities.input.includes("image") ? ["text", "image"] : ["text"],
|
||||
cost: {
|
||||
input: m.cost?.input || 0,
|
||||
output: m.cost?.output || 0,
|
||||
cacheRead: m.cost?.cache_read || 0,
|
||||
cacheWrite: m.cost?.cache_write || 0,
|
||||
},
|
||||
compat: NVIDIA_OPENAI_COMPAT,
|
||||
contextWindow: m.limit?.context || 4096,
|
||||
maxTokens: m.limit?.output || 4096,
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
// Process Together AI models
|
||||
const togetherProvider = data.together ?? data.togetherai ?? data["together-ai"];
|
||||
if (togetherProvider?.models) {
|
||||
|
||||
@@ -101,6 +101,7 @@ function getApiKeyEnvVars(provider: string): readonly string[] | undefined {
|
||||
const envMap: Record<string, string> = {
|
||||
openai: "OPENAI_API_KEY",
|
||||
"azure-openai-responses": "AZURE_OPENAI_API_KEY",
|
||||
nvidia: "NVIDIA_API_KEY",
|
||||
deepseek: "DEEPSEEK_API_KEY",
|
||||
google: "GEMINI_API_KEY",
|
||||
"google-vertex": "GOOGLE_CLOUD_API_KEY",
|
||||
|
||||
@@ -6335,6 +6335,369 @@ export const MODELS = {
|
||||
maxTokens: 262144,
|
||||
} satisfies Model<"openai-completions">,
|
||||
},
|
||||
"nvidia": {
|
||||
"meta/llama-3.1-70b-instruct": {
|
||||
id: "meta/llama-3.1-70b-instruct",
|
||||
name: "Llama 3.1 70b Instruct",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"meta/llama-3.1-8b-instruct": {
|
||||
id: "meta/llama-3.1-8b-instruct",
|
||||
name: "Llama 3.1 8B Instruct",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 16000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"meta/llama-3.2-11b-vision-instruct": {
|
||||
id: "meta/llama-3.2-11b-vision-instruct",
|
||||
name: "Llama 3.2 11b Vision Instruct",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: false,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"meta/llama-3.2-90b-vision-instruct": {
|
||||
id: "meta/llama-3.2-90b-vision-instruct",
|
||||
name: "Llama-3.2-90B-Vision-Instruct",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: false,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"meta/llama-3.3-70b-instruct": {
|
||||
id: "meta/llama-3.3-70b-instruct",
|
||||
name: "Llama 3.3 70b Instruct",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 4096,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"mistralai/mistral-large-3-675b-instruct-2512": {
|
||||
id: "mistralai/mistral-large-3-675b-instruct-2512",
|
||||
name: "Mistral Large 3 675B Instruct 2512",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: false,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 262144,
|
||||
maxTokens: 262144,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"mistralai/mistral-small-4-119b-2603": {
|
||||
id: "mistralai/mistral-small-4-119b-2603",
|
||||
name: "mistral-small-4-119b-2603",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 128000,
|
||||
maxTokens: 8192,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"moonshotai/kimi-k2.6": {
|
||||
id: "moonshotai/kimi-k2.6",
|
||||
name: "Kimi K2.6",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 262144,
|
||||
maxTokens: 262144,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"nvidia/llama-3.3-nemotron-super-49b-v1": {
|
||||
id: "nvidia/llama-3.3-nemotron-super-49b-v1",
|
||||
name: "Llama 3.3 Nemotron Super 49B v1",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 131072,
|
||||
maxTokens: 131072,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"nvidia/llama-3.3-nemotron-super-49b-v1.5": {
|
||||
id: "nvidia/llama-3.3-nemotron-super-49b-v1.5",
|
||||
name: "Llama 3.3 Nemotron Super 49B v1.5",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 131072,
|
||||
maxTokens: 131072,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"nvidia/nemotron-3-nano-30b-a3b": {
|
||||
id: "nvidia/nemotron-3-nano-30b-a3b",
|
||||
name: "nemotron-3-nano-30b-a3b",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 131072,
|
||||
maxTokens: 131072,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"nvidia/nemotron-3-nano-omni-30b-a3b-reasoning": {
|
||||
id: "nvidia/nemotron-3-nano-omni-30b-a3b-reasoning",
|
||||
name: "Nemotron 3 Nano Omni",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 256000,
|
||||
maxTokens: 65536,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"nvidia/nemotron-3-super-120b-a12b": {
|
||||
id: "nvidia/nemotron-3-super-120b-a12b",
|
||||
name: "Nemotron 3 Super",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0.2,
|
||||
output: 0.8,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 262144,
|
||||
maxTokens: 262144,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"nvidia/nvidia-nemotron-nano-9b-v2": {
|
||||
id: "nvidia/nvidia-nemotron-nano-9b-v2",
|
||||
name: "nvidia-nemotron-nano-9b-v2",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 131072,
|
||||
maxTokens: 131072,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"openai/gpt-oss-20b": {
|
||||
id: "openai/gpt-oss-20b",
|
||||
name: "GPT OSS 20B",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 131072,
|
||||
maxTokens: 32768,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"qwen/qwen3-coder-480b-a35b-instruct": {
|
||||
id: "qwen/qwen3-coder-480b-a35b-instruct",
|
||||
name: "Qwen3 Coder 480B A35B Instruct",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: false,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 262144,
|
||||
maxTokens: 66536,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"stepfun-ai/step-3.5-flash": {
|
||||
id: "stepfun-ai/step-3.5-flash",
|
||||
name: "Step 3.5 Flash",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 256000,
|
||||
maxTokens: 16384,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"stepfun-ai/step-3.7-flash": {
|
||||
id: "stepfun-ai/step-3.7-flash",
|
||||
name: "Step 3.7 Flash",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text", "image"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 256000,
|
||||
maxTokens: 16384,
|
||||
} satisfies Model<"openai-completions">,
|
||||
"z-ai/glm-5.1": {
|
||||
id: "z-ai/glm-5.1",
|
||||
name: "GLM-5.1",
|
||||
api: "openai-completions",
|
||||
provider: "nvidia",
|
||||
baseUrl: "https://integrate.api.nvidia.com/v1",
|
||||
headers: {"NVCF-POLL-SECONDS":"3600"},
|
||||
compat: {"supportsStore":false,"supportsDeveloperRole":false,"supportsReasoningEffort":false,"maxTokensField":"max_tokens","supportsStrictMode":false,"supportsLongCacheRetention":false},
|
||||
reasoning: true,
|
||||
input: ["text"],
|
||||
cost: {
|
||||
input: 0,
|
||||
output: 0,
|
||||
cacheRead: 0,
|
||||
cacheWrite: 0,
|
||||
},
|
||||
contextWindow: 131072,
|
||||
maxTokens: 131072,
|
||||
} satisfies Model<"openai-completions">,
|
||||
},
|
||||
"openai": {
|
||||
"gpt-4": {
|
||||
id: "gpt-4",
|
||||
|
||||
@@ -1078,8 +1078,10 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
|
||||
const isOpenRouter = provider === "openrouter" || baseUrl.includes("openrouter.ai");
|
||||
const isCloudflareWorkersAI = provider === "cloudflare-workers-ai" || baseUrl.includes("api.cloudflare.com");
|
||||
const isCloudflareAiGateway = provider === "cloudflare-ai-gateway" || baseUrl.includes("gateway.ai.cloudflare.com");
|
||||
const isNvidia = provider === "nvidia" || baseUrl.includes("integrate.api.nvidia.com");
|
||||
|
||||
const isNonStandard =
|
||||
isNvidia ||
|
||||
provider === "cerebras" ||
|
||||
baseUrl.includes("cerebras.ai") ||
|
||||
provider === "xai" ||
|
||||
@@ -1094,7 +1096,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
|
||||
isCloudflareWorkersAI ||
|
||||
isCloudflareAiGateway;
|
||||
|
||||
const useMaxTokens = baseUrl.includes("chutes.ai") || isMoonshot || isCloudflareAiGateway || isTogether;
|
||||
const useMaxTokens = baseUrl.includes("chutes.ai") || isMoonshot || isCloudflareAiGateway || isTogether || isNvidia;
|
||||
|
||||
const isGrok = provider === "xai" || baseUrl.includes("api.x.ai");
|
||||
const isDeepSeek = provider === "deepseek" || baseUrl.includes("deepseek.com");
|
||||
@@ -1103,7 +1105,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
|
||||
return {
|
||||
supportsStore: !isNonStandard,
|
||||
supportsDeveloperRole: !isNonStandard && !isOpenRouter,
|
||||
supportsReasoningEffort: !isGrok && !isZai && !isMoonshot && !isTogether && !isCloudflareAiGateway,
|
||||
supportsReasoningEffort: !isGrok && !isZai && !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia,
|
||||
supportsUsageInStreaming: true,
|
||||
maxTokensField: useMaxTokens ? "max_tokens" : "max_completion_tokens",
|
||||
requiresToolResultName: false,
|
||||
@@ -1122,10 +1124,10 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
|
||||
openRouterRouting: {},
|
||||
vercelGatewayRouting: {},
|
||||
zaiToolStream: false,
|
||||
supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway,
|
||||
supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia,
|
||||
cacheControlFormat,
|
||||
sendSessionAffinityHeaders: false,
|
||||
supportsLongCacheRetention: !(isTogether || isCloudflareWorkersAI || isCloudflareAiGateway),
|
||||
supportsLongCacheRetention: !(isTogether || isCloudflareWorkersAI || isCloudflareAiGateway || isNvidia),
|
||||
};
|
||||
}
|
||||
|
||||
|
||||
@@ -28,6 +28,7 @@ export type KnownProvider =
|
||||
| "openai"
|
||||
| "azure-openai-responses"
|
||||
| "openai-codex"
|
||||
| "nvidia"
|
||||
| "deepseek"
|
||||
| "github-copilot"
|
||||
| "xai"
|
||||
|
||||
@@ -787,6 +787,30 @@ describe("Generate E2E Tests", () => {
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.NVIDIA_API_KEY)("NVIDIA NIM Provider (Nemotron 3 Super via OpenAI Completions)", () => {
|
||||
const llm = getModel("nvidia", "nvidia/nemotron-3-super-120b-a12b");
|
||||
|
||||
it("should complete basic text generation", { retry: 3 }, async () => {
|
||||
await basicTextGeneration(llm);
|
||||
});
|
||||
|
||||
it("should handle tool calling", { retry: 3 }, async () => {
|
||||
await handleToolCall(llm);
|
||||
});
|
||||
|
||||
it("should handle streaming", { retry: 3 }, async () => {
|
||||
await handleStreaming(llm);
|
||||
});
|
||||
|
||||
it("should handle thinking mode", { retry: 3 }, async () => {
|
||||
await handleThinking(llm, { reasoningEffort: "high" });
|
||||
});
|
||||
|
||||
it("should handle multi-turn with thinking and tools", { retry: 3 }, async () => {
|
||||
await multiTurn(llm, { reasoningEffort: "high" });
|
||||
});
|
||||
});
|
||||
|
||||
describe.skipIf(!process.env.OPENROUTER_API_KEY)("OpenRouter Provider (glm-4.5v via OpenAI Completions)", () => {
|
||||
const llm = getModel("openrouter", "z-ai/glm-4.5v");
|
||||
|
||||
|
||||
@@ -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.
|
||||
|
||||
|
||||
@@ -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 |
|
||||
|
||||
|
||||
@@ -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": "..." },
|
||||
|
||||
@@ -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 |
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -2787,6 +2787,7 @@ export class AgentSession {
|
||||
customInstructions,
|
||||
replaceInstructions,
|
||||
reserveTokens: branchSummarySettings.reserveTokens,
|
||||
streamFn: this.agent.streamFn,
|
||||
});
|
||||
if (result.aborted) {
|
||||
return { cancelled: true, aborted: true };
|
||||
|
||||
@@ -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") {
|
||||
|
||||
@@ -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
97
packages/coding-agent/src/core/provider-attribution.ts
Normal 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;
|
||||
}
|
||||
@@ -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",
|
||||
|
||||
@@ -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) => {
|
||||
|
||||
@@ -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",
|
||||
|
||||
Reference in New Issue
Block a user