|
import { streamSimpleAnthropic } from "@earendil-works/pi-ai"; |
|
import type { |
|
ExtensionAPI, |
|
ExtensionContext, |
|
} from "@earendil-works/pi-coding-agent"; |
|
|
|
type AnthropicModel = Parameters<typeof streamSimpleAnthropic>[0]; |
|
|
|
const FAST_MODE_BETA = "fast-mode-2026-02-01"; |
|
const CLAUDE_CODE_BETAS = ["claude-code-20250219", "oauth-2025-04-20"]; |
|
|
|
// Model IDs that support fast mode, mapped to display names. |
|
const FAST_MODELS: Record<string, string> = { |
|
"claude-opus-4-8": "Opus 4.8", |
|
"claude-opus-4-7": "Opus 4.7", |
|
"claude-opus-4-6": "Opus 4.6", |
|
}; |
|
|
|
function isRecord(value: unknown): value is Record<string, unknown> { |
|
return typeof value === "object" && value !== null && !Array.isArray(value); |
|
} |
|
|
|
function isAnthropicProvider(ctx: ExtensionContext): boolean { |
|
return ctx.model?.provider === "anthropic"; |
|
} |
|
|
|
function isSupportedOpusFastModel(model: string): boolean { |
|
return ( |
|
FAST_MODELS[model] !== undefined || |
|
model.startsWith("claude-opus-4-8-") || |
|
model.startsWith("claude-opus-4-7-") || |
|
model.startsWith("claude-opus-4-6-") |
|
); |
|
} |
|
|
|
function getModelDisplayName(model: string): string { |
|
return ( |
|
FAST_MODELS[model] ?? |
|
(model.startsWith("claude-opus-4-8-") |
|
? "Opus 4.8" |
|
: model.startsWith("claude-opus-4-7-") |
|
? "Opus 4.7" |
|
: model.startsWith("claude-opus-4-6-") |
|
? "Opus 4.6" |
|
: model) |
|
); |
|
} |
|
|
|
function isOAuthToken(apiKey: string | undefined): boolean { |
|
return apiKey?.includes("sk-ant-oat") === true; |
|
} |
|
|
|
function getHeader( |
|
headers: Record<string, string> | undefined, |
|
name: string, |
|
): string | undefined { |
|
if (!headers) return undefined; |
|
const lowerName = name.toLowerCase(); |
|
const key = Object.keys(headers).find( |
|
(k) => k.toLowerCase() === lowerName, |
|
); |
|
return key ? headers[key] : undefined; |
|
} |
|
|
|
function withoutHeader( |
|
headers: Record<string, string> | undefined, |
|
name: string, |
|
): Record<string, string> { |
|
if (!headers) return {}; |
|
const lowerName = name.toLowerCase(); |
|
return Object.fromEntries( |
|
Object.entries(headers).filter( |
|
([key]) => key.toLowerCase() !== lowerName, |
|
), |
|
); |
|
} |
|
|
|
function appendBetas( |
|
...values: Array<string | undefined> |
|
): string { |
|
const betas = new Set<string>(); |
|
|
|
for (const value of values) { |
|
for (const beta of value?.split(",") ?? []) { |
|
const trimmed = beta.trim(); |
|
if (trimmed) betas.add(trimmed); |
|
} |
|
} |
|
|
|
betas.add(FAST_MODE_BETA); |
|
return [...betas].join(","); |
|
} |
|
|
|
export default function fastOpusExtension(pi: ExtensionAPI): void { |
|
// Override the default anthropic provider to inject fast-mode beta header. |
|
pi.registerProvider("anthropic", { |
|
api: "anthropic-messages", |
|
streamSimple(model, context, options) { |
|
const anthropicModel = model as AnthropicModel; |
|
|
|
if (!isSupportedOpusFastModel(model.id)) { |
|
return streamSimpleAnthropic(anthropicModel, context, options); |
|
} |
|
|
|
const incomingBeta = getHeader(options?.headers, "anthropic-beta"); |
|
const baseBeta = isOAuthToken(options?.apiKey) |
|
? CLAUDE_CODE_BETAS.join(",") |
|
: undefined; |
|
|
|
const headers = { |
|
...withoutHeader(options?.headers, "anthropic-beta"), |
|
"anthropic-beta": appendBetas(baseBeta, incomingBeta), |
|
}; |
|
|
|
return streamSimpleAnthropic(anthropicModel, context, { |
|
...options, |
|
headers, |
|
}); |
|
}, |
|
}); |
|
|
|
// Notify on model switch to a fast-mode supported model. |
|
pi.on("model_select", async (event, ctx) => { |
|
if (!ctx.hasUI) return; |
|
if (event.source === "restore") return; |
|
if (event.model.provider !== "anthropic") return; |
|
if (!isSupportedOpusFastModel(event.model.id)) return; |
|
|
|
const displayName = getModelDisplayName(event.model.id); |
|
ctx.ui.notify(`⚡ ${displayName}: Fast mode.`, "info"); |
|
}); |
|
|
|
// Inject speed=fast into the payload for supported models. |
|
pi.on("before_provider_request", (event, ctx) => { |
|
if (!isRecord(event.payload)) return; |
|
if (!isAnthropicProvider(ctx)) return; |
|
|
|
const model = event.payload.model; |
|
if (typeof model !== "string" || !isSupportedOpusFastModel(model)) { |
|
return; |
|
} |
|
if (Object.hasOwn(event.payload, "speed")) return; |
|
|
|
return { |
|
...event.payload, |
|
speed: "fast", |
|
}; |
|
}); |
|
} |