can1357
cba6641299
- Added `ApiKeyResolver`/`ApiKey` types and exported auth-retry helpers. - Changed stream and gateway auth retry handling to use resolver steps. - Added initial-key, force-refresh, and rotate credential retries for auth failures. - Updated agent and coding-agent integrations to use context-aware API-key resolvers.
1089 lines
39 KiB
TypeScript
1089 lines
39 KiB
TypeScript
import * as fs from "node:fs";
|
|
import * as os from "node:os";
|
|
import * as path from "node:path";
|
|
import { $env, $pickenv, extractHttpStatusFromError } from "@oh-my-pi/pi-utils";
|
|
import { getCustomApi } from "./api-registry";
|
|
import { AUTH_RETRY_STEPS, isApiKeyResolver, resolveRetryKey } from "./auth-retry";
|
|
import type { Effort } from "./effort";
|
|
import {
|
|
mapEffortToAnthropicAdaptiveEffort,
|
|
mapEffortToGoogleThinkingLevel,
|
|
modelOmitsReasoningEffort,
|
|
requireSupportedEffort,
|
|
} from "./model-thinking";
|
|
import type { BedrockOptions } from "./providers/amazon-bedrock";
|
|
import type { AnthropicOptions } from "./providers/anthropic";
|
|
import type { CursorOptions } from "./providers/cursor";
|
|
import { isGitLabDuoModel, streamGitLabDuo } from "./providers/gitlab-duo";
|
|
import type { GoogleOptions } from "./providers/google";
|
|
import { getVertexAccessToken } from "./providers/google-auth";
|
|
import type { GoogleGeminiCliOptions } from "./providers/google-gemini-cli";
|
|
import type { GoogleVertexOptions } from "./providers/google-vertex";
|
|
import { isKimiModel, streamKimi } from "./providers/kimi";
|
|
import type { OllamaChatOptions } from "./providers/ollama";
|
|
import type { OpenAICompletionsOptions } from "./providers/openai-completions";
|
|
import { streamPiNative } from "./providers/pi-native-client";
|
|
// Heavy provider stream functions are imported lazily via register-builtins,
|
|
// which wraps each provider module in a dynamic import. This keeps the
|
|
// AWS SDK, google-auth-library, @google/genai, @bufbuild/protobuf, and
|
|
// other provider SDKs out of the CLI startup parse graph. The
|
|
// gitlab-duo / kimi / synthetic providers stay eager because their modules
|
|
// export routing predicates (isGitLabDuoModel, isKimiModel, isSyntheticModel)
|
|
// that must be callable synchronously before streaming begins, and their
|
|
// modules are thin wrappers with no heavy SDK dependencies.
|
|
import {
|
|
streamAnthropic,
|
|
streamAzureOpenAIResponses,
|
|
streamBedrock,
|
|
streamCursor,
|
|
streamGoogle,
|
|
streamGoogleGeminiCli,
|
|
streamGoogleVertex,
|
|
streamOllama,
|
|
streamOpenAICodexResponses,
|
|
streamOpenAICompletions,
|
|
streamOpenAIResponses,
|
|
} from "./providers/register-builtins";
|
|
import { isSyntheticModel, streamSynthetic } from "./providers/synthetic";
|
|
import { streamXAIResponses } from "./providers/xai-responses";
|
|
import { isUsageLimitError } from "./rate-limit-utils";
|
|
import type {
|
|
Api,
|
|
AssistantMessage,
|
|
AssistantMessageEvent,
|
|
Context,
|
|
FetchImpl,
|
|
Model,
|
|
OptionsForApi,
|
|
SimpleStreamOptions,
|
|
StreamOptions,
|
|
ThinkingBudgets,
|
|
ToolChoice,
|
|
} from "./types";
|
|
import { AssistantMessageEventStream } from "./utils/event-stream";
|
|
import { isFoundryEnabled } from "./utils/foundry";
|
|
import { withRequestDebugFetch } from "./utils/request-debug";
|
|
|
|
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;
|
|
}
|
|
function isGoogleVertexAuthenticatedModel(model: Model<Api>): boolean {
|
|
return (
|
|
model.provider === "google-vertex" &&
|
|
((model.api === "openai-completions" && model.baseUrl.includes("/endpoints/openapi")) ||
|
|
(model.api === "anthropic-messages" && model.baseUrl.includes(":streamRawPredict")))
|
|
);
|
|
}
|
|
|
|
function createVertexAuthenticatedFetch(options: StreamOptions | undefined): FetchImpl {
|
|
const baseFetch = options?.fetch ?? fetch;
|
|
const vertexFetch = async (input: string | URL | Request, init?: RequestInit): Promise<Response> => {
|
|
const token = await getVertexAccessToken({ signal: options?.signal, fetch: baseFetch });
|
|
const headers = new Headers(init?.headers);
|
|
headers.set("Authorization", `Bearer ${token}`);
|
|
const rewritten = resolveVertexRequest(input);
|
|
const url = rewritten instanceof Request ? rewritten.url : rewritten.toString();
|
|
if (isVertexAnthropicRawPredict(url)) {
|
|
const bodyText = await readVertexRequestBody(rewritten, init);
|
|
const transformed = transformVertexAnthropicBody(bodyText);
|
|
return baseFetch(url, {
|
|
...init,
|
|
method: init?.method ?? (rewritten instanceof Request ? rewritten.method : "POST"),
|
|
headers,
|
|
body: transformed,
|
|
});
|
|
}
|
|
return baseFetch(rewritten, { ...init, headers });
|
|
};
|
|
return Object.assign(vertexFetch, baseFetch.preconnect ? { preconnect: baseFetch.preconnect } : {});
|
|
}
|
|
|
|
function isVertexAnthropicRawPredict(url: string): boolean {
|
|
return url.includes(":streamRawPredict") || url.includes(":rawPredict");
|
|
}
|
|
|
|
async function readVertexRequestBody(input: string | URL | Request, init: RequestInit | undefined): Promise<string> {
|
|
if (input instanceof Request) return input.clone().text();
|
|
const body = init?.body;
|
|
if (typeof body === "string") return body;
|
|
if (body instanceof Uint8Array) return new TextDecoder().decode(body);
|
|
if (body instanceof ArrayBuffer) return new TextDecoder().decode(body);
|
|
return "";
|
|
}
|
|
|
|
// Vertex Claude rejects the standard Anthropic body shape: the `model` field
|
|
// is encoded in the URL path and `anthropic_version: "vertex-2023-10-16"` is
|
|
// required in the JSON body instead of the `anthropic-version` HTTP header.
|
|
function transformVertexAnthropicBody(bodyText: string): string {
|
|
if (!bodyText) return bodyText;
|
|
try {
|
|
const payload = JSON.parse(bodyText) as Record<string, unknown>;
|
|
delete payload.model;
|
|
payload.anthropic_version = "vertex-2023-10-16";
|
|
return JSON.stringify(payload);
|
|
} catch {
|
|
return bodyText;
|
|
}
|
|
}
|
|
|
|
function resolveVertexRequest(input: string | URL | Request): string | URL | Request {
|
|
const project = $env.GOOGLE_CLOUD_PROJECT || $env.GCP_PROJECT || $env.GCLOUD_PROJECT;
|
|
const location = $env.GOOGLE_VERTEX_LOCATION || $env.GOOGLE_CLOUD_LOCATION || $env.VERTEX_LOCATION;
|
|
if (!project || !location) return input;
|
|
|
|
const rewriteUrl = (url: string): string => {
|
|
const hasPlaceholder =
|
|
url.includes("{project}") ||
|
|
url.includes("{location}") ||
|
|
url.includes("%7Bproject%7D") ||
|
|
url.includes("%7Blocation%7D");
|
|
const host = location === "global" ? "aiplatform.googleapis.com" : `${location}-aiplatform.googleapis.com`;
|
|
const rewritten = hasPlaceholder
|
|
? url
|
|
.replace("https://{location}-aiplatform.googleapis.com", `https://${host}`)
|
|
.replace("https://%7Blocation%7D-aiplatform.googleapis.com", `https://${host}`)
|
|
.replaceAll("{project}", encodeURIComponent(project))
|
|
.replaceAll("%7Bproject%7D", encodeURIComponent(project))
|
|
.replaceAll("{location}", encodeURIComponent(location))
|
|
.replaceAll("%7Blocation%7D", encodeURIComponent(location))
|
|
: url;
|
|
return rewritten.replace(":streamRawPredict/v1/messages", ":streamRawPredict");
|
|
};
|
|
|
|
if (input instanceof Request) {
|
|
const rewrittenUrl = rewriteUrl(input.url);
|
|
return rewrittenUrl === input.url ? input : new Request(rewrittenUrl, input);
|
|
}
|
|
if (input instanceof URL) {
|
|
const rewrittenUrl = rewriteUrl(input.toString());
|
|
return rewrittenUrl === input.toString() ? input : new URL(rewrittenUrl);
|
|
}
|
|
return rewriteUrl(input);
|
|
}
|
|
|
|
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",
|
|
"xai-oauth": () => $pickenv("XAI_OAUTH_TOKEN", "XAI_API_KEY"),
|
|
fireworks: "FIREWORKS_API_KEY",
|
|
firepass: "FIREPASS_API_KEY",
|
|
"wafer-pass": "WAFER_PASS_API_KEY",
|
|
"wafer-serverless": "WAFER_SERVERLESS_API_KEY",
|
|
openrouter: "OPENROUTER_API_KEY",
|
|
kilo: "KILO_API_KEY",
|
|
"vercel-ai-gateway": "AI_GATEWAY_API_KEY",
|
|
zai: "ZAI_API_KEY",
|
|
"zhipu-coding-plan": "ZHIPU_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.GCP_PROJECT || $env.GCLOUD_PROJECT);
|
|
const hasLocation = !!($env.GOOGLE_VERTEX_LOCATION || $env.GOOGLE_CLOUD_LOCATION || $env.VERTEX_LOCATION);
|
|
if (hasCredentials && hasProject && hasLocation) {
|
|
return "<authenticated>";
|
|
}
|
|
},
|
|
// Amazon Bedrock supports multiple credential sources:
|
|
// 1. AWS_BEARER_TOKEN_BEDROCK - Bedrock API keys (bearer token)
|
|
// 2. AWS_ACCESS_KEY_ID + AWS_SECRET_ACCESS_KEY - standard IAM keys
|
|
// 3. AWS_PROFILE - named profile from ~/.aws/credentials
|
|
// 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",
|
|
"xiaomi-token-plan-sgp": "XIAOMI_TOKEN_PLAN_SGP_API_KEY",
|
|
"xiaomi-token-plan-ams": "XIAOMI_TOKEN_PLAN_AMS_API_KEY",
|
|
"xiaomi-token-plan-cn": "XIAOMI_TOKEN_PLAN_CN_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?.();
|
|
}
|
|
|
|
/**
|
|
* Enumerate every provider that has an env-var fallback for `getEnvApiKey`.
|
|
* Used by `omp auth-broker migrate --include-env` to discover env-sourced keys
|
|
* that should be uploaded to the broker.
|
|
*/
|
|
export function listProvidersWithEnvKey(): string[] {
|
|
return Object.keys(serviceProviderMap);
|
|
}
|
|
|
|
export function stream<TApi extends Api>(
|
|
model: Model<TApi>,
|
|
context: Context,
|
|
options?: OptionsForApi<TApi>,
|
|
): AssistantMessageEventStream {
|
|
const requestOptions = withRequestDebugFetch(options as StreamOptions | undefined) as
|
|
| OptionsForApi<TApi>
|
|
| undefined;
|
|
|
|
// Check custom API registry first (extension-provided APIs like "vertex-claude-api")
|
|
const customApiProvider = getCustomApi(model.api);
|
|
if (customApiProvider) {
|
|
return customApiProvider.stream(model, context, requestOptions as StreamOptions);
|
|
}
|
|
|
|
if (isGitLabDuoModel(model)) {
|
|
const apiKey = (requestOptions as StreamOptions | undefined)?.apiKey || getEnvApiKey(model.provider);
|
|
if (!apiKey) {
|
|
throw new Error(`No API key for provider: ${model.provider}`);
|
|
}
|
|
return streamGitLabDuo(model, context, {
|
|
...(requestOptions 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, requestOptions 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,
|
|
(requestOptions || {}) as BedrockOptions,
|
|
);
|
|
}
|
|
|
|
const apiKey = requestOptions?.apiKey || getEnvApiKey(model.provider);
|
|
if (!apiKey) {
|
|
throw new Error(`No API key for provider: ${model.provider}`);
|
|
}
|
|
const providerOptions = isGoogleVertexAuthenticatedModel(model)
|
|
? {
|
|
...requestOptions,
|
|
apiKey: "vertex-adc",
|
|
fetch: createVertexAuthenticatedFetch(requestOptions as StreamOptions | undefined),
|
|
}
|
|
: { ...requestOptions, 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": {
|
|
if (model.provider === "xai-oauth") {
|
|
return streamXAIResponses(model as Model<"openai-responses">, context, providerOptions as any);
|
|
}
|
|
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();
|
|
}
|
|
|
|
type AuthRetryFailure = {
|
|
error: unknown;
|
|
bufferedEvents: AssistantMessageEvent[];
|
|
terminalEvent?: Extract<AssistantMessageEvent, { type: "error" }>;
|
|
};
|
|
|
|
function extractStatusFromAssistantError(message: AssistantMessage): number | undefined {
|
|
if (message.errorStatus !== undefined) return message.errorStatus;
|
|
if (!message.errorMessage) return undefined;
|
|
return extractHttpStatusFromError({ message: message.errorMessage });
|
|
}
|
|
|
|
function isRetryableUpstreamError(error: unknown, status: number | undefined, message: string | undefined): boolean {
|
|
// 401 means the credential is bad. Usage-limit phrasing (Codex's
|
|
// "You have hit your ChatGPT usage limit", Anthropic's "usage_limit_reached",
|
|
// Google's "resource_exhausted") means this account is parked but a
|
|
// sibling credential can usually pick the request up. Both are
|
|
// rotatable via `onAuthError` — the auth-gateway maps the former to
|
|
// `invalidateCredentialMatching` and the latter to `markUsageLimitReached`.
|
|
if (status === 401) return true;
|
|
void error;
|
|
return !!message && isUsageLimitError(message);
|
|
}
|
|
|
|
function createAssistantAuthError(message: AssistantMessage): Error & { status?: number } {
|
|
const error: Error & { status?: number } = new Error(message.errorMessage ?? "Provider authentication failed");
|
|
const status = extractStatusFromAssistantError(message);
|
|
if (status !== undefined) error.status = status;
|
|
return error;
|
|
}
|
|
|
|
function emitBufferedEvents(stream: AssistantMessageEventStream, events: AssistantMessageEvent[]): void {
|
|
for (const event of events) {
|
|
stream.push(event);
|
|
}
|
|
}
|
|
|
|
export function streamSimple<TApi extends Api>(
|
|
model: Model<TApi>,
|
|
context: Context,
|
|
options?: SimpleStreamOptions,
|
|
): AssistantMessageEventStream {
|
|
const requestOptions = withRequestDebugFetch(options);
|
|
const apiKeyResolver = isApiKeyResolver(requestOptions?.apiKey) ? requestOptions.apiKey : undefined;
|
|
if (apiKeyResolver) {
|
|
const outer = new AssistantMessageEventStream();
|
|
const signal = requestOptions?.signal;
|
|
// One inner attempt against a resolved string key. When
|
|
// `captureAuthFailure` is set, a retryable auth error that arrives before
|
|
// any replay-unsafe event is buffered and returned (so the caller can
|
|
// retry with a fresh key) instead of surfaced. The terminal attempt
|
|
// clears the flag and emits whatever it gets.
|
|
const runAttempt = async (apiKey: string, captureAuthFailure: boolean): Promise<AuthRetryFailure | undefined> => {
|
|
const bufferedEvents: AssistantMessageEvent[] = [];
|
|
let emittedReplayUnsafeEvent = false;
|
|
const flushBuffered = (): void => {
|
|
emitBufferedEvents(outer, bufferedEvents);
|
|
bufferedEvents.length = 0;
|
|
};
|
|
|
|
try {
|
|
const inner = streamSimple(model, context, { ...requestOptions, apiKey });
|
|
for await (const event of inner) {
|
|
if (!emittedReplayUnsafeEvent && event.type === "start") {
|
|
bufferedEvents.push(event);
|
|
continue;
|
|
}
|
|
if (
|
|
!emittedReplayUnsafeEvent &&
|
|
captureAuthFailure &&
|
|
event.type === "error" &&
|
|
isRetryableUpstreamError(
|
|
event.error,
|
|
extractStatusFromAssistantError(event.error),
|
|
event.error.errorMessage,
|
|
)
|
|
) {
|
|
return { error: createAssistantAuthError(event.error), bufferedEvents, terminalEvent: event };
|
|
}
|
|
flushBuffered();
|
|
emittedReplayUnsafeEvent = true;
|
|
outer.push(event);
|
|
if (outer.done) return undefined;
|
|
}
|
|
flushBuffered();
|
|
if (!outer.done) outer.end(await inner.result());
|
|
} catch (error) {
|
|
if (
|
|
!emittedReplayUnsafeEvent &&
|
|
captureAuthFailure &&
|
|
isRetryableUpstreamError(
|
|
error,
|
|
extractHttpStatusFromError(error),
|
|
error instanceof Error ? error.message : undefined,
|
|
)
|
|
) {
|
|
return { error, bufferedEvents };
|
|
}
|
|
flushBuffered();
|
|
outer.fail(error);
|
|
}
|
|
return undefined;
|
|
};
|
|
const emitFailure = (failure: AuthRetryFailure): void => {
|
|
emitBufferedEvents(outer, failure.bufferedEvents);
|
|
if (failure.terminalEvent) {
|
|
outer.push(failure.terminalEvent);
|
|
} else {
|
|
outer.fail(failure.error);
|
|
}
|
|
};
|
|
|
|
void (async () => {
|
|
let lastKey: string | undefined;
|
|
try {
|
|
lastKey = (await apiKeyResolver({ lastChance: false, error: undefined, signal })) || undefined;
|
|
} catch {
|
|
lastKey = undefined;
|
|
}
|
|
if (lastKey === undefined) {
|
|
outer.fail(new Error(`No API key for provider: ${model.provider}`));
|
|
return;
|
|
}
|
|
let failure = await runAttempt(lastKey, true);
|
|
if (!failure) return;
|
|
// a/b/c policy: refresh the same account (lastChance=false), then
|
|
// switch to a sibling (lastChance=true). A step is skipped when the
|
|
// resolver yields the same key it just tried or `undefined`; the
|
|
// final step's attempt clears the capture flag so it emits directly.
|
|
for (let step = 0; step < AUTH_RETRY_STEPS.length; step++) {
|
|
const nextKey = await resolveRetryKey(apiKeyResolver, AUTH_RETRY_STEPS[step]!, failure.error, signal);
|
|
if (nextKey === undefined || nextKey === lastKey) continue;
|
|
lastKey = nextKey;
|
|
const isLastStep = step === AUTH_RETRY_STEPS.length - 1;
|
|
const next = await runAttempt(nextKey, !isLastStep);
|
|
if (!next) return;
|
|
failure = next;
|
|
}
|
|
emitFailure(failure);
|
|
})();
|
|
return outer;
|
|
}
|
|
|
|
// Pi-native transport short-circuits the per-provider dispatch entirely:
|
|
// the gateway resolves provider + credential server-side, so we don't
|
|
// need an `apiKey` from `getEnvApiKey` here — `options.apiKey` carries
|
|
// the gateway bearer instead. Comes BEFORE the custom-API check so
|
|
// extension-registered APIs can't accidentally override a configured
|
|
// pi-native transport.
|
|
if (model.transport === "pi-native") {
|
|
return streamPiNative(model, context, requestOptions);
|
|
}
|
|
|
|
// Check custom API registry (extension-provided APIs)
|
|
const customApiProvider = getCustomApi(model.api);
|
|
if (customApiProvider) {
|
|
return customApiProvider.streamSimple(model, context, requestOptions);
|
|
}
|
|
|
|
// Vertex AI uses Application Default Credentials, not API keys
|
|
if (model.api === "google-vertex") {
|
|
const providerOptions = mapOptionsForApi(model, requestOptions, 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, requestOptions, undefined);
|
|
return stream(model, context, providerOptions);
|
|
}
|
|
|
|
// The resolver form is handled by the wrapper above; only a static string
|
|
// key reaches this point.
|
|
const apiKey =
|
|
(typeof requestOptions?.apiKey === "string" ? requestOptions.apiKey : undefined) || 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, {
|
|
...requestOptions,
|
|
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, {
|
|
...requestOptions,
|
|
apiKey,
|
|
format: requestOptions?.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, {
|
|
...requestOptions,
|
|
apiKey,
|
|
format: requestOptions?.syntheticApiFormat ?? "openai", // Default to OpenAI format
|
|
});
|
|
}
|
|
const providerOptions = mapOptionsForApi(model, requestOptions, 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;
|
|
}
|
|
|
|
export 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;
|
|
}
|
|
// Named-tool routing on Google: emit an `ANY`-mode allow-list of one entry,
|
|
// mirroring the Anthropic mapper that returns `{type: "tool", name}`.
|
|
if (choice.type === "tool") {
|
|
return choice.name ? { mode: "ANY", allowedFunctionNames: [choice.name] } : undefined;
|
|
}
|
|
if (choice.type === "function") {
|
|
const name = "function" in choice ? choice.function?.name : choice.name;
|
|
return name ? { mode: "ANY", allowedFunctionNames: [name] } : undefined;
|
|
}
|
|
return undefined;
|
|
}
|
|
|
|
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;
|
|
// Models with compat.supportsReasoningEffort: false reason natively but
|
|
// reject the wire effort param. The wire-side omitReasoningEffort gate
|
|
// (providers/xai-responses.ts:78) is the actual strip; returning
|
|
// undefined here avoids a redundant requireSupportedEffort throw that
|
|
// would defeat the gate and surface a confusing
|
|
// "Compaction failed: Thinking effort high is not supported by..." to
|
|
// the user.
|
|
if (modelOmitsReasoningEffort(model)) 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 ?? model.maxTokens,
|
|
signal: options?.signal,
|
|
apiKey: apiKey ?? (typeof options?.apiKey === "string" ? options.apiKey : undefined),
|
|
cacheRetention: options?.cacheRetention,
|
|
headers: options?.headers,
|
|
initiatorOverride: options?.initiatorOverride,
|
|
maxRetryDelayMs: options?.maxRetryDelayMs,
|
|
metadata: options?.metadata,
|
|
taskBudget: options?.taskBudget,
|
|
sessionId: options?.sessionId,
|
|
promptCacheKey: options?.promptCacheKey,
|
|
streamFirstEventTimeoutMs: options?.streamFirstEventTimeoutMs,
|
|
streamIdleTimeoutMs: options?.streamIdleTimeoutMs,
|
|
providerSessionState: options?.providerSessionState,
|
|
onPayload: options?.onPayload,
|
|
onResponse: options?.onResponse,
|
|
onSseEvent: options?.onSseEvent,
|
|
execHandlers: options?.execHandlers,
|
|
fetch: options?.fetch,
|
|
};
|
|
|
|
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),
|
|
thinkingDisplay: options?.hideThinkingSummary ? "omitted" : undefined,
|
|
serviceTier: options?.serviceTier,
|
|
});
|
|
}
|
|
|
|
let thinkingBudget = options.thinkingBudgets?.[reasoning] ?? ANTHROPIC_THINKING[reasoning];
|
|
if (thinkingBudget <= 0) {
|
|
return castApi<"anthropic-messages">({
|
|
...base,
|
|
thinkingEnabled: false,
|
|
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
|
|
thinkingDisplay: options?.hideThinkingSummary ? "omitted" : undefined,
|
|
serviceTier: options?.serviceTier,
|
|
});
|
|
}
|
|
|
|
// 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),
|
|
thinkingDisplay: options?.hideThinkingSummary ? "omitted" : undefined,
|
|
serviceTier: options?.serviceTier,
|
|
});
|
|
}
|
|
|
|
if (ANTHROPIC_USE_INTERLEAVED_THINKING) {
|
|
return castApi<"anthropic-messages">({
|
|
...base,
|
|
thinkingEnabled: true,
|
|
thinkingBudgetTokens: thinkingBudget,
|
|
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
|
|
thinkingDisplay: options?.hideThinkingSummary ? "omitted" : undefined,
|
|
serviceTier: options?.serviceTier,
|
|
});
|
|
}
|
|
|
|
// 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),
|
|
thinkingDisplay: options?.hideThinkingSummary ? "omitted" : undefined,
|
|
serviceTier: options?.serviceTier,
|
|
});
|
|
} else {
|
|
return castApi<"anthropic-messages">({
|
|
...base,
|
|
maxTokens,
|
|
thinkingEnabled: true,
|
|
thinkingBudgetTokens: thinkingBudget,
|
|
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
|
|
thinkingDisplay: options?.hideThinkingSummary ? "omitted" : undefined,
|
|
serviceTier: options?.serviceTier,
|
|
});
|
|
}
|
|
}
|
|
|
|
case "bedrock-converse-stream": {
|
|
const bedrockBase: BedrockOptions = {
|
|
...base,
|
|
reasoning: options?.reasoning,
|
|
thinkingBudgets: options?.thinkingBudgets,
|
|
toolChoice: mapAnthropicToolChoice(options?.toolChoice),
|
|
thinkingDisplay: options?.hideThinkingSummary ? "omitted" : undefined,
|
|
};
|
|
// 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),
|
|
disableReasoning: options?.disableReasoning,
|
|
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
|
|
serviceTier: options?.serviceTier,
|
|
openrouterVariant: options?.openrouterVariant,
|
|
});
|
|
|
|
case "openai-responses":
|
|
return castApi<"openai-responses">({
|
|
...base,
|
|
reasoning: resolveOpenAiReasoningEffort(model, options),
|
|
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
|
|
serviceTier: options?.serviceTier,
|
|
reasoningSummary: options?.hideThinkingSummary ? null : undefined,
|
|
});
|
|
|
|
case "azure-openai-responses":
|
|
return castApi<"azure-openai-responses">({
|
|
...base,
|
|
reasoning: resolveOpenAiReasoningEffort(model, options),
|
|
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
|
|
serviceTier: options?.serviceTier,
|
|
reasoningSummary: options?.hideThinkingSummary ? null : undefined,
|
|
});
|
|
|
|
case "openai-codex-responses":
|
|
return castApi<"openai-codex-responses">({
|
|
...base,
|
|
reasoning: resolveOpenAiReasoningEffort(model, options),
|
|
toolChoice: mapOpenAiToolChoice(options?.toolChoice),
|
|
serviceTier: options?.serviceTier,
|
|
preferWebsockets: options?.preferWebsockets,
|
|
reasoningSummary: options?.hideThinkingSummary ? null : undefined,
|
|
});
|
|
|
|
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;
|
|
}
|