alexion
4def357d09
Imported as a linear vendor commit so forge rebases do not replay the upstream root tree. git-subtree-dir: third_party/pi git-subtree-split: d981de1229ef899957bbe968bc8dcda02a21f477
96 lines
3.2 KiB
TypeScript
96 lines
3.2 KiB
TypeScript
import type {
|
|
Api,
|
|
Context,
|
|
Model,
|
|
SimpleStreamOptions,
|
|
StreamOptions,
|
|
ThinkingBudgets,
|
|
ThinkingLevel,
|
|
} from "../types.ts";
|
|
import { estimateContextTokens } from "../utils/estimate.ts";
|
|
|
|
const CONTEXT_SAFETY_TOKENS = 4096;
|
|
const MIN_MAX_TOKENS = 1;
|
|
|
|
export function clampMaxTokensToContext(model: Model<Api>, context: Context, maxTokens: number): number {
|
|
if (model.contextWindow <= 0) return Math.max(MIN_MAX_TOKENS, maxTokens);
|
|
const available = model.contextWindow - estimateContextTokens(context).tokens - CONTEXT_SAFETY_TOKENS;
|
|
return Math.min(maxTokens, Math.max(MIN_MAX_TOKENS, available));
|
|
}
|
|
|
|
export function buildBaseOptions(
|
|
model: Model<Api>,
|
|
context: Context,
|
|
options?: SimpleStreamOptions,
|
|
apiKey?: string,
|
|
): StreamOptions {
|
|
const samplingParams =
|
|
model.samplingParams || options?.samplingParams
|
|
? { ...model.samplingParams, ...options?.samplingParams }
|
|
: undefined;
|
|
return {
|
|
temperature: options?.temperature,
|
|
samplingParams,
|
|
maxTokens: clampMaxTokensToContext(model, context, options?.maxTokens ?? model.maxTokens),
|
|
signal: options?.signal,
|
|
telemetryContext: options?.telemetryContext,
|
|
apiKey: apiKey || options?.apiKey,
|
|
fetch: options?.fetch,
|
|
transport: options?.transport,
|
|
cacheRetention: options?.cacheRetention,
|
|
sessionId: options?.sessionId,
|
|
headers: options?.headers,
|
|
onPayload: options?.onPayload,
|
|
onResponse: options?.onResponse,
|
|
timeoutMs: options?.timeoutMs,
|
|
websocketConnectTimeoutMs: options?.websocketConnectTimeoutMs,
|
|
maxRetries: options?.maxRetries,
|
|
maxRetryDelayMs: options?.maxRetryDelayMs,
|
|
metadata: options?.metadata,
|
|
env: options?.env,
|
|
};
|
|
}
|
|
|
|
/** Tokens always left for the answer when a thinking budget shares the response ceiling. */
|
|
export const MIN_ANSWER_TOKENS = 1024;
|
|
|
|
export const DEFAULT_THINKING_BUDGETS: ThinkingBudgets = {
|
|
minimal: 1024,
|
|
low: 2048,
|
|
medium: 8192,
|
|
high: 16384,
|
|
};
|
|
|
|
export function clampReasoning(effort: ThinkingLevel | undefined): Exclude<ThinkingLevel, "xhigh" | "max"> | undefined {
|
|
return effort === "xhigh" || effort === "max" ? "high" : effort;
|
|
}
|
|
|
|
export function thinkingBudgetForLevel(reasoningLevel: ThinkingLevel, customBudgets?: ThinkingBudgets): number {
|
|
const budgets = { ...DEFAULT_THINKING_BUDGETS, ...customBudgets };
|
|
const level = clampReasoning(reasoningLevel)!;
|
|
return budgets[level]!;
|
|
}
|
|
|
|
/** Cap a thinking budget so at least MIN_ANSWER_TOKENS remain under a shared response ceiling. */
|
|
export function clampThinkingBudgetToAnswerRoom(thinkingBudget: number, ceiling: number): number {
|
|
return Math.min(thinkingBudget, Math.max(0, ceiling - MIN_ANSWER_TOKENS));
|
|
}
|
|
|
|
export function adjustMaxTokensForThinking(
|
|
// Undefined means no explicit caller cap. Use the model cap and fit thinking inside it.
|
|
baseMaxTokens: number | undefined,
|
|
modelMaxTokens: number,
|
|
reasoningLevel: ThinkingLevel,
|
|
customBudgets?: ThinkingBudgets,
|
|
): { maxTokens: number; thinkingBudget: number } {
|
|
let thinkingBudget = thinkingBudgetForLevel(reasoningLevel, customBudgets);
|
|
const maxTokens =
|
|
baseMaxTokens === undefined ? modelMaxTokens : Math.min(baseMaxTokens + thinkingBudget, modelMaxTokens);
|
|
|
|
if (maxTokens <= thinkingBudget) {
|
|
thinkingBudget = clampThinkingBudgetToAnswerRoom(thinkingBudget, maxTokens);
|
|
}
|
|
|
|
return { maxTokens, thinkingBudget };
|
|
}
|