ai-cli 0.0.12 → 0.1.0

package/src/lib/jobs.ts

@@ -0,0 +1,192 @@
+import {
+  supportsKittyGraphics,
+  displayImage,
+  displayVideoFrame,
+} from "./kitty.js";
+import type { OutputFormat } from "./output.js";
+import { writeOutput } from "./output.js";
+import { pMap } from "./p-map.js";
+import { Progress, MultiProgress, formatElapsed } from "./progress.js";
+
+export interface Job {
+  modelId: string;
+  label: string;
+  index: number;
+}
+
+export interface RunJobsOptions {
+  noun: string;
+  format: OutputFormat;
+  outputPath?: string;
+  quiet?: boolean;
+  json?: boolean;
+  concurrency: number;
+  display?: boolean;
+}
+
+export function buildJobs(models: string[], countPerModel: number): Job[] {
+  let jobIndex = 0;
+  return models.flatMap((modelId) =>
+    Array.from({ length: countPerModel }, (_, i) => ({
+      modelId,
+      label: models.length > 1 ? `${modelId} #${i + 1}` : `#${i + 1}`,
+      index: jobIndex++,
+    }))
+  );
+}
+
+export interface RunJobsResult {
+  total: number;
+  failed: number;
+}
+
+export async function runJobs(
+  jobs: Job[],
+  generate: (modelId: string) => Promise<Buffer | string>,
+  opts: RunJobsOptions
+): Promise<RunJobsResult> {
+  const { noun, format, outputPath, quiet, json, concurrency, display } = opts;
+
+  if (jobs.length === 1) {
+    const { modelId } = jobs[0];
+    const progress = new Progress(quiet);
+    const start = Date.now();
+    progress.start(`Generating ${noun} with ${modelId}`);
+
+    try {
+      const data = await generate(modelId);
+      const elapsed = Date.now() - start;
+      progress.stop(`Generated ${noun} with ${modelId}`);
+
+      if (json) {
+        const path = await writeOutput({
+          data,
+          format,
+          outputPath,
+          quiet: true,
+          display,
+        });
+        const meta = {
+          elapsed_ms: elapsed,
+          count: 1,
+          results: [
+            {
+              index: 1,
+              model: modelId,
+              elapsed_ms: elapsed,
+              success: true,
+              file: path,
+            },
+          ],
+        };
+        process.stdout.write(JSON.stringify(meta, null, 2) + "\n");
+      } else {
+        await writeOutput({ data, format, outputPath, quiet, display });
+      }
+    } catch (err) {
+      progress.stop();
+      throw err;
+    }
+    return { total: 1, failed: 0 };
+  }
+
+  const multi = new MultiProgress(quiet);
+  const start = Date.now();
+  const shouldDisplay =
+    display !== false &&
+    (format === "image" || format === "video") &&
+    process.stdout.isTTY &&
+    supportsKittyGraphics();
+
+  const lineIdxs = jobs.map((j) =>
+    multi.addLine(`Generating ${noun} ${j.label} with ${j.modelId}`)
+  );
+
+  const results: {
+    index: number;
+    model: string;
+    success: boolean;
+    elapsed_ms: number;
+    file: string | null;
+  }[] = [];
+  const pendingDisplayBuffers: Buffer[] = [];
+
+  await pMap(
+    jobs,
+    async (job, i) => {
+      multi.startLine(lineIdxs[i]);
+      const genStart = Date.now();
+      try {
+        const data = await generate(job.modelId);
+        const genElapsed = Date.now() - genStart;
+        const suffix = `${i + 1}`;
+        const path = await writeOutput({
+          data,
+          format,
+          outputPath,
+          suffix,
+          quiet: true,
+          display: false,
+        });
+        if (shouldDisplay && Buffer.isBuffer(data))
+          pendingDisplayBuffers.push(data);
+        const savedMsg = path
+          ? `Saved to ${path}`
+          : `${noun[0].toUpperCase()}${noun.slice(1)} ${job.label} written to stdout`;
+        multi.completeLine(
+          lineIdxs[i],
+          `${savedMsg} (${formatElapsed(genElapsed)})`
+        );
+        results.push({
+          index: i,
+          model: job.modelId,
+          success: true,
+          elapsed_ms: genElapsed,
+          file: path,
+        });
+      } catch (err: unknown) {
+        const genElapsed = Date.now() - genStart;
+        const msg = err instanceof Error ? err.message : String(err);
+        multi.completeLine(
+          lineIdxs[i],
+          `${noun[0].toUpperCase()}${noun.slice(1)} ${job.label} failed: ${msg} (${formatElapsed(genElapsed)})`
+        );
+        results.push({
+          index: i,
+          model: job.modelId,
+          success: false,
+          elapsed_ms: genElapsed,
+          file: null,
+        });
+      }
+    },
+    concurrency
+  );
+
+  if (json) {
+    const totalElapsed = Date.now() - start;
+    const meta = {
+      elapsed_ms: totalElapsed,
+      count: results.filter((r) => r.success).length,
+      results: results.map((r) => ({
+        index: r.index + 1,
+        model: r.model,
+        elapsed_ms: r.elapsed_ms,
+        success: r.success,
+        file: r.file,
+      })),
+    };
+    process.stdout.write(JSON.stringify(meta, null, 2) + "\n");
+  }
+
+  for (const buf of pendingDisplayBuffers) {
+    if (format === "video") {
+      await displayVideoFrame(buf);
+    } else {
+      displayImage(buf);
+    }
+  }
+
+  const failCount = results.filter((r) => !r.success).length;
+  return { total: results.length, failed: failCount };
+}

package/src/lib/kitty.ts

@@ -0,0 +1,55 @@
+const SUPPORTED_TERMS = new Set(["xterm-kitty"]);
+
+const SUPPORTED_TERM_PROGRAMS = new Set([
+  "kitty",
+  "ghostty",
+  "wezterm",
+  "warpterminal",
+]);
+
+export function supportsKittyGraphics(): boolean {
+  if (process.env.AI_CLI_PREVIEW === "0") return false;
+  if (process.env.AI_CLI_PREVIEW === "1") return true;
+
+  const term = process.env.TERM ?? "";
+  if (SUPPORTED_TERMS.has(term)) return true;
+
+  const termProgram = (process.env.TERM_PROGRAM ?? "").toLowerCase();
+  if (SUPPORTED_TERM_PROGRAMS.has(termProgram)) return true;
+
+  const lcTerminal = (process.env.LC_TERMINAL ?? "").toLowerCase();
+  if (lcTerminal === "iterm2") return true;
+
+  return false;
+}
+
+import { decodeIDR } from "./h264-wasm.js";
+import { extractKeyframe } from "./mp4.js";
+import { encodePNG } from "./png.js";
+
+const CHUNK_SIZE = 4096;
+
+export async function displayVideoFrame(buf: Buffer): Promise<void> {
+  try {
+    const kf = extractKeyframe(new Uint8Array(buf));
+    if (!kf) return;
+    const frame = await decodeIDR(kf.sps, kf.pps, kf.sliceData);
+    if (!frame) return;
+    const png = encodePNG(frame.yuv, frame.width, frame.height);
+    displayImage(png);
+  } catch {
+    // Preview is best-effort; skip silently on any failure
+  }
+}
+
+export function displayImage(buf: Buffer): void {
+  const encoded = buf.toString("base64");
+  for (let i = 0; i < encoded.length; i += CHUNK_SIZE) {
+    const chunk = encoded.slice(i, i + CHUNK_SIZE);
+    const isLast = i + CHUNK_SIZE >= encoded.length;
+    const control =
+      i === 0 ? `a=T,f=100,m=${isLast ? 0 : 1}` : `m=${isLast ? 0 : 1}`;
+    process.stderr.write(`\x1b_G${control};${chunk}\x1b\\`);
+  }
+  process.stderr.write("\n");
+}

package/src/lib/models.test.ts

@@ -0,0 +1,197 @@
+import { describe, expect, mock, test } from "bun:test";
+
+import {
+  resolveModels,
+  fetchGatewayModels,
+  FALLBACK_TEXT_MODELS,
+  FALLBACK_IMAGE_MODELS,
+  FALLBACK_VIDEO_MODELS,
+} from "./models.js";
+
+describe("resolveModels", () => {
+  test("returns default when no user model", () => {
+    expect(resolveModels("text")[0]).toContain("/");
+    expect(resolveModels("image")[0]).toContain("/");
+    expect(resolveModels("video")[0]).toContain("/");
+  });
+
+  test("returns fully-qualified model as-is", () => {
+    expect(resolveModels("text", "openai/gpt-4")).toEqual(["openai/gpt-4"]);
+    expect(resolveModels("image", "openai/gpt-image-1")).toEqual([
+      "openai/gpt-image-1",
+    ]);
+  });
+
+  test("expands short image model names", () => {
+    expect(resolveModels("image", "gpt-image-1")).toEqual([
+      "openai/gpt-image-1",
+    ]);
+    expect(resolveModels("image", "flux-2-pro")).toEqual(["bfl/flux-2-pro"]);
+  });
+
+  test("expands short video model names", () => {
+    expect(resolveModels("video", "seedance-2.0")).toEqual([
+      "bytedance/seedance-2.0",
+    ]);
+  });
+
+  test("expands short text model names", () => {
+    expect(resolveModels("text", "gpt-5.5")).toEqual(["openai/gpt-5.5"]);
+    expect(resolveModels("text", "o3")).toEqual(["openai/o3"]);
+  });
+
+  test("returns unknown short names as-is for text", () => {
+    expect(resolveModels("text", "gpt-image-1")).toEqual(["gpt-image-1"]);
+    expect(resolveModels("text", "my-model")).toEqual(["my-model"]);
+  });
+
+  test("returns unknown short names as-is for image/video", () => {
+    expect(resolveModels("image", "nonexistent-model")).toEqual([
+      "nonexistent-model",
+    ]);
+  });
+});
+
+describe("resolveModels multi", () => {
+  test("returns default when no user model", () => {
+    const result = resolveModels("text");
+    expect(result).toHaveLength(1);
+    expect(result[0]).toContain("/");
+  });
+
+  test("splits comma-separated models", () => {
+    const result = resolveModels("image", "openai/gpt-image-1,bfl/flux-2-pro");
+    expect(result).toEqual(["openai/gpt-image-1", "bfl/flux-2-pro"]);
+  });
+
+  test("trims whitespace around model names", () => {
+    const result = resolveModels(
+      "image",
+      "openai/gpt-image-1 , bfl/flux-2-pro"
+    );
+    expect(result).toEqual(["openai/gpt-image-1", "bfl/flux-2-pro"]);
+  });
+
+  test("expands short names in comma list", () => {
+    const result = resolveModels("image", "gpt-image-1,flux-2-pro");
+    expect(result).toEqual(["openai/gpt-image-1", "bfl/flux-2-pro"]);
+  });
+
+  test("filters empty segments from trailing comma", () => {
+    const result = resolveModels("image", "openai/gpt-image-1,");
+    expect(result).toEqual(["openai/gpt-image-1"]);
+  });
+
+  test("falls back to default when all segments are empty", () => {
+    const result = resolveModels("image", ",,,");
+    expect(result).toHaveLength(1);
+    expect(result[0]).toContain("/");
+  });
+});
+
+describe("fetchGatewayModels", () => {
+  test("partitions models by modelType", async () => {
+    const { gateway } = await import("ai");
+    const original = gateway.getAvailableModels;
+    gateway.getAvailableModels = mock(() =>
+      Promise.resolve({
+        models: [
+          { id: "openai/gpt-5", name: "GPT 5", modelType: "language" },
+          {
+            id: "openai/gpt-image-2",
+            name: "GPT Image 2",
+            description: "Image gen",
+            modelType: "image",
+          },
+          { id: "google/veo-3.0", name: "Veo 3", modelType: "video" },
+          {
+            id: "openai/text-embedding-3",
+            name: "Embedding",
+            modelType: "embedding",
+          },
+        ],
+      })
+    ) as unknown as typeof gateway.getAvailableModels;
+
+    try {
+      const result = await fetchGatewayModels();
+
+      expect(result.text).toEqual([
+        { id: "openai/gpt-5", name: "GPT 5", description: undefined },
+      ]);
+      expect(result.image).toEqual([
+        {
+          id: "openai/gpt-image-2",
+          name: "GPT Image 2",
+          description: "Image gen",
+        },
+      ]);
+      expect(result.video).toEqual([
+        { id: "google/veo-3.0", name: "Veo 3", description: undefined },
+      ]);
+    } finally {
+      gateway.getAvailableModels = original;
+    }
+  });
+
+  test("falls back to static lists on gateway error", async () => {
+    const { gateway } = await import("ai");
+    const original = gateway.getAvailableModels;
+    gateway.getAvailableModels = mock(() =>
+      Promise.reject(new Error("network error"))
+    ) as typeof gateway.getAvailableModels;
+
+    try {
+      const result = await fetchGatewayModels();
+
+      expect(result.text.map((m) => m.id)).toEqual(FALLBACK_TEXT_MODELS);
+      expect(result.image.map((m) => m.id)).toEqual(FALLBACK_IMAGE_MODELS);
+      expect(result.video.map((m) => m.id)).toEqual(FALLBACK_VIDEO_MODELS);
+    } finally {
+      gateway.getAvailableModels = original;
+    }
+  });
+
+  test("uses fallbacks when gateway returns no image/video models", async () => {
+    const { gateway } = await import("ai");
+    const original = gateway.getAvailableModels;
+    gateway.getAvailableModels = mock(() =>
+      Promise.resolve({
+        models: [{ id: "openai/gpt-5", name: "GPT 5", modelType: "language" }],
+      })
+    ) as unknown as typeof gateway.getAvailableModels;
+
+    try {
+      const result = await fetchGatewayModels();
+
+      expect(result.text).toHaveLength(1);
+      expect(result.text[0].id).toBe("openai/gpt-5");
+      expect(result.image.map((m) => m.id)).toEqual(FALLBACK_IMAGE_MODELS);
+      expect(result.video.map((m) => m.id)).toEqual(FALLBACK_VIDEO_MODELS);
+    } finally {
+      gateway.getAvailableModels = original;
+    }
+  });
+
+  test("uses text fallbacks when gateway returns no text models", async () => {
+    const { gateway } = await import("ai");
+    const original = gateway.getAvailableModels;
+    gateway.getAvailableModels = mock(() =>
+      Promise.resolve({
+        models: [
+          { id: "openai/gpt-image-2", name: "GPT Image 2", modelType: "image" },
+        ],
+      })
+    ) as unknown as typeof gateway.getAvailableModels;
+
+    try {
+      const result = await fetchGatewayModels();
+
+      expect(result.text.map((m) => m.id)).toEqual(FALLBACK_TEXT_MODELS);
+      expect(result.image).toHaveLength(1);
+      expect(result.video.map((m) => m.id)).toEqual(FALLBACK_VIDEO_MODELS);
+    } finally {
+      gateway.getAvailableModels = original;
+    }
+  });
+});

package/src/lib/models.ts

@@ -0,0 +1,163 @@
+import { gateway } from "ai";
+
+export type Modality = "text" | "image" | "video";
+
+const DEFAULTS: Record<Modality, string> = {
+  text: process.env.AI_CLI_TEXT_MODEL ?? "openai/gpt-5.5",
+  image: process.env.AI_CLI_IMAGE_MODEL ?? "openai/gpt-image-2",
+  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",
+];
+
+export interface ModelEntry {
+  id: string;
+  name?: string;
+  description?: string;
+}
+
+export interface GatewayModels {
+  text: ModelEntry[];
+  image: ModelEntry[];
+  video: ModelEntry[];
+}
+
+const MODEL_TYPE_TO_MODALITY: Record<string, Modality> = {
+  language: "text",
+  image: "image",
+  video: "video",
+};
+
+export async function fetchGatewayModels(): Promise<GatewayModels> {
+  const result: GatewayModels = { text: [], image: [], video: [] };
+
+  try {
+    const { models } = await gateway.getAvailableModels();
+    for (const m of models) {
+      const modality =
+        MODEL_TYPE_TO_MODALITY[(m as { modelType?: string }).modelType ?? ""];
+      if (!modality) continue;
+      result[modality].push({
+        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 }));
+  }
+
+  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;
+
+  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);
+
+  for (const list of knownLists) {
+    for (const fullId of list) {
+      const name = fullId.slice(fullId.indexOf("/") + 1);
+      if (name === input) return fullId;
+    }
+  }
+
+  return input;
+}
+
+export function resolveModels(
+  modality: Modality,
+  userModel?: string
+): string[] {
+  if (!userModel) return [DEFAULTS[modality]];
+  const models = userModel
+    .split(",")
+    .map((m) => m.trim())
+    .filter(Boolean)
+    .map((m) => expandModelId(m, modality));
+  return models.length > 0 ? models : [DEFAULTS[modality]];
+}