Azreen Jamal
362 lines
14 KiB
TypeScript
362 lines
14 KiB
TypeScript
/**
|
|
* Model Router — LLM-powered automatic model selection
|
|
*
|
|
* Uses Haiku (~200ms, ~$0.001/call) to classify every prompt's complexity,
|
|
* then routes to the optimal model + thinking level. No regex guessing.
|
|
*
|
|
* Tiers:
|
|
* ⚡ Haiku — Simple Q&A, file reads, small edits, quick lookups
|
|
* ⚖️ Sonnet — Code generation, debugging, refactoring, multi-file work
|
|
* 🧠 Opus — Architecture, complex reasoning, large-scale changes
|
|
*
|
|
* Mid-turn escalation:
|
|
* - 4+ edits/writes → escalate to Opus
|
|
* - 2+ consecutive errors → escalate (model struggling)
|
|
* - 6+ tool calls in one turn → escalate to at least Sonnet
|
|
*
|
|
* Usage: pi -e extensions/model-router.ts
|
|
*
|
|
* Commands:
|
|
* /router — Show routing state & stats
|
|
* /router lock — Lock current model (disable auto-routing)
|
|
* /router unlock — Re-enable auto-routing
|
|
* /router tier <n> — Force tier: 1=haiku, 2=sonnet, 3=opus
|
|
*/
|
|
|
|
import type { AssistantMessage } from "@mariozechner/pi-ai";
|
|
import { complete } from "@mariozechner/pi-ai";
|
|
import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
|
|
import { truncateToWidth, visibleWidth } from "@mariozechner/pi-tui";
|
|
import { basename } from "node:path";
|
|
import { applyExtensionDefaults } from "./themeMap.ts";
|
|
|
|
// ── Model Tiers ────────────────────────────────────────────────────────────
|
|
|
|
interface ModelTier {
|
|
name: string;
|
|
icon: string;
|
|
provider: string;
|
|
modelId: string;
|
|
thinking: "off" | "low" | "medium" | "high";
|
|
inputCost: number;
|
|
outputCost: number;
|
|
}
|
|
|
|
const TIERS: ModelTier[] = [
|
|
{ name: "Haiku", icon: "⚡", provider: "anthropic", modelId: "claude-haiku-4-5", thinking: "off", inputCost: 0.80, outputCost: 4.00 },
|
|
{ name: "Sonnet", icon: "⚖️", provider: "anthropic", modelId: "claude-sonnet-4-6", thinking: "low", inputCost: 3.00, outputCost: 15.00 },
|
|
{ name: "Opus", icon: "🧠", provider: "anthropic", modelId: "claude-opus-4-6", thinking: "high", inputCost: 15.00, outputCost: 75.00 },
|
|
];
|
|
|
|
// ── LLM Classifier ────────────────────────────────────────────────────────
|
|
|
|
const CLASSIFIER_PROMPT = `You are a task complexity classifier for a coding assistant. Given a user prompt, respond with ONLY a JSON object (no markdown, no explanation):
|
|
|
|
{"tier": <1|2|3>, "reason": "<5 words max>"}
|
|
|
|
TIER 1 (simple): Quick questions, yes/no, read/list/check files, small edits, conversational ("thanks", "ok"), simple lookups.
|
|
TIER 2 (medium): Write functions/classes, fix bugs, refactor single files, write tests, create endpoints, docker/CI config.
|
|
TIER 3 (complex): Architecture design, multi-file refactoring/migration, security audits, performance optimization, complex algorithms, system design from scratch.`;
|
|
|
|
interface Classification { tier: number; reason: string }
|
|
|
|
const cache = new Map<string, { result: Classification; ts: number }>();
|
|
const CACHE_TTL = 60_000;
|
|
|
|
async function classify(prompt: string, ctx: ExtensionContext): Promise<Classification> {
|
|
const key = prompt.trim().toLowerCase().slice(0, 200);
|
|
const c = cache.get(key);
|
|
if (c && Date.now() - c.ts < CACHE_TTL) return c.result;
|
|
|
|
const model = ctx.modelRegistry.find("anthropic", "claude-haiku-4-5");
|
|
if (!model) return { tier: 2, reason: "no classifier" };
|
|
const apiKey = await ctx.modelRegistry.getApiKey(model);
|
|
if (!apiKey) return { tier: 2, reason: "no key" };
|
|
|
|
try {
|
|
const resp = await complete(model, {
|
|
systemPrompt: CLASSIFIER_PROMPT,
|
|
messages: [{ role: "user" as const, content: prompt, timestamp: Date.now() }],
|
|
}, { reasoning: "off" });
|
|
|
|
const text = resp.content
|
|
.filter((c): c is { type: "text"; text: string } => c.type === "text")
|
|
.map(c => c.text).join("");
|
|
|
|
const json = text.replace(/```json?\n?/g, "").replace(/```/g, "").trim();
|
|
const parsed = JSON.parse(json);
|
|
const result: Classification = {
|
|
tier: Math.max(1, Math.min(3, parsed.tier || 2)),
|
|
reason: String(parsed.reason || "classified").slice(0, 50),
|
|
};
|
|
|
|
cache.set(key, { result, ts: Date.now() });
|
|
return result;
|
|
} catch {
|
|
return { tier: 2, reason: "classify failed" };
|
|
}
|
|
}
|
|
|
|
// ── Router State ───────────────────────────────────────────────────────────
|
|
|
|
interface RouterState {
|
|
currentTier: number;
|
|
locked: boolean;
|
|
totalSwitches: number;
|
|
turnsSinceLast: number;
|
|
consecutiveErrors: number;
|
|
toolCallsThisTurn: number;
|
|
editWriteThisTurn: number;
|
|
totalToolCalls: number;
|
|
totalErrors: number;
|
|
tierHistory: Array<{ tier: number; reason: string; turn: number }>;
|
|
turnCount: number;
|
|
savedVsAlwaysOpus: number;
|
|
classifierCost: number;
|
|
}
|
|
|
|
// ── Extension ──────────────────────────────────────────────────────────────
|
|
|
|
export default function modelRouter(pi: ExtensionAPI) {
|
|
const state: RouterState = {
|
|
currentTier: 1,
|
|
locked: false,
|
|
totalSwitches: 0,
|
|
turnsSinceLast: 0,
|
|
consecutiveErrors: 0,
|
|
toolCallsThisTurn: 0,
|
|
editWriteThisTurn: 0,
|
|
totalToolCalls: 0,
|
|
totalErrors: 0,
|
|
tierHistory: [],
|
|
turnCount: 0,
|
|
savedVsAlwaysOpus: 0,
|
|
classifierCost: 0,
|
|
};
|
|
|
|
async function switchToTier(tierIndex: number, reason: string, ctx?: ExtensionContext) {
|
|
const clamped = Math.max(0, Math.min(2, tierIndex));
|
|
if (clamped === state.currentTier) return;
|
|
const tier = TIERS[clamped];
|
|
const model = ctx?.modelRegistry.find(tier.provider, tier.modelId);
|
|
if (!model) return;
|
|
const success = await pi.setModel(model);
|
|
if (!success) return;
|
|
pi.setThinkingLevel(tier.thinking);
|
|
const prev = state.currentTier;
|
|
state.currentTier = clamped;
|
|
state.totalSwitches++;
|
|
state.turnsSinceLast = 0;
|
|
state.tierHistory.push({
|
|
tier: clamped,
|
|
reason: `${TIERS[prev].icon}→${tier.icon} ${reason}`,
|
|
turn: state.turnCount,
|
|
});
|
|
}
|
|
|
|
// ── LLM Classification on every prompt ─────────────────────────────────
|
|
|
|
pi.on("before_agent_start", async (event, ctx) => {
|
|
if (state.locked) return;
|
|
|
|
ctx.ui.setWorkingMessage("classifying...");
|
|
const result = await classify(event.prompt, ctx);
|
|
ctx.ui.setWorkingMessage();
|
|
|
|
// Track classifier cost (~180 tokens per call at Haiku rates)
|
|
state.classifierCost += (150 * 0.80 + 30 * 4.00) / 1_000_000;
|
|
|
|
let target = result.tier - 1; // 1-3 → 0-2
|
|
|
|
// Context-aware overrides
|
|
if (target === 0 && state.totalToolCalls > 10) target = 1;
|
|
if (state.consecutiveErrors >= 2 && target < 2) target = Math.min(target + 1, 2);
|
|
|
|
if (target !== state.currentTier) {
|
|
await switchToTier(target, result.reason, ctx);
|
|
}
|
|
});
|
|
|
|
// ── Turn & Tool Tracking ───────────────────────────────────────────────
|
|
|
|
pi.on("turn_start", async () => {
|
|
state.toolCallsThisTurn = 0;
|
|
state.editWriteThisTurn = 0;
|
|
state.turnCount++;
|
|
state.turnsSinceLast++;
|
|
});
|
|
|
|
pi.on("tool_execution_start", async (event) => {
|
|
state.toolCallsThisTurn++;
|
|
state.totalToolCalls++;
|
|
if (event.toolName === "edit" || event.toolName === "write") state.editWriteThisTurn++;
|
|
});
|
|
|
|
pi.on("tool_execution_end", async (event, ctx) => {
|
|
if (event.isError) {
|
|
state.consecutiveErrors++;
|
|
state.totalErrors++;
|
|
if (state.consecutiveErrors >= 2 && !state.locked && state.currentTier < 2) {
|
|
await switchToTier(state.currentTier + 1, `${state.consecutiveErrors} errors`, ctx);
|
|
}
|
|
} else {
|
|
state.consecutiveErrors = 0;
|
|
}
|
|
});
|
|
|
|
pi.on("turn_end", async (_event, ctx) => {
|
|
if (!state.locked) {
|
|
if (state.editWriteThisTurn >= 4 && state.currentTier < 2) {
|
|
await switchToTier(2, `heavy: ${state.editWriteThisTurn} edits`, ctx);
|
|
} else if (state.toolCallsThisTurn >= 6 && state.currentTier < 1) {
|
|
await switchToTier(1, `busy: ${state.toolCallsThisTurn} tools`, ctx);
|
|
}
|
|
}
|
|
const opusCost = (2000 * TIERS[2].inputCost + 1000 * TIERS[2].outputCost) / 1e6;
|
|
const actualCost = (2000 * TIERS[state.currentTier].inputCost + 1000 * TIERS[state.currentTier].outputCost) / 1e6;
|
|
state.savedVsAlwaysOpus += opusCost - actualCost;
|
|
});
|
|
|
|
// ── Session Start ──────────────────────────────────────────────────────
|
|
|
|
pi.on("session_start", async (_event, ctx) => {
|
|
applyExtensionDefaults(import.meta.url, ctx);
|
|
|
|
const initialTier = TIERS[1];
|
|
const model = ctx.modelRegistry.find(initialTier.provider, initialTier.modelId);
|
|
if (model) {
|
|
await pi.setModel(model);
|
|
pi.setThinkingLevel(initialTier.thinking);
|
|
}
|
|
|
|
ctx.ui.setFooter((tui, theme, footerData) => {
|
|
const unsub = footerData.onBranchChange(() => tui.requestRender());
|
|
return {
|
|
dispose: unsub,
|
|
invalidate() {},
|
|
render(width: number): string[] {
|
|
const tier = TIERS[state.currentTier];
|
|
const lockIcon = state.locked ? " 🔒" : "";
|
|
|
|
let tokIn = 0, tokOut = 0, cost = 0;
|
|
for (const entry of ctx.sessionManager.getBranch()) {
|
|
if (entry.type === "message" && entry.message.role === "assistant") {
|
|
const m = entry.message as AssistantMessage;
|
|
tokIn += m.usage.input;
|
|
tokOut += m.usage.output;
|
|
cost += m.usage.cost.total;
|
|
}
|
|
}
|
|
cost += state.classifierCost;
|
|
const fmt = (n: number) => n < 1000 ? `${n}` : `${(n / 1000).toFixed(1)}k`;
|
|
const usage = ctx.getContextUsage();
|
|
const pct = usage?.percent ?? 0;
|
|
const filled = Math.round(pct / 10) || 1;
|
|
|
|
const l1Left =
|
|
theme.fg("dim", " ") +
|
|
theme.fg("accent", `${tier.icon} ${tier.name}`) +
|
|
theme.fg("dim", ` [${tier.thinking}]${lockIcon} `) +
|
|
theme.fg("warning", "[") +
|
|
theme.fg("success", "#".repeat(filled)) +
|
|
theme.fg("dim", "-".repeat(10 - filled)) +
|
|
theme.fg("warning", "]") +
|
|
theme.fg("dim", " ") +
|
|
theme.fg("accent", `${Math.round(pct)}%`);
|
|
|
|
const l1Right =
|
|
theme.fg("success", `${fmt(tokIn)}`) +
|
|
theme.fg("dim", " in ") +
|
|
theme.fg("accent", `${fmt(tokOut)}`) +
|
|
theme.fg("dim", " out ") +
|
|
theme.fg("warning", `$${cost.toFixed(4)}`) +
|
|
theme.fg("dim", " ");
|
|
|
|
const pad1 = " ".repeat(Math.max(1, width - visibleWidth(l1Left) - visibleWidth(l1Right)));
|
|
const line1 = truncateToWidth(l1Left + pad1 + l1Right, width, "");
|
|
|
|
const dir = basename(ctx.cwd);
|
|
const branch = footerData.getGitBranch();
|
|
const lastReason = state.tierHistory.length > 0
|
|
? state.tierHistory[state.tierHistory.length - 1].reason
|
|
: "start: balanced";
|
|
|
|
const l2Left =
|
|
theme.fg("dim", ` ${dir}`) +
|
|
(branch ? theme.fg("dim", " ") + theme.fg("warning", "(") + theme.fg("success", branch) + theme.fg("warning", ")") : "") +
|
|
theme.fg("dim", " · ") +
|
|
theme.fg("accent", `${state.totalSwitches}`) +
|
|
theme.fg("dim", " switches") +
|
|
(state.savedVsAlwaysOpus > 0 ? theme.fg("dim", " · saved ") + theme.fg("success", `$${state.savedVsAlwaysOpus.toFixed(4)}`) : "");
|
|
|
|
const l2Right = theme.fg("dim", truncateToWidth(lastReason, Math.floor(width * 0.4), "…") + " ");
|
|
|
|
const pad2 = " ".repeat(Math.max(1, width - visibleWidth(l2Left) - visibleWidth(l2Right)));
|
|
const line2 = truncateToWidth(l2Left + pad2 + l2Right, width, "");
|
|
|
|
return [line1, line2];
|
|
},
|
|
};
|
|
});
|
|
|
|
ctx.ui.notify(
|
|
`Model Router active — LLM-powered (Haiku classifier)\n` +
|
|
`Start: ${TIERS[1].icon} ${TIERS[1].name} · Auto-routes per prompt\n` +
|
|
`/router — Status & controls`,
|
|
"info",
|
|
);
|
|
});
|
|
|
|
// ── /router Command ────────────────────────────────────────────────────
|
|
|
|
pi.registerCommand("router", {
|
|
description: "Model router status & controls. Usage: /router [lock|unlock|tier 1|2|3]",
|
|
async handler(args: string, ctx) {
|
|
const parts = args.trim().split(/\s+/);
|
|
const sub = parts[0]?.toLowerCase();
|
|
|
|
if (sub === "lock") { state.locked = true; ctx.ui.notify("🔒 Locked", "info"); return; }
|
|
if (sub === "unlock") { state.locked = false; ctx.ui.notify("🔓 Unlocked", "info"); return; }
|
|
if (sub === "tier") {
|
|
const n = parseInt(parts[1], 10);
|
|
if (n >= 1 && n <= 3) {
|
|
await switchToTier(n - 1, `manual /router tier ${n}`, ctx);
|
|
ctx.ui.notify(`${TIERS[n - 1].icon} ${TIERS[n - 1].name}`, "info");
|
|
return;
|
|
}
|
|
ctx.ui.notify("Usage: /router tier 1|2|3", "warning"); return;
|
|
}
|
|
|
|
const tier = TIERS[state.currentTier];
|
|
const history = state.tierHistory.length > 0
|
|
? state.tierHistory.slice(-5).map(h => ` Turn ${h.turn}: ${h.reason}`).join("\n")
|
|
: " (none)";
|
|
|
|
pi.sendMessage({
|
|
customType: "router-status",
|
|
content: [
|
|
``,
|
|
` ┌─── Model Router (LLM-powered) ───┐`,
|
|
` │ Current: ${tier.icon} ${tier.name.padEnd(22)}│`,
|
|
` │ Thinking: ${tier.thinking.padEnd(21)}│`,
|
|
` │ Locked: ${(state.locked ? "yes 🔒" : "no").padEnd(23)}│`,
|
|
` │ Switches: ${String(state.totalSwitches).padEnd(21)}│`,
|
|
` │ Turns: ${String(state.turnCount).padEnd(23)}│`,
|
|
` │ Tool calls: ${String(state.totalToolCalls).padEnd(18)}│`,
|
|
` │ Errors: ${String(state.totalErrors).padEnd(22)}│`,
|
|
` │ Saved vs Opus: $${state.savedVsAlwaysOpus.toFixed(4).padEnd(14)}│`,
|
|
` │ Classifier cost: $${state.classifierCost.toFixed(4).padEnd(12)}│`,
|
|
` └────────────────────────────────────┘`,
|
|
``,
|
|
` Recent routing:`,
|
|
history,
|
|
``,
|
|
` /router lock | unlock | tier 1|2|3`,
|
|
``,
|
|
].join("\n"),
|
|
display: "assistant",
|
|
});
|
|
},
|
|
});
|
|
}
|