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

@@ -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") {