diff --git a/CHANGELOG.md b/CHANGELOG.md index a0c317f0edd..de9f7609863 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -78,6 +78,7 @@ Docs: https://docs.openclaw.ai - Gateway/auth: reject gateway auth credentials that match published example placeholders at startup and secret reload, and keep cloud install snippets from publishing copy-paste gateway/keyring secrets. (#68404) Thanks @coygeek. - CLI/update: preserve macOS restart helper launchctl failures in the update restart log without letting log setup block the restart path. (#68492) Thanks @hclsys. - Slack/threads: keep file-only root messages as starter context so first thread replies can still hydrate starter media. (#68594) Thanks @martingarramon. +- Google/Antigravity: resolve forward-compatible Gemini 3.1 Pro custom-tools and Flash variants from the bundled Google plugin templates, so `google-antigravity/gemini-3.1-pro-preview-customtools` no longer falls through to an unknown-model error. Fixes #35512. ## 2026.4.15 diff --git a/extensions/google/api.ts b/extensions/google/api.ts index 86987cb070a..d09ff4b54c5 100644 --- a/extensions/google/api.ts +++ b/extensions/google/api.ts @@ -10,6 +10,19 @@ import { normalizeGoogleGenerativeAiBaseUrl, } from "./provider-policy.js"; export { normalizeAntigravityModelId, normalizeGoogleModelId } from "./model-id.js"; +export { + createGoogleThinkingPayloadWrapper, + createGoogleThinkingStreamWrapper, + isGoogleGemini3FlashModel, + isGoogleGemini3ProModel, + isGoogleGemini3ThinkingLevelModel, + isGoogleThinkingRequiredModel, + resolveGoogleGemini3ThinkingLevel, + sanitizeGoogleThinkingPayload, + stripInvalidGoogleThinkingBudget, + type GoogleThinkingInputLevel, + type GoogleThinkingLevel, +} from "./thinking-api.js"; export { DEFAULT_GOOGLE_API_BASE_URL, isGoogleGenerativeAiApi, diff --git a/extensions/google/provider-hooks.ts b/extensions/google/provider-hooks.ts index 6a2544fd27d..1775ecd95b8 100644 --- a/extensions/google/provider-hooks.ts +++ b/extensions/google/provider-hooks.ts @@ -1,9 +1,9 @@ import { buildProviderReplayFamilyHooks } from "openclaw/plugin-sdk/provider-model-shared"; -import { GOOGLE_THINKING_STREAM_HOOKS } from "openclaw/plugin-sdk/provider-stream-family"; +import { createGoogleThinkingStreamWrapper } from "./thinking-api.js"; export const GOOGLE_GEMINI_PROVIDER_HOOKS = { ...buildProviderReplayFamilyHooks({ family: "google-gemini", }), - ...GOOGLE_THINKING_STREAM_HOOKS, + wrapStreamFn: createGoogleThinkingStreamWrapper, }; diff --git a/extensions/google/provider-models.test.ts b/extensions/google/provider-models.test.ts index 71bcc8c1ec9..8633d4492a9 100644 --- a/extensions/google/provider-models.test.ts +++ b/extensions/google/provider-models.test.ts @@ -228,6 +228,112 @@ describe("resolveGoogleGeminiForwardCompatModel", () => { }); }); + it("resolves Antigravity Gemini 3.1 pro customtools from the low template", () => { + const model = resolveGoogleGeminiForwardCompatModel({ + providerId: "google-antigravity", + ctx: createContext({ + provider: "google-antigravity", + modelId: "gemini-3.1-pro-preview-customtools", + models: [ + createTemplateModel("google-antigravity", "gemini-3-pro-low", { + api: "openai-completions", + baseUrl: "https://antigravity.example/v1", + contextWindow: 1_048_576, + reasoning: true, + }), + ], + }), + }); + + expect(model).toMatchObject({ + provider: "google-antigravity", + id: "gemini-3.1-pro-preview-customtools", + api: "openai-completions", + baseUrl: "https://antigravity.example/v1", + contextWindow: 1_048_576, + reasoning: true, + }); + }); + + it("falls back to the Antigravity high template when the low template is unavailable", () => { + const model = resolveGoogleGeminiForwardCompatModel({ + providerId: "google-antigravity", + ctx: createContext({ + provider: "google-antigravity", + modelId: "gemini-3.1-pro-preview", + models: [ + createTemplateModel("google-antigravity", "gemini-3-pro-high", { + api: "openai-completions", + maxTokens: 65_536, + reasoning: true, + }), + ], + }), + }); + + expect(model).toMatchObject({ + provider: "google-antigravity", + id: "gemini-3.1-pro-preview", + api: "openai-completions", + maxTokens: 65_536, + reasoning: true, + }); + }); + + it("resolves Antigravity Gemini 3.1 flash variants from the flash template", () => { + const models = [ + createTemplateModel("google-antigravity", "gemini-3-flash", { + api: "openai-completions", + contextWindow: 1_048_576, + }), + ]; + + expect( + resolveGoogleGeminiForwardCompatModel({ + providerId: "google-antigravity", + ctx: createContext({ + provider: "google-antigravity", + modelId: "gemini-3.1-flash-preview", + models, + }), + }), + ).toMatchObject({ + provider: "google-antigravity", + id: "gemini-3.1-flash-preview", + api: "openai-completions", + contextWindow: 1_048_576, + }); + + expect( + resolveGoogleGeminiForwardCompatModel({ + providerId: "google-antigravity", + ctx: createContext({ + provider: "google-antigravity", + modelId: "gemini-3.1-flash-lite-preview", + models, + }), + }), + ).toMatchObject({ + provider: "google-antigravity", + id: "gemini-3.1-flash-lite-preview", + api: "openai-completions", + contextWindow: 1_048_576, + }); + }); + + it("returns undefined for Antigravity Gemini 3.1 models without a matching template", () => { + const model = resolveGoogleGeminiForwardCompatModel({ + providerId: "google-antigravity", + ctx: createContext({ + provider: "google-antigravity", + modelId: "gemini-3.1-pro-preview-customtools", + models: [createTemplateModel("google-antigravity", "claude-opus-4-6-thinking")], + }), + }); + + expect(model).toBeUndefined(); + }); + it("prefers the flash-lite template before the broader flash prefix", () => { const model = resolveGoogleGeminiForwardCompatModel({ providerId: "google-vertex", diff --git a/extensions/google/provider-models.ts b/extensions/google/provider-models.ts index 2aff4e2ad06..baa9fd662ed 100644 --- a/extensions/google/provider-models.ts +++ b/extensions/google/provider-models.ts @@ -6,6 +6,7 @@ import { cloneFirstTemplateModel } from "openclaw/plugin-sdk/provider-model-shar import { normalizeOptionalLowercaseString } from "openclaw/plugin-sdk/text-runtime"; const GOOGLE_GEMINI_CLI_PROVIDER_ID = "google-gemini-cli"; +const GOOGLE_ANTIGRAVITY_PROVIDER_ID = "google-antigravity"; const GEMINI_2_5_PRO_PREFIX = "gemini-2.5-pro"; const GEMINI_2_5_FLASH_LITE_PREFIX = "gemini-2.5-flash-lite"; const GEMINI_2_5_FLASH_PREFIX = "gemini-2.5-flash"; @@ -22,6 +23,8 @@ const GEMINI_2_5_FLASH_TEMPLATE_IDS = ["gemini-2.5-flash"] as const; const GEMINI_3_1_PRO_TEMPLATE_IDS = ["gemini-3-pro-preview"] as const; const GEMINI_3_1_FLASH_LITE_TEMPLATE_IDS = ["gemini-3.1-flash-lite-preview"] as const; const GEMINI_3_1_FLASH_TEMPLATE_IDS = ["gemini-3-flash-preview"] as const; +const GEMINI_3_PRO_ANTIGRAVITY_TEMPLATE_IDS = ["gemini-3-pro-low", "gemini-3-pro-high"] as const; +const GEMINI_3_FLASH_ANTIGRAVITY_TEMPLATE_IDS = ["gemini-3-flash"] as const; // Gemma uses the Gemini flash template as a forward-compat approximation // until a dedicated Gemma template is registered in the catalog. const GEMMA_TEMPLATE_IDS = GEMINI_3_1_FLASH_TEMPLATE_IDS; @@ -29,6 +32,7 @@ const GEMMA_TEMPLATE_IDS = GEMINI_3_1_FLASH_TEMPLATE_IDS; type GoogleForwardCompatFamily = { googleTemplateIds: readonly string[]; cliTemplateIds: readonly string[]; + antigravityTemplateIds?: readonly string[]; preferExternalFirstForCli?: boolean; }; @@ -61,13 +65,21 @@ function isGoogleGeminiCliProvider(providerId: string): boolean { return normalizeOptionalLowercaseString(providerId) === GOOGLE_GEMINI_CLI_PROVIDER_ID; } +function isGoogleAntigravityProvider(providerId: string): boolean { + return normalizeOptionalLowercaseString(providerId) === GOOGLE_ANTIGRAVITY_PROVIDER_ID; +} + function templateIdsForProvider( templateProviderId: string, family: GoogleForwardCompatFamily, ): readonly string[] { - return isGoogleGeminiCliProvider(templateProviderId) - ? family.cliTemplateIds - : family.googleTemplateIds; + if (isGoogleGeminiCliProvider(templateProviderId)) { + return family.cliTemplateIds; + } + if (isGoogleAntigravityProvider(templateProviderId)) { + return family.antigravityTemplateIds ?? family.googleTemplateIds; + } + return family.googleTemplateIds; } function buildGoogleTemplateSources(params: { @@ -135,6 +147,7 @@ export function resolveGoogleGeminiForwardCompatModel(params: { family = { googleTemplateIds: GEMINI_3_1_PRO_TEMPLATE_IDS, cliTemplateIds: GEMINI_3_1_PRO_TEMPLATE_IDS, + antigravityTemplateIds: GEMINI_3_PRO_ANTIGRAVITY_TEMPLATE_IDS, }; if (params.providerId === "google" || params.providerId === GOOGLE_GEMINI_CLI_PROVIDER_ID) { patch = { reasoning: true }; @@ -146,11 +159,13 @@ export function resolveGoogleGeminiForwardCompatModel(params: { family = { googleTemplateIds: GEMINI_3_1_FLASH_LITE_TEMPLATE_IDS, cliTemplateIds: GEMINI_3_1_FLASH_LITE_TEMPLATE_IDS, + antigravityTemplateIds: GEMINI_3_FLASH_ANTIGRAVITY_TEMPLATE_IDS, }; } else if (lower.startsWith(GEMINI_3_1_FLASH_PREFIX) || lower === GEMINI_FLASH_LATEST_ID) { family = { googleTemplateIds: GEMINI_3_1_FLASH_TEMPLATE_IDS, cliTemplateIds: GEMINI_3_1_FLASH_TEMPLATE_IDS, + antigravityTemplateIds: GEMINI_3_FLASH_ANTIGRAVITY_TEMPLATE_IDS, }; } else if (lower.startsWith(GEMMA_PREFIX)) { family = { diff --git a/extensions/google/runtime-api.ts b/extensions/google/runtime-api.ts index 8115792bfa9..e04bc297ba4 100644 --- a/extensions/google/runtime-api.ts +++ b/extensions/google/runtime-api.ts @@ -1,7 +1,17 @@ export { DEFAULT_GOOGLE_API_BASE_URL, + createGoogleThinkingPayloadWrapper, + createGoogleThinkingStreamWrapper, + isGoogleGemini3FlashModel, + isGoogleGemini3ProModel, + isGoogleGemini3ThinkingLevelModel, + isGoogleThinkingRequiredModel, normalizeGoogleApiBaseUrl, normalizeGoogleModelId, parseGeminiAuth, + resolveGoogleGemini3ThinkingLevel, resolveGoogleGenerativeAiHttpRequestConfig, + sanitizeGoogleThinkingPayload, + stripInvalidGoogleThinkingBudget, } from "./api.js"; +export type { GoogleThinkingInputLevel, GoogleThinkingLevel } from "./api.js"; diff --git a/extensions/google/thinking-api.ts b/extensions/google/thinking-api.ts new file mode 100644 index 00000000000..a13526e90cb --- /dev/null +++ b/extensions/google/thinking-api.ts @@ -0,0 +1,13 @@ +export { + createGoogleThinkingPayloadWrapper, + createGoogleThinkingStreamWrapper, + isGoogleGemini3FlashModel, + isGoogleGemini3ProModel, + isGoogleGemini3ThinkingLevelModel, + isGoogleThinkingRequiredModel, + resolveGoogleGemini3ThinkingLevel, + sanitizeGoogleThinkingPayload, + stripInvalidGoogleThinkingBudget, + type GoogleThinkingInputLevel, + type GoogleThinkingLevel, +} from "./thinking.js"; diff --git a/extensions/google/thinking.test.ts b/extensions/google/thinking.test.ts new file mode 100644 index 00000000000..d05f5cbc9e5 --- /dev/null +++ b/extensions/google/thinking.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, it } from "vitest"; +import { + resolveGoogleGemini3ThinkingLevel, + sanitizeGoogleThinkingPayload, +} from "./thinking-api.js"; + +describe("google thinking policy", () => { + it.each([ + ["off", "LOW"], + ["minimal", "LOW"], + ["low", "LOW"], + ["medium", "HIGH"], + ["adaptive", "HIGH"], + ["high", "HIGH"], + ["xhigh", "HIGH"], + ] as const)("maps Gemini 3 Pro thinking level %s to %s", (thinkingLevel, expected) => { + expect( + resolveGoogleGemini3ThinkingLevel({ + modelId: "gemini-3.1-pro-preview", + thinkingLevel, + }), + ).toBe(expected); + }); + + it.each([ + [0, "LOW"], + [2048, "LOW"], + [2049, "HIGH"], + ] as const)("maps Gemini 3 Pro budget %s to %s", (thinkingBudget, expected) => { + expect( + resolveGoogleGemini3ThinkingLevel({ + modelId: "gemini-pro-latest", + thinkingBudget, + }), + ).toBe(expected); + }); + + it.each([ + ["off", "MINIMAL"], + ["minimal", "MINIMAL"], + ["low", "LOW"], + ["medium", "MEDIUM"], + ["adaptive", "MEDIUM"], + ["high", "HIGH"], + ["xhigh", "HIGH"], + ] as const)("maps Gemini 3 Flash thinking level %s to %s", (thinkingLevel, expected) => { + expect( + resolveGoogleGemini3ThinkingLevel({ + modelId: "gemini-flash-latest", + thinkingLevel, + }), + ).toBe(expected); + }); + + it.each([ + [-1, "MINIMAL"], + [0, "MINIMAL"], + [2048, "LOW"], + [8192, "MEDIUM"], + [8193, "HIGH"], + ] as const)("maps Gemini 3 Flash budget %s to %s", (thinkingBudget, expected) => { + expect( + resolveGoogleGemini3ThinkingLevel({ + modelId: "gemini-3.1-flash-lite-preview", + thinkingBudget, + }), + ).toBe(expected); + }); + + it("removes thinkingBudget=0 for Gemini 2.5 Pro", () => { + const payload = { + config: { + thinkingConfig: { thinkingBudget: 0 }, + }, + }; + + sanitizeGoogleThinkingPayload({ payload, modelId: "google/gemini-2.5-pro-preview" }); + + expect(payload.config).not.toHaveProperty("thinkingConfig"); + }); + + it("rewrites Gemini 3 thinking budgets to thinkingLevel", () => { + const payload = { + generationConfig: { + thinkingConfig: { thinkingBudget: 8193, includeThoughts: true }, + }, + }; + + sanitizeGoogleThinkingPayload({ + payload, + modelId: "gemini-3-flash-preview", + thinkingLevel: "medium", + }); + + expect(payload.generationConfig.thinkingConfig).toEqual({ + includeThoughts: true, + thinkingLevel: "MEDIUM", + }); + }); + + it("maps Gemma 4 thinking mode without sending thinkingBudget", () => { + const payload = { + config: { + thinkingConfig: { thinkingBudget: 4096 }, + }, + }; + + sanitizeGoogleThinkingPayload({ + payload, + modelId: "gemma-4-26b-a4b-it", + thinkingLevel: "high", + }); + + expect(payload.config.thinkingConfig).toEqual({ thinkingLevel: "HIGH" }); + }); +}); diff --git a/extensions/google/thinking.ts b/extensions/google/thinking.ts new file mode 100644 index 00000000000..792ace94036 --- /dev/null +++ b/extensions/google/thinking.ts @@ -0,0 +1,263 @@ +import type { ProviderWrapStreamFnContext } from "openclaw/plugin-sdk/plugin-entry"; +import { createPayloadPatchStreamWrapper } from "openclaw/plugin-sdk/provider-stream-shared"; +import { normalizeLowercaseStringOrEmpty } from "openclaw/plugin-sdk/text-runtime"; + +export type GoogleThinkingLevel = "MINIMAL" | "LOW" | "MEDIUM" | "HIGH"; +export type GoogleThinkingInputLevel = + | "off" + | "minimal" + | "low" + | "medium" + | "adaptive" + | "high" + | "xhigh"; + +export function isGoogleThinkingRequiredModel(modelId: string): boolean { + return normalizeLowercaseStringOrEmpty(modelId).includes("gemini-2.5-pro"); +} + +export function isGoogleGemini3ProModel(modelId: string): boolean { + const normalized = normalizeLowercaseStringOrEmpty(modelId); + return /(?:^|\/)gemini-(?:3(?:\.\d+)?-pro|pro-latest)(?:-|$)/.test(normalized); +} + +export function isGoogleGemini3FlashModel(modelId: string): boolean { + const normalized = normalizeLowercaseStringOrEmpty(modelId); + return /(?:^|\/)gemini-(?:3(?:\.\d+)?-flash|flash(?:-lite)?-latest)(?:-|$)/.test(normalized); +} + +export function isGoogleGemini3ThinkingLevelModel(modelId: string): boolean { + return isGoogleGemini3ProModel(modelId) || isGoogleGemini3FlashModel(modelId); +} + +export function resolveGoogleGemini3ThinkingLevel(params: { + modelId?: string; + thinkingLevel?: GoogleThinkingInputLevel; + thinkingBudget?: number; +}): GoogleThinkingLevel | undefined { + if (typeof params.modelId !== "string") { + return undefined; + } + if (isGoogleGemini3ProModel(params.modelId)) { + switch (params.thinkingLevel) { + case "off": + case "minimal": + case "low": + return "LOW"; + case "medium": + case "adaptive": + case "high": + case "xhigh": + return "HIGH"; + } + if (typeof params.thinkingBudget === "number") { + return params.thinkingBudget <= 2048 ? "LOW" : "HIGH"; + } + return undefined; + } + if (!isGoogleGemini3FlashModel(params.modelId)) { + return undefined; + } + switch (params.thinkingLevel) { + case "off": + case "minimal": + return "MINIMAL"; + case "low": + return "LOW"; + case "medium": + case "adaptive": + return "MEDIUM"; + case "high": + case "xhigh": + return "HIGH"; + } + if (typeof params.thinkingBudget !== "number") { + return undefined; + } + if (params.thinkingBudget <= 0) { + return "MINIMAL"; + } + if (params.thinkingBudget <= 2048) { + return "LOW"; + } + if (params.thinkingBudget <= 8192) { + return "MEDIUM"; + } + return "HIGH"; +} + +export function stripInvalidGoogleThinkingBudget(params: { + thinkingConfig: Record; + modelId?: string; +}): boolean { + if ( + params.thinkingConfig.thinkingBudget !== 0 || + typeof params.modelId !== "string" || + !isGoogleThinkingRequiredModel(params.modelId) + ) { + return false; + } + delete params.thinkingConfig.thinkingBudget; + return true; +} + +function isGemma4Model(modelId: string): boolean { + return normalizeLowercaseStringOrEmpty(modelId).startsWith("gemma-4"); +} + +function mapThinkLevelToGemma4ThinkingLevel( + thinkingLevel?: GoogleThinkingInputLevel, +): "MINIMAL" | "HIGH" | undefined { + switch (thinkingLevel) { + case "off": + return undefined; + case "minimal": + case "low": + return "MINIMAL"; + case "medium": + case "adaptive": + case "high": + case "xhigh": + return "HIGH"; + default: + return undefined; + } +} + +function normalizeGemma4ThinkingLevel(value: unknown): "MINIMAL" | "HIGH" | undefined { + if (typeof value !== "string") { + return undefined; + } + switch (value.trim().toUpperCase()) { + case "MINIMAL": + case "LOW": + return "MINIMAL"; + case "MEDIUM": + case "HIGH": + return "HIGH"; + default: + return undefined; + } +} + +export function sanitizeGoogleThinkingPayload(params: { + payload: unknown; + modelId?: string; + thinkingLevel?: GoogleThinkingInputLevel; +}): void { + if (!params.payload || typeof params.payload !== "object") { + return; + } + const payloadObj = params.payload as Record; + sanitizeGoogleThinkingConfigContainer({ + container: payloadObj.config, + modelId: params.modelId, + thinkingLevel: params.thinkingLevel, + }); + sanitizeGoogleThinkingConfigContainer({ + container: payloadObj.generationConfig, + modelId: params.modelId, + thinkingLevel: params.thinkingLevel, + }); +} + +function sanitizeGoogleThinkingConfigContainer(params: { + container: unknown; + modelId?: string; + thinkingLevel?: GoogleThinkingInputLevel; +}): void { + if (!params.container || typeof params.container !== "object") { + return; + } + const configObj = params.container as Record; + const thinkingConfig = configObj.thinkingConfig; + if (!thinkingConfig || typeof thinkingConfig !== "object") { + return; + } + const thinkingConfigObj = thinkingConfig as Record; + + if (typeof params.modelId === "string" && isGemma4Model(params.modelId)) { + const normalizedThinkingLevel = normalizeGemma4ThinkingLevel(thinkingConfigObj.thinkingLevel); + const explicitMappedLevel = mapThinkLevelToGemma4ThinkingLevel(params.thinkingLevel); + const disabledViaBudget = + typeof thinkingConfigObj.thinkingBudget === "number" && thinkingConfigObj.thinkingBudget <= 0; + const hadThinkingBudget = thinkingConfigObj.thinkingBudget !== undefined; + delete thinkingConfigObj.thinkingBudget; + + if ( + params.thinkingLevel === "off" || + (disabledViaBudget && explicitMappedLevel === undefined && !normalizedThinkingLevel) + ) { + delete thinkingConfigObj.thinkingLevel; + if (Object.keys(thinkingConfigObj).length === 0) { + delete configObj.thinkingConfig; + } + return; + } + + const mappedLevel = + explicitMappedLevel ?? normalizedThinkingLevel ?? (hadThinkingBudget ? "MINIMAL" : undefined); + + if (mappedLevel) { + thinkingConfigObj.thinkingLevel = mappedLevel; + } + return; + } + + const thinkingBudget = thinkingConfigObj.thinkingBudget; + + if (typeof params.modelId === "string" && isGoogleGemini3ThinkingLevelModel(params.modelId)) { + const mappedLevel = resolveGoogleGemini3ThinkingLevel({ + modelId: params.modelId, + thinkingLevel: params.thinkingLevel, + thinkingBudget: typeof thinkingBudget === "number" ? thinkingBudget : undefined, + }); + delete thinkingConfigObj.thinkingBudget; + if (mappedLevel) { + thinkingConfigObj.thinkingLevel = mappedLevel; + } + if (Object.keys(thinkingConfigObj).length === 0) { + delete configObj.thinkingConfig; + } + return; + } + + if ( + stripInvalidGoogleThinkingBudget({ thinkingConfig: thinkingConfigObj, modelId: params.modelId }) + ) { + if (Object.keys(thinkingConfigObj).length === 0) { + delete configObj.thinkingConfig; + } + return; + } + + if (typeof thinkingBudget !== "number" || thinkingBudget >= 0) { + return; + } + + delete thinkingConfigObj.thinkingBudget; + if (Object.keys(thinkingConfigObj).length === 0) { + delete configObj.thinkingConfig; + } +} + +export function createGoogleThinkingPayloadWrapper( + baseStreamFn: ProviderWrapStreamFnContext["streamFn"], + thinkingLevel?: GoogleThinkingInputLevel, +): NonNullable { + return createPayloadPatchStreamWrapper(baseStreamFn, ({ payload, model }) => { + if (model.api === "google-generative-ai") { + sanitizeGoogleThinkingPayload({ + payload, + modelId: model.id, + thinkingLevel, + }); + } + }); +} + +export function createGoogleThinkingStreamWrapper( + ctx: ProviderWrapStreamFnContext, +): NonNullable { + return createGoogleThinkingPayloadWrapper(ctx.streamFn, ctx.thinkingLevel); +} diff --git a/src/agents/google-thinking-compat.ts b/src/agents/google-thinking-compat.ts index 9d8f91f7601..1173d8cf739 100644 --- a/src/agents/google-thinking-compat.ts +++ b/src/agents/google-thinking-compat.ts @@ -1,102 +1,10 @@ -import { normalizeLowercaseStringOrEmpty } from "../shared/string-coerce.js"; - -export type GoogleThinkingLevel = "MINIMAL" | "LOW" | "MEDIUM" | "HIGH"; -export type GoogleThinkingInputLevel = - | "off" - | "minimal" - | "low" - | "medium" - | "adaptive" - | "high" - | "xhigh"; - -// Gemini 2.5 Pro only works in thinking mode and rejects thinkingBudget=0 with -// "Budget 0 is invalid. This model only works in thinking mode." -export function isGoogleThinkingRequiredModel(modelId: string): boolean { - return normalizeLowercaseStringOrEmpty(modelId).includes("gemini-2.5-pro"); -} - -export function isGoogleGemini3ProModel(modelId: string): boolean { - const normalized = normalizeLowercaseStringOrEmpty(modelId); - return /(?:^|\/)gemini-(?:3(?:\.\d+)?-pro|pro-latest)(?:-|$)/.test(normalized); -} - -export function isGoogleGemini3FlashModel(modelId: string): boolean { - const normalized = normalizeLowercaseStringOrEmpty(modelId); - return /(?:^|\/)gemini-(?:3(?:\.\d+)?-flash|flash(?:-lite)?-latest)(?:-|$)/.test(normalized); -} - -export function isGoogleGemini3ThinkingLevelModel(modelId: string): boolean { - return isGoogleGemini3ProModel(modelId) || isGoogleGemini3FlashModel(modelId); -} - -export function resolveGoogleGemini3ThinkingLevel(params: { - modelId?: string; - thinkingLevel?: GoogleThinkingInputLevel; - thinkingBudget?: number; -}): GoogleThinkingLevel | undefined { - if (typeof params.modelId !== "string") { - return undefined; - } - if (isGoogleGemini3ProModel(params.modelId)) { - switch (params.thinkingLevel) { - case "off": - case "minimal": - case "low": - return "LOW"; - case "medium": - case "adaptive": - case "high": - case "xhigh": - return "HIGH"; - } - if (typeof params.thinkingBudget === "number") { - return params.thinkingBudget <= 2048 ? "LOW" : "HIGH"; - } - return undefined; - } - if (!isGoogleGemini3FlashModel(params.modelId)) { - return undefined; - } - switch (params.thinkingLevel) { - case "off": - case "minimal": - return "MINIMAL"; - case "low": - return "LOW"; - case "medium": - case "adaptive": - return "MEDIUM"; - case "high": - case "xhigh": - return "HIGH"; - } - if (typeof params.thinkingBudget !== "number") { - return undefined; - } - if (params.thinkingBudget <= 0) { - return "MINIMAL"; - } - if (params.thinkingBudget <= 2048) { - return "LOW"; - } - if (params.thinkingBudget <= 8192) { - return "MEDIUM"; - } - return "HIGH"; -} - -export function stripInvalidGoogleThinkingBudget(params: { - thinkingConfig: Record; - modelId?: string; -}): boolean { - if ( - params.thinkingConfig.thinkingBudget !== 0 || - typeof params.modelId !== "string" || - !isGoogleThinkingRequiredModel(params.modelId) - ) { - return false; - } - delete params.thinkingConfig.thinkingBudget; - return true; -} +export { + isGoogleGemini3FlashModel, + isGoogleGemini3ProModel, + isGoogleGemini3ThinkingLevelModel, + isGoogleThinkingRequiredModel, + resolveGoogleGemini3ThinkingLevel, + stripInvalidGoogleThinkingBudget, + type GoogleThinkingInputLevel, + type GoogleThinkingLevel, +} from "../../extensions/google/thinking-api.js"; diff --git a/src/agents/pi-embedded-runner/google-stream-wrappers.ts b/src/agents/pi-embedded-runner/google-stream-wrappers.ts index cd25414e534..4d8e9bf9be5 100644 --- a/src/agents/pi-embedded-runner/google-stream-wrappers.ts +++ b/src/agents/pi-embedded-runner/google-stream-wrappers.ts @@ -1,170 +1,4 @@ -import type { StreamFn } from "@mariozechner/pi-agent-core"; -import { streamSimple } from "@mariozechner/pi-ai"; -import type { ThinkLevel } from "../../auto-reply/thinking.js"; -import { normalizeLowercaseStringOrEmpty } from "../../shared/string-coerce.js"; -import { - isGoogleGemini3ThinkingLevelModel, - resolveGoogleGemini3ThinkingLevel, - stripInvalidGoogleThinkingBudget, -} from "../google-thinking-compat.js"; -import { streamWithPayloadPatch } from "./stream-payload-utils.js"; - -function isGemma4Model(modelId: string): boolean { - return normalizeLowercaseStringOrEmpty(modelId).startsWith("gemma-4"); -} - -function mapThinkLevelToGemma4ThinkingLevel( - thinkingLevel?: ThinkLevel, -): "MINIMAL" | "HIGH" | undefined { - switch (thinkingLevel) { - case "off": - return undefined; - case "minimal": - case "low": - return "MINIMAL"; - case "medium": - case "adaptive": - case "high": - case "xhigh": - return "HIGH"; - default: - return undefined; - } -} - -function normalizeGemma4ThinkingLevel(value: unknown): "MINIMAL" | "HIGH" | undefined { - if (typeof value !== "string") { - return undefined; - } - switch (value.trim().toUpperCase()) { - case "MINIMAL": - case "LOW": - return "MINIMAL"; - case "MEDIUM": - case "HIGH": - return "HIGH"; - default: - return undefined; - } -} - -export function sanitizeGoogleThinkingPayload(params: { - payload: unknown; - modelId?: string; - thinkingLevel?: ThinkLevel; -}): void { - if (!params.payload || typeof params.payload !== "object") { - return; - } - const payloadObj = params.payload as Record; - sanitizeGoogleThinkingConfigContainer({ - container: payloadObj.config, - modelId: params.modelId, - thinkingLevel: params.thinkingLevel, - }); - sanitizeGoogleThinkingConfigContainer({ - container: payloadObj.generationConfig, - modelId: params.modelId, - thinkingLevel: params.thinkingLevel, - }); -} - -function sanitizeGoogleThinkingConfigContainer(params: { - container: unknown; - modelId?: string; - thinkingLevel?: ThinkLevel; -}): void { - if (!params.container || typeof params.container !== "object") { - return; - } - const configObj = params.container as Record; - const thinkingConfig = configObj.thinkingConfig; - if (!thinkingConfig || typeof thinkingConfig !== "object") { - return; - } - const thinkingConfigObj = thinkingConfig as Record; - - if (typeof params.modelId === "string" && isGemma4Model(params.modelId)) { - const normalizedThinkingLevel = normalizeGemma4ThinkingLevel(thinkingConfigObj.thinkingLevel); - const explicitMappedLevel = mapThinkLevelToGemma4ThinkingLevel(params.thinkingLevel); - const disabledViaBudget = - typeof thinkingConfigObj.thinkingBudget === "number" && thinkingConfigObj.thinkingBudget <= 0; - const hadThinkingBudget = thinkingConfigObj.thinkingBudget !== undefined; - delete thinkingConfigObj.thinkingBudget; - - if ( - params.thinkingLevel === "off" || - (disabledViaBudget && explicitMappedLevel === undefined && !normalizedThinkingLevel) - ) { - delete thinkingConfigObj.thinkingLevel; - if (Object.keys(thinkingConfigObj).length === 0) { - delete configObj.thinkingConfig; - } - return; - } - - const mappedLevel = - explicitMappedLevel ?? normalizedThinkingLevel ?? (hadThinkingBudget ? "MINIMAL" : undefined); - - if (mappedLevel) { - thinkingConfigObj.thinkingLevel = mappedLevel; - } - return; - } - - const thinkingBudget = thinkingConfigObj.thinkingBudget; - - if (typeof params.modelId === "string" && isGoogleGemini3ThinkingLevelModel(params.modelId)) { - const mappedLevel = resolveGoogleGemini3ThinkingLevel({ - modelId: params.modelId, - thinkingLevel: params.thinkingLevel, - thinkingBudget: typeof thinkingBudget === "number" ? thinkingBudget : undefined, - }); - delete thinkingConfigObj.thinkingBudget; - if (mappedLevel) { - thinkingConfigObj.thinkingLevel = mappedLevel; - } - if (Object.keys(thinkingConfigObj).length === 0) { - delete configObj.thinkingConfig; - } - return; - } - - if ( - stripInvalidGoogleThinkingBudget({ thinkingConfig: thinkingConfigObj, modelId: params.modelId }) - ) { - if (Object.keys(thinkingConfigObj).length === 0) { - delete configObj.thinkingConfig; - } - return; - } - - if (typeof thinkingBudget !== "number" || thinkingBudget >= 0) { - return; - } - - // pi-ai can emit thinkingBudget=-1 for some Google model IDs; a negative budget - // is invalid for Google-compatible backends and can lead to malformed handling. - delete thinkingConfigObj.thinkingBudget; - if (Object.keys(thinkingConfigObj).length === 0) { - delete configObj.thinkingConfig; - } -} - -export function createGoogleThinkingPayloadWrapper( - baseStreamFn: StreamFn | undefined, - thinkingLevel?: ThinkLevel, -): StreamFn { - const underlying = baseStreamFn ?? streamSimple; - return (model, context, options) => { - return streamWithPayloadPatch(underlying, model, context, options, (payload) => { - if (model.api === "google-generative-ai") { - sanitizeGoogleThinkingPayload({ - payload, - modelId: model.id, - thinkingLevel, - }); - } - }); - }; -} +export { + createGoogleThinkingPayloadWrapper, + sanitizeGoogleThinkingPayload, +} from "../../../extensions/google/thinking-api.js"; diff --git a/src/plugin-sdk/provider-stream-shared.ts b/src/plugin-sdk/provider-stream-shared.ts index 97ec6d7e0b1..ceabdcea590 100644 --- a/src/plugin-sdk/provider-stream-shared.ts +++ b/src/plugin-sdk/provider-stream-shared.ts @@ -1,5 +1,6 @@ import type { StreamFn } from "@mariozechner/pi-agent-core"; import { streamSimple } from "@mariozechner/pi-ai"; +import { streamWithPayloadPatch } from "../agents/pi-embedded-runner/stream-payload-utils.js"; export type ProviderStreamWrapperFactory = | ((streamFn: StreamFn | undefined) => StreamFn | undefined) @@ -132,6 +133,20 @@ export function createHtmlEntityToolCallArgumentDecodingWrapper( }; } +export function createPayloadPatchStreamWrapper( + baseStreamFn: StreamFn | undefined, + patchPayload: (params: { + payload: Record; + model: Parameters[0]; + }) => void, +): StreamFn { + const underlying = baseStreamFn ?? streamSimple; + return (model, context, options) => + streamWithPayloadPatch(underlying, model, context, options, (payload) => + patchPayload({ payload, model }), + ); +} + export { applyAnthropicPayloadPolicyToParams, resolveAnthropicPayloadPolicy, @@ -149,7 +164,7 @@ export { createMoonshotThinkingWrapper, resolveMoonshotThinkingType, } from "../agents/pi-embedded-runner/moonshot-thinking-stream-wrappers.js"; -export { streamWithPayloadPatch } from "../agents/pi-embedded-runner/stream-payload-utils.js"; +export { streamWithPayloadPatch }; export { createToolStreamWrapper, createZaiToolStreamWrapper,