From 858a3f72fab0e397aa0c3916acb5b1643da3e63f Mon Sep 17 00:00:00 2001 From: Peter Steinberger Date: Sat, 18 Apr 2026 19:34:07 +0100 Subject: [PATCH] fix(agents): keep google compat facades in core --- src/agents/google-thinking-compat.ts | 112 +++++++++++- .../google-stream-wrappers.ts | 173 +++++++++++++++++- 2 files changed, 271 insertions(+), 14 deletions(-) diff --git a/src/agents/google-thinking-compat.ts b/src/agents/google-thinking-compat.ts index 1173d8cf739..9d8f91f7601 100644 --- a/src/agents/google-thinking-compat.ts +++ b/src/agents/google-thinking-compat.ts @@ -1,10 +1,102 @@ -export { - isGoogleGemini3FlashModel, - isGoogleGemini3ProModel, - isGoogleGemini3ThinkingLevelModel, - isGoogleThinkingRequiredModel, - resolveGoogleGemini3ThinkingLevel, - stripInvalidGoogleThinkingBudget, - type GoogleThinkingInputLevel, - type GoogleThinkingLevel, -} from "../../extensions/google/thinking-api.js"; +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; +} diff --git a/src/agents/pi-embedded-runner/google-stream-wrappers.ts b/src/agents/pi-embedded-runner/google-stream-wrappers.ts index 4d8e9bf9be5..c64269a1405 100644 --- a/src/agents/pi-embedded-runner/google-stream-wrappers.ts +++ b/src/agents/pi-embedded-runner/google-stream-wrappers.ts @@ -1,4 +1,169 @@ -export { - createGoogleThinkingPayloadWrapper, - sanitizeGoogleThinkingPayload, -} from "../../../extensions/google/thinking-api.js"; +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) => + streamWithPayloadPatch(underlying, model, context, options, (payload) => { + if (model.api === "google-generative-ai") { + sanitizeGoogleThinkingPayload({ + payload, + modelId: model.id, + thinkingLevel, + }); + } + }); +}