ai-cli 0.1.1 → 0.2.0

package/src/lib/models.ts

@@ -1,5 +1,3 @@
-import { gateway } from "ai";
-
 export type Modality = "text" | "image" | "video";
 
 const DEFAULTS: Record<Modality, string> = {
@@ -8,156 +6,168 @@ const DEFAULTS: Record<Modality, string> = {
   video: process.env.AI_CLI_VIDEO_MODEL ?? "bytedance/seedance-2.0",
 };
 
-export const FALLBACK_TEXT_MODELS = [
-  "anthropic/claude-sonnet-4",
-  "google/gemini-2.5-pro",
-  "meta/llama-4-maverick",
-  "openai/gpt-4.1",
-  "openai/gpt-4.1-mini",
-  "openai/gpt-4.1-nano",
-  "openai/gpt-5.5",
-  "openai/o3",
-  "openai/o4-mini",
-  "xai/grok-3",
-];
-
-export const FALLBACK_IMAGE_MODELS = [
-  "bfl/flux-2-flex",
-  "bfl/flux-2-klein-4b",
-  "bfl/flux-2-klein-9b",
-  "bfl/flux-2-max",
-  "bfl/flux-2-pro",
-  "bfl/flux-kontext-max",
-  "bfl/flux-kontext-pro",
-  "bfl/flux-pro-1.0-fill",
-  "bfl/flux-pro-1.1",
-  "bfl/flux-pro-1.1-ultra",
-  "bytedance/seedream-4.0",
-  "bytedance/seedream-4.5",
-  "bytedance/seedream-5.0-lite",
-  "google/imagen-4.0-fast-generate-001",
-  "google/imagen-4.0-generate-001",
-  "google/imagen-4.0-ultra-generate-001",
-  "openai/gpt-image-1",
-  "openai/gpt-image-1-mini",
-  "openai/gpt-image-1.5",
-  "openai/gpt-image-2",
-  "prodia/flux-fast-schnell",
-  "recraft/recraft-v2",
-  "recraft/recraft-v3",
-  "recraft/recraft-v4",
-  "recraft/recraft-v4-pro",
-  "xai/grok-imagine-image",
-  "xai/grok-imagine-image-pro",
-];
-
-export const FALLBACK_VIDEO_MODELS = [
-  "alibaba/wan-v2.5-t2v-preview",
-  "alibaba/wan-v2.6-i2v",
-  "alibaba/wan-v2.6-i2v-flash",
-  "alibaba/wan-v2.6-r2v",
-  "alibaba/wan-v2.6-r2v-flash",
-  "alibaba/wan-v2.6-t2v",
-  "bytedance/seedance-2.0",
-  "bytedance/seedance-2.0-fast",
-  "bytedance/seedance-v1.0-lite-i2v",
-  "bytedance/seedance-v1.0-lite-t2v",
-  "bytedance/seedance-v1.0-pro",
-  "bytedance/seedance-v1.0-pro-fast",
-  "bytedance/seedance-v1.5-pro",
-  "google/veo-3.0-fast-generate-001",
-  "google/veo-3.0-generate-001",
-  "google/veo-3.1-fast-generate-001",
-  "google/veo-3.1-generate-001",
-  "klingai/kling-v2.5-turbo-i2v",
-  "klingai/kling-v2.5-turbo-t2v",
-  "klingai/kling-v2.6-i2v",
-  "klingai/kling-v2.6-motion-control",
-  "klingai/kling-v2.6-t2v",
-  "klingai/kling-v3.0-i2v",
-  "klingai/kling-v3.0-t2v",
-  "xai/grok-imagine-video",
-];
+const GATEWAY_MODELS_URL = "https://ai-gateway.vercel.sh/v1/models";
+const GATEWAY_TIMEOUT_MS = 5_000;
+
+export interface ModelPricing {
+  input?: string;
+  output?: string;
+  image?: string;
+}
 
 export interface ModelEntry {
   id: string;
   name?: string;
   description?: string;
+  creator: string;
+  capabilities: Modality[];
+  pricing?: ModelPricing;
 }
 
 export interface GatewayModels {
   text: ModelEntry[];
   image: ModelEntry[];
   video: ModelEntry[];
+  all: ModelEntry[];
+  languageImageModelIds: Set<string>;
 }
 
-const MODEL_TYPE_TO_MODALITY: Record<string, Modality> = {
-  language: "text",
-  image: "image",
-  video: "video",
-};
+interface RawGatewayModel {
+  id: string;
+  name?: string;
+  description?: string;
+  owned_by?: string;
+  type?: string;
+  tags?: string[];
+  pricing?: {
+    input?: string;
+    output?: string;
+    image?: string;
+  };
+}
 
-export async function fetchGatewayModels(): Promise<GatewayModels> {
-  const result: GatewayModels = { text: [], image: [], video: [] };
+let cached: Promise<GatewayModels> | null = null;
+
+export function fetchGatewayModels(): Promise<GatewayModels> {
+  if (!cached) {
+    cached = doFetch().catch((err) => {
+      cached = null;
+      throw err;
+    });
+  }
+  return cached;
+}
+
+export function resetGatewayCache(): void {
+  cached = null;
+}
+
+async function doFetch(): Promise<GatewayModels> {
+  const result: GatewayModels = {
+    text: [],
+    image: [],
+    video: [],
+    all: [],
+    languageImageModelIds: new Set(),
+  };
 
   try {
-    const { models } = await gateway.getAvailableModels();
+    const res = await fetch(GATEWAY_MODELS_URL, {
+      signal: AbortSignal.timeout(GATEWAY_TIMEOUT_MS),
+    });
+    if (!res.ok) throw new Error(`HTTP ${res.status}`);
+    const json = (await res.json()) as { data?: RawGatewayModel[] };
+    const models = json.data ?? [];
+
+    const entryMap = new Map<string, ModelEntry>();
+
     for (const m of models) {
-      const modality =
-        MODEL_TYPE_TO_MODALITY[(m as { modelType?: string }).modelType ?? ""];
-      if (!modality) continue;
-      result[modality].push({
+      const tags = m.tags ?? [];
+      const isImageGen = tags.includes("image-generation");
+      const capabilities: Modality[] = [];
+
+      switch (m.type) {
+        case "language":
+          capabilities.push("text");
+          if (isImageGen) capabilities.push("image");
+          break;
+        case "image":
+          capabilities.push("image");
+          break;
+        case "video":
+          capabilities.push("video");
+          break;
+        default:
+          continue;
+      }
+
+      const creator =
+        m.owned_by ??
+        (m.id.slice(0, Math.max(0, m.id.indexOf("/"))) || "other");
+
+      const pricing: ModelPricing | undefined =
+        m.pricing?.input || m.pricing?.output || m.pricing?.image
+          ? {
+              ...(m.pricing.input ? { input: m.pricing.input } : {}),
+              ...(m.pricing.output ? { output: m.pricing.output } : {}),
+              ...(m.pricing.image ? { image: m.pricing.image } : {}),
+            }
+          : undefined;
+
+      const entry: ModelEntry = {
         id: m.id,
         name: m.name,
-        description: m.description ?? undefined,
-      });
-    }
-  } catch {
-    result.text = FALLBACK_TEXT_MODELS.map((id) => ({ id }));
-    result.image = FALLBACK_IMAGE_MODELS.map((id) => ({ id }));
-    result.video = FALLBACK_VIDEO_MODELS.map((id) => ({ id }));
-  }
+        description: m.description,
+        creator,
+        capabilities,
+        pricing,
+      };
 
-  if (result.text.length === 0) {
-    result.text = FALLBACK_TEXT_MODELS.map((id) => ({ id }));
-  }
-  if (result.image.length === 0) {
-    result.image = FALLBACK_IMAGE_MODELS.map((id) => ({ id }));
-  }
-  if (result.video.length === 0) {
-    result.video = FALLBACK_VIDEO_MODELS.map((id) => ({ id }));
-  }
-
-  return result;
-}
-
-function expandModelId(input: string, modality: Modality): string {
-  if (input.includes("/")) return input;
+      entryMap.set(m.id, entry);
 
-  const knownLists: string[][] = [];
-  if (modality === "text") knownLists.push(FALLBACK_TEXT_MODELS);
-  else if (modality === "image") knownLists.push(FALLBACK_IMAGE_MODELS);
-  else if (modality === "video") knownLists.push(FALLBACK_VIDEO_MODELS);
+      if (capabilities.includes("text")) result.text.push(entry);
+      if (capabilities.includes("image")) result.image.push(entry);
+      if (capabilities.includes("video")) result.video.push(entry);
 
-  for (const list of knownLists) {
-    for (const fullId of list) {
-      const name = fullId.slice(fullId.indexOf("/") + 1);
-      if (name === input) return fullId;
+      if (m.type === "language" && isImageGen) {
+        result.languageImageModelIds.add(m.id);
+      }
     }
+
+    result.all = [...entryMap.values()];
+  } catch {
+    cached = null;
+    process.stderr.write("Warning: could not fetch models from AI Gateway\n");
   }
 
-  return input;
+  return result;
 }
 
 export function resolveModels(
   modality: Modality,
-  userModel?: string
+  userModel?: string,
+  knownModels?: Pick<ModelEntry, "id">[]
 ): string[] {
   if (!userModel) return [DEFAULTS[modality]];
   const models = userModel
     .split(",")
     .map((m) => m.trim())
     .filter(Boolean)
-    .map((m) => expandModelId(m, modality));
+    .map((m) => expandModelId(m, knownModels));
   return models.length > 0 ? models : [DEFAULTS[modality]];
 }
+
+function expandModelId(
+  input: string,
+  knownModels?: Pick<ModelEntry, "id">[]
+): string {
+  if (input.includes("/")) return input;
+  if (!knownModels) return input;
+
+  for (const m of knownModels) {
+    const name = m.id.slice(m.id.indexOf("/") + 1);
+    if (name === input) return m.id;
+  }
+
+  return input;
+}

package/src/commands/completions.ts

@@ -1,296 +0,0 @@
-import type { Command } from "commander";
-
-import {
-  FALLBACK_TEXT_MODELS,
-  FALLBACK_IMAGE_MODELS,
-  FALLBACK_VIDEO_MODELS,
-} from "../lib/models.js";
-
-export function registerCompletionsCommand(program: Command) {
-  program
-    .command("completions")
-    .description("Output shell completion script")
-    .argument("<shell>", "Shell type: zsh, bash, fish")
-    .action((shell: string) => {
-      switch (shell.toLowerCase()) {
-        case "zsh":
-          process.stdout.write(generateZsh());
-          break;
-        case "bash":
-          process.stdout.write(generateBash());
-          break;
-        case "fish":
-          process.stdout.write(generateFish());
-          break;
-        default:
-          process.stderr.write(
-            `Unknown shell: ${shell}. Supported: zsh, bash, fish\n`
-          );
-          process.exit(1);
-      }
-    });
-}
-
-const SUBCOMMANDS = ["text", "image", "video", "models", "completions", "help"];
-const GLOBAL_FLAGS = [
-  "--model",
-  "--output",
-  "--count",
-  "--concurrency",
-  "--quiet",
-  "--json",
-  "--help",
-  "--version",
-];
-const TEXT_FLAGS = ["--format", "--system", "--max-tokens", "--temperature"];
-const IMAGE_FLAGS = [
-  "--size",
-  "--aspect-ratio",
-  "--quality",
-  "--style",
-  "--no-preview",
-];
-const VIDEO_FLAGS = ["--aspect-ratio", "--duration", "--no-preview"];
-const MODEL_FLAGS = ["--type", "--provider", "--json", "--help"];
-
-const ALL_MODELS = [
-  ...FALLBACK_TEXT_MODELS,
-  ...FALLBACK_IMAGE_MODELS,
-  ...FALLBACK_VIDEO_MODELS,
-];
-const MODEL_NAMES = ALL_MODELS.map((m) => m.slice(m.indexOf("/") + 1));
-
-function generateZsh(): string {
-  return `#compdef ai
-
-_ai() {
-  local -a subcommands
-  subcommands=(
-    'text:Generate text from a prompt'
-    'image:Generate an image from a prompt'
-    'video:Generate a video from a prompt'
-    'models:List available models'
-    'completions:Output shell completion script'
-    'help:Display help'
-  )
-
-  local -a models
-  models=(${ALL_MODELS.join(" ")})
-
-  local -a model_names
-  model_names=(${MODEL_NAMES.join(" ")})
-
-  _arguments -C \\
-    '1:command:->cmd' \\
-    '*::arg:->args'
-
-  case $state in
-    cmd)
-      _describe 'command' subcommands
-      ;;
-    args)
-      case $words[1] in
-        text)
-          _arguments \\
-            '-m[Model ID]:model:($models $model_names)' \\
-            '--model[Model ID]:model:($models $model_names)' \\
-            '-o[Output path]:file:_files' \\
-            '--output[Output path]:file:_files' \\
-            '-f[Format]:format:(md txt)' \\
-            '--format[Format]:format:(md txt)' \\
-            '-n[Count]:count:' \\
-            '--count[Count]:count:' \\
-            '-p[Concurrency]:concurrency:' \\
-            '--concurrency[Concurrency]:concurrency:' \\
-            '-s[System prompt]:system:' \\
-            '--system[System prompt]:system:' \\
-            '--max-tokens[Max tokens]:tokens:' \\
-            '-t[Temperature]:temp:' \\
-            '--temperature[Temperature]:temp:' \\
-            '-q[Quiet]' \\
-            '--quiet[Quiet]' \\
-            '--json[JSON output]' \\
-            '*:prompt:'
-          ;;
-        image)
-          _arguments \\
-            '-m[Model ID]:model:($models $model_names)' \\
-            '--model[Model ID]:model:($models $model_names)' \\
-            '-o[Output path]:file:_files' \\
-            '--output[Output path]:file:_files' \\
-            '-n[Count]:count:' \\
-            '--count[Count]:count:' \\
-            '-p[Concurrency]:concurrency:' \\
-            '--concurrency[Concurrency]:concurrency:' \\
-            '--size[Size]:size:' \\
-            '--aspect-ratio[Aspect ratio]:ratio:' \\
-            '--quality[Quality]:quality:(standard hd)' \\
-            '--style[Style]:style:(vivid natural)' \\
-            '--no-preview[Disable inline image preview]' \\
-            '-q[Quiet]' \\
-            '--quiet[Quiet]' \\
-            '--json[JSON output]' \\
-            '*:prompt:'
-          ;;
-        video)
-          _arguments \\
-            '-m[Model ID]:model:($models $model_names)' \\
-            '--model[Model ID]:model:($models $model_names)' \\
-            '-o[Output path]:file:_files' \\
-            '--output[Output path]:file:_files' \\
-            '-n[Count]:count:' \\
-            '--count[Count]:count:' \\
-            '-p[Concurrency]:concurrency:' \\
-            '--concurrency[Concurrency]:concurrency:' \\
-            '--aspect-ratio[Aspect ratio]:ratio:' \\
-            '--duration[Duration]:seconds:' \\
-            '--no-preview[Disable inline video frame preview]' \\
-            '-q[Quiet]' \\
-            '--quiet[Quiet]' \\
-            '--json[JSON output]' \\
-            '*:prompt:'
-          ;;
-        models)
-          _arguments \\
-            '--type[Filter by type]:type:(text image video)' \\
-            '--provider[Filter by provider]:provider:' \\
-            '--json[JSON output]'
-          ;;
-        completions)
-          _arguments '1:shell:(zsh bash fish)'
-          ;;
-      esac
-      ;;
-  esac
-}
-
-_ai "$@"
-`;
-}
-
-function generateBash(): string {
-  return `_ai_completions() {
-  local cur prev subcmd
-  COMPREPLY=()
-  cur="\${COMP_WORDS[COMP_CWORD]}"
-  prev="\${COMP_WORDS[COMP_CWORD-1]}"
-  subcmd="\${COMP_WORDS[1]}"
-
-  if [[ \${COMP_CWORD} -eq 1 ]]; then
-    COMPREPLY=($(compgen -W "${SUBCOMMANDS.join(" ")}" -- "$cur"))
-    return
-  fi
-
-  case "$prev" in
-    -m|--model)
-      COMPREPLY=($(compgen -W "${ALL_MODELS.join(" ")} ${MODEL_NAMES.join(" ")}" -- "$cur"))
-      return
-      ;;
-    -o|--output)
-      COMPREPLY=($(compgen -f -- "$cur"))
-      return
-      ;;
-    -f|--format)
-      COMPREPLY=($(compgen -W "md txt" -- "$cur"))
-      return
-      ;;
-    --quality)
-      COMPREPLY=($(compgen -W "standard hd" -- "$cur"))
-      return
-      ;;
-    --style)
-      COMPREPLY=($(compgen -W "vivid natural" -- "$cur"))
-      return
-      ;;
-    --type)
-      COMPREPLY=($(compgen -W "text image video" -- "$cur"))
-      return
-      ;;
-  esac
-
-  case "$subcmd" in
-    text)
-      COMPREPLY=($(compgen -W "${[...GLOBAL_FLAGS, ...TEXT_FLAGS].join(" ")}" -- "$cur"))
-      ;;
-    image)
-      COMPREPLY=($(compgen -W "${[...GLOBAL_FLAGS, ...IMAGE_FLAGS].join(" ")}" -- "$cur"))
-      ;;
-    video)
-      COMPREPLY=($(compgen -W "${[...GLOBAL_FLAGS, ...VIDEO_FLAGS].join(" ")}" -- "$cur"))
-      ;;
-    models)
-      COMPREPLY=($(compgen -W "${MODEL_FLAGS.join(" ")}" -- "$cur"))
-      ;;
-    completions)
-      COMPREPLY=($(compgen -W "zsh bash fish" -- "$cur"))
-      ;;
-  esac
-}
-
-complete -F _ai_completions ai
-`;
-}
-
-function generateFish(): string {
-  const lines: string[] = [];
-  lines.push("# ai completions for fish");
-  lines.push("");
-
-  for (const sub of SUBCOMMANDS) {
-    lines.push(`complete -c ai -n '__fish_use_subcommand' -a '${sub}'`);
-  }
-  lines.push("");
-
-  const SHORT_FLAG_MAP: Record<string, string> = {
-    "--model": "m",
-    "--output": "o",
-    "--count": "n",
-    "--concurrency": "p",
-    "--quiet": "q",
-    "--format": "f",
-    "--system": "s",
-    "--temperature": "t",
-  };
-
-  const addFlags = (sub: string, flags: string[]) => {
-    for (const flag of flags) {
-      const name = flag.replace(/^--/, "");
-      const short = SHORT_FLAG_MAP[flag];
-      const shortPart = short ? ` -s ${short}` : "";
-      lines.push(
-        `complete -c ai -n '__fish_seen_subcommand_from ${sub}'${shortPart} -l '${name}'`
-      );
-    }
-  };
-
-  addFlags("text", [...GLOBAL_FLAGS, ...TEXT_FLAGS]);
-  addFlags("image", [...GLOBAL_FLAGS, ...IMAGE_FLAGS]);
-  addFlags("video", [...GLOBAL_FLAGS, ...VIDEO_FLAGS]);
-  addFlags("models", MODEL_FLAGS);
-
-  lines.push("");
-  lines.push(
-    `complete -c ai -n '__fish_seen_subcommand_from completions' -a 'zsh bash fish'`
-  );
-  lines.push("");
-
-  const modelCompletions = ALL_MODELS.concat(MODEL_NAMES);
-  lines.push(
-    `complete -c ai -n '__fish_seen_subcommand_from text image video' -s m -l model -a '${modelCompletions.join(" ")}'`
-  );
-  lines.push(
-    `complete -c ai -n '__fish_seen_subcommand_from text' -s f -l format -a 'md txt'`
-  );
-  lines.push(
-    `complete -c ai -n '__fish_seen_subcommand_from image' -l quality -a 'standard hd'`
-  );
-  lines.push(
-    `complete -c ai -n '__fish_seen_subcommand_from image' -l style -a 'vivid natural'`
-  );
-  lines.push(
-    `complete -c ai -n '__fish_seen_subcommand_from models' -l type -a 'text image video'`
-  );
-  lines.push("");
-
-  return lines.join("\n");
-}