From f3ea30f83c9526efa9d732db67a3e288ed06c6d0 Mon Sep 17 00:00:00 2001 From: RiteshTiwari1 Date: Tue, 2 Jun 2026 00:49:16 +0530 Subject: [PATCH] feat: auto model routing for SkillFlows --- src/exports.ts | 3 + src/loader.ts | 6 ++ src/model-routing.ts | 115 +++++++++++++++++++++++++++++++++++++ src/skills.ts | 2 + src/workflows.ts | 5 +- test/model-routing.test.ts | 111 +++++++++++++++++++++++++++++++++++ 6 files changed, 241 insertions(+), 1 deletion(-) create mode 100644 src/model-routing.ts create mode 100644 test/model-routing.test.ts diff --git a/src/exports.ts b/src/exports.ts index 787d60f..955edb0 100644 --- a/src/exports.ts +++ b/src/exports.ts @@ -89,6 +89,9 @@ export { saveFlowDefinition, deleteFlowDefinition, } from "./workflows.js"; +// Model routing — consumed by @open-gitagent/voice to pick a per-step model. +export { resolveRoutedModel } from "./model-routing.js"; +export type { ModelTier, RoutingConfig, RouteInput, RouteResult } from "./model-routing.js"; export { discoverSchedules, saveSchedule, diff --git a/src/loader.ts b/src/loader.ts index b8d193e..3a912d6 100644 --- a/src/loader.ts +++ b/src/loader.ts @@ -42,6 +42,12 @@ export interface AgentManifest { top_k?: number; stop_sequences?: string[]; }; + routing?: { + enabled?: boolean; + lightweight?: string; + reasoning?: string; + rules?: Array<{ tier: "lightweight" | "reasoning"; match: string[] }>; + }; }; tools: string[]; skills?: string[]; diff --git a/src/model-routing.ts b/src/model-routing.ts new file mode 100644 index 0000000..473b11c --- /dev/null +++ b/src/model-routing.ts @@ -0,0 +1,115 @@ +// Classifies each SkillFlow step by complexity and resolves the model it should +// run on: lightweight tasks (summarize/extract/classify/transform) to a cheap +// model, reasoning-heavy tasks to the configured reasoning model. Explicit +// per-step / per-skill settings win; anything unresolved falls back to primary. + +export type ModelTier = "lightweight" | "reasoning"; + +export interface RoutingConfig { + /** Defaults to true when a routing block is present. */ + enabled?: boolean; + /** Model id for lightweight tasks, e.g. "openai:gpt-4o-mini". */ + lightweight?: string; + /** Model id for reasoning tasks, e.g. "openai:gpt-4o". */ + reasoning?: string; + /** Classification overrides — first matching rule wins. */ + rules?: Array<{ tier: ModelTier; match: string[] }>; +} + +export interface RouteInput { + /** Explicit per-step model (highest priority); alias or model id. */ + stepModel?: string; + /** Per-skill default from SKILL.md frontmatter; alias or model id. */ + skillModel?: string; + /** Text used to classify the task (skill name + step prompt). */ + classifyText: string; + routing?: RoutingConfig; + /** The agent's preferred model — the ultimate fallback. */ + primaryModel?: string; +} + +export interface RouteResult { + /** Resolved "provider:model" (undefined → let the runtime decide). */ + model?: string; + /** Tier, when the model came from automatic classification. */ + tier: ModelTier | null; + source: "step" | "skill" | "auto" | "fallback"; +} + +// Matched against word starts, so "summarize"/"summary"/"summarization" all hit +// "summ" without "already" matching "read". +const DEFAULT_LIGHTWEIGHT = [ + "summ", "extract", "classif", "transform", "format", "convert", + "parse", "fetch", "read", "load", "lookup", "normaliz", "translat", + "rephrase", "rewrite", "tag", "label", "render", +]; +const DEFAULT_REASONING = [ + "search", "analy", "plan", "decid", "decision", "orchestrat", "solve", + "reason", "validat", "evaluat", "review", "audit", "diagnos", "debug", + "architect", "design", "strateg", "investigat", "assess", "judge", + "verify", "critique", "infer", "deduc", +]; + +function matchesAny(text: string, keywords: string[]): boolean { + for (const kw of keywords) { + const re = new RegExp(`\\b${kw.replace(/[.*+?^${}()|[\]\\]/g, "\\$&")}`, "i"); + if (re.test(text)) return true; + } + return false; +} + +/** + * Classify a task into a complexity tier. User rules take precedence over the + * built-in defaults. A task that matches neither — or both — resolves to + * "reasoning", so cost optimization never silently degrades quality. + */ +export function classifyTaskTier( + classifyText: string, + rules?: Array<{ tier: ModelTier; match: string[] }>, +): ModelTier { + const text = classifyText || ""; + + if (rules) { + for (const rule of rules) { + if (Array.isArray(rule.match) && matchesAny(text, rule.match)) { + return rule.tier; + } + } + } + + if (matchesAny(text, DEFAULT_REASONING)) return "reasoning"; + if (matchesAny(text, DEFAULT_LIGHTWEIGHT)) return "lightweight"; + return "reasoning"; +} + +/** Resolve a tier alias ("lightweight"/"reasoning") or pass a model id through. */ +export function resolveModelAlias(ref: string | undefined, routing?: RoutingConfig): string | undefined { + if (!ref) return undefined; + if (ref === "lightweight") return routing?.lightweight || undefined; + if (ref === "reasoning") return routing?.reasoning || undefined; + return ref; +} + +/** + * Decide which model a task runs on, in precedence order: explicit per-step + * model, per-skill model, automatic classification (only when a routing block + * is present and enabled), then the primary model. + */ +export function resolveRoutedModel(input: RouteInput): RouteResult { + const { stepModel, skillModel, classifyText, routing, primaryModel } = input; + + const fromStep = resolveModelAlias(stepModel, routing); + if (fromStep) return { model: fromStep, tier: null, source: "step" }; + + const fromSkill = resolveModelAlias(skillModel, routing); + if (fromSkill) return { model: fromSkill, tier: null, source: "skill" }; + + const autoEnabled = !!routing && routing.enabled !== false && !!(routing.lightweight || routing.reasoning); + if (autoEnabled) { + const tier = classifyTaskTier(classifyText, routing!.rules); + const model = tier === "lightweight" ? routing!.lightweight : routing!.reasoning; + if (model) return { model, tier, source: "auto" }; + } + + return { model: primaryModel, tier: null, source: "fallback" }; +} diff --git a/src/skills.ts b/src/skills.ts index 74879de..fab47ce 100644 --- a/src/skills.ts +++ b/src/skills.ts @@ -11,6 +11,7 @@ export interface SkillMetadata { usage_count?: number; success_count?: number; failure_count?: number; + model?: string; } export interface ParsedSkill extends SkillMetadata { @@ -96,6 +97,7 @@ export async function discoverSkills(agentDir: string): Promise if (typeof frontmatter.usage_count === "number") meta.usage_count = frontmatter.usage_count; if (typeof frontmatter.success_count === "number") meta.success_count = frontmatter.success_count; if (typeof frontmatter.failure_count === "number") meta.failure_count = frontmatter.failure_count; + if (typeof frontmatter.model === "string") meta.model = frontmatter.model; skills.push(meta); } diff --git a/src/workflows.ts b/src/workflows.ts index 03fa300..1ed13ab 100644 --- a/src/workflows.ts +++ b/src/workflows.ts @@ -7,6 +7,7 @@ export interface SkillFlowStep { skill: string; prompt: string; channel?: string; + model?: string; } export interface SkillFlowDefinition { @@ -68,6 +69,7 @@ export async function discoverWorkflows(agentDir: string): Promise ({ skill: s.skill, prompt: s.prompt, ...(s.channel ? { channel: s.channel } : {}) })), + steps: flow.steps.map((s) => ({ skill: s.skill, prompt: s.prompt, ...(s.channel ? { channel: s.channel } : {}), ...(s.model ? { model: s.model } : {}) })), }, { lineWidth: 120 }); await writeFile(filePath, content, "utf-8"); return filePath; diff --git a/test/model-routing.test.ts b/test/model-routing.test.ts new file mode 100644 index 0000000..0e60f13 --- /dev/null +++ b/test/model-routing.test.ts @@ -0,0 +1,111 @@ +// Tests for the resolution priority chain (step > skill > auto > fallback) and +// the classifier's safety default (ambiguous tasks resolve to reasoning). + +import test from "node:test"; +import assert from "node:assert/strict"; + +import { + classifyTaskTier, + resolveModelAlias, + resolveRoutedModel, + type RoutingConfig, +} from "../src/model-routing.ts"; + +const routing: RoutingConfig = { + lightweight: "openai:gpt-4o-mini", + reasoning: "openai:gpt-4o", +}; + +// ── classifyTaskTier ─────────────────────────────────────────────────── + +test("classifies lightweight task types as lightweight", () => { + assert.equal(classifyTaskTier("summarize the pull request diff"), "lightweight"); + assert.equal(classifyTaskTier("extract the linked issue number"), "lightweight"); + assert.equal(classifyTaskTier("format the report as markdown"), "lightweight"); +}); + +test("classifies reasoning task types as reasoning", () => { + assert.equal(classifyTaskTier("analyze the security implications"), "reasoning"); + assert.equal(classifyTaskTier("plan a multi-step remediation"), "reasoning"); + assert.equal(classifyTaskTier("validate the truth score"), "reasoning"); +}); + +test("unknown tasks default to reasoning (never silently downgrade quality)", () => { + assert.equal(classifyTaskTier("frobnicate the widget"), "reasoning"); + assert.equal(classifyTaskTier(""), "reasoning"); +}); + +test("a task matching both tiers resolves to reasoning", () => { + // "summarize" (lightweight) + "analyze" (reasoning) → reasoning wins. + assert.equal(classifyTaskTier("summarize and analyze the results"), "reasoning"); +}); + +test("user rules take precedence over the built-in defaults", () => { + // "analyze" would default to reasoning, but a user rule forces lightweight. + const rules = [{ tier: "lightweight" as const, match: ["analyze"] }]; + assert.equal(classifyTaskTier("analyze the log lines", rules), "lightweight"); +}); + +// ── resolveModelAlias ────────────────────────────────────────────────── + +test("resolves tier aliases and passes literal model ids through", () => { + assert.equal(resolveModelAlias("lightweight", routing), "openai:gpt-4o-mini"); + assert.equal(resolveModelAlias("reasoning", routing), "openai:gpt-4o"); + assert.equal(resolveModelAlias("anthropic:claude-sonnet-4-5", routing), "anthropic:claude-sonnet-4-5"); + assert.equal(resolveModelAlias(undefined, routing), undefined); +}); + +// ── resolveRoutedModel priority chain ────────────────────────────────── + +test("explicit per-step model wins over everything", () => { + const r = resolveRoutedModel({ + stepModel: "openai:gpt-4o", + skillModel: "openai:gpt-4o-mini", + classifyText: "summarize the diff", + routing, + primaryModel: "openai:gpt-5-reasoning", + }); + assert.equal(r.model, "openai:gpt-4o"); + assert.equal(r.source, "step"); +}); + +test("per-skill model wins when no step model is set", () => { + const r = resolveRoutedModel({ + skillModel: "lightweight", + classifyText: "analyze the diff", + routing, + primaryModel: "openai:gpt-5-reasoning", + }); + assert.equal(r.model, "openai:gpt-4o-mini"); + assert.equal(r.source, "skill"); +}); + +test("auto classification routes lightweight tasks to the cheap model", () => { + const r = resolveRoutedModel({ + classifyText: "summarize the pull request", + routing, + primaryModel: "openai:gpt-5-reasoning", + }); + assert.equal(r.model, "openai:gpt-4o-mini"); + assert.equal(r.tier, "lightweight"); + assert.equal(r.source, "auto"); +}); + +test("routing stays opt-in — no routing block falls back to primary", () => { + const r = resolveRoutedModel({ + classifyText: "summarize the pull request", + primaryModel: "openai:gpt-5-reasoning", + }); + assert.equal(r.model, "openai:gpt-5-reasoning"); + assert.equal(r.source, "fallback"); +}); + +test("disabled routing falls back to primary", () => { + const r = resolveRoutedModel({ + classifyText: "summarize the pull request", + routing: { ...routing, enabled: false }, + primaryModel: "openai:gpt-5-reasoning", + }); + assert.equal(r.model, "openai:gpt-5-reasoning"); + assert.equal(r.source, "fallback"); +});