diff --git a/src/exports.ts b/src/exports.ts index 787d60f..5347a6a 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, RouteQuery, RouteDeps } 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..9a4ff79 --- /dev/null +++ b/src/model-routing.ts @@ -0,0 +1,185 @@ +// Resolves which model each SkillFlow step runs on. Explicit per-step / per-skill +// settings win; otherwise the step is classified by one call to the lightweight +// model via an injected `query` (RouteDeps), falling back to a keyword heuristic +// when no `query` is supplied. + +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, before the LLM/keyword step. */ + 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; +} + +/** Minimal shape of the SDK `query()` used to classify — injected so core stays decoupled and testable. */ +export type RouteQuery = (opts: { + prompt: string; + model?: string; + dir?: string; + env?: string; + systemPrompt?: string; + maxTurns?: number; + tools?: []; + replaceBuiltinTools?: boolean; +}) => AsyncIterable<{ type: string; content?: string }>; + +export interface RouteDeps { + /** When present (and routing is enabled) classification runs one LLM call on the lightweight model. */ + query?: RouteQuery; + dir?: string; + env?: 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"; +} + +// Keyword fallback, used only when no `query` is injected. Ambiguous verbs like +// "search" that appear in both trivial and complex prompts are deliberately left +// out so they don't systematically overpay. +const DEFAULT_LIGHTWEIGHT = [ + "summ", "extract", "classif", "transform", "format", "convert", + "parse", "fetch", "read", "load", "lookup", "normaliz", "translat", + "rephrase", "rewrite", "tag", "label", "render", +]; +const DEFAULT_REASONING = [ + "analy", "plan", "decid", "decision", "orchestrat", "solve", + "reason", "validat", "evaluat", "review", "audit", "diagnos", "debug", + "architect", "design", "strateg", "investigat", "research", "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; +} + +function matchRules(text: string, rules?: Array<{ tier: ModelTier; match: string[] }>): ModelTier | null { + if (rules) { + for (const rule of rules) { + if (Array.isArray(rule.match) && matchesAny(text, rule.match)) return rule.tier; + } + } + return null; +} + +/** + * Keyword-based tier classification — the offline fallback. User rules win; + * otherwise a task that matches neither list (or both) resolves to "reasoning", + * so cost optimization never silently degrades quality. + */ +export function classifyByKeywords( + classifyText: string, + rules?: Array<{ tier: ModelTier; match: string[] }>, +): ModelTier { + const text = classifyText || ""; + const ruled = matchRules(text, rules); + if (ruled) return ruled; + if (matchesAny(text, DEFAULT_REASONING)) return "reasoning"; + if (matchesAny(text, DEFAULT_LIGHTWEIGHT)) return "lightweight"; + return "reasoning"; +} + +const CLASSIFY_INSTRUCTIONS = + `Classify the following agent task as exactly one word: "lightweight" or "reasoning".\n` + + `- lightweight: mechanical work with a known shape (summarize, extract, format, fetch, read, look up).\n` + + `- reasoning: needs analysis, planning, multi-step logic, or judgment.\n` + + `If unsure, answer "reasoning". Reply with only the single word.\n\nTask: `; + +/** One classification call on the lightweight model. Returns null if the call fails or is unparseable. */ +async function classifyViaLLM(text: string, model: string, deps: RouteDeps): Promise { + try { + // Constrained one-shot: no tools, single turn — keeps it a single cheap call. + const result = deps.query!({ + prompt: CLASSIFY_INSTRUCTIONS + text.trim(), + model, + dir: deps.dir, + env: deps.env, + systemPrompt: "You are a task classifier. Reply with exactly one word.", + maxTurns: 1, + tools: [], + replaceBuiltinTools: true, + }); + let out = ""; + for await (const msg of result) { + if (msg.type === "assistant" && msg.content) out += msg.content; + } + const answer = out.toLowerCase(); + if (answer.includes("lightweight")) return "lightweight"; + if (answer.includes("reason")) return "reasoning"; + return null; + } catch { + return null; + } +} + +/** + * Resolve a tier alias ("lightweight"/"reasoning") or pass a model id through. + * Warns when an alias is requested but the routing block has no model for that + * tier, so silently falling through to the fallback stays observable. + */ +export function resolveModelAlias(ref: string | undefined, routing?: RoutingConfig): string | undefined { + if (!ref) return undefined; + if (ref === "lightweight" || ref === "reasoning") { + const configured = routing?.[ref]; + if (!configured) { + console.warn(`[routing] tier alias '${ref}' requested but routing.${ref} is not configured; falling through`); + return undefined; + } + return configured; + } + 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. Automatic classification + * uses the injected `query` fn when available, else the keyword fallback. + */ +export async function resolveRoutedModel(input: RouteInput, deps: RouteDeps = {}): Promise { + 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) { + let tier = matchRules(classifyText || "", routing!.rules); + if (!tier && deps.query && routing!.lightweight) { + tier = await classifyViaLLM(classifyText, routing!.lightweight, deps); + } + if (!tier) tier = classifyByKeywords(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..eaa9901 --- /dev/null +++ b/test/model-routing.test.ts @@ -0,0 +1,181 @@ +// Tests for the resolution priority chain (step > skill > auto > fallback), the +// LLM classifier (via an injected fake query), the keyword fallback, and a +// save -> load roundtrip proving step `model` survives the YAML plumbing. + +import test from "node:test"; +import assert from "node:assert/strict"; +import { mkdtemp, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; + +import { + classifyByKeywords, + resolveModelAlias, + resolveRoutedModel, + type RoutingConfig, + type RouteQuery, +} from "../src/model-routing.ts"; +import { saveFlowDefinition, loadFlowDefinition } from "../src/workflows.ts"; + +const routing: RoutingConfig = { + lightweight: "openai:gpt-4o-mini", + reasoning: "openai:gpt-4o", +}; + +// A fake injected query that always "answers" with the given text. +function fakeQuery(answer: string): RouteQuery { + return () => (async function* () { + yield { type: "assistant", content: answer }; + })(); +} + +// ── classifyByKeywords (offline fallback) ────────────────────────────── + +test("keyword fallback classifies lightweight and reasoning verbs", () => { + assert.equal(classifyByKeywords("summarize the pull request diff"), "lightweight"); + assert.equal(classifyByKeywords("analyze the security implications"), "reasoning"); +}); + +test("unknown tasks default to reasoning (never silently downgrade quality)", () => { + assert.equal(classifyByKeywords("frobnicate the widget"), "reasoning"); + assert.equal(classifyByKeywords(""), "reasoning"); +}); + +test("bare 'search' no longer forces reasoning, but 'research' still does", () => { + // Previously "search" was in the reasoning list, so this whole step went + // expensive; now the lightweight "summarize" verb wins. + assert.equal(classifyByKeywords("search the logs and summarize them"), "lightweight"); + assert.equal(classifyByKeywords("research the root cause"), "reasoning"); + assert.equal(classifyByKeywords("investigate the regression"), "reasoning"); +}); + +test("user rules take precedence over defaults", () => { + const rules = [{ tier: "lightweight" as const, match: ["analyze"] }]; + assert.equal(classifyByKeywords("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); +}); + +test("alias with no configured tier model warns and returns undefined", () => { + const warnings: string[] = []; + const orig = console.warn; + console.warn = (m?: any) => { warnings.push(String(m)); }; + try { + assert.equal(resolveModelAlias("lightweight", { reasoning: "openai:gpt-4o" }), undefined); + } finally { + console.warn = orig; + } + assert.equal(warnings.length, 1); + assert.match(warnings[0], /routing\.lightweight is not configured/); +}); + +// ── resolveRoutedModel priority chain ────────────────────────────────── + +test("explicit per-step model wins over everything", async () => { + const r = await 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", async () => { + const r = await 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 routing uses the injected LLM classifier", async () => { + const r = await resolveRoutedModel( + { classifyText: "do the thing", routing, primaryModel: "openai:gpt-5-reasoning" }, + { query: fakeQuery("lightweight") }, + ); + assert.equal(r.model, "openai:gpt-4o-mini"); + assert.equal(r.tier, "lightweight"); + assert.equal(r.source, "auto"); +}); + +test("LLM classifier answering 'reasoning' routes to the reasoning model", async () => { + const r = await resolveRoutedModel( + { classifyText: "do the thing", routing, primaryModel: "openai:gpt-5-reasoning" }, + { query: fakeQuery("reasoning") }, + ); + assert.equal(r.model, "openai:gpt-4o"); + assert.equal(r.tier, "reasoning"); +}); + +test("unparseable LLM output falls back to the keyword heuristic", async () => { + const r = await resolveRoutedModel( + { classifyText: "xyzzy", routing, primaryModel: "openai:gpt-5-reasoning" }, + { query: fakeQuery("¯\\_(ツ)_/¯") }, + ); + assert.equal(r.tier, "reasoning"); // keyword fallback → unknown defaults to reasoning + assert.equal(r.source, "auto"); +}); + +test("no injected query → keyword fallback still routes", async () => { + const r = await resolveRoutedModel({ + classifyText: "summarize the pull request", + routing, + primaryModel: "openai:gpt-5-reasoning", + }); + assert.equal(r.model, "openai:gpt-4o-mini"); + assert.equal(r.source, "auto"); +}); + +test("routing stays opt-in — no routing block falls back to primary", async () => { + const r = await 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", async () => { + const r = await resolveRoutedModel( + { classifyText: "summarize the pull request", routing: { ...routing, enabled: false }, primaryModel: "openai:gpt-5-reasoning" }, + { query: fakeQuery("lightweight") }, + ); + assert.equal(r.model, "openai:gpt-5-reasoning"); + assert.equal(r.source, "fallback"); +}); + +// ── Integration: step `model` survives the YAML roundtrip ────────────── + +test("step model survives saveFlowDefinition -> loadFlowDefinition", async () => { + const dir = await mkdtemp(join(tmpdir(), "gitagent-flow-")); + try { + const filePath = await saveFlowDefinition(dir, { + name: "route-test", + description: "roundtrip", + steps: [ + { skill: "summarize", prompt: "summarize {input}", model: "openai:gpt-4o-mini" }, + { skill: "plan", prompt: "plan the fix", model: "reasoning" }, + { skill: "noop", prompt: "no model here" }, + ], + }); + const loaded = await loadFlowDefinition(filePath); + assert.equal(loaded.steps[0].model, "openai:gpt-4o-mini"); + assert.equal(loaded.steps[1].model, "reasoning"); + assert.equal(loaded.steps[2].model, undefined); + } finally { + await rm(dir, { recursive: true, force: true }); + } +});