93 lines
2.7 KiB
TypeScript
93 lines
2.7 KiB
TypeScript
import type { StreamFn } from "@mariozechner/pi-agent-core";
|
|
import { streamSimple } from "@mariozechner/pi-ai";
|
|
import type { ThinkLevel } from "../../auto-reply/thinking.js";
|
|
|
|
function isGemini31Model(modelId: string): boolean {
|
|
const normalized = modelId.toLowerCase();
|
|
return normalized.includes("gemini-3.1-pro") || normalized.includes("gemini-3.1-flash");
|
|
}
|
|
|
|
function mapThinkLevelToGoogleThinkingLevel(
|
|
thinkingLevel: ThinkLevel,
|
|
): "MINIMAL" | "LOW" | "MEDIUM" | "HIGH" | undefined {
|
|
switch (thinkingLevel) {
|
|
case "minimal":
|
|
return "MINIMAL";
|
|
case "low":
|
|
return "LOW";
|
|
case "medium":
|
|
case "adaptive":
|
|
return "MEDIUM";
|
|
case "high":
|
|
case "xhigh":
|
|
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<string, unknown>;
|
|
const config = payloadObj.config;
|
|
if (!config || typeof config !== "object") {
|
|
return;
|
|
}
|
|
const configObj = config as Record<string, unknown>;
|
|
const thinkingConfig = configObj.thinkingConfig;
|
|
if (!thinkingConfig || typeof thinkingConfig !== "object") {
|
|
return;
|
|
}
|
|
const thinkingConfigObj = thinkingConfig as Record<string, unknown>;
|
|
const thinkingBudget = thinkingConfigObj.thinkingBudget;
|
|
if (typeof thinkingBudget !== "number" || thinkingBudget >= 0) {
|
|
return;
|
|
}
|
|
|
|
// pi-ai can emit thinkingBudget=-1 for some Gemini 3.1 IDs; a negative budget
|
|
// is invalid for Google-compatible backends and can lead to malformed handling.
|
|
delete thinkingConfigObj.thinkingBudget;
|
|
|
|
if (
|
|
typeof params.modelId === "string" &&
|
|
isGemini31Model(params.modelId) &&
|
|
params.thinkingLevel &&
|
|
params.thinkingLevel !== "off" &&
|
|
thinkingConfigObj.thinkingLevel === undefined
|
|
) {
|
|
const mappedLevel = mapThinkLevelToGoogleThinkingLevel(params.thinkingLevel);
|
|
if (mappedLevel) {
|
|
thinkingConfigObj.thinkingLevel = mappedLevel;
|
|
}
|
|
}
|
|
}
|
|
|
|
export function createGoogleThinkingPayloadWrapper(
|
|
baseStreamFn: StreamFn | undefined,
|
|
thinkingLevel?: ThinkLevel,
|
|
): StreamFn {
|
|
const underlying = baseStreamFn ?? streamSimple;
|
|
return (model, context, options) => {
|
|
const onPayload = options?.onPayload;
|
|
return underlying(model, context, {
|
|
...options,
|
|
onPayload: (payload) => {
|
|
if (model.api === "google-generative-ai") {
|
|
sanitizeGoogleThinkingPayload({
|
|
payload,
|
|
modelId: model.id,
|
|
thinkingLevel,
|
|
});
|
|
}
|
|
return onPayload?.(payload, model);
|
|
},
|
|
});
|
|
};
|
|
}
|