Files
oh-my-pi/packages/ai/src/stream.ts
T

743 lines
24 KiB
TypeScript

import * as fs from "node:fs";
import * as os from "node:os";
import * as path from "node:path";
import { $env, $pickenv } from "@oh-my-pi/pi-utils";
import { getCustomApi } from "./api-registry";
import type { Effort } from "./model-thinking";
import {
mapEffortToAnthropicAdaptiveEffort,
mapEffortToGoogleThinkingLevel,
requireSupportedEffort,
} from "./model-thinking";
import { type BedrockOptions, streamBedrock } from "./providers/amazon-bedrock";
import { type AnthropicOptions, streamAnthropic } from "./providers/anthropic";
import { streamAzureOpenAIResponses } from "./providers/azure-openai-responses";
import { type CursorOptions, streamCursor } from "./providers/cursor";
import { isGitLabDuoModel, streamGitLabDuo } from "./providers/gitlab-duo";
import { type GoogleOptions, streamGoogle } from "./providers/google";
import { type GoogleGeminiCliOptions, streamGoogleGeminiCli } from "./providers/google-gemini-cli";
import { type GoogleVertexOptions, streamGoogleVertex } from "./providers/google-vertex";
import { isKimiModel, streamKimi } from "./providers/kimi";
import { type OllamaChatOptions, streamOllama } from "./providers/ollama";
import { streamOpenAICodexResponses } from "./providers/openai-codex-responses";
import { type OpenAICompletionsOptions, streamOpenAICompletions } from "./providers/openai-completions";
import { streamOpenAIResponses } from "./providers/openai-responses";
import { isSyntheticModel, streamSynthetic } from "./providers/synthetic";
import type {
Api,
AssistantMessage,
AssistantMessageEventStream,
Context,
Model,
OptionsForApi,
SimpleStreamOptions,
StreamOptions,
ThinkingBudgets,
ToolChoice,
} from "./types";
import { isFoundryEnabled } from "./utils/foundry";
let cachedVertexAdcCredentialsExists: boolean | null = null;
function hasVertexAdcCredentials(): boolean {
if (cachedVertexAdcCredentialsExists === null) {
const gacPath = $env.GOOGLE_APPLICATION_CREDENTIALS;
if (gacPath) {
cachedVertexAdcCredentialsExists = fs.existsSync(gacPath);
} else {
cachedVertexAdcCredentialsExists = fs.existsSync(
path.join(os.homedir(), ".config", "gcloud", "application_default_credentials.json"),
);
}
}
return cachedVertexAdcCredentialsExists;
}
type KeyResolver = string | (() => string | undefined);
const serviceProviderMap: Record<string, KeyResolver> = {
"alibaba-coding-plan": "ALIBABA_CODING_PLAN_API_KEY",
openai: "OPENAI_API_KEY",
google: "GEMINI_API_KEY",
groq: "GROQ_API_KEY",
cerebras: "CEREBRAS_API_KEY",
xai: "XAI_API_KEY",
fireworks: "FIREWORKS_API_KEY",
openrouter: "OPENROUTER_API_KEY",
kilo: "KILO_API_KEY",
"vercel-ai-gateway": "AI_GATEWAY_API_KEY",
zai: "ZAI_API_KEY",
mistral: "MISTRAL_API_KEY",
minimax: "MINIMAX_API_KEY",
"minimax-code": "MINIMAX_CODE_API_KEY",
"minimax-code-cn": "MINIMAX_CODE_CN_API_KEY",
"opencode-go": "OPENCODE_API_KEY",
"opencode-zen": "OPENCODE_API_KEY",
cursor: "CURSOR_ACCESS_TOKEN",
deepseek: "DEEPSEEK_API_KEY",
"openai-codex": "OPENAI_CODEX_OAUTH_TOKEN",
"azure-openai-responses": "AZURE_OPENAI_API_KEY",
exa: "EXA_API_KEY",
jina: "JINA_API_KEY",
brave: "BRAVE_API_KEY",
perplexity: "PERPLEXITY_API_KEY",
tavily: "TAVILY_API_KEY",
parallel: "PARALLEL_API_KEY",
kagi: "KAGI_API_KEY",
// GitHub Copilot uses GitHub personal access token
"github-copilot": () => $pickenv("COPILOT_GITHUB_TOKEN", "GH_TOKEN", "GITHUB_TOKEN"),
// Foundry mode optionally switches Anthropic auth to enterprise gateway credentials.
anthropic: () =>
isFoundryEnabled()
? $pickenv("ANTHROPIC_FOUNDRY_API_KEY", "ANTHROPIC_OAUTH_TOKEN", "ANTHROPIC_API_KEY")
: $pickenv("ANTHROPIC_OAUTH_TOKEN", "ANTHROPIC_API_KEY"),
"gitlab-duo": "GITLAB_TOKEN",
// Vertex AI supports either GOOGLE_CLOUD_API_KEY or Application Default Credentials.
"google-vertex": () => {
if ($env.GOOGLE_CLOUD_API_KEY) {
return $env.GOOGLE_CLOUD_API_KEY;
}
const hasCredentials = hasVertexAdcCredentials();
const hasProject = !!($env.GOOGLE_CLOUD_PROJECT || $env.GCLOUD_PROJECT);
const hasLocation = !!$env.GOOGLE_CLOUD_LOCATION;
if (hasCredentials && hasProject && hasLocation) {
return "<authenticated>";
}
},
// Amazon Bedrock supports multiple credential sources:
// 1. AWS_PROFILE - named profile from ~/.aws/credentials
// 2. AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY - standard IAM keys
// 3. AWS_BEARER_TOKEN_BEDROCK - Bedrock API keys (bearer token)
// 4. AWS_CONTAINER_CREDENTIALS_* - ECS/Task IAM role credentials
// 5. AWS_WEB_IDENTITY_TOKEN_FILE + AWS_ROLE_ARN - IRSA (EKS) web identity
"amazon-bedrock": () => {
const hasEcsCredentials =
!!$env.AWS_CONTAINER_CREDENTIALS_RELATIVE_URI || !!$env.AWS_CONTAINER_CREDENTIALS_FULL_URI;
const hasWebIdentity = !!$env.AWS_WEB_IDENTITY_TOKEN_FILE && !!$env.AWS_ROLE_ARN;
if (
$env.AWS_PROFILE ||
($env.AWS_ACCESS_KEY_ID && $env.AWS_SECRET_ACCESS_KEY) ||
$env.AWS_BEARER_TOKEN_BEDROCK ||
hasEcsCredentials ||
hasWebIdentity
) {
return "<authenticated>";
}
},
synthetic: "SYNTHETIC_API_KEY",
"cloudflare-ai-gateway": "CLOUDFLARE_AI_GATEWAY_API_KEY",
huggingface: () => $pickenv("HUGGINGFACE_HUB_TOKEN", "HF_TOKEN"),
litellm: "LITELLM_API_KEY",
moonshot: "MOONSHOT_API_KEY",
nvidia: "NVIDIA_API_KEY",
nanogpt: "NANO_GPT_API_KEY",
"lm-studio": "LM_STUDIO_API_KEY",
ollama: "OLLAMA_API_KEY",
"ollama-cloud": "OLLAMA_CLOUD_API_KEY",
"llama.cpp": "LLAMA_CPP_API_KEY",
qianfan: "QIANFAN_API_KEY",
"qwen-portal": () => $pickenv("QWEN_OAUTH_TOKEN", "QWEN_PORTAL_API_KEY"),
together: "TOGETHER_API_KEY",
zenmux: "ZENMUX_API_KEY",
venice: "VENICE_API_KEY",
vllm: "VLLM_API_KEY",
xiaomi: "XIAOMI_API_KEY",
};
/**
* Get API key for provider from known environment variables, e.g. OPENAI_API_KEY.
*
* Will not return API keys for providers that require OAuth tokens.
* Checks Bun.env, then cwd/.env, then ~/.env.
*/
export function getEnvApiKey(provider: string): string | undefined {
const resolver = serviceProviderMap[provider];
if (typeof resolver === "string") {
return $env[resolver];
}
return resolver?.();
}
export function stream<TApi extends Api>(
model: Model<TApi>,
context: Context,
options?: OptionsForApi<TApi>,
): AssistantMessageEventStream {
// Check custom API registry first (extension-provided APIs like "vertex-claude-api")
const customApiProvider = getCustomApi(model.api);
if (customApiProvider) {
return customApiProvider.stream(model, context, options as StreamOptions);
}
if (isGitLabDuoModel(model)) {
const apiKey = (options as StreamOptions | undefined)?.apiKey || getEnvApiKey(model.provider);
if (!apiKey) {
throw new Error(`No API key for provider: ${model.provider}`);
}
return streamGitLabDuo(model, context, {
...(options as SimpleStreamOptions | undefined),
apiKey,
});
}
// Vertex AI uses Application Default Credentials, not API keys
if (model.api === "google-vertex") {
return streamGoogleVertex(model as Model<"google-vertex">, context, options as GoogleVertexOptions);
} else if (model.api === "bedrock-converse-stream") {
// Bedrock doesn't have any API keys instead it sources credentials from standard AWS env variables or from given AWS profile.
return streamBedrock(model as Model<"bedrock-converse-stream">, context, (options || {}) as BedrockOptions);
}
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
if (!apiKey) {
throw new Error(`No API key for provider: ${model.provider}`);
}
const providerOptions = { ...options, apiKey };
const api: Api = model.api;
switch (api) {
case "anthropic-messages": {
const anthropicOptions = providerOptions as AnthropicOptions;
return streamAnthropic(model as Model<"anthropic-messages">, context, {
...anthropicOptions,
isOAuth: anthropicOptions.isOAuth ?? model.isOAuth,
});
}
case "openai-completions":
return streamOpenAICompletions(model as Model<"openai-completions">, context, providerOptions as any);
case "openai-responses":
return streamOpenAIResponses(model as Model<"openai-responses">, context, providerOptions as any);
case "azure-openai-responses":
return streamAzureOpenAIResponses(model as Model<"azure-openai-responses">, context, providerOptions as any);
case "openai-codex-responses":
return streamOpenAICodexResponses(model as Model<"openai-codex-responses">, context, providerOptions as any);
case "google-generative-ai":
return streamGoogle(model as Model<"google-generative-ai">, context, providerOptions);
case "google-gemini-cli":
return streamGoogleGeminiCli(
model as Model<"google-gemini-cli">,
context,
providerOptions as GoogleGeminiCliOptions,
);
case "ollama-chat":
return streamOllama(model as Model<"ollama-chat">, context, providerOptions as OllamaChatOptions);
case "cursor-agent":
return streamCursor(model as Model<"cursor-agent">, context, providerOptions as CursorOptions);
default:
throw new Error(`Unhandled API: ${api}`);
}
}
export async function complete<TApi extends Api>(
model: Model<TApi>,
context: Context,
options?: OptionsForApi<TApi>,
): Promise<AssistantMessage> {
const s = stream(model, context, options);
return s.result();
}
export function streamSimple<TApi extends Api>(
model: Model<TApi>,
context: Context,
options?: SimpleStreamOptions,
): AssistantMessageEventStream {
// Check custom API registry first (extension-provided APIs)
const customApiProvider = getCustomApi(model.api);
if (customApiProvider) {
return customApiProvider.streamSimple(model, context, options);
}
// Vertex AI uses Application Default Credentials, not API keys
if (model.api === "google-vertex") {
const providerOptions = mapOptionsForApi(model, options, undefined);
return stream(model, context, providerOptions);
} else if (model.api === "bedrock-converse-stream") {
// Bedrock doesn't have any API keys instead it sources credentials from standard AWS env variables or from given AWS profile.
const providerOptions = mapOptionsForApi(model, options, undefined);
return stream(model, context, providerOptions);
}
const apiKey = options?.apiKey || getEnvApiKey(model.provider);
if (!apiKey) {
throw new Error(`No API key for provider: ${model.provider}`);
}
// GitLab Duo - wraps Anthropic/OpenAI behind GitLab AI Gateway direct access tokens
if (isGitLabDuoModel(model)) {
return streamGitLabDuo(model, context, {
...options,
apiKey,
});
}
// Kimi Code - route to dedicated handler that wraps OpenAI or Anthropic API
if (isKimiModel(model)) {
// Pass raw SimpleStreamOptions - streamKimi handles mapping internally
return streamKimi(model as Model<"openai-completions">, context, {
...options,
apiKey,
format: options?.kimiApiFormat ?? "anthropic",
});
}
// Synthetic - route to dedicated handler that wraps OpenAI or Anthropic API
if (isSyntheticModel(model)) {
// Pass raw SimpleStreamOptions - streamSynthetic handles mapping internally
return streamSynthetic(model as Model<"openai-completions">, context, {
...options,
apiKey,
format: options?.syntheticApiFormat ?? "openai", // Default to OpenAI format
});
}
const providerOptions = mapOptionsForApi(model, options, apiKey);
return stream(model, context, providerOptions);
}
export async function completeSimple<TApi extends Api>(
model: Model<TApi>,
context: Context,
options?: SimpleStreamOptions,
): Promise<AssistantMessage> {
const s = streamSimple(model, context, options);
return s.result();
}
const MIN_OUTPUT_TOKENS = 1024;
export const OUTPUT_FALLBACK_BUFFER = 4000;
const ANTHROPIC_USE_INTERLEAVED_THINKING = Bun.env.PI_NO_INTERLEAVED_THINKING !== "1";
export const ANTHROPIC_THINKING: Record<Effort, number> = {
minimal: 1024,
low: 4096,
medium: 8192,
high: 16384,
xhigh: 32768,
};
const GOOGLE_THINKING: Record<Effort, number> = {
minimal: 1024,
low: 4096,
medium: 8192,
high: 16384,
xhigh: 24575,
};
const BEDROCK_CLAUDE_THINKING: Record<Effort, number> = {
minimal: 1024,
low: 2048,
medium: 8192,
high: 16384,
xhigh: 16384,
};
function resolveBedrockThinkingBudget(
model: Model<"bedrock-converse-stream">,
options?: SimpleStreamOptions,
): { budget: number; level: Effort } | null {
if (!options?.reasoning || !model.reasoning) return null;
const level = requireSupportedEffort(model, options.reasoning);
const budget = options.thinkingBudgets?.[level] ?? BEDROCK_CLAUDE_THINKING[level];
return { budget, level };
}
export function mapAnthropicToolChoice(choice?: ToolChoice): AnthropicOptions["toolChoice"] {
if (!choice) return undefined;
if (typeof choice === "string") {
if (choice === "required") return "any";
if (choice === "auto" || choice === "none" || choice === "any") return choice;
return undefined;
}
if (choice.type === "tool") {
return choice.name ? { type: "tool", name: choice.name } : undefined;
}
if (choice.type === "function") {
const name = "function" in choice ? choice.function?.name : choice.name;
return name ? { type: "tool", name } : undefined;
}
return undefined;
}
function mapGoogleToolChoice(
choice?: ToolChoice,
): GoogleOptions["toolChoice"] | GoogleGeminiCliOptions["toolChoice"] | GoogleVertexOptions["toolChoice"] {
if (!choice) return undefined;
if (typeof choice === "string") {
if (choice === "required") return "any";
if (choice === "auto" || choice === "none" || choice === "any") return choice;
return undefined;
}
return "any";
}
function mapOpenAiToolChoice(choice?: ToolChoice): OpenAICompletionsOptions["toolChoice"] {
if (!choice) return undefined;
if (typeof choice === "string") {
if (choice === "any") return "required";
if (choice === "auto" || choice === "none" || choice === "required") return choice;
return undefined;
}
if (choice.type === "tool") {
return choice.name ? { type: "function", function: { name: choice.name } } : undefined;
}
if (choice.type === "function") {
const name = "function" in choice ? choice.function?.name : choice.name;
return name ? { type: "function", function: { name } } : undefined;
}
return undefined;
}
function resolveOpenAiReasoningEffort<TApi extends Api>(
model: Model<TApi>,
options?: SimpleStreamOptions,
): Effort | undefined {
const reasoning = options?.reasoning;
if (!reasoning || !model.reasoning) return undefined;
return requireSupportedEffort(model, reasoning);
}
const castApi = <TApi extends Api>(api: OptionsForApi<TApi>): OptionsForApi<Api> => api as OptionsForApi<Api>;
function mapOptionsForApi<TApi extends Api>(
model: Model<TApi>,
options?: SimpleStreamOptions,
apiKey?: string,
): OptionsForApi<TApi> {
const base = {
temperature: options?.temperature,
topP: options?.topP,
topK: options?.topK,
minP: options?.minP,
presencePenalty: options?.presencePenalty,
repetitionPenalty: options?.repetitionPenalty,
maxTokens: options?.maxTokens || Math.min(model.maxTokens, 32000),
signal: options?.signal,
apiKey: apiKey || options?.apiKey,
cacheRetention: options?.cacheRetention,
headers: options?.headers,
initiatorOverride: options?.initiatorOverride,
maxRetryDelayMs: options?.maxRetryDelayMs,
metadata: options?.metadata,
sessionId: options?.sessionId,
providerSessionState: options?.providerSessionState,
onPayload: options?.onPayload,
onResponse: options?.onResponse,
execHandlers: options?.execHandlers,
};
switch (model.api) {
case "anthropic-messages": {
// Explicitly disable thinking when reasoning is not specified or model doesn't support it
const reasoning = options?.reasoning;
if (!reasoning || !model.reasoning) {
return castApi<"anthropic-messages">({
...base,
thinkingEnabled: false,
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
});
}
let thinkingBudget = options.thinkingBudgets?.[reasoning] ?? ANTHROPIC_THINKING[reasoning];
if (thinkingBudget <= 0) {
return castApi<"anthropic-messages">({
...base,
thinkingEnabled: false,
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
});
}
// For Opus 4.6+ and Sonnet 4.6+: use adaptive thinking with effort level
// For older models: use budget-based thinking
if (model.thinking?.mode === "anthropic-adaptive") {
const effort = mapEffortToAnthropicAdaptiveEffort(model, reasoning);
return castApi<"anthropic-messages">({
...base,
thinkingEnabled: true,
effort,
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
});
}
if (ANTHROPIC_USE_INTERLEAVED_THINKING) {
return castApi<"anthropic-messages">({
...base,
thinkingEnabled: true,
thinkingBudgetTokens: thinkingBudget,
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
});
}
// Caller's maxTokens is the desired output; add thinking budget on top, capped at model limit
const maxTokens = Math.min((base.maxTokens || 0) + thinkingBudget, model.maxTokens);
// If not enough room for thinking + output, reduce thinking budget
if (maxTokens <= thinkingBudget) {
thinkingBudget = maxTokens - MIN_OUTPUT_TOKENS;
}
// If thinking budget is too low, disable thinking
if (thinkingBudget <= 0) {
return castApi<"anthropic-messages">({
...base,
thinkingEnabled: false,
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
});
} else {
return castApi<"anthropic-messages">({
...base,
maxTokens,
thinkingEnabled: true,
thinkingBudgetTokens: thinkingBudget,
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
});
}
}
case "bedrock-converse-stream": {
const bedrockBase: BedrockOptions = {
...base,
reasoning: options?.reasoning,
thinkingBudgets: options?.thinkingBudgets,
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
};
// Adaptive mode sends effort directly, no budget_tokens — skip budget inflation.
if (model.thinking?.mode === "anthropic-adaptive") {
return castApi<"bedrock-converse-stream">(bedrockBase);
}
const budgetInfo = resolveBedrockThinkingBudget(model as Model<"bedrock-converse-stream">, options);
if (!budgetInfo) return bedrockBase as OptionsForApi<TApi>;
let maxTokens = bedrockBase.maxTokens ?? model.maxTokens;
let thinkingBudgets = bedrockBase.thinkingBudgets;
if (maxTokens <= budgetInfo.budget) {
const desiredMaxTokens = Math.min(model.maxTokens, budgetInfo.budget + MIN_OUTPUT_TOKENS);
if (desiredMaxTokens > maxTokens) {
maxTokens = desiredMaxTokens;
}
}
if (maxTokens <= budgetInfo.budget) {
const adjustedBudget = Math.max(0, maxTokens - MIN_OUTPUT_TOKENS);
thinkingBudgets = { ...(thinkingBudgets ?? {}), [budgetInfo.level]: adjustedBudget };
}
return castApi<"bedrock-converse-stream">({ ...bedrockBase, maxTokens, thinkingBudgets });
}
case "openai-completions":
return castApi<"openai-completions">({
...base,
reasoning: resolveOpenAiReasoningEffort(model, options),
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
serviceTier: options?.serviceTier,
});
case "openai-responses":
return castApi<"openai-responses">({
...base,
reasoning: resolveOpenAiReasoningEffort(model, options),
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
serviceTier: options?.serviceTier,
});
case "azure-openai-responses":
return castApi<"azure-openai-responses">({
...base,
reasoning: resolveOpenAiReasoningEffort(model, options),
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
serviceTier: options?.serviceTier,
});
case "openai-codex-responses":
return castApi<"openai-codex-responses">({
...base,
reasoning: resolveOpenAiReasoningEffort(model, options),
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
serviceTier: options?.serviceTier,
preferWebsockets: options?.preferWebsockets,
});
case "google-generative-ai": {
// Explicitly disable thinking when reasoning is not specified or model doesn't support it
// This is needed because Gemini has "dynamic thinking" enabled by default
const reasoning = options?.reasoning;
if (!reasoning || !model.reasoning) {
return castApi<"google-generative-ai">({
...base,
thinking: { enabled: false },
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
const googleModel = model as Model<"google-generative-ai">;
const effort = requireSupportedEffort(googleModel, reasoning);
// Gemini 3+ models use thinkingLevel exclusively instead of thinkingBudget.
// https://ai.google.dev/gemini-api/docs/thinking#set-budget
if (googleModel.thinking?.mode === "google-level") {
return castApi<"google-generative-ai">({
...base,
thinking: {
enabled: true,
level: mapEffortToGoogleThinkingLevel(googleModel, effort),
},
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
return castApi<"google-gemini-cli">({
...base,
thinking: {
enabled: true,
budgetTokens: getGoogleBudget(googleModel, effort, options?.thinkingBudgets),
},
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
case "google-gemini-cli": {
const reasoning = options?.reasoning;
if (!reasoning || !model.reasoning) {
return castApi<"google-gemini-cli">({
...base,
thinking: { enabled: false },
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
const effort = requireSupportedEffort(model, reasoning);
// Gemini 3+ models use thinkingLevel instead of thinkingBudget
if (model.thinking?.mode === "google-level") {
return castApi<"google-gemini-cli">({
...base,
thinking: {
enabled: true,
level: mapEffortToGoogleThinkingLevel(model, effort),
},
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
let thinkingBudget = options.thinkingBudgets?.[effort] ?? GOOGLE_THINKING[effort];
// Caller's maxTokens is the desired output; add thinking budget on top, capped at model limit
const maxTokens = Math.min((base.maxTokens || 0) + thinkingBudget, model.maxTokens);
// If not enough room for thinking + output, reduce thinking budget
if (maxTokens <= thinkingBudget) {
thinkingBudget = Math.max(0, maxTokens - MIN_OUTPUT_TOKENS) ?? 0;
}
// If thinking budget is too low, disable thinking
if (thinkingBudget <= 0) {
return castApi<"google-gemini-cli">({
...base,
thinking: { enabled: false },
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
} else {
return castApi<"google-gemini-cli">({
...base,
maxTokens,
thinking: { enabled: true, budgetTokens: thinkingBudget },
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
}
case "google-vertex": {
// Explicitly disable thinking when reasoning is not specified or model doesn't support it
const reasoning = options?.reasoning;
if (!reasoning || !model.reasoning) {
return castApi<"google-vertex">({
...base,
thinking: { enabled: false },
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
const vertexModel = model as Model<"google-vertex">;
const effort = requireSupportedEffort(vertexModel, reasoning);
const geminiModel = vertexModel as unknown as Model<"google-generative-ai">;
if (geminiModel.thinking?.mode === "google-level") {
return castApi<"google-vertex">({
...base,
thinking: {
enabled: true,
level: mapEffortToGoogleThinkingLevel(geminiModel, effort),
},
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
return castApi<"google-vertex">({
...base,
thinking: {
enabled: true,
budgetTokens: getGoogleBudget(geminiModel, effort, options?.thinkingBudgets),
},
toolChoice: mapGoogleToolChoice(options?.toolChoice),
});
}
case "ollama-chat":
return castApi<"ollama-chat">({
...base,
reasoning: resolveOpenAiReasoningEffort(model, options),
toolChoice: options?.toolChoice,
});
case "cursor-agent": {
const execHandlers = options?.cursorExecHandlers ?? options?.execHandlers;
const onToolResult = options?.cursorOnToolResult ?? execHandlers?.onToolResult;
return castApi<"cursor-agent">({
...base,
execHandlers,
onToolResult,
});
}
default:
throw new Error(`Unhandled API in mapOptionsForApi: ${model.api}`);
}
}
function getGoogleBudget(
model: Model<"google-generative-ai">,
effort: Effort,
customBudgets?: ThinkingBudgets,
): number {
requireSupportedEffort(model, effort);
// Custom budgets take precedence if provided for this level
if (customBudgets?.[effort] !== undefined) {
return customBudgets[effort]!;
}
// See https://ai.google.dev/gemini-api/docs/thinking#set-budget
if (model.id.includes("2.5-")) {
switch (effort) {
case "minimal":
return 128;
case "low":
return 2048;
case "medium":
return 8192;
default:
return model.id.includes("2.5-flash") ? 24576 : 32768;
}
}
// Unknown model - use dynamic
return -1;
}