Something went wrong. Try again.
fork of prime agent with changes i want but are not upstreamable probqbly (and some fixes)
Something went wrong. Try again.
22 kB · 639 lines
TypeScript
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640/** * Model resolution, scoping, and initial selection */
import type { ThinkingLevel } from "@earendil-works/pi-agent-core";import { type Api, getLogger, type KnownProvider, type Model, modelsAreEqual } from "@earendil-works/pi-ai";import chalk from "chalk";import { minimatch } from "minimatch";import { isValidThinkingLevel } from "../cli/args.js";import { APP_NAME } from "../config.js";import { DEFAULT_THINKING_LEVEL } from "./defaults.js";import type { ModelRegistry } from "./model-registry.js";import { isPrivatePrimeInferenceModel } from "./prime-inference-models.js";
const log = getLogger("coding-agent.model-resolver");
export const PRIME_INFERENCE_DEFAULT_MODEL_ID = "z-ai/glm-5.2";
/** Default model IDs for each known provider */export const defaultModelPerProvider: Record<KnownProvider, string> = { "amazon-bedrock": "us.anthropic.claude-opus-4-6-v1", anthropic: "claude-opus-4-7", openai: "gpt-5.4", "azure-openai-responses": "gpt-5.4", "openai-codex": "gpt-5.5", "prime-inference": PRIME_INFERENCE_DEFAULT_MODEL_ID, deepseek: "deepseek-v4-pro", google: "gemini-3.1-pro-preview", "google-vertex": "gemini-3.1-pro-preview", "github-copilot": "gpt-5.4", openrouter: "moonshotai/kimi-k2.6", "vercel-ai-gateway": "zai/glm-5.1", xai: "grok-4.20-0309-reasoning", groq: "openai/gpt-oss-120b", cerebras: "gpt-oss-120b", zai: "glm-5.3", mistral: "devstral-medium-latest", minimax: "MiniMax-M2.7", "minimax-cn": "MiniMax-M2.7", moonshotai: "kimi-k2.6", "moonshotai-cn": "kimi-k2.6", huggingface: "moonshotai/Kimi-K2.6", fireworks: "accounts/fireworks/models/kimi-k2p6", opencode: "kimi-k2.6", "opencode-go": "kimi-k2.6", "kimi-coding": "kimi-for-coding", "cloudflare-workers-ai": "@cf/moonshotai/kimi-k2.6", "cloudflare-ai-gateway": "claude-sonnet-4.5", xiaomi: "mimo-v2.5-pro", "xiaomi-token-plan-cn": "mimo-v2.5-pro", "xiaomi-token-plan-ams": "mimo-v2.5-pro", "xiaomi-token-plan-sgp": "mimo-v2.5-pro",};
export interface ScopedModel { model: Model<Api>; /** Thinking level if explicitly specified in pattern (e.g., "model:high"), undefined otherwise */ thinkingLevel?: ThinkingLevel;}
/** * Helper to check if a model ID looks like an alias (no date suffix) * Dates are typically in format: -20241022 or -20250929 */function isAlias(id: string): boolean { if (id.endsWith("-latest")) return true; const datePattern = /-\d{8}$/; return !datePattern.test(id);}
/** * Find an exact model reference match. * Supports either a bare model id or a canonical provider/modelId reference. * When matching by bare id, ambiguous matches across providers are rejected. */export function findExactModelReferenceMatch( modelReference: string, availableModels: Model<Api>[],): Model<Api> | undefined { const trimmedReference = modelReference.trim(); if (!trimmedReference) { return undefined; }
const normalizedReference = trimmedReference.toLowerCase();
const canonicalMatches = availableModels.filter( (model) => `${model.provider}/${model.id}`.toLowerCase() === normalizedReference, ); if (canonicalMatches.length === 1) { return canonicalMatches[0]; } if (canonicalMatches.length > 1) { return undefined; }
const slashIndex = trimmedReference.indexOf("/"); if (slashIndex !== -1) { const provider = trimmedReference.substring(0, slashIndex).trim(); const modelId = trimmedReference.substring(slashIndex + 1).trim(); if (provider && modelId) { const providerMatches = availableModels.filter( (model) => model.provider.toLowerCase() === provider.toLowerCase() && model.id.toLowerCase() === modelId.toLowerCase(), ); if (providerMatches.length === 1) { return providerMatches[0]; } if (providerMatches.length > 1) { return undefined; } } }
const idMatches = availableModels.filter((model) => model.id.toLowerCase() === normalizedReference); return idMatches.length === 1 ? idMatches[0] : undefined;}
/** * Try to match a pattern to a model from the available models list. * Returns the matched model or undefined if no match found. */function tryMatchModel(modelPattern: string, availableModels: Model<Api>[]): Model<Api> | undefined { const exactMatch = findExactModelReferenceMatch(modelPattern, availableModels); if (exactMatch) { return exactMatch; } const matches = availableModels.filter( (m) => m.id.toLowerCase().includes(modelPattern.toLowerCase()) || m.name?.toLowerCase().includes(modelPattern.toLowerCase()), );
if (matches.length === 0) { return undefined; } const aliases = matches.filter((m) => isAlias(m.id)); const datedVersions = matches.filter((m) => !isAlias(m.id));
if (aliases.length > 0) { aliases.sort((a, b) => b.id.localeCompare(a.id)); return aliases[0]; } else { datedVersions.sort((a, b) => b.id.localeCompare(a.id)); return datedVersions[0]; }}
interface ParsedModelResult { model: Model<Api> | undefined; /** Thinking level if explicitly specified in pattern, undefined otherwise */ thinkingLevel?: ThinkingLevel; warning: string | undefined;}
function buildFallbackModel(provider: string, modelId: string, availableModels: Model<Api>[]): Model<Api> | undefined { const providerModels = availableModels.filter((m) => m.provider === provider); if (providerModels.length === 0) return undefined;
const defaultId = defaultModelPerProvider[provider as KnownProvider]; const baseModel = defaultId ? (providerModels.find((m) => m.id === defaultId) ?? providerModels[0]) : providerModels[0];
return { ...baseModel, id: modelId, name: modelId, };}
function findPreferredDefaultModel(availableModels: Model<Api>[]): Model<Api> | undefined { const primeInferenceDefault = availableModels.find( (model) => model.provider === "prime-inference" && model.id === PRIME_INFERENCE_DEFAULT_MODEL_ID, ); if (primeInferenceDefault) { return primeInferenceDefault; }
for (const provider of Object.keys(defaultModelPerProvider) as KnownProvider[]) { const defaultId = defaultModelPerProvider[provider]; const match = availableModels.find((model) => model.provider === provider && model.id === defaultId); if (match) { return match; } }
return undefined;}
/** * Parse a pattern to extract model and thinking level. * Handles models with colons in their IDs (e.g., OpenRouter's :exacto suffix). * * Algorithm: * 1. Try to match full pattern as a model * 2. If found, return it with "off" thinking level * 3. If not found and has colons, split on last colon: * - If suffix is valid thinking level, use it and recurse on prefix * - If suffix is invalid, warn and recurse on prefix with "off" */function parseModelPattern( pattern: string, availableModels: Model<Api>[], options?: { allowInvalidThinkingLevelFallback?: boolean },): ParsedModelResult { const exactMatch = tryMatchModel(pattern, availableModels); if (exactMatch) { return { model: exactMatch, thinkingLevel: undefined, warning: undefined }; } const lastColonIndex = pattern.lastIndexOf(":"); if (lastColonIndex === -1) { return { model: undefined, thinkingLevel: undefined, warning: undefined }; }
const prefix = pattern.substring(0, lastColonIndex); const suffix = pattern.substring(lastColonIndex + 1);
if (isValidThinkingLevel(suffix)) { const result = parseModelPattern(prefix, availableModels, options); if (result.model) { return { model: result.model, thinkingLevel: result.warning ? undefined : suffix, warning: result.warning, }; } return result; } else { const allowFallback = options?.allowInvalidThinkingLevelFallback ?? true; if (!allowFallback) { // In strict mode (CLI --model parsing), treat it as part of the model id and fail. // This avoids accidentally resolving to a different model. return { model: undefined, thinkingLevel: undefined, warning: undefined }; } const result = parseModelPattern(prefix, availableModels, options); if (result.model) { return { model: result.model, thinkingLevel: undefined, warning: `Invalid thinking level "${suffix}" in pattern "${pattern}". Using default instead.`, }; } return result; }}
/** * Resolve model patterns to actual Model objects with optional thinking levels * Format: "pattern:level" where :level is optional * For each pattern, finds all matching models and picks the best version: * 1. Prefer alias (e.g., claude-sonnet-4-5) over dated versions (claude-sonnet-4-5-20250929) * 2. If no alias, pick the latest dated version * * Supports models with colons in their IDs (e.g., OpenRouter's model:exacto). * The algorithm tries to match the full pattern first, then progressively * strips colon-suffixes to find a match. */export function resolveModelScopeFromModels(patterns: string[], availableModels: Model<Api>[]): ScopedModel[] { const scopedModels: ScopedModel[] = [];
for (const pattern of patterns) { if (pattern.includes("*") || pattern.includes("?") || pattern.includes("[")) { const colonIdx = pattern.lastIndexOf(":"); let globPattern = pattern; let thinkingLevel: ThinkingLevel | undefined;
if (colonIdx !== -1) { const suffix = pattern.substring(colonIdx + 1); if (isValidThinkingLevel(suffix)) { thinkingLevel = suffix; globPattern = pattern.substring(0, colonIdx); } } // This allows "*sonnet*" to match without requiring "anthropic/*sonnet*" const matchingModels = availableModels.filter((m) => { const fullId = `${m.provider}/${m.id}`; return minimatch(fullId, globPattern, { nocase: true }) || minimatch(m.id, globPattern, { nocase: true }); });
if (matchingModels.length === 0) { log.warn("no models match pattern", { pattern }); console.warn(chalk.yellow(`Warning: No models match pattern "${pattern}"`)); continue; }
for (const model of matchingModels) { if (!scopedModels.find((sm) => modelsAreEqual(sm.model, model))) { scopedModels.push({ model, thinkingLevel }); } } continue; }
const { model, thinkingLevel, warning } = parseModelPattern(pattern, availableModels);
if (warning) { log.warn(warning, { pattern }); console.warn(chalk.yellow(`Warning: ${warning}`)); }
if (!model) { log.warn("no models match pattern", { pattern }); console.warn(chalk.yellow(`Warning: No models match pattern "${pattern}"`)); continue; } if (!scopedModels.find((sm) => modelsAreEqual(sm.model, model))) { scopedModels.push({ model, thinkingLevel }); } }
return scopedModels;}
export async function resolveModelScope(patterns: string[], modelRegistry: ModelRegistry): Promise<ScopedModel[]> { const availableModels = await modelRegistry.refreshAvailableModels(); return resolveModelScopeFromModels(patterns, availableModels);}
export interface ResolveCliModelResult { model: Model<Api> | undefined; thinkingLevel?: ThinkingLevel; warning: string | undefined; /** * Error message suitable for CLI display. * When set, model will be undefined. */ error: string | undefined;}
/** * Resolve a single model from CLI flags. * * Supports: * - --provider <provider> --model <pattern> * - --model <provider>/<pattern> * - Fuzzy matching (same rules as model scoping: exact id, then partial id/name) * * Note: This does not apply the thinking level by itself, but it may *parse* and * return a thinking level from "<pattern>:<thinking>" so the caller can apply it. */export function resolveCliModel(options: { cliProvider?: string; cliModel?: string; modelRegistry: ModelRegistry;}): ResolveCliModelResult { const { cliProvider, cliModel, modelRegistry } = options;
if (!cliModel) { return { model: undefined, warning: undefined, error: undefined }; }
// Important: use *all* models here, not just models with pre-configured auth. // This allows "--api-key" to be used for first-time setup. const availableModels = modelRegistry.getAll(); if (availableModels.length === 0) { return { model: undefined, warning: undefined, error: "No models available. Check your installation or add models to models.json.", }; } const providerMap = new Map<string, string>(); for (const m of availableModels) { providerMap.set(m.provider.toLowerCase(), m.provider); }
let provider = cliProvider ? providerMap.get(cliProvider.toLowerCase()) : undefined; if (cliProvider && !provider) { return { model: undefined, warning: undefined, error: `Unknown provider "${cliProvider}". Use "${APP_NAME} model list" to see available providers/models.`, }; }
// If no explicit --provider, try to interpret "provider/model" format first. // When the prefix before the first slash matches a known provider, prefer that // interpretation over matching models whose IDs literally contain slashes // (e.g. "zai/glm-5" should resolve to provider=zai, model=glm-5, not to a // vercel-ai-gateway model with id "zai/glm-5"). let pattern = cliModel; let inferredProvider = false;
if (!provider) { const slashIndex = cliModel.indexOf("/"); if (slashIndex !== -1) { const maybeProvider = cliModel.substring(0, slashIndex); const canonical = providerMap.get(maybeProvider.toLowerCase()); if (canonical) { provider = canonical; pattern = cliModel.substring(slashIndex + 1); inferredProvider = true; } } }
// If no provider was inferred from the slash, try exact matches without provider inference. // This handles models whose IDs naturally contain slashes (e.g. OpenRouter-style IDs). if (!provider) { const lower = cliModel.toLowerCase(); const exact = availableModels.find( (m) => m.id.toLowerCase() === lower || `${m.provider}/${m.id}`.toLowerCase() === lower, ); if (exact) { return { model: exact, warning: undefined, thinkingLevel: undefined, error: undefined }; } }
if (cliProvider && provider) { // If both were provided, tolerate --model <provider>/<pattern> by stripping the provider prefix const prefix = `${provider}/`; if (cliModel.toLowerCase().startsWith(prefix.toLowerCase())) { pattern = cliModel.substring(prefix.length); } }
const candidates = provider ? availableModels.filter((m) => m.provider === provider) : availableModels; const { model, thinkingLevel, warning } = parseModelPattern(pattern, candidates, { allowInvalidThinkingLevelFallback: false, });
if (model) { return { model, thinkingLevel, warning, error: undefined }; }
// If we inferred a provider from the slash but found no match within that provider, // fall back to matching the full input as a raw model id across all models. // This handles OpenRouter-style IDs like "openai/gpt-4o:extended" where "openai" // looks like a provider but the full string is actually a model id on openrouter. if (inferredProvider) { const lower = cliModel.toLowerCase(); const exact = availableModels.find( (m) => m.id.toLowerCase() === lower || `${m.provider}/${m.id}`.toLowerCase() === lower, ); if (exact) { return { model: exact, warning: undefined, thinkingLevel: undefined, error: undefined }; } const fallback = parseModelPattern(cliModel, availableModels, { allowInvalidThinkingLevelFallback: false, }); if (fallback.model) { return { model: fallback.model, thinkingLevel: fallback.thinkingLevel, warning: fallback.warning, error: undefined, }; } }
if (provider) { const fallbackModel = buildFallbackModel(provider, pattern, availableModels); if (fallbackModel) { const fallbackWarning = warning ? `${warning} Model "${pattern}" not found for provider "${provider}". Using custom model id.` : `Model "${pattern}" not found for provider "${provider}". Using custom model id.`; return { model: fallbackModel, thinkingLevel: undefined, warning: fallbackWarning, error: undefined }; } }
const display = provider ? `${provider}/${pattern}` : cliModel; return { model: undefined, thinkingLevel: undefined, warning, error: `Model "${display}" not found. Use "${APP_NAME} model list" to see available models.`, };}
export interface InitialModelResult { model: Model<Api> | undefined; thinkingLevel: ThinkingLevel; fallbackMessage: string | undefined;}
/** * Find the initial model to use based on priority: * 1. CLI args (provider + model) * 2. First model from scoped models (if not continuing/resuming) * 3. Restored from session (if continuing/resuming) * 4. Saved default from settings * 5. First available model with valid API key */export async function findInitialModel(options: { cliProvider?: string; cliModel?: string; scopedModels: ScopedModel[]; isContinuing: boolean; defaultProvider?: string; defaultModelId?: string; defaultThinkingLevel?: ThinkingLevel; modelRegistry: ModelRegistry;}): Promise<InitialModelResult> { const { cliProvider, cliModel, scopedModels, isContinuing, defaultProvider, defaultModelId, defaultThinkingLevel, modelRegistry, } = options;
let model: Model<Api> | undefined; let thinkingLevel: ThinkingLevel = DEFAULT_THINKING_LEVEL; let cachedAvailableModels: Model<Api>[] | undefined; const getAvailableModels = async (): Promise<Model<Api>[]> => { cachedAvailableModels ??= await modelRegistry.refreshAvailableModels(); return cachedAvailableModels; }; if (cliProvider && cliModel) { const resolved = resolveCliModel({ cliProvider, cliModel, modelRegistry, }); if (resolved.error) { log.error(resolved.error, { cliProvider, cliModel }); console.error(chalk.red(resolved.error)); process.exit(1); } const resolvedModel = resolved.model; if (resolvedModel) { if (isPrivatePrimeInferenceModel(resolvedModel)) { const availableModel = (await getAvailableModels()).find((candidate) => modelsAreEqual(candidate, resolvedModel), ); if (!availableModel) { const error = `Model "${resolvedModel.provider}/${resolvedModel.id}" is not available for the current Prime team.`; log.error(error, { cliProvider, cliModel }); console.error(chalk.red(error)); process.exit(1); } return { model: availableModel, thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined }; } return { model: resolvedModel, thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined }; } } if (scopedModels.length > 0 && !isContinuing) { return { model: scopedModels[0].model, thinkingLevel: scopedModels[0].thinkingLevel ?? defaultThinkingLevel ?? DEFAULT_THINKING_LEVEL, fallbackMessage: undefined, }; } const availableModels = await getAvailableModels(); if (defaultProvider && defaultModelId) { // Rebuild from the provider template when the saved id is missing from this // build's snapshot (e.g. prime-inference catalog churn), so it survives updates. const found = availableModels.find( (candidate) => candidate.provider === defaultProvider && candidate.id === defaultModelId, ) ?? (!isPrivatePrimeInferenceModel({ provider: defaultProvider, id: defaultModelId }) ? buildFallbackModel(defaultProvider, defaultModelId, availableModels) : undefined); if (found) { model = found; if (defaultThinkingLevel) { thinkingLevel = defaultThinkingLevel; } return { model, thinkingLevel, fallbackMessage: undefined }; } } if (availableModels.length > 0) { const defaultModel = findPreferredDefaultModel(availableModels); if (defaultModel) { return { model: defaultModel, thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined }; } return { model: availableModels[0], thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined }; } return { model: undefined, thinkingLevel: DEFAULT_THINKING_LEVEL, fallbackMessage: undefined };}
/** * Restore model from session, with fallback to available models */export async function restoreModelFromSession( savedProvider: string, savedModelId: string, currentModel: Model<Api> | undefined, shouldPrintMessages: boolean, modelRegistry: ModelRegistry,): Promise<{ model: Model<Api> | undefined; fallbackMessage: string | undefined }> { const availableModels = await modelRegistry.refreshAvailableModels(); const restoredModel = availableModels.find( (candidate) => candidate.provider === savedProvider && candidate.id === savedModelId, );
if (restoredModel) { if (shouldPrintMessages) { console.log(chalk.dim(`Restored model: ${savedProvider}/${savedModelId}`)); } return { model: restoredModel, fallbackMessage: undefined }; } const registeredModel = modelRegistry.find(savedProvider, savedModelId); const reason = !registeredModel ? "model no longer exists" : !modelRegistry.hasConfiguredAuth(registeredModel) ? "no auth configured" : "model is not available"; log.warn("could not restore model", { provider: savedProvider, model: savedModelId, reason });
if (shouldPrintMessages) { console.error(chalk.yellow(`Warning: Could not restore model ${savedProvider}/${savedModelId} (${reason}).`)); } const availableCurrentModel = currentModel ? availableModels.find((candidate) => modelsAreEqual(candidate, currentModel)) : undefined; const fallbackCurrentModel = currentModel && (!isPrivatePrimeInferenceModel(currentModel) || availableCurrentModel) ? (availableCurrentModel ?? currentModel) : undefined; if (fallbackCurrentModel) { if (shouldPrintMessages) { console.log(chalk.dim(`Falling back to: ${fallbackCurrentModel.provider}/${fallbackCurrentModel.id}`)); } return { model: fallbackCurrentModel, fallbackMessage: `Could not restore model ${savedProvider}/${savedModelId} (${reason}). Using ${fallbackCurrentModel.provider}/${fallbackCurrentModel.id}.`, }; } if (availableModels.length > 0) { const fallbackModel = findPreferredDefaultModel(availableModels) ?? availableModels[0];
if (shouldPrintMessages) { console.log(chalk.dim(`Falling back to: ${fallbackModel.provider}/${fallbackModel.id}`)); }
return { model: fallbackModel, fallbackMessage: `Could not restore model ${savedProvider}/${savedModelId} (${reason}). Using ${fallbackModel.provider}/${fallbackModel.id}.`, }; } return { model: undefined, fallbackMessage: undefined };}