stack/apps/gateway/src/agent/model-capabilities.ts

import type { ModelCapability } from '@mosaic/types';

/**
 * Comprehensive capability matrix for all target models.
 * Cost fields are optional and will be filled in when real pricing data is available.
 */
export const MODEL_CAPABILITIES: ModelCapability[] = [
  {
    id: 'claude-opus-4-6',
    provider: 'anthropic',
    displayName: 'Claude Opus 4.6',
    tier: 'premium',
    contextWindow: 200000,
    maxOutputTokens: 32000,
    capabilities: {
      tools: true,
      vision: true,
      streaming: true,
      reasoning: true,
      embedding: false,
    },
  },
  {
    id: 'claude-sonnet-4-6',
    provider: 'anthropic',
    displayName: 'Claude Sonnet 4.6',
    tier: 'standard',
    contextWindow: 200000,
    maxOutputTokens: 16000,
    capabilities: {
      tools: true,
      vision: true,
      streaming: true,
      reasoning: true,
      embedding: false,
    },
  },
  {
    id: 'claude-haiku-4-5',
    provider: 'anthropic',
    displayName: 'Claude Haiku 4.5',
    tier: 'cheap',
    contextWindow: 200000,
    maxOutputTokens: 8192,
    capabilities: {
      tools: true,
      vision: true,
      streaming: true,
      reasoning: false,
      embedding: false,
    },
  },
  {
    id: 'codex-gpt-5.4',
    provider: 'openai',
    displayName: 'Codex gpt-5.4',
    tier: 'premium',
    contextWindow: 128000,
    maxOutputTokens: 16384,
    capabilities: {
      tools: true,
      vision: true,
      streaming: true,
      reasoning: true,
      embedding: false,
    },
  },
  {
    id: 'glm-5',
    provider: 'zai',
    displayName: 'GLM-5',
    tier: 'standard',
    contextWindow: 128000,
    maxOutputTokens: 8192,
    capabilities: {
      tools: true,
      vision: false,
      streaming: true,
      reasoning: false,
      embedding: false,
    },
  },
  {
    id: 'llama3.2',
    provider: 'ollama',
    displayName: 'llama3.2',
    tier: 'local',
    contextWindow: 128000,
    maxOutputTokens: 8192,
    capabilities: {
      tools: true,
      vision: false,
      streaming: true,
      reasoning: false,
      embedding: false,
    },
  },
  {
    id: 'codellama',
    provider: 'ollama',
    displayName: 'codellama',
    tier: 'local',
    contextWindow: 16000,
    maxOutputTokens: 4096,
    capabilities: {
      tools: true,
      vision: false,
      streaming: true,
      reasoning: false,
      embedding: false,
    },
  },
  {
    id: 'mistral',
    provider: 'ollama',
    displayName: 'mistral',
    tier: 'local',
    contextWindow: 32000,
    maxOutputTokens: 8192,
    capabilities: {
      tools: true,
      vision: false,
      streaming: true,
      reasoning: false,
      embedding: false,
    },
  },
  {
    id: 'nomic-embed-text',
    provider: 'ollama',
    displayName: 'nomic-embed-text',
    tier: 'local',
    contextWindow: 8192,
    maxOutputTokens: 0,
    capabilities: {
      tools: false,
      vision: false,
      streaming: false,
      reasoning: false,
      embedding: true,
    },
  },
  {
    id: 'mxbai-embed-large',
    provider: 'ollama',
    displayName: 'mxbai-embed-large',
    tier: 'local',
    contextWindow: 8192,
    maxOutputTokens: 0,
    capabilities: {
      tools: false,
      vision: false,
      streaming: false,
      reasoning: false,
      embedding: true,
    },
  },
];

/**
 * Look up a model by its ID.
 * Returns undefined if the model is not found.
 */
export function getModelCapability(modelId: string): ModelCapability | undefined {
  return MODEL_CAPABILITIES.find((m) => m.id === modelId);
}

/**
 * Find models matching a partial capability filter.
 * All provided filter keys must match for a model to be included.
 */
export function findModelsByCapability(
  filter: Partial<Pick<ModelCapability, 'tier' | 'provider'>> & {
    capabilities?: Partial<ModelCapability['capabilities']>;
  },
): ModelCapability[] {
  return MODEL_CAPABILITIES.filter((model) => {
    if (filter.tier !== undefined && model.tier !== filter.tier) return false;
    if (filter.provider !== undefined && model.provider !== filter.provider) return false;
    if (filter.capabilities) {
      for (const [key, value] of Object.entries(filter.capabilities) as [
        keyof ModelCapability['capabilities'],
        boolean,
      ][]) {
        if (model.capabilities[key] !== value) return false;
      }
    }
    return true;
  });
}

/**
 * Get all models for a specific provider.
 */
export function getModelsByProvider(provider: string): ModelCapability[] {
  return MODEL_CAPABILITIES.filter((m) => m.provider === provider);
}

/**
 * Get the full list of all known models.
 */
export function getAllModels(): ModelCapability[] {
  return MODEL_CAPABILITIES;
}