ton/pi_harness
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ai): clamp streamSimple max tokens

Browse Source 
Clamps streamSimple max-token defaults against estimated context, addressing #5595.

closes #6061
This commit is contained in:
Mario Zechner
2026-06-25 14:10:05 +02:00
parent 49956a7cd0
commit 09f1059575
13 changed files with 198 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
packages/ai/src/api/anthropic-messages.ts
+6 -4
View File
@@ -36,7 +36,7 @@ import { getProviderEnvValue } from "../utils/provider-env.ts";
import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts";
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.ts";
import { adjustMaxTokensForThinking, buildBaseOptions } from "./simple-options.ts";
import { adjustMaxTokensForThinking, buildBaseOptions, clampMaxTokensToContext } from "./simple-options.ts";
import { transformMessages } from "./transform-messages.ts";
/**
@@ -771,7 +771,7 @@ export const streamSimple: StreamFunction<"anthropic-messages", SimpleStreamOpti
): AssistantMessageEventStream => {
assertRequestAuth(model.provider, options?.apiKey, options?.headers);
const base = buildBaseOptions(model, options, options?.apiKey);
const base = buildBaseOptions(model, context, options, options?.apiKey);
if (!options?.reasoning) {
return stream(model, context, { ...base, thinkingEnabled: false } satisfies AnthropicOptions);
}
@@ -796,11 +796,13 @@ export const streamSimple: StreamFunction<"anthropic-messages", SimpleStreamOpti
options.thinkingBudgets,
);
const maxTokens = clampMaxTokensToContext(model, context, adjusted.maxTokens);
return stream(model, context, {
...base,
maxTokens: adjusted.maxTokens,
maxTokens,
thinkingEnabled: true,
thinkingBudgetTokens: adjusted.thinkingBudget,
thinkingBudgetTokens: Math.min(adjusted.thinkingBudget, Math.max(0, maxTokens - 1024)),
} satisfies AnthropicOptions);
};
packages/ai/src/api/azure-openai-responses.ts
+1 -1
View File
@@ -160,7 +160,7 @@ export const streamSimple: StreamFunction<"azure-openai-responses", SimpleStream
throw new Error(`No API key for provider: ${model.provider}`);
}
const base = buildBaseOptions(model, options, apiKey);
const base = buildBaseOptions(model, context, options, apiKey);
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning;
packages/ai/src/api/bedrock-converse-stream.ts
+11 -4
View File
@@ -53,7 +53,12 @@ import { parseStreamingJson } from "../utils/json-parse.ts";
import { resolveHttpProxyUrlForTarget } from "../utils/node-http-proxy.ts";
import { getProviderEnvValue } from "../utils/provider-env.ts";
import { sanitizeSurrogates } from "../utils/sanitize-unicode.ts";
import { adjustMaxTokensForThinking, buildBaseOptions, clampReasoning } from "./simple-options.ts";
import {
adjustMaxTokensForThinking,
buildBaseOptions,
clampMaxTokensToContext,
clampReasoning,
} from "./simple-options.ts";
import { transformMessages } from "./transform-messages.ts";
export type BedrockThinkingDisplay = "summarized" | "omitted";
@@ -374,7 +379,7 @@ export const streamSimple: StreamFunction<"bedrock-converse-stream", SimpleStrea
context: Context,
options?: SimpleStreamOptions,
): AssistantMessageEventStream => {
const base = buildBaseOptions(model, options, undefined);
const base = buildBaseOptions(model, context, options, undefined);
if (!options?.reasoning) {
return stream(model, context, { ...base, reasoning: undefined } satisfies BedrockOptions);
}
@@ -397,13 +402,15 @@ export const streamSimple: StreamFunction<"bedrock-converse-stream", SimpleStrea
options.thinkingBudgets,
);
const maxTokens = clampMaxTokensToContext(model, context, adjusted.maxTokens);
return stream(model, context, {
...base,
maxTokens: adjusted.maxTokens,
maxTokens,
reasoning: options.reasoning,
thinkingBudgets: {
...(options.thinkingBudgets || {}),
[clampReasoning(options.reasoning)!]: adjusted.thinkingBudget,
[clampReasoning(options.reasoning)!]: Math.min(adjusted.thinkingBudget, Math.max(0, maxTokens - 1024)),
},
} satisfies BedrockOptions);
}
packages/ai/src/api/google-generative-ai.ts
+1 -1
View File
@@ -290,7 +290,7 @@ export const streamSimple: StreamFunction<"google-generative-ai", SimpleStreamOp
throw new Error(`No API key for provider: ${model.provider}`);
}
const base = buildBaseOptions(model, options, apiKey);
const base = buildBaseOptions(model, context, options, apiKey);
if (!options?.reasoning) {
return stream(model, context, { ...base, thinking: { enabled: false } } satisfies GoogleOptions);
}
packages/ai/src/api/google-vertex.ts
+1 -1
View File
@@ -302,7 +302,7 @@ export const streamSimple: StreamFunction<"google-vertex", SimpleStreamOptions>
context: Context,
options?: SimpleStreamOptions,
): AssistantMessageEventStream => {
const base = buildBaseOptions(model, options, undefined);
const base = buildBaseOptions(model, context, options, undefined);
if (!options?.reasoning) {
return stream(model, context, {
...base,
packages/ai/src/api/mistral-conversations.ts
+1 -1
View File
@@ -117,7 +117,7 @@ export const streamSimple: StreamFunction<"mistral-conversations", SimpleStreamO
throw new Error(`No API key for provider: ${model.provider}`);
}
const base = buildBaseOptions(model, options, apiKey);
const base = buildBaseOptions(model, context, options, apiKey);
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
const reasoning = clampedReasoning === "off" ? undefined : clampedReasoning;
const shouldUseReasoning = model.reasoning && reasoning !== undefined;
packages/ai/src/api/openai-codex-responses.ts
+1 -1
View File
@@ -430,7 +430,7 @@ export const streamSimple: StreamFunction<"openai-codex-responses", SimpleStream
throw new Error(`No API key for provider: ${model.provider}`);
}
const base = buildBaseOptions(model, options, apiKey);
const base = buildBaseOptions(model, context, options, apiKey);
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning;
packages/ai/src/api/openai-completions.ts
+1 -1
View File
@@ -482,7 +482,7 @@ export const streamSimple: StreamFunction<"openai-completions", SimpleStreamOpti
): AssistantMessageEventStream => {
getClientApiKey(model.provider, options?.apiKey, options?.headers);
const base = buildBaseOptions(model, options, options?.apiKey);
const base = buildBaseOptions(model, context, options, options?.apiKey);
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning;
const toolChoice = (options as OpenAICompletionsOptions | undefined)?.toolChoice;
packages/ai/src/api/openai-responses.ts
+1 -1
View File
@@ -180,7 +180,7 @@ export const streamSimple: StreamFunction<"openai-responses", SimpleStreamOption
): AssistantMessageEventStream => {
getClientApiKey(model.provider, options?.apiKey, options?.headers);
const base = buildBaseOptions(model, options, options?.apiKey);
const base = buildBaseOptions(model, context, options, options?.apiKey);
const clampedReasoning = options?.reasoning ? clampThinkingLevel(model, options.reasoning) : undefined;
const reasoningEffort = clampedReasoning === "off" ? undefined : clampedReasoning;
packages/ai/src/api/simple-options.ts
+26 -3
View File
@@ -1,9 +1,32 @@
import type { Api, Model, SimpleStreamOptions, StreamOptions, ThinkingBudgets, ThinkingLevel } from "../types.ts";
import type {
Api,
Context,
Model,
SimpleStreamOptions,
StreamOptions,
ThinkingBudgets,
ThinkingLevel,
} from "../types.ts";
import { estimateContextTokens } from "../utils/estimate.ts";
export function buildBaseOptions(_model: Model<Api>, options?: SimpleStreamOptions, apiKey?: string): StreamOptions {
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 {
return {
temperature: options?.temperature,
maxTokens: options?.maxTokens,
maxTokens: clampMaxTokensToContext(model, context, options?.maxTokens ?? model.maxTokens),
signal: options?.signal,
apiKey: apiKey || options?.apiKey,
transport: options?.transport,