From 32aef469a63d8635e2e7f98edcdacf6fc497386d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sebastian=20M=C3=BCller?= Date: Wed, 22 Jul 2026 10:53:39 +0200 Subject: [PATCH] fix(ai): align Vertex Flash-Lite thinking budget (#463) * fix(ai): align Vertex Flash-Lite thinking budget * test(ai): decouple Vertex preview coverage from catalog --- packages/ai/CHANGELOG.md | 2 + packages/ai/src/providers/google-shared.ts | 28 ++++++- packages/ai/src/providers/google-vertex.ts | 36 +-------- packages/ai/src/providers/google.ts | 46 +----------- .../google-vertex-thinking-budget.test.ts | 73 +++++++++++++++++++ 5 files changed, 106 insertions(+), 79 deletions(-) create mode 100644 packages/ai/test/google-vertex-thinking-budget.test.ts diff --git a/packages/ai/CHANGELOG.md b/packages/ai/CHANGELOG.md index 8381d78a1..e80a0edbd 100644 --- a/packages/ai/CHANGELOG.md +++ b/packages/ai/CHANGELOG.md @@ -2,6 +2,8 @@ ## [Unreleased] +- Fixed Vertex Gemini 2.5 Flash-Lite minimal reasoning requests to use the supported 512-token thinking budget. + ## [0.3.2] - 2026-07-20 - Added Kimi K3 to the Prime Inference and OpenRouter model catalogs with multimodal input and mandatory max reasoning. diff --git a/packages/ai/src/providers/google-shared.ts b/packages/ai/src/providers/google-shared.ts index 9f0ab1806..e9d05e37a 100644 --- a/packages/ai/src/providers/google-shared.ts +++ b/packages/ai/src/providers/google-shared.ts @@ -3,7 +3,7 @@ */ import { type Content, FinishReason, FunctionCallingConfigMode, type Part } from "@google/genai"; -import type { Context, ImageContent, Model, StopReason, TextContent, Tool } from "../types.js"; +import type { Context, ImageContent, Model, StopReason, TextContent, ThinkingBudgets, Tool } from "../types.js"; import { sanitizeSurrogates } from "../utils/sanitize-unicode.js"; import { transformMessages } from "./transform-messages.js"; @@ -15,6 +15,32 @@ type GoogleApiType = "google-generative-ai" | "google-vertex"; */ export type GoogleThinkingLevel = "THINKING_LEVEL_UNSPECIFIED" | "MINIMAL" | "LOW" | "MEDIUM" | "HIGH"; +type GoogleBudgetThinkingLevel = "minimal" | "low" | "medium" | "high"; + +export function getGoogleThinkingBudget( + modelId: string, + effort: GoogleBudgetThinkingLevel, + customBudgets?: ThinkingBudgets, +): number { + if (customBudgets?.[effort] !== undefined) { + return customBudgets[effort]!; + } + + if (modelId.includes("2.5-pro")) { + return { minimal: 128, low: 2048, medium: 8192, high: 32768 }[effort]; + } + + if (modelId.includes("2.5-flash-lite")) { + return { minimal: 512, low: 2048, medium: 8192, high: 24576 }[effort]; + } + + if (modelId.includes("2.5-flash")) { + return { minimal: 128, low: 2048, medium: 8192, high: 24576 }[effort]; + } + + return -1; +} + /** * Determines whether a streamed Gemini `Part` should be treated as "thinking". * diff --git a/packages/ai/src/providers/google-vertex.ts b/packages/ai/src/providers/google-vertex.ts index 9bf70809d..b6e0bbe84 100644 --- a/packages/ai/src/providers/google-vertex.ts +++ b/packages/ai/src/providers/google-vertex.ts @@ -18,7 +18,6 @@ import type { StreamFunction, StreamOptions, TextContent, - ThinkingBudgets, ThinkingContent, ToolCall, } from "../types.js"; @@ -33,6 +32,7 @@ import type { GoogleThinkingLevel } from "./google-shared.js"; import { convertMessages, convertTools, + getGoogleThinkingBudget, isThinkingPart, mapStopReason, mapToolChoice, @@ -332,7 +332,7 @@ export const streamSimpleGoogleVertex: StreamFunction<"google-vertex", SimpleStr ...base, thinking: { enabled: true, - budgetTokens: getGoogleBudget(geminiModel, effort, options.thinkingBudgets), + budgetTokens: getGoogleThinkingBudget(geminiModel.id, effort, options.thinkingBudgets), }, } satisfies GoogleVertexOptions); }; @@ -543,35 +543,3 @@ function getGemini3ThinkingLevel( return "HIGH"; } } - -function getGoogleBudget( - model: Model<"google-generative-ai">, - effort: ClampedThinkingLevel, - customBudgets?: ThinkingBudgets, -): number { - if (customBudgets?.[effort] !== undefined) { - return customBudgets[effort]!; - } - - if (model.id.includes("2.5-pro")) { - const budgets: Record = { - minimal: 128, - low: 2048, - medium: 8192, - high: 32768, - }; - return budgets[effort]; - } - - if (model.id.includes("2.5-flash")) { - const budgets: Record = { - minimal: 128, - low: 2048, - medium: 8192, - high: 24576, - }; - return budgets[effort]; - } - - return -1; -} diff --git a/packages/ai/src/providers/google.ts b/packages/ai/src/providers/google.ts index 68664e089..05914f687 100644 --- a/packages/ai/src/providers/google.ts +++ b/packages/ai/src/providers/google.ts @@ -15,7 +15,6 @@ import type { StreamFunction, StreamOptions, TextContent, - ThinkingBudgets, ThinkingContent, ThinkingLevel, ToolCall, @@ -31,6 +30,7 @@ import type { GoogleThinkingLevel } from "./google-shared.js"; import { convertMessages, convertTools, + getGoogleThinkingBudget, isThinkingPart, mapStopReason, mapToolChoice, @@ -317,7 +317,7 @@ export const streamSimpleGoogle: StreamFunction<"google-generative-ai", SimpleSt ...base, thinking: { enabled: true, - budgetTokens: getGoogleBudget(googleModel, effort, options.thinkingBudgets), + budgetTokens: getGoogleThinkingBudget(googleModel.id, effort, options.thinkingBudgets), }, } satisfies GoogleOptions); }; @@ -466,45 +466,3 @@ function getThinkingLevel(effort: ClampedThinkingLevel, model: Model<"google-gen return "HIGH"; } } - -function getGoogleBudget( - model: Model<"google-generative-ai">, - effort: ClampedThinkingLevel, - customBudgets?: ThinkingBudgets, -): number { - if (customBudgets?.[effort] !== undefined) { - return customBudgets[effort]!; - } - - if (model.id.includes("2.5-pro")) { - const budgets: Record = { - minimal: 128, - low: 2048, - medium: 8192, - high: 32768, - }; - return budgets[effort]; - } - - if (model.id.includes("2.5-flash-lite")) { - const budgets: Record = { - minimal: 512, - low: 2048, - medium: 8192, - high: 24576, - }; - return budgets[effort]; - } - - if (model.id.includes("2.5-flash")) { - const budgets: Record = { - minimal: 128, - low: 2048, - medium: 8192, - high: 24576, - }; - return budgets[effort]; - } - - return -1; -} diff --git a/packages/ai/test/google-vertex-thinking-budget.test.ts b/packages/ai/test/google-vertex-thinking-budget.test.ts new file mode 100644 index 000000000..885138821 --- /dev/null +++ b/packages/ai/test/google-vertex-thinking-budget.test.ts @@ -0,0 +1,73 @@ +import type * as GoogleGenAi from "@google/genai"; +import type { GenerateContentParameters } from "@google/genai"; +import { describe, expect, it, vi } from "vitest"; + +vi.mock("@google/genai", async (importOriginal) => { + const actual = await importOriginal(); + class GoogleGenAI { + models = { + generateContentStream: async function* () { + yield { + candidates: [{ content: { parts: [{ text: "ok" }] }, finishReason: "STOP" }], + usageMetadata: { promptTokenCount: 1, candidatesTokenCount: 1, totalTokenCount: 2 }, + }; + }, + }; + } + + return { + ...actual, + GoogleGenAI, + ResourceScope: { COLLECTION: "COLLECTION" }, + ThinkingLevel: { + THINKING_LEVEL_UNSPECIFIED: "THINKING_LEVEL_UNSPECIFIED", + MINIMAL: "MINIMAL", + LOW: "LOW", + MEDIUM: "MEDIUM", + HIGH: "HIGH", + }, + }; +}); + +import { getModel } from "../src/models.js"; +import { streamSimpleGoogleVertex } from "../src/providers/google-vertex.js"; +import type { Context } from "../src/types.js"; + +const context: Context = { + messages: [{ role: "user", content: "hello", timestamp: Date.now() }], +}; + +const stableFlashLite = getModel("google-vertex", "gemini-2.5-flash-lite"); +const flashLiteModels = [stableFlashLite, { ...stableFlashLite, id: "gemini-2.5-flash-lite-preview" }] as const; + +async function captureMinimalReasoningPayload( + model: (typeof flashLiteModels)[number], +): Promise { + let capturedPayload: GenerateContentParameters | undefined; + const stream = streamSimpleGoogleVertex(model, context, { + apiKey: "fake-key", + reasoning: "minimal", + onPayload: (payload) => { + capturedPayload = payload as GenerateContentParameters; + return payload; + }, + }); + + await stream.result(); + + if (!capturedPayload) { + throw new Error("Expected Vertex payload to be captured"); + } + return capturedPayload; +} + +describe("Google Vertex thinking budget payload", () => { + it.each(flashLiteModels)("uses the supported minimal budget for $id", async (model) => { + const payload = await captureMinimalReasoningPayload(model); + + expect(payload.config?.thinkingConfig).toEqual({ + includeThoughts: true, + thinkingBudget: 512, + }); + }); +}); -- 2.51.2