Luigit
repositories / pi-ext

pi-ext

bugabingas pi extensions

owned by admin

extensions/klaus/src/models.ts

Raw
import type {
	AnthropicMessagesCompat,
	Api,
	Model,
} from "@earendil-works/pi-ai";
import {
	getBuiltinModel,
	getBuiltinModels,
} from "@earendil-works/pi-ai/providers/all";
import type {
	ModelRegistry,
	ProviderModelConfig,
} from "@earendil-works/pi-coding-agent";
import { dbg } from "./debug.js";
import { KLAUS_API } from "./protocol.js";

type ProviderChatModelConfig = Extract<
	ProviderModelConfig,
	{ reasoning: boolean }
>;

const FABLE_FALLBACK_IDS = ["claude-opus-4-8", "claude-opus-5"] as const;

export const KLAUS_PROVIDER = "klaus";
export const KLAUS_BASE_URL = "https://klaus.invalid";

/** Dated snapshot ids (e.g. `claude-opus-4-5-20251101`) stay hidden; the
 * current id of each family is always projected. */
export function isProjectableClaudeModel(model: Model<Api>): boolean {
	return (
		model.provider === "anthropic" &&
		!/-\d{8}$/.test(model.id) &&
		!model.id.endsWith("-latest")
	);
}

/** Pi catalog entries may be tuned for the direct Anthropic API, where request
 * headers unlock capacity that the Claude Agent SDK path does not provide.
 * Klaus drops those headers, so it must not advertise more than the builtin
 * catalog promises for the same model. */
/** Typed escape hatch: catalog lookups accept only literal ids at the type
 * level, but klaus passes runtime catalog ids from enumerations. */
function builtinAnthropic(modelId: string): Model<Api> | undefined {
	return getBuiltinModel(
		"anthropic",
		modelId as Parameters<typeof getBuiltinModel>[1],
	) as Model<Api> | undefined;
}

function clampToBaseline(
	field: "contextWindow" | "maxTokens",
	value: number,
	baseline: Model<Api> | undefined,
): number {
	const ceiling = baseline?.[field];
	if (ceiling === undefined || value <= ceiling) return value;
	dbg?.("models.clamp");
	return ceiling;
}

function fableCompat(model: Model<Api>): Model<Api>["compat"] {
	return {
		...(structuredClone(model.compat) as AnthropicMessagesCompat),
		allowedFallbackModels: FABLE_FALLBACK_IDS.flatMap((id) => {
			const fallback = getBuiltinModel("anthropic", id);
			return fallback
				? [
						{
							provider: "anthropic",
							model: id,
							cost: structuredClone(fallback.cost),
						},
					]
				: [];
		}),
	} satisfies AnthropicMessagesCompat;
}

/** Project one Pi Anthropic model onto the Klaus provider. Transport-owned
 * fields (`api`, `baseUrl`, `headers`, sampling parameters) are never
 * inherited, because the Claude Agent SDK owns the request. */
export function projectModel(model: Model<Api>): ProviderChatModelConfig {
	dbg?.("models.projectModel");
	const baseline = builtinAnthropic(model.id);
	return {
		id: model.id,
		name: model.name,
		reasoning: model.reasoning,
		thinkingLevelMap: model.thinkingLevelMap,
		input: [...model.input],
		cost: structuredClone(model.cost),
		compat:
			model.id === "claude-fable-5-1"
				? fableCompat(model)
				: structuredClone(model.compat),
		contextWindow: clampToBaseline(
			"contextWindow",
			model.contextWindow,
			baseline,
		),
		maxTokens: clampToBaseline("maxTokens", model.maxTokens, baseline),
	};
}

/** Project every projectable Claude model from a catalog enumeration. */
export function projectModels(models: Model<Api>[]): ProviderChatModelConfig[] {
	const projected = models.filter(isProjectableClaudeModel).map(projectModel);
	dbg?.("models.projectModels", { count: projected.length });
	return projected;
}

/** Startup projection from Pi's builtin catalog before any session exists. */
export function projectBuiltinModels(): ProviderChatModelConfig[] {
	return projectModels(getBuiltinModels("anthropic"));
}

/** Attach the transport identity Pi needs for a registered provider model. */
function toKlausModel(config: ProviderChatModelConfig): Model<Api> {
	return {
		id: config.id,
		name: config.name,
		api: KLAUS_API,
		provider: KLAUS_PROVIDER,
		baseUrl: KLAUS_BASE_URL,
		reasoning: config.reasoning,
		...(config.thinkingLevelMap !== undefined
			? { thinkingLevelMap: config.thinkingLevelMap }
			: {}),
		input: config.input,
		cost: config.cost,
		...(config.compat !== undefined ? { compat: config.compat } : {}),
		contextWindow: config.contextWindow,
		maxTokens: config.maxTokens,
	};
}

/** On-demand Klaus catalog for `Provider.getModels()`. Pi's composed Anthropic
 * provider carries models.json overrides and is fresh even while Pi rebuilds
 * its snapshot, so it outranks the builtin catalog whenever it still knows
 * Claude models. Must not throw: Pi treats a throwing `getModels()` as empty. */
export function klausModels(registry: ModelRegistry | undefined): Model<Api>[] {
	let projected: ProviderChatModelConfig[] = [];
	try {
		const anthropic = registry?.getProvider("anthropic")?.getModels();
		if (anthropic) projected = projectModels([...anthropic]);
	} catch {
		projected = [];
	}
	const source = projected.length ? "registry" : "builtin";
	if (!projected.length) projected = projectBuiltinModels();
	dbg?.("models.klausModels", { source, count: projected.length });
	return projected.map(toKlausModel);
}