sscoble/pi-mono
1
0
Fork 0
mirror of https://github.com/badlogic/pi-mono.git synced 2026-11-05 23:37:25 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 4

feat(ai): generalize openai-completions thinking token budget fields (#8275)

Browse source 
#7638 hardcoded vLLM's thinking_token_budget. Map thinkingBudgets through
thinkingTokenBudgetField, keep the boolean as a vLLM alias, clamp so the
answer keeps 1024 tokens, and document it.

fixes #8274

Co-authored-by: Cursor <cursoragent@cursor.com>
Co-authored-by: Mario Zechner <badlogicgames@gmail.com>
This commit is contained in:
Ben 2026-08-18 17:35:35 +03:00 • committed by GitHub
parent a6c6f80180
commit b237412699
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 218 additions and 76 deletions
Download patch file Download diff file Expand all files Collapse all files

6
packages/ai/README.md
View file

@ -1183,8 +1183,10 @@ interface OpenAICompletionsCompat {
requiresThinkingAsText?: boolean; // Whether thinking blocks must be converted to text (default: false)
requiresReasoningContentOnAssistantMessages?: boolean; // Whether all replayed assistant messages must include empty reasoning_content when reasoning is enabled (default: auto-detected for DeepSeek)
thinkingFormat?: 'openai' | 'openrouter' | 'deepseek' | 'together' | 'baseten' | 'zai' | 'qwen' | 'chat-template' | 'qwen-chat-template' | 'string-thinking' | 'ant-ling'; // Format for reasoning param: 'openai' uses reasoning_effort, 'openrouter' uses reasoning: { effort }, 'deepseek' uses thinking: { type } plus reasoning_effort when supported, 'together' uses reasoning: { enabled } plus reasoning_effort when supported, 'baseten' uses configurable chat_template_args plus reasoning_effort when supported, 'zai' uses thinking: { type }, 'qwen' uses enable_thinking, 'chat-template' uses configurable chat_template_kwargs, 'qwen-chat-template' uses chat_template_kwargs.enable_thinking and preserve_thinking, 'string-thinking' uses top-level thinking, 'ant-ling' uses reasoning: { effort } only for mapped efforts (default: openai)
chatTemplateKwargs?: Record<string, string | number | boolean | null | { '$var': 'thinking.enabled' | 'thinking.effort'; omitWhenOff?: boolean }>; // chat_template_kwargs values; use $var for pi-controlled thinking values
chatTemplateArgs?: Record<string, string | number | boolean | null | { '$var': 'thinking.enabled' | 'thinking.effort'; omitWhenOff?: boolean }>; // chat_template_args values for thinkingFormat: 'baseten'; use $var for pi-controlled thinking values
chatTemplateKwargs?: Record<string, string | number | boolean | null | { '$var': 'thinking.enabled' | 'thinking.effort' | 'thinking.budget'; omitWhenOff?: boolean }>; // chat_template_kwargs values; use $var for pi-controlled thinking values
chatTemplateArgs?: Record<string, string | number | boolean | null | { '$var': 'thinking.enabled' | 'thinking.effort' | 'thinking.budget'; omitWhenOff?: boolean }>; // chat_template_args values for thinkingFormat: 'baseten'; use $var for pi-controlled thinking values
thinkingTokenBudgetField?: 'thinking_token_budget' | 'thinking_budget' | 'thinking_budget_tokens'; // Top-level field that caps reasoning tokens from thinkingBudgets (vLLM / Qwen / llama.cpp). Off by default.
supportsThinkingTokenBudget?: boolean; // Alias for thinkingTokenBudgetField: 'thinking_token_budget' (vLLM). Prefer thinkingTokenBudgetField. Default: false.
cacheControlFormat?: 'anthropic'; // Anthropic-style cache_control on system prompt, last tool, and last user/assistant text content
openRouterRouting?: OpenRouterRouting; // OpenRouter routing preferences (default: {})
vercelGatewayRouting?: VercelGatewayRouting; // Vercel AI Gateway routing preferences (default: {})

7
packages/ai/scripts/generate-models.ts
View file

@ -594,7 +594,12 @@ const OPENAI_COMPLETIONS_DEFAULT_COMPAT = {
supportsOpenAIGrammarTools: false,
sendSessionAffinityHeaders: false,
supportsLongCacheRetention: true,
} satisfies Required<Omit<OpenAICompletionsCompat, "cacheControlFormat" | "deferredToolsMode">> & {
} satisfies Required<
Omit<
OpenAICompletionsCompat,
"cacheControlFormat" | "deferredToolsMode" | "supportsThinkingTokenBudget" | "thinkingTokenBudgetField"
>
> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
};

71
packages/ai/src/api/openai-completions.ts
View file

@ -30,6 +30,7 @@ import type {
TextContent,
ThinkingBudgets,
ThinkingContent,
ThinkingTokenBudgetField,
Tool,
ToolCall,
ToolResultMessage,
@ -54,7 +55,7 @@ import {
} from "./constrained-sampling.ts";
import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.ts";
import { clampOpenAIPromptCacheKey } from "./openai-prompt-cache.ts";
import { buildBaseOptions, clampReasoning, MIN_ANSWER_TOKENS } from "./simple-options.ts";
import { buildBaseOptions, clampThinkingBudgetToAnswerRoom, thinkingBudgetForLevel } from "./simple-options.ts";
import { transformMessages } from "./transform-messages.ts";
/**
@ -144,7 +145,7 @@ function isEncryptedReasoningDetail(detail: unknown): detail is OpenAIEncryptedR
export interface OpenAICompletionsOptions extends StreamOptions {
toolChoice?: OpenAI.Chat.Completions.ChatCompletionToolChoiceOption;
reasoningEffort?: "minimal" | "low" | "medium" | "high" | "xhigh" | "max";
/** Token budgets per thinking level. Only used when `compat.supportsThinkingTokenBudget` is set. */
/** Token budgets per thinking level. Used when `compat.thinkingTokenBudgetField` or `compat.supportsThinkingTokenBudget` is set, or by `{ "$var": "thinking.budget" }`. */
thinkingBudgets?: ThinkingBudgets;
}
@ -159,11 +160,12 @@ interface OpenAICompatCacheControl {
type ResolvedOpenAICompletionsCompat = Omit<
Required<OpenAICompletionsCompat>,
"cacheControlFormat" | "deferredToolsMode" | "supportsThinkingTokenBudget"
"cacheControlFormat" | "deferredToolsMode" | "supportsThinkingTokenBudget" | "thinkingTokenBudgetField"
> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
supportsThinkingTokenBudget?: OpenAICompletionsCompat["supportsThinkingTokenBudget"];
thinkingTokenBudgetField?: OpenAICompletionsCompat["thinkingTokenBudgetField"];
};
type ResolvedChatTemplateKwargValue = string | number | boolean | null;
@ -752,6 +754,9 @@ function buildParams(
params.tool_choice = options.toolChoice;
}
const thinkingTokenBudgetField = resolveThinkingTokenBudgetField(compat);
const thinkingBudget = resolveClampedThinkingBudget(model, options, params);
if (compat.thinkingFormat === "zai" && model.reasoning) {
const zaiParams = params as Omit<typeof params, "reasoning_effort"> & {
thinking?: { type: "enabled" | "disabled"; clear_thinking?: boolean };
@ -779,7 +784,7 @@ function buildParams(
preserve_thinking: true,
};
} else if (compat.thinkingFormat === "chat-template" && model.reasoning) {
const chatTemplateKwargs = buildChatTemplateValues(model, options, compat.chatTemplateKwargs);
const chatTemplateKwargs = buildChatTemplateValues(model, options, compat.chatTemplateKwargs, thinkingBudget);
if (chatTemplateKwargs) {
(params as any).chat_template_kwargs = chatTemplateKwargs;
}
@ -788,7 +793,7 @@ function buildParams(
chat_template_args?: Record<string, ResolvedChatTemplateKwargValue>;
reasoning_effort?: string;
};
const chatTemplateArgs = buildChatTemplateValues(model, options, compat.chatTemplateArgs);
const chatTemplateArgs = buildChatTemplateValues(model, options, compat.chatTemplateArgs, thinkingBudget);
if (chatTemplateArgs) {
basetenParams.chat_template_args = chatTemplateArgs;
}
@ -851,25 +856,12 @@ function buildParams(
}
}
// vLLM caps reasoning with a top-level thinking_token_budget. Independent of
// thinkingFormat: the same server can serve zai, qwen or chat-template models.
// Reasoning and the answer share max_tokens here, so an uncapped reasoning
// phase can consume the whole response and leave no answer and no tool call.
if (compat.supportsThinkingTokenBudget && options?.reasoningEffort && model.reasoning) {
const level = clampReasoning(options.reasoningEffort)!;
const budgets: ThinkingBudgets = {
minimal: 1024,
low: 2048,
medium: 8192,
high: 16384,
...options.thinkingBudgets,
};
const ceiling = (params as { max_tokens?: number }).max_tokens ?? params.max_completion_tokens ?? model.maxTokens;
// Always leave room for the answer, otherwise the budget recreates the bug it prevents.
const budget = Math.min(budgets[level]!, Math.max(0, ceiling - MIN_ANSWER_TOKENS));
if (budget > 0) {
(params as { thinking_token_budget?: number }).thinking_token_budget = budget;
}
// Cap reasoning with a top-level budget field. Independent of thinkingFormat: the
// same server can serve zai, qwen or chat-template models. Reasoning and the answer
// share max_tokens here, so an uncapped reasoning phase can consume the whole
// response and leave no answer and no tool call.
if (thinkingTokenBudgetField && thinkingBudget !== undefined) {
Object.assign(params, { [thinkingTokenBudgetField]: thinkingBudget });
}
// OpenRouter provider routing preferences
@ -896,15 +888,38 @@ function buildParams(
return params;
}
function resolveThinkingTokenBudgetField(
compat: Pick<OpenAICompletionsCompat, "thinkingTokenBudgetField" | "supportsThinkingTokenBudget">,
): ThinkingTokenBudgetField | undefined {
if (compat.thinkingTokenBudgetField) return compat.thinkingTokenBudgetField;
if (compat.supportsThinkingTokenBudget) return "thinking_token_budget";
return undefined;
}
function resolveClampedThinkingBudget(
model: Model<"openai-completions">,
options: OpenAICompletionsOptions | undefined,
params: { max_tokens?: number | null; max_completion_tokens?: number | null },
): number | undefined {
if (!options?.reasoningEffort || !model.reasoning) return undefined;
const ceiling = params.max_tokens ?? params.max_completion_tokens ?? model.maxTokens;
const budget = clampThinkingBudgetToAnswerRoom(
thinkingBudgetForLevel(options.reasoningEffort, options.thinkingBudgets),
ceiling,
);
return budget > 0 ? budget : undefined;
}
function buildChatTemplateValues(
model: Model<"openai-completions">,
options: OpenAICompletionsOptions | undefined,
values: Record<string, ChatTemplateKwargValue>,
thinkingBudget?: number,
): Record<string, ResolvedChatTemplateKwargValue> | undefined {
const resolvedValues: Record<string, ResolvedChatTemplateKwargValue> = {};
for (const [key, value] of Object.entries(values)) {
const resolved = resolveChatTemplateKwargValue(model, options, value);
const resolved = resolveChatTemplateKwargValue(model, options, value, thinkingBudget);
if (resolved !== undefined) {
resolvedValues[key] = resolved;
}
@ -917,6 +932,7 @@ function resolveChatTemplateKwargValue(
model: Model<"openai-completions">,
options: OpenAICompletionsOptions | undefined,
value: ChatTemplateKwargValue,
thinkingBudget?: number,
): ResolvedChatTemplateKwargValue | undefined {
if (typeof value !== "object" || value === null) {
return value;
@ -929,6 +945,9 @@ function resolveChatTemplateKwargValue(
if (value.$var === "thinking.enabled") {
return !!reasoningEffort;
}
if (value.$var === "thinking.budget") {
return thinkingBudget;
}
const mappedValue = reasoningEffort ? model.thinkingLevelMap?.[reasoningEffort] : model.thinkingLevelMap?.off;
return mappedValue === undefined ? reasoningEffort : typeof mappedValue === "string" ? mappedValue : undefined;
@ -1532,6 +1551,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet
chatTemplateArgs: {},
zaiToolStream: false,
supportsThinkingTokenBudget: false,
thinkingTokenBudgetField: undefined,
supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia,
supportsOpenAIGrammarTools: false,
cacheControlFormat,
@ -1577,6 +1597,7 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion
chatTemplateArgs: model.compat.chatTemplateArgs ?? detected.chatTemplateArgs,
zaiToolStream: model.compat.zaiToolStream ?? detected.zaiToolStream,
supportsThinkingTokenBudget: model.compat.supportsThinkingTokenBudget ?? detected.supportsThinkingTokenBudget,
thinkingTokenBudgetField: model.compat.thinkingTokenBudgetField ?? detected.thinkingTokenBudgetField,
supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode,
supportsOpenAIGrammarTools: model.compat.supportsOpenAIGrammarTools ?? detected.supportsOpenAIGrammarTools,
cacheControlFormat: model.compat.cacheControlFormat ?? detected.cacheControlFormat,

31
packages/ai/src/api/simple-options.ts
View file

@ -54,10 +54,28 @@ export function buildBaseOptions(
/** 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,
@ -65,21 +83,12 @@ export function adjustMaxTokensForThinking(
reasoningLevel: ThinkingLevel,
customBudgets?: ThinkingBudgets,
): { maxTokens: number; thinkingBudget: number } {
const defaultBudgets: ThinkingBudgets = {
minimal: 1024,
low: 2048,
medium: 8192,
high: 16384,
};
const budgets = { ...defaultBudgets, ...customBudgets };
const level = clampReasoning(reasoningLevel)!;
let thinkingBudget = budgets[level]!;
let thinkingBudget = thinkingBudgetForLevel(reasoningLevel, customBudgets);
const maxTokens =
baseMaxTokens === undefined ? modelMaxTokens : Math.min(baseMaxTokens + thinkingBudget, modelMaxTokens);
if (maxTokens <= thinkingBudget) {
thinkingBudget = Math.max(0, maxTokens - MIN_ANSWER_TOKENS);
thinkingBudget = clampThinkingBudgetToAnswerRoom(thinkingBudget, maxTokens);
}
return { maxTokens, thinkingBudget };

19
packages/ai/src/types.ts
View file

@ -89,10 +89,13 @@ export type ChatTemplateKwargValue =
| boolean
| null
| {
$var: "thinking.enabled" | "thinking.effort";
$var: "thinking.enabled" | "thinking.effort" | "thinking.budget";
omitWhenOff?: boolean;
};
/** Top-level request field used to cap reasoning tokens on OpenAI-compatible servers. */
export type ThinkingTokenBudgetField = "thinking_token_budget" | "thinking_budget" | "thinking_budget_tokens";
/** Token budgets for each thinking level (token-based providers only) */
export interface ThinkingBudgets {
minimal?: number;
@ -583,9 +586,9 @@ export interface OpenAICompletionsCompat {
| "qwen-chat-template"
| "string-thinking"
| "ant-ling";
/** Kwargs to send as `chat_template_kwargs` when `thinkingFormat` is `chat-template`. Use `{ "$var": "thinking.enabled" }` or `{ "$var": "thinking.effort" }` for pi-controlled thinking values. */
/** Kwargs to send as `chat_template_kwargs` when `thinkingFormat` is `chat-template`. Use `{ "$var": "thinking.enabled" }`, `{ "$var": "thinking.effort" }`, or `{ "$var": "thinking.budget" }` for pi-controlled thinking values. */
chatTemplateKwargs?: Record<string, ChatTemplateKwargValue>;
/** Arguments to send as `chat_template_args` when `thinkingFormat` is `baseten`. Use `{ "$var": "thinking.enabled" }` or `{ "$var": "thinking.effort" }` for pi-controlled thinking values. */
/** Arguments to send as `chat_template_args` when `thinkingFormat` is `baseten`. Use `{ "$var": "thinking.enabled" }`, `{ "$var": "thinking.effort" }`, or `{ "$var": "thinking.budget" }` for pi-controlled thinking values. */
chatTemplateArgs?: Record<string, ChatTemplateKwargValue>;
/** OpenRouter-compatible routing preferences sent as the `provider` request field. */
openRouterRouting?: OpenRouterRouting;
@ -593,7 +596,15 @@ export interface OpenAICompletionsCompat {
vercelGatewayRouting?: VercelGatewayRouting;
/** Whether z.ai supports top-level `tool_stream: true` for streaming tool call deltas. Default: false. */
zaiToolStream?: boolean;
/** Whether the provider supports top-level `thinking_token_budget` to cap reasoning tokens (vLLM). Reasoning and the answer share `max_tokens` on these endpoints, so without a budget a reasoning-heavy turn can consume the whole response and emit no answer. Default: false. */
/**
* Top-level request field used to cap reasoning tokens from `thinkingBudgets`.
* Reasoning and the answer share `max_tokens` on these endpoints, so without a budget a
* reasoning-heavy turn can consume the whole response and emit no answer.
* `"thinking_token_budget"` is vLLM, `"thinking_budget"` is Qwen/DashScope/SGLang,
* `"thinking_budget_tokens"` is llama.cpp. Off by default; not set on the generated catalog.
*/
thinkingTokenBudgetField?: ThinkingTokenBudgetField;
/** Alias for `thinkingTokenBudgetField: "thinking_token_budget"` (vLLM). Prefer `thinkingTokenBudgetField`. Default: false. */
supportsThinkingTokenBudget?: boolean;
/** Whether the provider supports OpenAI custom tools with Lark/regex grammar formats. When false, grammar-constrained tools fall back to normal function tools. Default: false; the generated model catalog enables it for capable models. */
supportsOpenAIGrammarTools?: boolean;

7
packages/ai/test/openai-completions-thinking-as-text.test.ts
View file

@ -39,15 +39,20 @@ const compat = {
chatTemplateArgs: {},
zaiToolStream: false,
supportsThinkingTokenBudget: false,
thinkingTokenBudgetField: undefined,
supportsStrictMode: true,
supportsOpenAIGrammarTools: false,
cacheControlFormat: undefined,
sendSessionAffinityHeaders: false,
sessionAffinityFormat: "openai",
supportsLongCacheRetention: true,
} satisfies Omit<Required<OpenAICompletionsCompat>, "cacheControlFormat" | "deferredToolsMode"> & {
} satisfies Omit<
Required<OpenAICompletionsCompat>,
"cacheControlFormat" | "deferredToolsMode" | "thinkingTokenBudgetField"
> & {
cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"];
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
thinkingTokenBudgetField?: OpenAICompletionsCompat["thinkingTokenBudgetField"];
};
function buildModel(baseUrl = "http://127.0.0.1:1"): Model<"openai-completions"> {

130
packages/ai/test/openai-completions-thinking-token-budget.test.ts
View file

@ -41,21 +41,35 @@ vi.mock("openai", () => {
return { default: FakeOpenAI };
});
// vLLM-served reasoning model: reasoning and the answer share max_tokens.
const vllmModel: Model<"openai-completions"> = {
id: "zai-org/glm-5.2",
name: "GLM 5.2 (local vLLM)",
api: "openai-completions",
provider: "local-vllm",
baseUrl: "http://localhost:8000/v1",
reasoning: true,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 262144,
maxTokens: 16384,
compat: { thinkingFormat: "zai", supportsThinkingTokenBudget: true },
type CapturedParams = {
thinking_token_budget?: number;
thinking_budget?: number;
thinking_budget_tokens?: number;
thinking?: unknown;
chat_template_kwargs?: Record<string, unknown>;
};
function vllmModel(
compat: Model<"openai-completions">["compat"] = {
thinkingFormat: "zai",
supportsThinkingTokenBudget: true,
},
): Model<"openai-completions"> {
return {
id: "zai-org/glm-5.2",
name: "GLM 5.2 (local vLLM)",
api: "openai-completions",
provider: "local-vllm",
baseUrl: "http://localhost:8000/v1",
reasoning: true,
input: ["text"],
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0 },
contextWindow: 262144,
maxTokens: 16384,
compat,
};
}
async function capture(
model: Model<"openai-completions">,
options?: {
@ -63,7 +77,7 @@ async function capture(
thinkingBudgets?: ThinkingBudgets;
maxTokens?: number;
},
): Promise<{ thinking_token_budget?: number; thinking?: unknown }> {
): Promise<CapturedParams> {
let payload: unknown;
await streamSimple(
@ -80,45 +94,109 @@ async function capture(
},
).result();
return (payload ?? mockState.lastParams) as { thinking_token_budget?: number; thinking?: unknown };
return (payload ?? mockState.lastParams) as CapturedParams;
}
describe("openai-completions thinking_token_budget", () => {
describe("openai-completions thinking token budget", () => {
beforeEach(() => {
mockState.lastParams = undefined;
});
it("sends the configured budget for the requested level", async () => {
const params = await capture(vllmModel, { reasoning: "medium", thinkingBudgets: { medium: 4096 } });
const params = await capture(vllmModel(), { reasoning: "medium", thinkingBudgets: { medium: 4096 } });
expect(params.thinking_token_budget).toBe(4096);
});
it("omits the budget when the compat flag is not set", async () => {
const model = { ...vllmModel, compat: { thinkingFormat: "zai" } } as Model<"openai-completions">;
const params = await capture(model, { reasoning: "medium", thinkingBudgets: { medium: 4096 } });
it("omits the budget when neither the field nor the alias is set", async () => {
const params = await capture(vllmModel({ thinkingFormat: "zai" }), {
reasoning: "medium",
thinkingBudgets: { medium: 4096 },
});
expect(params.thinking_token_budget).toBeUndefined();
expect(params.thinking_budget).toBeUndefined();
expect(params.thinking_budget_tokens).toBeUndefined();
});
it("omits the budget when thinking is off", async () => {
const params = await capture(vllmModel, { reasoning: undefined, thinkingBudgets: { high: 8192 } });
const params = await capture(vllmModel(), { reasoning: undefined, thinkingBudgets: { high: 8192 } });
expect(params.thinking_token_budget).toBeUndefined();
});
it("clamps xhigh and max to the high budget", async () => {
const xhigh = await capture(vllmModel, { reasoning: "xhigh", thinkingBudgets: { high: 8192 } });
const max = await capture(vllmModel, { reasoning: "max", thinkingBudgets: { high: 8192 } });
const xhigh = await capture(vllmModel(), { reasoning: "xhigh", thinkingBudgets: { high: 8192 } });
const max = await capture(vllmModel(), { reasoning: "max", thinkingBudgets: { high: 8192 } });
expect(xhigh.thinking_token_budget).toBe(8192);
expect(max.thinking_token_budget).toBe(8192);
});
it("leaves room for the answer when the budget meets the response ceiling", async () => {
// Default high budget (16384) equals the model ceiling, which would leave no answer.
const params = await capture(vllmModel, { reasoning: "high" });
const params = await capture(vllmModel(), { reasoning: "high" });
expect(params.thinking_token_budget).toBe(16384 - 1024);
});
it("uses the caller max_tokens as the ceiling when it is lower than the model cap", async () => {
const params = await capture(vllmModel, { reasoning: "high", thinkingBudgets: { high: 8192 }, maxTokens: 4096 });
const params = await capture(vllmModel(), {
reasoning: "high",
thinkingBudgets: { high: 8192 },
maxTokens: 4096,
});
expect(params.thinking_token_budget).toBe(4096 - 1024);
});
it.each(["thinking_budget", "thinking_budget_tokens"] as const)(
"sends %s when thinkingTokenBudgetField is set",
async (field) => {
const params = await capture(vllmModel({ thinkingFormat: "qwen", thinkingTokenBudgetField: field }), {
reasoning: "medium",
thinkingBudgets: { medium: 4096 },
});
expect(params[field]).toBe(4096);
expect(params.thinking_token_budget).toBeUndefined();
},
);
it("lets thinkingTokenBudgetField win over the boolean alias", async () => {
const params = await capture(
vllmModel({
thinkingFormat: "zai",
supportsThinkingTokenBudget: true,
thinkingTokenBudgetField: "thinking_budget",
}),
{ reasoning: "medium", thinkingBudgets: { medium: 4096 } },
);
expect(params.thinking_budget).toBe(4096);
expect(params.thinking_token_budget).toBeUndefined();
});
it("puts the clamped budget in chat_template_kwargs when $var is thinking.budget", async () => {
const params = await capture(
vllmModel({
thinkingFormat: "chat-template",
chatTemplateKwargs: {
enable_thinking: { $var: "thinking.enabled" },
thinking_budget: { $var: "thinking.budget" },
},
}),
{ reasoning: "high" },
);
expect(params.chat_template_kwargs).toEqual({
enable_thinking: true,
thinking_budget: 16384 - 1024,
});
expect(params.thinking_token_budget).toBeUndefined();
});
it("omits thinking.budget from chat_template_kwargs when thinking is off", async () => {
const params = await capture(
vllmModel({
thinkingFormat: "chat-template",
chatTemplateKwargs: {
enable_thinking: { $var: "thinking.enabled" },
thinking_budget: { $var: "thinking.budget" },
},
}),
{ reasoning: undefined },
);
expect(params.chat_template_kwargs).toEqual({ enable_thinking: false });
});
});

4
packages/ai/test/openai-completions-tool-result-images.test.ts
View file

@ -19,8 +19,9 @@ const emptyUsage: Usage = {
cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
};
const compat: Omit<Required<OpenAICompletionsCompat>, "deferredToolsMode"> & {
const compat: Omit<Required<OpenAICompletionsCompat>, "deferredToolsMode" | "thinkingTokenBudgetField"> & {
deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"];
thinkingTokenBudgetField?: OpenAICompletionsCompat["thinkingTokenBudgetField"];
} = {
supportsStore: true,
supportsDeveloperRole: true,
@ -39,6 +40,7 @@ const compat: Omit<Required<OpenAICompletionsCompat>, "deferredToolsMode"> & {
chatTemplateArgs: {},
zaiToolStream: false,
supportsThinkingTokenBudget: false,
thinkingTokenBudgetField: undefined,
supportsStrictMode: true,
supportsOpenAIGrammarTools: false,
cacheControlFormat: "anthropic",

7
packages/coding-agent/docs/custom-provider.md
View file

@ -752,8 +752,10 @@ interface ProviderModelConfig {
requiresThinkingAsText?: boolean;
requiresReasoningContentOnAssistantMessages?: boolean;
thinkingFormat?: "openai" | "openrouter" | "deepseek" | "together" | "baseten" | "zai" | "qwen" | "chat-template" | "qwen-chat-template" | "string-thinking" | "ant-ling";
chatTemplateKwargs?: Record<string, string | number | boolean | null | { "$var": "thinking.enabled" | "thinking.effort"; omitWhenOff?: boolean }>;
chatTemplateArgs?: Record<string, string | number | boolean | null | { "$var": "thinking.enabled" | "thinking.effort"; omitWhenOff?: boolean }>;
chatTemplateKwargs?: Record<string, string | number | boolean | null | { "$var": "thinking.enabled" | "thinking.effort" | "thinking.budget"; omitWhenOff?: boolean }>;
chatTemplateArgs?: Record<string, string | number | boolean | null | { "$var": "thinking.enabled" | "thinking.effort" | "thinking.budget"; omitWhenOff?: boolean }>;
thinkingTokenBudgetField?: "thinking_token_budget" | "thinking_budget" | "thinking_budget_tokens";
supportsThinkingTokenBudget?: boolean;
cacheControlFormat?: "anthropic";
sessionAffinityFormat?: "openai" | "openai-nosession" | "openrouter";
sendSessionAffinityHeaders?: boolean;
@ -771,4 +773,5 @@ interface ProviderModelConfig {
```
`openrouter` sends `reasoning: { effort }`. `deepseek` sends `thinking: { type: "enabled" | "disabled" }` and `reasoning_effort` when enabled. `together` sends `reasoning: { enabled }` and also `reasoning_effort` when `supportsReasoningEffort` is enabled. `qwen` is for DashScope-style top-level `enable_thinking`. Use `qwen-chat-template` for local Qwen-compatible servers that read `chat_template_kwargs.enable_thinking` and need `preserve_thinking`. Use `chat-template` for configurable `chat_template_kwargs`, for example DeepSeek V3.x behind vLLM with `chatTemplateKwargs: { "thinking": { "$var": "thinking.enabled" } }`. Use `thinkingFormat: "baseten"` with `chatTemplateArgs` when the provider expects toggle values under `chat_template_args` and optionally supports top-level `reasoning_effort`.
`thinkingTokenBudgetField` sends a clamped per-level thinking budget as a top-level request field (`thinking_token_budget` on vLLM, `thinking_budget` on Qwen/SGLang, `thinking_budget_tokens` on llama.cpp). `supportsThinkingTokenBudget: true` is an alias for the vLLM field name. Do not combine it with `reasoning_effort` on DashScope Qwen models.
`cacheControlFormat: "anthropic"` applies Anthropic-style `cache_control` markers to the system prompt, last tool definition, and last user, assistant, or tool-result text content.

10
packages/coding-agent/docs/models.md
View file

@ -254,6 +254,8 @@ Current behavior:
Only OpenAI-compatible APIs apply it (`openai-completions`, `openai-responses`, `azure-openai-responses`); other APIs ignore it. Keys override pi's named request fields (for example a `temperature` key here beats the request-level temperature), so prefer it as the single source of sampling truth for a model. In `modelOverrides`, `samplingParams` merges per key with the base model's value.
A constant thinking-token cap can go here too, but it will not follow `thinkingBudgets` or leave room for the answer. Prefer `compat.thinkingTokenBudgetField` (or the `supportsThinkingTokenBudget` alias) for that.
### Thinking Level Map
Use `thinkingLevelMap` on a model to describe model-specific thinking controls. Keys are pi thinking levels: `off`, `minimal`, `low`, `medium`, `high`, `xhigh`, `max`. Maps may contain holes; for example, a model can expose `high` and `max` without exposing `xhigh`.
@ -467,8 +469,10 @@ For providers with partial OpenAI compatibility, use the `compat` field.
| `requiresThinkingAsText` | Convert thinking blocks to plain text |
| `requiresReasoningContentOnAssistantMessages` | Include empty `reasoning_content` on all replayed assistant messages when reasoning is enabled |
| `thinkingFormat` | Use `reasoning_effort`, `openrouter`, `deepseek`, `together`, `baseten`, `zai`, `qwen`, `chat-template`, or `qwen-chat-template` thinking parameters |
| `chatTemplateKwargs` | `chat_template_kwargs` values for `thinkingFormat: "chat-template"`; use `{ "$var": "thinking.enabled" }` or `{ "$var": "thinking.effort" }` for pi-controlled thinking values |
| `chatTemplateArgs` | `chat_template_args` values for `thinkingFormat: "baseten"`; use `{ "$var": "thinking.enabled" }` or `{ "$var": "thinking.effort" }` for pi-controlled thinking values |
| `chatTemplateKwargs` | `chat_template_kwargs` values for `thinkingFormat: "chat-template"`; use `{ "$var": "thinking.enabled" }`, `{ "$var": "thinking.effort" }`, or `{ "$var": "thinking.budget" }` for pi-controlled thinking values |
| `chatTemplateArgs` | `chat_template_args` values for `thinkingFormat: "baseten"`; use `{ "$var": "thinking.enabled" }`, `{ "$var": "thinking.effort" }`, or `{ "$var": "thinking.budget" }` for pi-controlled thinking values |
| `thinkingTokenBudgetField` | Top-level request field used to cap reasoning tokens from `thinkingBudgets`, clamped so at least 1024 tokens remain for the answer. `"thinking_token_budget"` (vLLM), `"thinking_budget"` (Qwen/DashScope/SGLang), `"thinking_budget_tokens"` (llama.cpp). Off by default; not set on the generated catalog. |
| `supportsThinkingTokenBudget` | Alias for `thinkingTokenBudgetField: "thinking_token_budget"` (vLLM). Prefer `thinkingTokenBudgetField`. Default: `false`. |
| `cacheControlFormat` | Use Anthropic-style `cache_control` markers on the system prompt, last tool definition, and last user, assistant, or tool-result text content. Currently only `anthropic` is supported. |
| `sendSessionAffinityHeaders` | For `openai-completions`, send session-affinity headers from the session id when caching is enabled. Default: `false`. |
| `sessionAffinityFormat` | For `openai-completions` and `openai-responses`, the session-affinity header format: `openai` sends `session_id`/`x-client-request-id` (completions also `x-session-affinity`), `openai-nosession` omits the underscore-containing `session_id` header, `openrouter` sends `x-session-id`. Does not affect the `prompt_cache_key` body param. Default: auto-detected. |
@ -481,6 +485,8 @@ For providers with partial OpenAI compatibility, use the `compat` field.
`openrouter` uses `reasoning: { effort }`. `together` uses `reasoning: { enabled }` and also `reasoning_effort` when `supportsReasoningEffort` is enabled. `qwen` uses top-level `enable_thinking`. Use `qwen-chat-template` for local Qwen-compatible servers that require `chat_template_kwargs.enable_thinking` and `preserve_thinking`. Use `chat-template` for vLLM/Hugging Face chat templates that need configurable `chat_template_kwargs`, such as `chatTemplateKwargs: { "thinking": { "$var": "thinking.enabled" } }` for DeepSeek V3.x templates. Use `thinkingFormat: "baseten"` with `chatTemplateArgs` for providers that expose toggle controls through `chat_template_args` and optionally support top-level `reasoning_effort`.
`thinkingTokenBudgetField` is independent of `thinkingFormat`. Do not enable it on the generated Qwen catalog: those models already send `reasoning_effort`, and DashScope rejects `thinking_budget` together with `reasoning_effort`.
`cacheControlFormat: "anthropic"` is for OpenAI-compatible providers that expose Anthropic-style prompt caching through `cache_control` markers on text content and tool definitions.
Example:

2
packages/coding-agent/docs/settings.md
View file

@ -32,7 +32,7 @@ Use `/trust` in interactive mode to save a project trust decision for future ses
| `defaultThinkingLevel` | string | - | `"off"`, `"minimal"`, `"low"`, `"medium"`, `"high"`, `"xhigh"`, `"max"` |
| `hideThinkingBlock` | boolean | `false` | Hide thinking blocks in output |
| `showCacheMissNotices` | boolean | `false` | Show transcript notices for significant prompt-cache misses and compaction or branch-summary usage |
| `thinkingBudgets` | object | - | Custom token budgets per thinking level |
| `thinkingBudgets` | object | - | Custom token budgets per thinking level. Anthropic, Google, and Bedrock use these natively. OpenAI-compatible models use them when `compat.thinkingTokenBudgetField` (or `supportsThinkingTokenBudget`) is set. |
#### thinkingBudgets