boocode/packages/ion/src/engine/model-validation.ts

/**
 * Provider/model resolution for the Ion workflow engine.
 *
 * Maps model references (aliases, tier presets, or literal specs)
 * to concrete AI model configurations.
 */

import type { ProviderConfig } from './deps.js';

// ---------------------------------------------------------------------------
// Types
// ---------------------------------------------------------------------------

/** A concrete model specification with all fields resolved. */
export interface LiteralModelSpec {
  /** The AI provider (e.g. "openai", "anthropic"). */
  provider: string;
  /** The model identifier (e.g. "gpt-4o", "claude-sonnet-4-20250514"). */
  model: string;
  /** Optional effort level (e.g. "low", "medium", "high"). */
  effort?: string;
  /** Optional thinking/reasoning configuration. */
  thinking?: {
    type: string;
    budgetTokens?: number;
  };
}

/** A preset that maps an alias to a concrete model configuration. */
export interface ModelAliasPreset {
  /** The provider for this preset. */
  provider: string;
  /** The model identifier. */
  model: string;
  /** Optional effort level. */
  effort?: string;
  /** Optional thinking configuration. */
  thinking?: {
    type: string;
    budgetTokens?: number;
  };
}

/** Tier definitions for an AI profile. */
export interface AiProfileTiers {
  /** Fast/cheap tier. */
  fast?: ModelAliasPreset;
  /** Balanced tier. */
  balanced?: ModelAliasPreset;
  /** Powerful/expensive tier. */
  powerful?: ModelAliasPreset;
}

/** An AI profile with tiers and named aliases. */
export interface AiProfile {
  /** The default provider. */
  defaultProvider: string;
  /** Named provider configurations. */
  providers: Record<string, ProviderConfig>;
  /** Tier presets. */
  tiers: AiProfileTiers;
  /** Named aliases mapping to presets. */
  aliases: Record<string, ModelAliasPreset>;
}

/** Options for building an AI profile. */
export interface BuildAiProfileOptions {
  /** The default assistant/provider id. */
  assistant: string;
  /** Named provider configurations. */
  assistants: Record<string, ProviderConfig>;
  /** Optional model overrides from workflow config. */
  modelOverrides?: Record<string, ModelAliasPreset>;
}

// ---------------------------------------------------------------------------
// Type guards
// ---------------------------------------------------------------------------

/**
 * Check if a model spec is a literal (fully resolved) spec.
 *
 * A literal spec has a `provider` and `model` field directly.
 */
export function isLiteralSpec(
  spec: unknown,
): spec is LiteralModelSpec {
  if (typeof spec !== 'object' || spec === null) {
    return false;
  }
  const obj = spec as Record<string, unknown>;
  return typeof obj['provider'] === 'string' && typeof obj['model'] === 'string';
}

// ---------------------------------------------------------------------------
// Profile builder
// ---------------------------------------------------------------------------

/** Default tier presets for common providers. */
const DEFAULT_TIERS: Record<string, AiProfileTiers> = {
  openai: {
    fast: { provider: 'openai', model: 'gpt-4o-mini' },
    balanced: { provider: 'openai', model: 'gpt-4o' },
    powerful: { provider: 'openai', model: 'o1' },
  },
  anthropic: {
    fast: { provider: 'anthropic', model: 'claude-haiku-4-20250414' },
    balanced: { provider: 'anthropic', model: 'claude-sonnet-4-20250514' },
    powerful: {
      provider: 'anthropic',
      model: 'claude-opus-4-20250514',
      thinking: { type: 'enabled', budgetTokens: 10000 },
    },
  },
};

/**
 * Build an AI profile from workflow configuration.
 *
 * Merges default tier presets with any model overrides from the config.
 */
export function buildAiProfile(
  opts: BuildAiProfileOptions,
): AiProfile {
  const providers = { ...opts.assistants };

  // Determine the default provider from the assistant config.
  const defaultProviderConfig = providers[opts.assistant];
  const defaultProvider = defaultProviderConfig?.provider ?? opts.assistant;

  // Start with default tiers for the default provider.
  const baseTiers = DEFAULT_TIERS[defaultProvider] ?? {};

  // Apply model overrides if provided.
  const tiers: AiProfileTiers = { ...baseTiers };
  if (opts.modelOverrides) {
    for (const [key, preset] of Object.entries(opts.modelOverrides)) {
      if (key === 'fast' || key === 'balanced' || key === 'powerful') {
        tiers[key] = preset;
      }
    }
  }

  // Build aliases from overrides and tiers.
  const aliases: Record<string, ModelAliasPreset> = {};

  // Tier-based aliases.
  if (tiers.fast) aliases['fast'] = tiers.fast;
  if (tiers.balanced) aliases['balanced'] = tiers.balanced;
  if (tiers.powerful) aliases['powerful'] = tiers.powerful;

  // Custom overrides as aliases.
  if (opts.modelOverrides) {
    for (const [key, preset] of Object.entries(opts.modelOverrides)) {
      if (key !== 'fast' && key !== 'balanced' && key !== 'powerful') {
        aliases[key] = preset;
      }
    }
  }

  return {
    defaultProvider,
    providers,
    tiers,
    aliases,
  };
}

// ---------------------------------------------------------------------------
// Model resolution
// ---------------------------------------------------------------------------

/**
 * Resolve a model reference to a literal model spec.
 *
 * A model reference can be:
 * - A literal spec (has `provider` and `model` fields) → returned as-is
 * - A tier name ("fast", "balanced", "powerful") → resolved from profile tiers
 * - A named alias → resolved from profile aliases
 * - A provider-prefixed reference ("openai/gpt-4o") → parsed into a spec
 * - A bare model name → resolved using the default provider
 *
 * Throws if the reference cannot be resolved.
 */
export function resolveModelSpec(
  profile: AiProfile,
  modelRef: string | LiteralModelSpec,
): LiteralModelSpec {
  // Already a literal spec.
  if (typeof modelRef !== 'string') {
    if (isLiteralSpec(modelRef)) {
      return modelRef;
    }
    throw new Error(`Invalid model spec: ${JSON.stringify(modelRef)}`);
  }

  // Check aliases first (includes tier aliases).
  const alias = profile.aliases[modelRef];
  if (alias) {
    return {
      provider: alias.provider,
      model: alias.model,
      effort: alias.effort,
      thinking: alias.thinking,
    };
  }

  // Provider-prefixed reference: "provider/model"
  if (modelRef.includes('/')) {
    const slashIndex = modelRef.indexOf('/');
    const provider = modelRef.slice(0, slashIndex)!;
    const model = modelRef.slice(slashIndex + 1);

    if (!provider || !model) {
      throw new Error(
        `Invalid provider-prefixed model reference: "${modelRef}". Expected format "provider/model".`,
      );
    }

    return { provider, model };
  }

  // Bare model name — use default provider.
  return {
    provider: profile.defaultProvider,
    model: modelRef,
  };
}