diff --git a/packages/ai/src/api/openai-completions.ts b/packages/ai/src/api/openai-completions.ts index 890745b7e..20da4e05d 100644 --- a/packages/ai/src/api/openai-completions.ts +++ b/packages/ai/src/api/openai-completions.ts @@ -28,6 +28,7 @@ import type { StreamFunction, StreamOptions, TextContent, + ThinkingBudgets, ThinkingContent, Tool, ToolCall, @@ -51,7 +52,7 @@ import { } from "./constrained-sampling.ts"; import { buildCopilotDynamicHeaders, hasCopilotVisionInput } from "./github-copilot-headers.ts"; import { clampOpenAIPromptCacheKey } from "./openai-prompt-cache.ts"; -import { buildBaseOptions } from "./simple-options.ts"; +import { buildBaseOptions, clampReasoning, MIN_ANSWER_TOKENS } from "./simple-options.ts"; import { transformMessages } from "./transform-messages.ts"; /** @@ -141,6 +142,8 @@ 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. */ + thinkingBudgets?: ThinkingBudgets; } export interface ConvertCompletionsMessagesOptions { @@ -154,10 +157,11 @@ interface OpenAICompatCacheControl { type ResolvedOpenAICompletionsCompat = Omit< Required, - "cacheControlFormat" | "deferredToolsMode" + "cacheControlFormat" | "deferredToolsMode" | "supportsThinkingTokenBudget" > & { cacheControlFormat?: OpenAICompletionsCompat["cacheControlFormat"]; deferredToolsMode?: OpenAICompletionsCompat["deferredToolsMode"]; + supportsThinkingTokenBudget?: OpenAICompletionsCompat["supportsThinkingTokenBudget"]; }; type ResolvedChatTemplateKwargValue = string | number | boolean | null; @@ -625,6 +629,7 @@ export const streamSimple: StreamFunction<"openai-completions", SimpleStreamOpti ...base, reasoningEffort, toolChoice, + thinkingBudgets: options?.thinkingBudgets, } satisfies OpenAICompletionsOptions); }; @@ -839,6 +844,27 @@ 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; + } + } + // OpenRouter provider routing preferences if (model.compat?.openRouterRouting) { (params as any).provider = model.compat.openRouterRouting; @@ -1492,6 +1518,7 @@ function detectCompat(model: Model<"openai-completions">): ResolvedOpenAIComplet chatTemplateKwargs: {}, chatTemplateArgs: {}, zaiToolStream: false, + supportsThinkingTokenBudget: false, supportsStrictMode: !isMoonshot && !isTogether && !isCloudflareAiGateway && !isNvidia, supportsOpenAIGrammarTools: false, cacheControlFormat, @@ -1536,6 +1563,7 @@ function getCompat(model: Model<"openai-completions">): ResolvedOpenAICompletion chatTemplateKwargs: model.compat.chatTemplateKwargs ?? detected.chatTemplateKwargs, chatTemplateArgs: model.compat.chatTemplateArgs ?? detected.chatTemplateArgs, zaiToolStream: model.compat.zaiToolStream ?? detected.zaiToolStream, + supportsThinkingTokenBudget: model.compat.supportsThinkingTokenBudget ?? detected.supportsThinkingTokenBudget, supportsStrictMode: model.compat.supportsStrictMode ?? detected.supportsStrictMode, supportsOpenAIGrammarTools: model.compat.supportsOpenAIGrammarTools ?? detected.supportsOpenAIGrammarTools, cacheControlFormat: model.compat.cacheControlFormat ?? detected.cacheControlFormat, diff --git a/packages/ai/src/api/simple-options.ts b/packages/ai/src/api/simple-options.ts index 42499aa68..067d01e59 100644 --- a/packages/ai/src/api/simple-options.ts +++ b/packages/ai/src/api/simple-options.ts @@ -51,6 +51,9 @@ export function buildBaseOptions( }; } +/** Tokens always left for the answer when a thinking budget shares the response ceiling. */ +export const MIN_ANSWER_TOKENS = 1024; + export function clampReasoning(effort: ThinkingLevel | undefined): Exclude | undefined { return effort === "xhigh" || effort === "max" ? "high" : effort; } @@ -70,14 +73,13 @@ export function adjustMaxTokensForThinking( }; const budgets = { ...defaultBudgets, ...customBudgets }; - const minOutputTokens = 1024; const level = clampReasoning(reasoningLevel)!; let thinkingBudget = budgets[level]!; const maxTokens = baseMaxTokens === undefined ? modelMaxTokens : Math.min(baseMaxTokens + thinkingBudget, modelMaxTokens); if (maxTokens <= thinkingBudget) { - thinkingBudget = Math.max(0, maxTokens - minOutputTokens); + thinkingBudget = Math.max(0, maxTokens - MIN_ANSWER_TOKENS); } return { maxTokens, thinkingBudget }; diff --git a/packages/ai/src/types.ts b/packages/ai/src/types.ts index e2777ed3e..82a3a9e1e 100644 --- a/packages/ai/src/types.ts +++ b/packages/ai/src/types.ts @@ -578,6 +578,8 @@ 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. */ + 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; /** Whether the provider supports the `strict` field in tool definitions. Default: true. */ diff --git a/packages/ai/test/openai-completions-thinking-as-text.test.ts b/packages/ai/test/openai-completions-thinking-as-text.test.ts index 1078e3f5c..2e451f6b7 100644 --- a/packages/ai/test/openai-completions-thinking-as-text.test.ts +++ b/packages/ai/test/openai-completions-thinking-as-text.test.ts @@ -38,6 +38,7 @@ const compat = { chatTemplateKwargs: {}, chatTemplateArgs: {}, zaiToolStream: false, + supportsThinkingTokenBudget: false, supportsStrictMode: true, supportsOpenAIGrammarTools: false, cacheControlFormat: undefined, diff --git a/packages/ai/test/openai-completions-thinking-token-budget.test.ts b/packages/ai/test/openai-completions-thinking-token-budget.test.ts new file mode 100644 index 000000000..2ca0c3ce0 --- /dev/null +++ b/packages/ai/test/openai-completions-thinking-token-budget.test.ts @@ -0,0 +1,124 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; +import { streamSimple } from "../src/compat.ts"; +import type { Model, SimpleStreamOptions, ThinkingBudgets } from "../src/types.ts"; + +const mockState = vi.hoisted(() => ({ + lastParams: undefined as unknown, +})); + +vi.mock("openai", () => { + class FakeOpenAI { + chat = { + completions: { + create: (params: unknown) => { + mockState.lastParams = params; + const stream = { + async *[Symbol.asyncIterator]() { + yield { + choices: [{ delta: {}, finish_reason: "stop" }], + usage: { + prompt_tokens: 1, + completion_tokens: 1, + prompt_tokens_details: { cached_tokens: 0 }, + completion_tokens_details: { reasoning_tokens: 0 }, + }, + }; + }, + }; + const promise = Promise.resolve(stream) as Promise & { + withResponse: () => Promise<{ data: typeof stream; response: { status: number; headers: Headers } }>; + }; + promise.withResponse = async () => ({ + data: stream, + response: { status: 200, headers: new Headers() }, + }); + return promise; + }, + }, + }; + } + + 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 }, +}; + +async function capture( + model: Model<"openai-completions">, + options?: { + reasoning?: SimpleStreamOptions["reasoning"]; + thinkingBudgets?: ThinkingBudgets; + maxTokens?: number; + }, +): Promise<{ thinking_token_budget?: number; thinking?: unknown }> { + let payload: unknown; + + await streamSimple( + model, + { messages: [{ role: "user", content: "Hi", timestamp: Date.now() }] }, + { + apiKey: "test", + reasoning: options?.reasoning, + thinkingBudgets: options?.thinkingBudgets, + maxTokens: options?.maxTokens, + onPayload: (params: unknown) => { + payload = params; + }, + }, + ).result(); + + return (payload ?? mockState.lastParams) as { thinking_token_budget?: number; thinking?: unknown }; +} + +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 } }); + 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 } }); + expect(params.thinking_token_budget).toBeUndefined(); + }); + + it("omits the budget when thinking is off", async () => { + 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 } }); + 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" }); + 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 }); + expect(params.thinking_token_budget).toBe(4096 - 1024); + }); +}); diff --git a/packages/ai/test/openai-completions-tool-result-images.test.ts b/packages/ai/test/openai-completions-tool-result-images.test.ts index 6b977714d..6beb7f284 100644 --- a/packages/ai/test/openai-completions-tool-result-images.test.ts +++ b/packages/ai/test/openai-completions-tool-result-images.test.ts @@ -38,6 +38,7 @@ const compat: Omit, "deferredToolsMode"> & { chatTemplateKwargs: {}, chatTemplateArgs: {}, zaiToolStream: false, + supportsThinkingTokenBudget: false, supportsStrictMode: true, supportsOpenAIGrammarTools: false, cacheControlFormat: "anthropic",